diff --git a/.gitattributes b/.gitattributes index 37ab2e466a..65a6b946ab 100644 --- a/.gitattributes +++ b/.gitattributes @@ -9,3 +9,4 @@ *.mp4 filter=lfs diff=lfs merge=lfs -text *.pt filter=lfs diff=lfs merge=lfs -text *.jit filter=lfs diff=lfs merge=lfs -text +*.hdf5 filter=lfs diff=lfs merge=lfs -text diff --git a/.github/workflows/docs.yaml b/.github/workflows/docs.yaml index 8665194686..213bfc8858 100644 --- a/.github/workflows/docs.yaml +++ b/.github/workflows/docs.yaml @@ -1,42 +1,65 @@ name: Build & deploy docs -on: [push] +on: + push: + branches: + - main + pull_request: + types: [opened, synchronize, reopened] + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true jobs: + check-secrets: + name: Check secrets + runs-on: ubuntu-latest + outputs: + trigger-deploy: ${{ steps.trigger-deploy.outputs.defined }} + steps: + - id: trigger-deploy + env: + REPO_NAME: ${{ secrets.REPO_NAME }} + BRANCH_REF: ${{ secrets.BRANCH_REF }} + if: "${{ github.repository == env.REPO_NAME && github.ref == env.BRANCH_REF }}" + run: echo "defined=true" >> "$GITHUB_OUTPUT" + build-docs: name: Build Docs runs-on: ubuntu-latest + needs: [check-secrets] steps: - - name: Checkout code - uses: actions/checkout@v2 + - name: Checkout code + uses: actions/checkout@v2 - - name: Setup python - uses: actions/setup-python@v2 - with: - python-version: "3.10" - architecture: x64 + - name: Setup python + uses: actions/setup-python@v2 + with: + python-version: "3.10" + architecture: x64 - - name: Install dev requirements - working-directory: ./docs - run: pip install -r requirements.txt + - name: Install dev requirements + working-directory: ./docs + run: pip install -r requirements.txt - - name: Generate docs - working-directory: ./docs - run: make html + - name: Check branch docs building + working-directory: ./docs + if: needs.check-secrets.outputs.trigger-deploy != 'true' + run: make current-docs - check-secrets: - name: Check secrets - runs-on: ubuntu-latest - outputs: - trigger-deploy: ${{ steps.trigger-deploy.outputs.defined }} - steps: - - id: trigger-deploy - env: - REPO_NAME: ${{ secrets.REPO_NAME }} - BRANCH_REF: ${{ secrets.BRANCH_REF }} - if: "${{ github.repository == env.REPO_NAME && github.ref == env.BRANCH_REF }}" - run: echo "defined=true" >> "$GITHUB_OUTPUT" + - name: Generate multi-version docs + working-directory: ./docs + run: | + git fetch --prune --unshallow --tags + make multi-docs + + - name: Upload docs artifact + uses: actions/upload-artifact@v4 + with: + name: docs-html + path: ./docs/_build deploy-docs-preview: name: Deploy Docs Preview @@ -55,8 +78,14 @@ jobs: if: needs.check-secrets.outputs.trigger-deploy == 'true' steps: - - name: Deploy to gh-pages - uses: peaceiris/actions-gh-pages@v3 - with: - github_token: ${{ secrets.GITHUB_TOKEN }} - publish_dir: ./docs/_build/html + - name: Download docs artifact + uses: actions/download-artifact@v4 + with: + name: docs-html + path: ./docs/_build + + - name: Deploy to gh-pages + uses: peaceiris/actions-gh-pages@v3 + with: + github_token: ${{ secrets.GITHUB_TOKEN }} + publish_dir: ./docs/_build diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 9dd7c6d9b2..8bc6718254 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -38,7 +38,7 @@ repos: - id: pyupgrade args: ["--py310-plus"] # FIXME: This is a hack because Pytorch does not like: torch.Tensor | dict aliasing - exclude: "source/extensions/omni.isaac.lab/omni/isaac/lab/envs/common.py" + exclude: "source/extensions/omni.isaac.lab/omni/isaac/lab/envs/common.py|source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/image_plot.py" - repo: https://github.com/codespell-project/codespell rev: v2.2.6 hooks: diff --git a/.vscode/tools/settings.template.json b/.vscode/tools/settings.template.json index 41e7130cda..c794051784 100644 --- a/.vscode/tools/settings.template.json +++ b/.vscode/tools/settings.template.json @@ -45,7 +45,6 @@ "teleoperation", "xform", "numpy", - "tensordict", "flatcache", "physx", "dpad", diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index b75e4e5bc1..b837c0df92 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -5,4 +5,4 @@ the framework more mature and useful for everyone. These may happen in forms of design proposals and more. For general information on how to contribute see -. +. diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 244b910786..5ae174987b 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -6,7 +6,7 @@ To see the full list of contributors, please check the revision history in the s Guidelines for modifications: -* Please keep the lists sorted alphabetically. +* Please keep the **lists sorted alphabetically**. * Names should be added to this file as: *individual names* or *organizations*. * E-mail addresses are tracked elsewhere to avoid spam. @@ -19,6 +19,7 @@ Guidelines for modifications: --- +* Antonio Serrano-Muñoz * David Hoeller * Farbod Farshidian * Hunter Hansen @@ -28,6 +29,7 @@ Guidelines for modifications: * Mayank Mittal * Nikita Rudin * Pascal Roth +* Sheikh Dawood ## Contributors @@ -35,40 +37,53 @@ Guidelines for modifications: * Amr Mousa * Andrej Orsula * Anton Bjørndahl Mortensen -* Antonio Serrano-Muñoz * Arjun Bhardwaj * Brayden Zhang * Calvin Yu * Chenyu Yang +* CY (Chien-Ying) Chen * David Yang +* Dorsa Rohani +* Felix Yu * Gary Lvov * Giulio Romualdi +* Haoran Zhou * HoJin Jeon +* Hongwei Xiong +* Iretiayo Akinola +* Jan Kerner * Jean Tampon * Jia Lin Yuan +* Jinghuan Shang * Jingzhou Liu * Johnson Sun * Kaixi Bao * Kourosh Darvish * Lionel Gulich +* Louis Le Lay * Lorenz Wellhausen * Masoud Moghani * Michael Gussert +* Michael Noseworthy * Muhong Guo * Nuralem Abizov +* Oyindamola Omotuyi * Özhan Özen +* Peter Du +* Qian Wan * Qinxi Yu * René Zurbrügg * Ritvik Singh * Rosario Scalise +* Ryley McCarroll * Shafeef Omar * Vladimir Fokow * Wei Yang * Xavier Nal * Yang Jin +* Yujian Zhang * Zhengyu Zhang * Ziqi Fan -* Qian Wan ## Acknowledgements @@ -78,3 +93,4 @@ Guidelines for modifications: * Gavriel State * Hammad Mazhar * Marco Hutter +* Yashraj Narang diff --git a/README.md b/README.md index 590fd5aeb7..a46cd89d8f 100644 --- a/README.md +++ b/README.md @@ -13,23 +13,38 @@ [![License](https://img.shields.io/badge/license-BSD--3-yellow.svg)](https://opensource.org/licenses/BSD-3-Clause) -**Isaac Lab** is a unified and modular framework for robot learning that aims to simplify common workflows -in robotics research (such as RL, learning from demonstrations, and motion planning). It is built upon -[NVIDIA Isaac Sim](https://docs.omniverse.nvidia.com/isaacsim/latest/overview.html) to leverage the latest -simulation capabilities for photo-realistic scenes and fast and accurate simulation. +**Isaac Lab** is a GPU-accelerated, open-source framework designed to unify and simplify robotics research workflows, such as reinforcement learning, imitation learning, and motion planning. Built on [NVIDIA Isaac Sim](https://docs.omniverse.nvidia.com/isaacsim/latest/overview.html), it combines fast and accurate physics and sensor simulation, making it an ideal choice for sim-to-real transfer in robotics. + +Isaac Lab provides developers with a range of essential features for accurate sensor simulation, such as RTX-based cameras, LIDAR, or contact sensors. The framework's GPU acceleration enables users to run complex simulations and computations faster, which is key for iterative processes like reinforcement learning and data-intensive tasks. Moreover, Isaac Lab can run locally or be distributed across the cloud, offering flexibility for large-scale deployments. + +## Key Features + +Isaac Lab offers a comprehensive set of tools and environments designed to facilitate robot learning: +- **Robots**: A diverse collection of robots, from manipulators, quadrupeds, to humanoids, with 16 commonly available models. +- **Environments**: Ready-to-train implementations of more than 30 environments, which can be trained with popular reinforcement learning frameworks such as RSL RL, SKRL, RL Games, or Stable Baselines. We also support multi-agent reinforcement learning. +- **Physics**: Rigid bodies, articulated systems, deformable objects +- **Sensors**: RGB/depth/segmentation cameras, camera annotations, IMU, contact sensors, ray casters. + + +## Getting Started + +Our [documentation page](https://isaac-sim.github.io/IsaacLab) provides everything you need to get started, including detailed tutorials and step-by-step guides. Follow these links to learn more about: + +- [Installation steps](https://isaac-sim.github.io/IsaacLab/main/source/setup/installation/index.html#local-installation) +- [Reinforcement learning](https://isaac-sim.github.io/IsaacLab/main/source/overview/reinforcement-learning/rl_existing_scripts.html) +- [Tutorials](https://isaac-sim.github.io/IsaacLab/main/source/tutorials/index.html) +- [Available environments](https://isaac-sim.github.io/IsaacLab/main/source/overview/environments.html) -Please refer to our [documentation page](https://isaac-sim.github.io/IsaacLab) to learn more about the -installation steps, features, tutorials, and how to set up your project with Isaac Lab. ## Contributing to Isaac Lab We wholeheartedly welcome contributions from the community to make this framework mature and useful for everyone. These may happen as bug reports, feature requests, or code contributions. For details, please check our -[contribution guidelines](https://isaac-sim.github.io/IsaacLab/source/refs/contributing.html). +[contribution guidelines](https://isaac-sim.github.io/IsaacLab/main/source/refs/contributing.html). ## Troubleshooting -Please see the [troubleshooting](https://isaac-sim.github.io/IsaacLab/source/refs/troubleshooting.html) section for +Please see the [troubleshooting](https://isaac-sim.github.io/IsaacLab/main/source/refs/troubleshooting.html) section for common fixes or [submit an issue](https://github.com/isaac-sim/IsaacLab/issues). For issues related to Isaac Sim, we recommend checking its [documentation](https://docs.omniverse.nvidia.com/app_isaacsim/app_isaacsim/overview.html) diff --git a/VERSION b/VERSION index 26aaba0e86..f0bb29e763 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -1.2.0 +1.3.0 diff --git a/docker/cluster/cluster_interface.sh b/docker/cluster/cluster_interface.sh index fd325e4ca4..6684d835c9 100755 --- a/docker/cluster/cluster_interface.sh +++ b/docker/cluster/cluster_interface.sh @@ -16,8 +16,18 @@ SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" #== # Functions #== +# Function to display warnings in red +display_warning() { + echo -e "\033[31mWARNING: $1\033[0m" +} + +# Helper function to compare version numbers +version_gte() { + # Returns 0 if the first version is greater than or equal to the second, otherwise 1 + [ "$(printf '%s\n' "$1" "$2" | sort -V | head -n 1)" == "$2" ] +} + # Function to check docker versions -# If docker version is more than 25, the script errors out. check_docker_version() { # check if docker is installed if ! command -v docker &> /dev/null; then @@ -28,12 +38,17 @@ check_docker_version() { docker_version=$(docker --version | awk '{ print $3 }') apptainer_version=$(apptainer --version | awk '{ print $3 }') - # Check if version is above 25.xx - if [ "$(echo "${docker_version}" | cut -d '.' -f 1)" -ge 25 ]; then - echo "[ERROR]: Docker version ${docker_version} is not compatible with Apptainer version ${apptainer_version}. Exiting." - exit 1 + # Check if Docker version is exactly 24.0.7 or Apptainer version is exactly 1.2.5 + if [ "$docker_version" = "24.0.7" ] && [ "$apptainer_version" = "1.2.5" ]; then + echo "[INFO]: Docker version ${docker_version} and Apptainer version ${apptainer_version} are tested and compatible." + + # Check if Docker version is >= 27.0.0 and Apptainer version is >= 1.3.4 + elif version_gte "$docker_version" "27.0.0" && version_gte "$apptainer_version" "1.3.4"; then + echo "[INFO]: Docker version ${docker_version} and Apptainer version ${apptainer_version} are tested and compatible." + + # Else, display a warning for non-tested versions else - echo "[INFO]: Building singularity with docker version: ${docker_version} and Apptainer version: ${apptainer_version}." + display_warning "Docker version ${docker_version} and Apptainer version ${apptainer_version} are non-tested versions. There could be issues, please try to update them. More info: https://isaac-sim.github.io/IsaacLab/source/deployment/cluster.html" fi } @@ -139,7 +154,7 @@ case $command in fi # Check if Docker image exists check_image_exists isaac-lab-$profile:latest - # Check if Docker version is greater than 25 + # Check docker and apptainer version check_docker_version # source env file to get cluster login and path information source $SCRIPT_DIR/.env.cluster @@ -162,7 +177,7 @@ case $command in job) if [ $# -ge 1 ]; then passed_profile=$1 - if [ -f ".env.$passed_profile" ]; then + if [ -f "$SCRIPT_DIR/../.env.$passed_profile" ]; then profile=$passed_profile shift fi diff --git a/docs/Makefile b/docs/Makefile index d4bb2cbb9e..ce33dad503 100644 --- a/docs/Makefile +++ b/docs/Makefile @@ -8,13 +8,11 @@ SPHINXBUILD ?= sphinx-build SOURCEDIR = . BUILDDIR = _build -# Put it first so that "make" without argument is like "make help". -help: - @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) +.PHONY: multi-docs +multi-docs: + @sphinx-multiversion "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) + @cp _redirect/index.html $(BUILDDIR)/index.html -.PHONY: help Makefile - -# Catch-all target: route all unknown targets to Sphinx using the new -# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). -%: Makefile - @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) +.PHONY: current-docs +current-docs: + @$(SPHINXBUILD) "$(SOURCEDIR)" "$(BUILDDIR)/current" $(SPHINXOPTS) diff --git a/docs/README.md b/docs/README.md index c154e0ad0f..69a77a48d9 100644 --- a/docs/README.md +++ b/docs/README.md @@ -1,30 +1,75 @@ # Building Documentation -We use [Sphinx](https://www.sphinx-doc.org/en/master/) with the [Book Theme](https://sphinx-book-theme.readthedocs.io/en/stable/) for maintaining the documentation. +We use [Sphinx](https://www.sphinx-doc.org/en/master/) with the [Book Theme](https://sphinx-book-theme.readthedocs.io/en/stable/) for maintaining and generating our documentation. -> **Note:** To build the documentation, we recommend creating a virtual environment to avoid any conflicts with system installed dependencies. +> **Note:** To avoid dependency conflicts, we strongly recommend using a Python virtual environment to isolate the required dependencies from your system's global Python environment. -Execute the following instructions to build the documentation (assumed from the top of the repository): +## Current-Version Documentation -1. Install the dependencies for [Sphinx](https://www.sphinx-doc.org/en/master/): +This section describes how to build the documentation for the current version of the project. - ```bash - # enter the location where this readme exists - cd docs - # install dependencies - pip install -r requirements.txt - ``` +
+Linux -2. Generate the documentation file via: +```bash +# 1. Navigate to the docs directory and install dependencies +cd docs +pip install -r requirements.txt - ```bash - # make the html version - make html - ``` +# 2. Build the current documentation +make current-docs -3. The documentation is now available at `docs/_build/html/index.html`: +# 3. Open the current docs +xdg-open _build/current/index.html +``` +
- ```bash - # open on default browser - xdg-open _build/html/index.html - ``` +
Windows + +```batch +:: 1. Navigate to the docs directory and install dependencies +cd docs +pip install -r requirements.txt + +:: 2. Build the current documentation +make current-docs + +:: 3. Open the current docs +start _build\current\index.html +``` +
+ + +## Multi-Version Documentation + +This section describes how to build the multi-version documentation, which includes previous tags and the main branch. + +
Linux + +```bash +# 1. Navigate to the docs directory and install dependencies +cd docs +pip install -r requirements.txt + +# 2. Build the multi-version documentation +make multi-docs + +# 3. Open the multi-version docs +xdg-open _build/index.html +``` +
+ +
Windows + +```batch +:: 1. Navigate to the docs directory and install dependencies +cd docs +pip install -r requirements.txt + +:: 2. Build the multi-version documentation +make multi-docs + +:: 3. Open the multi-version docs +start _build\index.html +``` +
diff --git a/docs/_redirect/index.html b/docs/_redirect/index.html new file mode 100644 index 0000000000..5208597ed1 --- /dev/null +++ b/docs/_redirect/index.html @@ -0,0 +1,8 @@ + + + + Redirecting to the latest Isaac Lab documentation + + + + diff --git a/docs/_templates/versioning.html b/docs/_templates/versioning.html new file mode 100644 index 0000000000..eb67be60e1 --- /dev/null +++ b/docs/_templates/versioning.html @@ -0,0 +1,21 @@ +{% if versions %} + +{% endif %} diff --git a/docs/conf.py b/docs/conf.py index a30e673280..0fccd611de 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -54,6 +54,8 @@ "sphinxcontrib.icon", "sphinx_copybutton", "sphinx_design", + "sphinx_tabs.tabs", # backwards compatibility for building docs on v1.0.0 + "sphinx_multiversion", ] # mathjax hacks @@ -115,7 +117,7 @@ # List of patterns, relative to source directory, that match files and # directories to ignore when looking for source files. # This pattern also affects html_static_path and html_extra_path. -exclude_patterns = ["_build", "Thumbs.db", ".DS_Store", "README.md", "licenses/*"] +exclude_patterns = ["_build", "_redirect", "_templates", "Thumbs.db", ".DS_Store", "README.md", "licenses/*"] # Mock out modules that are not available on RTD autodoc_mock_imports = [ @@ -190,7 +192,7 @@ import sphinx_book_theme -html_title = "Isaac Lab documentation" +html_title = "Isaac Lab Documentation" html_theme_path = [sphinx_book_theme.get_html_theme_path()] html_theme = "sphinx_book_theme" html_favicon = "source/_static/favicon.ico" @@ -213,7 +215,7 @@ "show_toc_level": 1, "use_sidenotes": True, "logo": { - "text": "Isaac Lab documentation", + "text": "Isaac Lab Documentation", "image_light": "source/_static/NVIDIA-logo-white.png", "image_dark": "source/_static/NVIDIA-logo-black.png", }, @@ -240,7 +242,19 @@ "icon_links_label": "Quick Links", } -html_sidebars = {"**": ["navbar-logo.html", "icon-links.html", "search-field.html", "sbt-sidebar-nav.html"]} +templates_path = [ + "_templates", +] + +# Whitelist pattern for remotes +smv_remote_whitelist = r"^.*$" +# Whitelist pattern for branches (set to None to ignore all branches) +smv_branch_whitelist = os.getenv("SMV_BRANCH_WHITELIST", r"^(main|devel)$") +# Whitelist pattern for tags (set to None to ignore all tags) +smv_tag_whitelist = os.getenv("SMV_TAG_WHITELIST", r"^v[1-9]\d*\.\d+\.\d+$") +html_sidebars = { + "**": ["navbar-logo.html", "versioning.html", "icon-links.html", "search-field.html", "sbt-sidebar-nav.html"] +} # -- Advanced configuration ------------------------------------------------- @@ -248,7 +262,7 @@ def skip_member(app, what, name, obj, skip, options): # List the names of the functions you want to skip here - exclusions = ["from_dict", "to_dict", "replace", "copy", "__post_init__"] + exclusions = ["from_dict", "to_dict", "replace", "copy", "validate", "__post_init__"] if name in exclusions: return True return None diff --git a/docs/index.rst b/docs/index.rst index 873ca677cd..4af2072b12 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -74,7 +74,9 @@ Table of Contents :maxdepth: 2 :caption: Getting Started + source/setup/ecosystem source/setup/installation/index + source/setup/installation/cloud_installation source/setup/faq .. toctree:: @@ -84,6 +86,7 @@ Table of Contents source/overview/developer-guide/index source/overview/core-concepts/index + source/overview/sensors/index source/overview/environments source/overview/reinforcement-learning/index source/overview/teleop_imitation @@ -96,9 +99,9 @@ Table of Contents source/features/hydra source/features/multi_gpu - source/features/tiled_rendering + Tiled Rendering + source/features/ray source/features/reproducibility - .. source/features/motion_generators .. toctree:: :maxdepth: 1 @@ -128,6 +131,7 @@ Table of Contents :maxdepth: 1 :caption: References + source/refs/reference_architecture/index source/refs/additional_resources source/refs/contributing source/refs/troubleshooting diff --git a/docs/licenses/dependencies/einops-license.txt b/docs/licenses/dependencies/einops-license.txt new file mode 100644 index 0000000000..3a654e9066 --- /dev/null +++ b/docs/licenses/dependencies/einops-license.txt @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2018 Alex Rogozhnikov + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/docs/licenses/dependencies/sphinx-multiversion-license.txt b/docs/licenses/dependencies/sphinx-multiversion-license.txt new file mode 100644 index 0000000000..172d6b3f5d --- /dev/null +++ b/docs/licenses/dependencies/sphinx-multiversion-license.txt @@ -0,0 +1,25 @@ +BSD 2-Clause License + +Copyright (c) 2020, Jan Holthuis +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/docs/licenses/dependencies/transformers-license.txt b/docs/licenses/dependencies/transformers-license.txt new file mode 100644 index 0000000000..68b7d66c97 --- /dev/null +++ b/docs/licenses/dependencies/transformers-license.txt @@ -0,0 +1,203 @@ +Copyright 2018- The Hugging Face team. All rights reserved. + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/docs/make.bat b/docs/make.bat index 2119f51099..cdaf22f257 100644 --- a/docs/make.bat +++ b/docs/make.bat @@ -2,34 +2,63 @@ pushd %~dp0 -REM Command file for Sphinx documentation +REM Command file to build Sphinx documentation -if "%SPHINXBUILD%" == "" ( - set SPHINXBUILD=sphinx-build -) set SOURCEDIR=. set BUILDDIR=_build -if "%1" == "" goto help - -%SPHINXBUILD% >NUL 2>NUL -if errorlevel 9009 ( - echo. - echo.The 'sphinx-build' command was not found. Make sure you have Sphinx - echo.installed, then set the SPHINXBUILD environment variable to point - echo.to the full path of the 'sphinx-build' executable. Alternatively you - echo.may add the Sphinx directory to PATH. - echo. - echo.If you don't have Sphinx installed, grab it from - echo.http://sphinx-doc.org/ - exit /b 1 +REM Check if a specific target was passed +if "%1" == "multi-docs" ( + REM Check if SPHINXBUILD is set, if not default to sphinx-multiversion + if "%SPHINXBUILD%" == "" ( + set SPHINXBUILD=sphinx-multiversion + ) + %SPHINXBUILD% >NUL 2>NUL + if errorlevel 9009 ( + echo. + echo.The 'sphinx-multiversion' command was not found. Make sure you have Sphinx + echo.installed, then set the SPHINXBUILD environment variable to point + echo.to the full path of the 'sphinx-multiversion' executable. Alternatively you + echo.may add the Sphinx directory to PATH. + echo. + echo.If you don't have Sphinx installed, grab it from + echo.http://sphinx-doc.org/ + exit /b 1 + ) + %SPHINXBUILD% %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% + + REM Copy the redirect index.html to the build directory + copy _redirect\index.html %BUILDDIR%\index.html + goto end ) -%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% -goto end +if "%1" == "current-docs" ( + REM Check if SPHINXBUILD is set, if not default to sphinx-build + if "%SPHINXBUILD%" == "" ( + set SPHINXBUILD=sphinx-build + ) + %SPHINXBUILD% >NUL 2>NUL + if errorlevel 9009 ( + echo. + echo.The 'sphinx-build' command was not found. Make sure you have Sphinx + echo.installed, then set the SPHINXBUILD environment variable to point + echo.to the full path of the 'sphinx-build' executable. Alternatively you + echo.may add the Sphinx directory to PATH. + echo. + echo.If you don't have Sphinx installed, grab it from + echo.http://sphinx-doc.org/ + exit /b 1 + ) + %SPHINXBUILD% %SOURCEDIR% %BUILDDIR%\current %SPHINXOPTS% %O% + goto end +) -:help -%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% +REM If no valid target is passed, show usage instructions +echo. +echo.Usage: +echo. make.bat multi-docs - To build the multi-version documentation. +echo. make.bat current-docs - To build the current documentation. +echo. :end popd diff --git a/docs/requirements.txt b/docs/requirements.txt index 33917bd712..13b2bfe9d6 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -7,6 +7,8 @@ sphinx-copybutton sphinx-icon sphinx_design sphinxemoji +sphinx-tabs # backwards compatibility for building docs on v1.0.0 +sphinx-multiversion==0.2.4 # basic python numpy diff --git a/docs/source/_static/demos/multi_asset.jpg b/docs/source/_static/demos/multi_asset.jpg index c7124d20e6..a59388532b 100644 Binary files a/docs/source/_static/demos/multi_asset.jpg and b/docs/source/_static/demos/multi_asset.jpg differ diff --git a/docs/source/_static/overview/overview_sensors_contact_diagram.png b/docs/source/_static/overview/overview_sensors_contact_diagram.png new file mode 100644 index 0000000000..91b05c51fa Binary files /dev/null and b/docs/source/_static/overview/overview_sensors_contact_diagram.png differ diff --git a/docs/source/_static/overview/overview_sensors_depth.png b/docs/source/_static/overview/overview_sensors_depth.png new file mode 100644 index 0000000000..5437e0bf8d Binary files /dev/null and b/docs/source/_static/overview/overview_sensors_depth.png differ diff --git a/docs/source/_static/overview/overview_sensors_frame_transformer.png b/docs/source/_static/overview/overview_sensors_frame_transformer.png new file mode 100644 index 0000000000..c16944060e Binary files /dev/null and b/docs/source/_static/overview/overview_sensors_frame_transformer.png differ diff --git a/docs/source/_static/overview/overview_sensors_ft_visualizer.png b/docs/source/_static/overview/overview_sensors_ft_visualizer.png new file mode 100644 index 0000000000..c16ea4aa05 Binary files /dev/null and b/docs/source/_static/overview/overview_sensors_ft_visualizer.png differ diff --git a/docs/source/_static/overview/overview_sensors_instance.png b/docs/source/_static/overview/overview_sensors_instance.png new file mode 100644 index 0000000000..8683e71946 Binary files /dev/null and b/docs/source/_static/overview/overview_sensors_instance.png differ diff --git a/docs/source/_static/overview/overview_sensors_instanceID.png b/docs/source/_static/overview/overview_sensors_instanceID.png new file mode 100644 index 0000000000..2a3d4d9036 Binary files /dev/null and b/docs/source/_static/overview/overview_sensors_instanceID.png differ diff --git a/docs/source/_static/overview/overview_sensors_normals.png b/docs/source/_static/overview/overview_sensors_normals.png new file mode 100644 index 0000000000..6e3ab7f3cd Binary files /dev/null and b/docs/source/_static/overview/overview_sensors_normals.png differ diff --git a/docs/source/_static/overview/overview_sensors_rc_patterns.png b/docs/source/_static/overview/overview_sensors_rc_patterns.png new file mode 100644 index 0000000000..fc3c001b14 Binary files /dev/null and b/docs/source/_static/overview/overview_sensors_rc_patterns.png differ diff --git a/docs/source/_static/overview/overview_sensors_rc_visualizer.png b/docs/source/_static/overview/overview_sensors_rc_visualizer.png new file mode 100644 index 0000000000..ac4a06595e Binary files /dev/null and b/docs/source/_static/overview/overview_sensors_rc_visualizer.png differ diff --git a/docs/source/_static/overview/overview_sensors_rgb.png b/docs/source/_static/overview/overview_sensors_rgb.png new file mode 100644 index 0000000000..2f17e81039 Binary files /dev/null and b/docs/source/_static/overview/overview_sensors_rgb.png differ diff --git a/docs/source/_static/overview/overview_sensors_semantic.png b/docs/source/_static/overview/overview_sensors_semantic.png new file mode 100644 index 0000000000..aa14f16e15 Binary files /dev/null and b/docs/source/_static/overview/overview_sensors_semantic.png differ diff --git a/docs/source/_static/reference-architecture/deployment-dark.svg b/docs/source/_static/reference-architecture/deployment-dark.svg new file mode 100644 index 0000000000..ba8c739724 --- /dev/null +++ b/docs/source/_static/reference-architecture/deployment-dark.svg @@ -0,0 +1,3 @@ + + +
Robot Hardware
NVIDIA Isaac Perceptor
Custom Estimator
Isaac ROS Packages
State Estimator
Extract Observation
Actions Controller
Trained Model (.onnx, .pt)
Commanded Actions
Model Inference Runtime
Scaled Actions
1
2
3
4
5
6
7
 diff --git a/docs/source/_static/reference-architecture/deployment-light.svg b/docs/source/_static/reference-architecture/deployment-light.svg new file mode 100644 index 0000000000..33857dbdbe --- /dev/null +++ b/docs/source/_static/reference-architecture/deployment-light.svg @@ -0,0 +1,3 @@ + + +
Robot Hardware
NVIDIA Isaac Perceptor
Custom Estimator
Isaac ROS Packages
State Estimator
Extract Observation
Actions Controller
Trained Model (.onnx, .pt)
Commanded Actions
Model Inference Runtime
Scaled Actions
1
2
3
4
5
6
7
 diff --git a/docs/source/_static/reference-architecture/isaac-lab-ra-dark.svg b/docs/source/_static/reference-architecture/isaac-lab-ra-dark.svg new file mode 100644 index 0000000000..90ed29a93e --- /dev/null +++ b/docs/source/_static/reference-architecture/isaac-lab-ra-dark.svg @@ -0,0 +1,3 @@ + + +
Scene Assets
Robot Assets (.usd, .urdf)
Asset Input
Design Robot Learning Task
Scene Configuration
Asset Configuration
Configuration
Register Environment with Gymnasium
Learning Framework Wrapper
Video Wrapper
Wrap Environment
Wrapper API
Test Model
Multi-Node Training
Single GPU Training
Cloud-based Training
Multi-GPU Training
Run Training
 diff --git a/docs/source/_static/reference-architecture/isaac-lab-ra-light.svg b/docs/source/_static/reference-architecture/isaac-lab-ra-light.svg new file mode 100644 index 0000000000..a58a6af1ec --- /dev/null +++ b/docs/source/_static/reference-architecture/isaac-lab-ra-light.svg @@ -0,0 +1,3 @@ + + +
Scene Assets
Robot Assets (.usd, .urdf)
Asset Input
Design Robot Learning Task
Scene Configuration
Asset Configuration
Configuration
Register Environment with Gymnasium
Learning Framework Wrapper
Video Wrapper
Wrap Environment
Wrapper API
Test Model
Multi-Node Training
Single GPU Training
Cloud-based Training
Multi-GPU Training
Run Training
 diff --git a/docs/source/_static/reference-architecture/multi-gpu-training-dark.svg b/docs/source/_static/reference-architecture/multi-gpu-training-dark.svg new file mode 100644 index 0000000000..dd63d5769c --- /dev/null +++ b/docs/source/_static/reference-architecture/multi-gpu-training-dark.svg @@ -0,0 +1,3 @@ + + +
Environment 0
Gradients
Global Network
Updated Model
GPU 0
Learner 0
Environment 1
GPU 1
Learner 1
Environment N
GPU N
Learner N
 diff --git a/docs/source/_static/reference-architecture/multi-gpu-training-light.svg b/docs/source/_static/reference-architecture/multi-gpu-training-light.svg new file mode 100644 index 0000000000..a279bff5e8 --- /dev/null +++ b/docs/source/_static/reference-architecture/multi-gpu-training-light.svg @@ -0,0 +1,3 @@ + + +
Environment 0
Gradients
Global Network
Updated Model
GPU 0
Learner 0
Environment 1
GPU 1
Learner 1
Environment N
GPU N
Learner N
 diff --git a/docs/source/_static/reference-architecture/single-gpu-training-dark.svg b/docs/source/_static/reference-architecture/single-gpu-training-dark.svg new file mode 100644 index 0000000000..69f8e0f3c0 --- /dev/null +++ b/docs/source/_static/reference-architecture/single-gpu-training-dark.svg @@ -0,0 +1,3 @@ + + +
Isaac Lab
2
States
Add Noise
Assets
Rendering
Trained Model (.pt, .onnx)
Isaac Sim
1
Physics Sim
4
5
RL Libraries
3
Policy
Actions
Observations
6
 diff --git a/docs/source/_static/reference-architecture/single-gpu-training-light.svg b/docs/source/_static/reference-architecture/single-gpu-training-light.svg new file mode 100644 index 0000000000..7463c470d8 --- /dev/null +++ b/docs/source/_static/reference-architecture/single-gpu-training-light.svg @@ -0,0 +1,3 @@ + + +
Isaac Lab
2
States
Add Noise
Assets
Rendering
Trained Model (.pt, .onnx)
Isaac Sim
1
Physics Sim
4
5
RL Libraries
3
Policy
Actions
Observations
6
 diff --git a/docs/source/_static/refs.bib b/docs/source/_static/refs.bib index e7f82d17da..c87c90c076 100644 --- a/docs/source/_static/refs.bib +++ b/docs/source/_static/refs.bib @@ -139,3 +139,18 @@ @article{mittal2023orbit pages={3740-3747}, doi={10.1109/LRA.2023.3270034} } + +@article{shang2024theia, + title={Theia: Distilling diverse vision foundation models for robot learning}, + author={Shang, Jinghuan and Schmeckpeper, Karl and May, Brandon B and Minniti, Maria Vittoria and Kelestemur, Tarik and Watkins, David and Herlant, Laura}, + journal={arXiv preprint arXiv:2407.20179}, + year={2024} +} + +@inproceedings{he2016deep, + title={Deep residual learning for image recognition}, + author={He, Kaiming and Zhang, Xiangyu and Ren, Shaoqing and Sun, Jian}, + booktitle={Proceedings of the IEEE conference on computer vision and pattern recognition}, + pages={770--778}, + year={2016} +} diff --git a/docs/source/_static/setup/ecosystem-dark.jpg b/docs/source/_static/setup/ecosystem-dark.jpg new file mode 100644 index 0000000000..9057aeb48b Binary files /dev/null and b/docs/source/_static/setup/ecosystem-dark.jpg differ diff --git a/docs/source/_static/setup/ecosystem-light.jpg b/docs/source/_static/setup/ecosystem-light.jpg new file mode 100644 index 0000000000..5c2bbf5fe1 Binary files /dev/null and b/docs/source/_static/setup/ecosystem-light.jpg differ diff --git a/docs/source/_static/setup/verify_install.jpg b/docs/source/_static/setup/verify_install.jpg new file mode 100644 index 0000000000..166840dc70 Binary files /dev/null and b/docs/source/_static/setup/verify_install.jpg differ diff --git a/docs/source/_static/task-workflows/direct-based-dark.svg b/docs/source/_static/task-workflows/direct-based-dark.svg new file mode 100644 index 0000000000..2709dae140 --- /dev/null +++ b/docs/source/_static/task-workflows/direct-based-dark.svg @@ -0,0 +1,3 @@ + + +
Sensors
Observations
Apply Actions
Define Task
Define Step
Environment Scripting
NVIDIA
Isaac Sim
Perform Resets
Rewards
Commands
Learning Agent
Scene Creation
Articulation
Objects
Compute Signals
Actions
Perform
Randomization
 diff --git a/docs/source/_static/task-workflows/direct-based-light.svg b/docs/source/_static/task-workflows/direct-based-light.svg new file mode 100644 index 0000000000..b39bffa353 --- /dev/null +++ b/docs/source/_static/task-workflows/direct-based-light.svg @@ -0,0 +1,3 @@ + + +
Sensors
Observations
Apply Actions
Define Task
Define Step
Environment Scripting
NVIDIA
Isaac Sim
Perform Resets
Rewards
Commands
Learning Agent
Scene Creation
Articulation
Objects
Compute Signals
Actions
Perform
Randomization
 diff --git a/docs/source/_static/task-workflows/manager-based-dark.svg b/docs/source/_static/task-workflows/manager-based-dark.svg new file mode 100644 index 0000000000..b66ba1bfee --- /dev/null +++ b/docs/source/_static/task-workflows/manager-based-dark.svg @@ -0,0 +1,3 @@ + + +
Action Manager
Task-space
Joint-space
Termination Manager
Reward Manager
Curriculum Manager
Learning
Agent
Custom
Command Manager
Custom
Velocity
Pose
Observation Manager
Proprioception
Custom
Exteroception
NVIDIA
Isaac Sim
Articulation
Objects
Sensors
Interactive Scene
Event Manager
External Disturbances
Domain Randomization
 diff --git a/docs/source/_static/task-workflows/manager-based-light.svg b/docs/source/_static/task-workflows/manager-based-light.svg new file mode 100644 index 0000000000..60fe6119d4 --- /dev/null +++ b/docs/source/_static/task-workflows/manager-based-light.svg @@ -0,0 +1,3 @@ + + +
Action Manager
Task-space
Joint-space
Termination Manager
Reward Manager
Curriculum Manager
Learning
Agent
Custom
Command Manager
Custom
Velocity
Pose
Observation Manager
Proprioception
Custom
Exteroception
NVIDIA
Isaac Sim
Articulation
Objects
Sensors
Interactive Scene
Event Manager
External Disturbances
Domain Randomization
 diff --git a/docs/source/_static/tasks/factory/gear_mesh.jpg b/docs/source/_static/tasks/factory/gear_mesh.jpg new file mode 100644 index 0000000000..af26e5818a Binary files /dev/null and b/docs/source/_static/tasks/factory/gear_mesh.jpg differ diff --git a/docs/source/_static/tasks/factory/nut_thread.jpg b/docs/source/_static/tasks/factory/nut_thread.jpg new file mode 100644 index 0000000000..9a3155e6c4 Binary files /dev/null and b/docs/source/_static/tasks/factory/nut_thread.jpg differ diff --git a/docs/source/_static/tasks/factory/peg_insert.jpg b/docs/source/_static/tasks/factory/peg_insert.jpg new file mode 100644 index 0000000000..7aaa2c1450 Binary files /dev/null and b/docs/source/_static/tasks/factory/peg_insert.jpg differ diff --git a/docs/source/_static/tasks/manipulation/franka_stack.jpg b/docs/source/_static/tasks/manipulation/franka_stack.jpg new file mode 100644 index 0000000000..1da36113ec Binary files /dev/null and b/docs/source/_static/tasks/manipulation/franka_stack.jpg differ diff --git a/docs/source/_static/tutorials/tutorial_operational_space_controller.jpg b/docs/source/_static/tutorials/tutorial_operational_space_controller.jpg new file mode 100644 index 0000000000..e356eee4ff Binary files /dev/null and b/docs/source/_static/tutorials/tutorial_operational_space_controller.jpg differ diff --git a/docs/source/api/lab/omni.isaac.lab.assets.rst b/docs/source/api/lab/omni.isaac.lab.assets.rst index d7dcfbac19..ba139a2395 100644 --- a/docs/source/api/lab/omni.isaac.lab.assets.rst +++ b/docs/source/api/lab/omni.isaac.lab.assets.rst @@ -12,6 +12,9 @@ RigidObject RigidObjectData RigidObjectCfg + RigidObjectCollection + RigidObjectCollectionData + RigidObjectCollectionCfg Articulation ArticulationData ArticulationCfg @@ -51,6 +54,26 @@ Rigid Object :show-inheritance: :exclude-members: __init__, class_type +Rigid Object Collection +----------------------- + +.. autoclass:: RigidObjectCollection + :members: + :inherited-members: + :show-inheritance: + +.. autoclass:: RigidObjectCollectionData + :members: + :inherited-members: + :show-inheritance: + :exclude-members: __init__ + +.. autoclass:: RigidObjectCollectionCfg + :members: + :inherited-members: + :show-inheritance: + :exclude-members: __init__, class_type + Articulation ------------ diff --git a/docs/source/api/lab/omni.isaac.lab.controllers.rst b/docs/source/api/lab/omni.isaac.lab.controllers.rst index d56efdd05a..67136f0a3a 100644 --- a/docs/source/api/lab/omni.isaac.lab.controllers.rst +++ b/docs/source/api/lab/omni.isaac.lab.controllers.rst @@ -9,6 +9,8 @@ DifferentialIKController DifferentialIKControllerCfg + OperationalSpaceController + OperationalSpaceControllerCfg Differential Inverse Kinematics ------------------------------- @@ -23,3 +25,17 @@ Differential Inverse Kinematics :inherited-members: :show-inheritance: :exclude-members: __init__, class_type + +Operational Space controllers +----------------------------- + +.. autoclass:: OperationalSpaceController + :members: + :inherited-members: + :show-inheritance: + +.. autoclass:: OperationalSpaceControllerCfg + :members: + :inherited-members: + :show-inheritance: + :exclude-members: __init__, class_type diff --git a/docs/source/api/lab/omni.isaac.lab.sensors.rst b/docs/source/api/lab/omni.isaac.lab.sensors.rst index 9eac25b797..3a82a42904 100644 --- a/docs/source/api/lab/omni.isaac.lab.sensors.rst +++ b/docs/source/api/lab/omni.isaac.lab.sensors.rst @@ -31,6 +31,8 @@ RayCasterCfg RayCasterCamera RayCasterCameraCfg + Imu + ImuCfg Sensor Base ----------- @@ -150,3 +152,17 @@ Ray-Cast Camera :inherited-members: :show-inheritance: :exclude-members: __init__, class_type + +Inertia Measurement Unit +------------------------ + +.. autoclass:: Imu + :members: + :inherited-members: + :show-inheritance: + +.. autoclass:: ImuCfg + :members: + :inherited-members: + :show-inheritance: + :exclude-members: __init__, class_type diff --git a/docs/source/api/lab/omni.isaac.lab.sim.rst b/docs/source/api/lab/omni.isaac.lab.sim.rst index 4e37e33b4a..b2b582c68b 100644 --- a/docs/source/api/lab/omni.isaac.lab.sim.rst +++ b/docs/source/api/lab/omni.isaac.lab.sim.rst @@ -19,6 +19,7 @@ SimulationContext SimulationCfg PhysxCfg + RenderCfg .. rubric:: Functions @@ -46,6 +47,11 @@ Simulation Configuration :show-inheritance: :exclude-members: __init__ +.. autoclass:: RenderCfg + :members: + :show-inheritance: + :exclude-members: __init__ + Simulation Context Builder -------------------------- diff --git a/docs/source/api/lab/omni.isaac.lab.utils.rst b/docs/source/api/lab/omni.isaac.lab.utils.rst index 2143ecfc28..c14ae19b6c 100644 --- a/docs/source/api/lab/omni.isaac.lab.utils.rst +++ b/docs/source/api/lab/omni.isaac.lab.utils.rst @@ -18,6 +18,7 @@ noise string timer + types warp .. Rubric:: Functions @@ -123,6 +124,13 @@ Timer operations :members: :show-inheritance: +Type operations +~~~~~~~~~~~~~~~ + +.. automodule:: omni.isaac.lab.utils.types + :members: + :show-inheritance: + Warp operations ~~~~~~~~~~~~~~~ diff --git a/docs/source/deployment/cluster.rst b/docs/source/deployment/cluster.rst index beae7ef438..467fda90f4 100644 --- a/docs/source/deployment/cluster.rst +++ b/docs/source/deployment/cluster.rst @@ -45,15 +45,13 @@ development machine and the cluster. Such a connection will simplify the file tr the user cluster password from being requested multiple times. .. attention:: - The workflow has been tested with ``apptainer version 1.2.5-1.el7`` and ``docker version 24.0.7``. + The workflow has been tested with: - - ``apptainer``: - There have been reported binding issues with previous versions (such as ``apptainer version 1.1.3-1.el7``). Please - ensure that you are using the latest version. - - ``Docker``: - The latest versions (``25.x``) cannot be used as they are not compatible yet with apptainer/ singularity. + - ``apptainer version 1.2.5-1.el7`` and ``docker version 24.0.7`` + - ``apptainer version 1.3.4`` and ``docker version 27.3.1`` + + In the case of issues, please try to switch to those versions. - We are waiting for an update from the apptainer team. To track this issue, please check the `forum post`_. Configuring the cluster parameters ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -211,4 +209,3 @@ The above will, in addition, also render videos of the training progress and sto .. _apptainer: https://apptainer.org/ .. _documentation: https://www.apptainer.org/docs/admin/main/installation.html#install-ubuntu-packages .. _SLURM documentation: https://www.slurm.schedmd.com/sbatch.html -.. _forum post: https://forums.docker.com/t/trouble-after-upgrade-to-docker-ce-25-0-1-on-debian-12/139613 diff --git a/docs/source/features/hydra.rst b/docs/source/features/hydra.rst index 2fa5d2d706..577eb7a6ec 100644 --- a/docs/source/features/hydra.rst +++ b/docs/source/features/hydra.rst @@ -89,8 +89,8 @@ Elements in dictionaries are handled as a parameters in the hierarchy. For examp .. literalinclude:: ../../../source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/cartpole_env_cfg.py :language: python - :lines: 99-111 - :emphasize-lines: 10 + :lines: 90-114 + :emphasize-lines: 11 the ``position_range`` parameter can be modified with ``env.events.reset_cart_position.params.position_range="[-2.0, 2.0]"``. This example shows two noteworthy points: @@ -112,10 +112,10 @@ For example, for the configuration of the Cartpole camera depth environment: :language: python :start-at: class CartpoleDepthCameraEnvCfg :end-at: tiled_camera.width - :emphasize-lines: 16 + :emphasize-lines: 10, 15 If the user were to modify the width of the camera, i.e. ``env.tiled_camera.width=128``, then the parameter -``env.num_observations=10240`` (1*80*128) must be updated and given as input as well. +``env.observation_space=[80,128,1]`` must be updated and given as input as well. Similarly, the ``__post_init__`` method is not updated with the command line inputs. In the ``LocomotionVelocityRoughEnvCfg``, for example, the post init update is as follows: diff --git a/docs/source/features/ray.rst b/docs/source/features/ray.rst new file mode 100644 index 0000000000..86fdf48e5d --- /dev/null +++ b/docs/source/features/ray.rst @@ -0,0 +1,380 @@ +=========================== +Ray Job Dispatch and Tuning +=========================== + +.. currentmodule:: omni.isaac.lab + +Isaac Lab supports `Ray `_ for streamlining dispatching multiple training jobs (in parallel and in series), +and hyperparameter tuning, both on local and remote configurations. + +This `independent community contributed walkthrough video `_ +demonstrates some of the core functionality of the Ray integration covered in this overview. Although there may be some +differences in the codebase (such as file names being shortened) since the creation of the video, +the general workflow is the same. + +.. attention:: + + This functionality is experimental, and has been tested only on Linux. + +.. contents:: Table of Contents + :depth: 3 + :local: + +Overview +-------- + +The Ray integration is useful for the following: + +- Dispatching several training jobs in parallel or sequentially with minimal interaction +- Tuning hyperparameters; in parallel or sequentially with support for multiple GPUs and/or multiple GPU Nodes +- Using the same training setup everywhere (on cloud and local) with minimal overhead +- Resource Isolation for training jobs + +The core functionality of the Ray workflow consists of two main scripts that enable the orchestration +of resource-wrapped and tuning aggregate jobs. These scripts facilitate the decomposition of +aggregate jobs (overarching experiments) into individual jobs, which are discrete commands +executed on the cluster. An aggregate job can include multiple individual jobs. +For clarity, this guide refers to the jobs one layer below the topmost aggregate level as sub-jobs. + +Both resource-wrapped and tuning aggregate jobs dispatch individual jobs to a designated Ray +cluster, which leverages the cluster's resources (e.g., a single workstation node or multiple nodes) +to execute these jobs with workers in parallel and/or sequentially. By default, aggregate jobs use all \ +available resources on each available GPU-enabled node for each sub-job worker. This can be changed through +specifying the ``--num_workers`` argument, especially critical for parallel aggregate +job processing on local or virtual multi-GPU machines + +In resource-wrapped aggregate jobs, each sub-job and its +resource requirements are defined manually, enabling resource isolation. +For tuning aggregate jobs, individual jobs are generated automatically based on a hyperparameter +sweep configuration. This assumes homogeneous node resource composition for nodes with GPUs. + +.. dropdown:: source/standalone/workflows/ray/wrap_resources.py + :icon: code + + .. literalinclude:: ../../../source/standalone/workflows/ray/wrap_resources.py + :language: python + :emphasize-lines: 14-66 + +.. dropdown:: source/standalone/workflows/ray/tuner.py + :icon: code + + .. literalinclude:: ../../../source/standalone/workflows/ray/tuner.py + :language: python + :emphasize-lines: 18-53 + + +The following script can be used to submit aggregate +jobs to one or more Ray cluster(s), which can be used for +running jobs on a remote cluster or simultaneous jobs with heterogeneous +resource requirements: + +.. dropdown:: source/standalone/workflows/ray/submit_job.py + :icon: code + + .. literalinclude:: ../../../source/standalone/workflows/ray/submit_job.py + :language: python + :emphasize-lines: 12-53 + +The following script can be used to extract KubeRay Cluster information for aggregate job submission. + +.. dropdown:: source/standalone/workflows/ray/grok_cluster_with_kubectl.py + :icon: code + + .. literalinclude:: ../../../source/standalone/workflows/ray/grok_cluster_with_kubectl.py + :language: python + :emphasize-lines: 14-26 + +The following script can be used to easily create clusters on Google GKE. + +.. dropdown:: source/standalone/workflows/ray/launch.py + :icon: code + + .. literalinclude:: ../../../source/standalone/workflows/ray/launch.py + :language: python + :emphasize-lines: 16-37 + +**Installation** +---------------- + +The Ray functionality requires additional dependencies be installed. + +To use Ray without Kubernetes, like on a local computer or VM, +``kubectl`` is not required. For use on Kubernetes clusters with KubeRay, +such as Google Kubernetes Engine or Amazon Elastic Kubernetes Service, ``kubectl`` is required, and can +be installed via the `Kubernetes website `_ + +The pythonic dependencies can be installed with: + +.. code-block:: bash + + # For multi-run support and resource isolation + ./isaaclab.sh -p -m pip install ray[default]==2.31.0 + # For hyperparameter tuning + ./isaaclab.sh -p -m pip install ray[tune]==2.31.0 + ./isaaclab.sh -p -m pip install optuna bayesian-optimization + # MLFlow is needed only for fetching logs on clusters, not needed for local + ./isaaclab.sh -p -m pip install mlflow + +If using KubeRay clusters on Google GKE with the batteries-included cluster launch file, +the following dependencies are also needed. + +.. code-block:: bash + + ./isaaclab.sh -p -m pip install kubernetes Jinja2 + +**Setup Overview: Cluster Configuration** +----------------------------------------- + +Select one of the following methods to create a Ray Cluster to accept and execute dispatched jobs. + +Single-Node Ray Cluster (Recommended for Beginners) +''''''''''''''''''''''''''''''''''''''''''''''''''' +For use on a single machine (node) such as a local computer or VM, the +following command can be used start a ray server. This is compatible with +multiple-GPU machines. This Ray server will run indefinitely until it is stopped with ``CTRL + C`` + +.. code-block:: bash + + echo "import ray; ray.init(); import time; [time.sleep(10) for _ in iter(int, 1)]" | ./isaaclab.sh -p + +KubeRay Clusters +'''''''''''''''' +.. attention:: + The ``ray`` command should be modified to use Isaac python, which could be achieved in a fashion similar to + ``sed -i "1i $(echo "#!/workspace/isaaclab/_isaac_sim/python.sh")" \ + /isaac-sim/kit/python/bin/ray && ln -s /isaac-sim/kit/python/bin/ray /usr/local/bin/ray``. + +Google Cloud is currently the only platform tested, although +any cloud provider should work if one configures the following: + +- An container registry (NGC, GCS artifact registry, AWS ECR, etc) with + an Isaac Lab image configured to support Ray. See ``cluster_configs/Dockerfile`` to see how to modify the ``isaac-lab-base`` + container for Ray compatibility. Ray should use the isaac sim python shebang, and ``nvidia-smi`` + should work within the container. Be careful with the setup here as + paths need to be configured correctly for everything to work. It's likely that + the example dockerfile will work out of the box and can be pushed to the registry, as + long as the base image has already been built as in the container guide +- A Kubernetes setup with available NVIDIA RTX (likely ``l4`` or ``l40`` or ``tesla-t4`` or ``a10``) GPU-passthrough node-pool resources, + that has access to your container registry/storage bucket and has the Ray operator enabled with correct IAM + permissions. This can be easily achieved with services such as Google GKE or AWS EKS, + provided that your account or organization has been granted a GPU-budget. It is recommended + to use manual kubernetes services as opposed to "autopilot" services for cost-effective + experimentation as this way clusters can be completely shut down when not in use, although + this may require installing the `Nvidia GPU Operator `_ +- An MLFlow server that your cluster has access to. +- A ``kuberay.yaml.ninja`` file that describes how to allocate resources (already included for + Google Cloud, which can be referenced for the format and MLFlow integration) + +Ray Clusters (Without Kubernetes) +''''''''''''''''''''''''''''''''' +.. attention:: + Modify the Ray command to use Isaac Python like in KubeRay Clusters, and follow the same + steps for creating an image/cluster permissions/bucket access. + +See the `Ray Clusters Overview `_ or +`Anyscale `_ for more information + + +**Dispatching Jobs and Tuning** +------------------------------- + +Select one of the following guides that matches your desired Cluster configuration. + +Simple Ray Cluster (Local/VM) +''''''''''''''''''''''''''''' + +This guide assumes that there is a Ray cluster already running, and that this script is run locally on the cluster, or +that the cluster job submission address is known. + +1.) Testing that the cluster works can be done as follows. + +.. code-block:: bash + + ./isaaclab.sh -p source/standalone/workflows/ray/wrap_resources.py --test + +2.) Submitting resource-wrapped sub-jobs can be done as described in the following file: + +.. dropdown:: source/standalone/workflows/ray/wrap_resources.py + :icon: code + + .. literalinclude:: ../../../source/standalone/workflows/ray/wrap_resources.py + :language: python + :emphasize-lines: 14-66 + +3.) For tuning jobs, specify the hyperparameter sweep similar to the following two files. + +.. dropdown:: source/standalone/workflows/ray/hyperparameter_tuning/vision_cfg.py + :icon: code + + .. literalinclude:: ../../../source/standalone/workflows/ray/hyperparameter_tuning/vision_cfg.py + :language: python + +.. dropdown:: source/standalone/workflows/ray/hyperparameter_tuning/vision_cartpole_cfg.py + :icon: code + + .. literalinclude:: ../../../source/standalone/workflows/ray/hyperparameter_tuning/vision_cartpole_cfg.py + :language: python + +Then, see the local examples in the following file to see how to start a tuning run. + +.. dropdown:: source/standalone/workflows/ray/tuner.py + :icon: code + + .. literalinclude:: ../../../source/standalone/workflows/ray/tuner.py + :language: python + :emphasize-lines: 18-53 + + + +To view the logs, simply run ``tensorboard --logdir=`` + +Remote Ray Cluster Setup and Use +''''''''''''''''''''''''''''''''' +This guide assumes that one desires to create a cluster on a remote host or server. This +guide includes shared steps, and KubeRay or Ray specific steps. Follow all shared steps (part I and II), and then +only the KubeRay or Ray steps depending on your desired configuration, in order of shared steps part I, then +the configuration specific steps, then shared steps part II. + +Shared Steps Between KubeRay and Pure Ray Part I +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +1.) Build the Isaac Ray image, and upload it to your container registry of choice. + +.. code-block:: bash + + # Login with NGC (nvcr.io) registry first, see docker steps in repo. + ./isaaclab.sh -p docker/container.py start + # Build the special Isaac Lab Ray Image + docker build -t -f source/standalone/workflows/ray/cluster_configs/Dockerfile . + # Push the image to your registry of choice. + docker push + +KubeRay Specific +~~~~~~~~~~~~~~~~ +`k9s `_ is a great tool for monitoring your clusters that can +easily be installed with ``snap install k9s --devmode``. + +1.) Verify cluster access, and that the correct operators are installed. + +.. code-block:: bash + + # Verify cluster access + kubectl cluster-info + # If using a manually managed cluster (not Autopilot or the like) + # verify that there are node pools + kubectl get nodes + # Check that the ray operator is installed on the cluster + # should list rayclusters.ray.io , rayjobs.ray.io , and rayservices.ray.io + kubectl get crds | grep ray + # Check that the NVIDIA Driver Operator is installed on the cluster + # should list clusterpolicies.nvidia.com + kubectl get crds | grep nvidia + +2.) Create the KubeRay cluster and an MLFlow server for receiving logs +that your cluster has access to. +This can be done automatically for Google GKE, +where instructions are included in the following creation file. More than once cluster +can be created at once. Each cluster can have heterogeneous resources if so desired, +although only +For other cloud services, the ``kuberay.yaml.ninja`` will be similar to that of +Google's. + +.. dropdown:: source/standalone/workflows/ray/launch.py + :icon: code + + .. literalinclude:: ../../../source/standalone/workflows/ray/launch.py + :language: python + :emphasize-lines: 15-37 + +3.) Fetch the KubeRay cluster IP addresses, and the MLFLow Server IP. +This can be done automatically for KubeRay clusters, +where instructions are included in the following fetching file. +The KubeRay clusters are saved to a file, but the MLFLow Server IP is +printed. + +.. dropdown:: source/standalone/workflows/ray/grok_cluster_with_kubectl.py + :icon: code + + .. literalinclude:: ../../../source/standalone/workflows/ray/grok_cluster_with_kubectl.py + :language: python + :emphasize-lines: 14-26 + +Ray Specific +~~~~~~~~~~~~ + + +1.) Verify cluster access. + +2.) Create a ``~/.cluster_config`` file, where ``name: address: http://:`` is on +a new line for each unique cluster. For one cluster, there should only be one line in this file. + +3.) Start an MLFLow Server to receive the logs that the ray cluster has access to, +and determine the server URI. + +Shared Steps Between KubeRay and Pure Ray Part II +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +1.) Test that your cluster is operational with the following. + +.. code-block:: bash + + # Test that NVIDIA GPUs are visible and that Ray is operation with the following command: + ./isaaclab.sh -p source/standalone/workflows/ray/wrap_resources.py + --jobs wrap_resources.py --test + +2.) Submitting Jobs can be done in the following manner, with the following script. + +.. dropdown:: source/standalone/workflows/ray/submit_job.py + :icon: code + + .. literalinclude:: ../../../source/standalone/workflows/ray/submit_job.py + :language: python + :emphasize-lines: 12-53 + +3.) For tuning jobs, specify the hyperparameter sweep similar to :class:`RLGamesCameraJobCfg` in the following file: + +.. dropdown:: source/standalone/workflows/ray/tuner.py + :icon: code + + .. literalinclude:: ../../../source/standalone/workflows/ray/tuner.py + :language: python + :emphasize-lines: 18-53 + +For example, see the Cartpole Example configurations. + +.. dropdown:: source/standalone/workflows/ray/hyperparameter_tuning/vision_cartpole_cfg.py + :icon: code + + .. literalinclude:: ../../../source/standalone/workflows/ray/hyperparameter_tuning/vision_cartpole_cfg.py + :language: python + + +Tuning jobs can also be submitted via ``submit_job.py`` + +To view the tuning results, view the MLFlow dashboard of the server that you created. +For KubeRay, this can be done through port forwarding the MLFlow dashboard, with + +``kubectl port-forward service/isaacray-mlflow 5000:5000`` + +and visiting the following address in a browser. + +``localhost:5000`` + +If the MLFlow port is forwarded like above, it can be converted into tensorboard logs with +this following command. + +``./isaaclab.sh -p source/standalone/workflows/ray/mlflow_to_local_tensorboard.py \ +--uri http://localhost:5000 --experiment-name IsaacRay--tune --download-dir test`` + + +**Cluster Cleanup** +''''''''''''''''''' + +For the sake of conserving resources, and potentially freeing precious GPU resources for other people to use +on shared compute platforms, please destroy the Ray cluster after use. They can be easily +recreated! For KubeRay clusters, this can be done as follows. + +.. code-block:: bash + + kubectl get raycluster | egrep 'isaacray' | awk '{print $1}' | xargs kubectl delete raycluster && + kubectl get deployments | egrep 'mlflow' | awk '{print $1}' | xargs kubectl delete deployment && + kubectl get services | egrep 'mlflow' | awk '{print $1}' | xargs kubectl delete service diff --git a/docs/source/how-to/multi_asset_spawning.rst b/docs/source/how-to/multi_asset_spawning.rst index 9f74e39f6b..fe141981fc 100644 --- a/docs/source/how-to/multi_asset_spawning.rst +++ b/docs/source/how-to/multi_asset_spawning.rst @@ -1,16 +1,23 @@ + Spawning Multiple Assets ======================== .. currentmodule:: omni.isaac.lab -Typical, spawning configurations (introduced in the :ref:`tutorial-spawn-prims` tutorial) copy the same +Typical spawning configurations (introduced in the :ref:`tutorial-spawn-prims` tutorial) copy the same asset (or USD primitive) across the different resolved prim paths from the expressions. For instance, if the user specifies to spawn the asset at "/World/Table\_.*/Object", the same asset is created at the paths "/World/Table_0/Object", "/World/Table_1/Object" and so on. -However, at times, it might be desirable to spawn different assets under the prim paths to -ensure a diversity in the simulation. This guide describes how to create different assets under -each prim path using the spawning functionality. +However, we also support multi-asset spawning with two mechanisms: + +1. Rigid object collections. This allows the user to spawn multiple rigid objects in each environment and access/modify + them with a unified API, improving performance. + +2. Spawning different assets under the same prim path. This allows the user to create diverse simulations, where each + environment has a different asset. + +This guide describes how to use these two mechanisms. The sample script ``multi_asset.py`` is used as a reference, located in the ``IsaacLab/source/standalone/demos`` directory. @@ -20,20 +27,41 @@ The sample script ``multi_asset.py`` is used as a reference, located in the .. literalinclude:: ../../../source/standalone/demos/multi_asset.py :language: python - :emphasize-lines: 101-123, 130-149 + :emphasize-lines: 109-131, 135-179, 184-203 :linenos: -This script creates multiple environments, where each environment has a rigid object that is either a cone, -a cube, or a sphere, and an articulation that is either the ANYmal-C or ANYmal-D robot. +This script creates multiple environments, where each environment has: + +* a rigid object collection containing a cone, a cube, and a sphere +* a rigid object that is either a cone, a cube, or a sphere, chosen at random +* an articulation that is either the ANYmal-C or ANYmal-D robot, chosen at random .. image:: ../_static/demos/multi_asset.jpg :width: 100% :alt: result of multi_asset.py -Using Multi-Asset Spawning Functions ------------------------------------- -It is possible to spawn different assets and USDs in each environment using the spawners +Rigid Object Collections +------------------------ + +Multiple rigid objects can be spawned in each environment and accessed/modified with a unified ``(env_ids, obj_ids)`` API. +While the user could also create multiple rigid objects by spawning them individually, the API is more user-friendly and +more efficient since it uses a single physics view under the hood to handle all the objects. + +.. literalinclude:: ../../../source/standalone/demos/multi_asset.py + :language: python + :lines: 135-179 + :dedent: + +The configuration :class:`~assets.RigidObjectCollectionCfg` is used to create the collection. It's attribute :attr:`~assets.RigidObjectCollectionCfg.rigid_objects` +is a dictionary containing :class:`~assets.RigidObjectCfg` objects. The keys serve as unique identifiers for each +rigid object in the collection. + + +Spawning different assets under the same prim path +-------------------------------------------------- + +It is possible to spawn different assets and USDs under the same prim path in each environment using the spawners :class:`~sim.spawners.wrappers.MultiAssetSpawnerCfg` and :class:`~sim.spawners.wrappers.MultiUsdFileCfg`: * We set the spawn configuration in :class:`~assets.RigidObjectCfg` to be @@ -41,7 +69,7 @@ It is possible to spawn different assets and USDs in each environment using the .. literalinclude:: ../../../source/standalone/demos/multi_asset.py :language: python - :lines: 99-125 + :lines: 107-133 :dedent: This function allows you to define a list of different assets that can be spawned as rigid objects. @@ -53,14 +81,14 @@ It is possible to spawn different assets and USDs in each environment using the .. literalinclude:: ../../../source/standalone/demos/multi_asset.py :language: python - :lines: 128-161 + :lines: 182-215 :dedent: Similar to before, this configuration allows the selection of different USD files representing articulated assets. Things to Note --------------- +~~~~~~~~~~~~~~ Similar asset structuring ~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -85,7 +113,7 @@ anymore. Hence the flag :attr:`scene.InteractiveScene.replicate_physics` must be .. literalinclude:: ../../../source/standalone/demos/multi_asset.py :language: python - :lines: 221-224 + :lines: 280-283 :dedent: The Code Execution diff --git a/docs/source/migration/migrating_from_isaacgymenvs.rst b/docs/source/migration/migrating_from_isaacgymenvs.rst index c903d83c3d..22bc33c9a5 100644 --- a/docs/source/migration/migrating_from_isaacgymenvs.rst +++ b/docs/source/migration/migrating_from_isaacgymenvs.rst @@ -45,9 +45,9 @@ Below is an example skeleton of a task config class: # env decimation = 2 episode_length_s = 5.0 - num_actions = 1 - num_observations = 4 - num_states = 0 + action_space = 1 + observation_space = 4 + state_space = 0 # task-specific parameters ... @@ -135,9 +135,9 @@ The following parameters must be set for each environment config: decimation = 2 episode_length_s = 5.0 - num_actions = 1 - num_observations = 4 - num_states = 0 + action_space = 1 + observation_space = 4 + state_space = 0 Note that the maximum episode length parameter (now ``episode_length_s``) is in seconds instead of steps as it was in IsaacGymEnvs. To convert between step count to seconds, use the equation: @@ -196,7 +196,7 @@ adding any other optional objects into the scene, such as lights. | self.sim = super().create_sim(self.device_id, self.graphics_device_id, | # clone, filter, and replicate | | self.physics_engine, self.sim_params) | self.scene.clone_environments(copy_from_source=False) | | self._create_ground_plane() | self.scene.filter_collisions(global_prim_paths=[]) | -| self._create_envs(self.num_envs, self.cfg["env"]['envSpacing'], | # add articultion to scene | +| self._create_envs(self.num_envs, self.cfg["env"]['envSpacing'], | # add articulation to scene | | int(np.sqrt(self.num_envs))) | self.scene.articulations["cartpole"] = self.cartpole | | | # add lights | | | light_cfg = sim_utils.DomeLightCfg(intensity=2000.0) | @@ -569,9 +569,9 @@ Task Config | | decimation = 2 | | asset: | episode_length_s = 5.0 | | assetRoot: "../../assets" | action_scale = 100.0 # [N] | -| assetFileName: "urdf/cartpole.urdf" | num_actions = 1 | -| | num_observations = 4 | -| enableCameraSensors: False | num_states = 0 | +| assetFileName: "urdf/cartpole.urdf" | action_space = 1 | +| | observation_space = 4 | +| enableCameraSensors: False | state_space = 0 | | | # reset | | sim: | max_cart_pos = 3.0 | | dt: 0.0166 # 1/60 s | initial_pole_angle_range = [-0.25, 0.25] | @@ -660,7 +660,7 @@ the need to set simulation parameters for actors in the task implementation. | self._create_ground_plane() | copy_from_source=False) | | self._create_envs(self.num_envs, | self.scene.filter_collisions( | | self.cfg["env"]['envSpacing'], | global_prim_paths=[]) | -| int(np.sqrt(self.num_envs))) | # add articultion to scene | +| int(np.sqrt(self.num_envs))) | # add articulation to scene | | | self.scene.articulations["cartpole"] = self.cartpole | | def _create_ground_plane(self): | # add lights | | plane_params = gymapi.PlaneParams() | light_cfg = sim_utils.DomeLightCfg( | @@ -814,9 +814,9 @@ The ``progress_buf`` variable has also been renamed to ``episode_length_buf``. | | | | | self.joint_pos[env_ids] = joint_pos | | | | -| | self.cartpole.write_root_pose_to_sim( | +| | self.cartpole.write_root_link_pose_to_sim( | | | default_root_state[:, :7], env_ids) | -| | self.cartpole.write_root_velocity_to_sim( | +| | self.cartpole.write_root_com_velocity_to_sim( | | | default_root_state[:, 7:], env_ids) | | | self.cartpole.write_joint_state_to_sim( | | | joint_pos, joint_vel, None, env_ids) | diff --git a/docs/source/migration/migrating_from_omniisaacgymenvs.rst b/docs/source/migration/migrating_from_omniisaacgymenvs.rst index 50f9d5b9d6..629be58bd7 100644 --- a/docs/source/migration/migrating_from_omniisaacgymenvs.rst +++ b/docs/source/migration/migrating_from_omniisaacgymenvs.rst @@ -46,9 +46,9 @@ Below is an example skeleton of a task config class: # env decimation = 2 episode_length_s = 5.0 - num_actions = 1 - num_observations = 4 - num_states = 0 + action_space = 1 + observation_space = 4 + state_space = 0 # task-specific parameters ... @@ -158,9 +158,9 @@ The following parameters must be set for each environment config: decimation = 2 episode_length_s = 5.0 - num_actions = 1 - num_observations = 4 - num_states = 0 + action_space = 1 + observation_space = 4 + state_space = 0 RL Config Setup @@ -219,7 +219,7 @@ will automatically be created for the actor. This avoids the need to separately | self._cartpoles = ArticulationView( | # clone, filter, and replicate | | prim_paths_expr="/World/envs/.*/Cartpole", | self.scene.clone_environments(copy_from_source=False) | | name="cartpole_view", reset_xform_properties=False | self.scene.filter_collisions(global_prim_paths=[]) | -| ) | # add articultion to scene | +| ) | # add articulation to scene | | scene.add(self._cartpoles) | self.scene.articulations["cartpole"] = self.cartpole | | | # add lights | | | light_cfg = sim_utils.DomeLightCfg(intensity=2000.0) | @@ -364,8 +364,8 @@ In Isaac Lab, ``root_pose`` and ``root_velocity`` have been combined into single .. code-block::python - self.cartpole.write_root_pose_to_sim(default_root_state[:, :7], env_ids) - self.cartpole.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) + self.cartpole.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids) + self.cartpole.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids) Creating a New Environment @@ -501,9 +501,9 @@ Task config in Isaac Lab can be split into the main task configuration class and | clipObservations: 5.0 | decimation = 2 | | clipActions: 1.0 | episode_length_s = 5.0 | | controlFrequencyInv: 2 # 60 Hz | action_scale = 100.0 # [N] | -| | num_actions = 1 | -| sim: | num_observations = 4 | -| | num_states = 0 | +| | action_space = 1 | +| sim: | observation_space = 4 | +| | state_space = 0 | | dt: 0.0083 # 1/120 s | # reset | | use_gpu_pipeline: ${eq:${...pipeline},"gpu"} | max_cart_pos = 3.0 | | gravity: [0.0, 0.0, -9.81] | initial_pole_angle_range = [-0.25, 0.25] | @@ -638,7 +638,7 @@ Adding actors to the scene has been replaced by ``self.scene.articulations["cart | reset_xform_properties=False | copy_from_source=False) | | ) | self.scene.filter_collisions( | | scene.add(self._cartpoles) | global_prim_paths=[]) | -| return | # add articultion to scene | +| return | # add articulation to scene | | | self.scene.articulations["cartpole"] = self.cartpole | | def get_cartpole(self): | | | cartpole = Cartpole( | # add lights | @@ -738,9 +738,9 @@ reset the ``episode_length_buf`` buffer. | 1.0 - 2.0 * torch.rand(num_resets, device=self._device)) | self.joint_pos[env_ids] = joint_pos | | | self.joint_vel[env_ids] = joint_vel | | # apply resets | | -| indices = env_ids.to(dtype=torch.int32) | self.cartpole.write_root_pose_to_sim( | +| indices = env_ids.to(dtype=torch.int32) | self.cartpole.write_root_link_pose_to_sim( | | self._cartpoles.set_joint_positions(dof_pos, indices=indices) | default_root_state[:, :7], env_ids) | -| self._cartpoles.set_joint_velocities(dof_vel, indices=indices) | self.cartpole.write_root_velocity_to_sim( | +| self._cartpoles.set_joint_velocities(dof_vel, indices=indices) | self.cartpole.write_root_com_velocity_to_sim( | | | default_root_state[:, 7:], env_ids) | | # bookkeeping | self.cartpole.write_joint_state_to_sim( | | self.reset_buf[env_ids] = 0 | joint_pos, joint_vel, None, env_ids) | diff --git a/docs/source/overview/core-concepts/actuators.rst b/docs/source/overview/core-concepts/actuators.rst index 7513709b49..15bde0cdb8 100644 --- a/docs/source/overview/core-concepts/actuators.rst +++ b/docs/source/overview/core-concepts/actuators.rst @@ -74,4 +74,4 @@ The following figure shows the actuator groups for a legged mobile manipulator: .. seealso:: We provide implementations for various explicit actuator models. These are detailed in - `omni.isaac.lab.actuators <../api/lab/omni.isaac.lab.actuators.html>`_ sub-package. + `omni.isaac.lab.actuators <../../api/lab/omni.isaac.lab.actuators.html>`_ sub-package. diff --git a/docs/source/overview/core-concepts/task_workflows.rst b/docs/source/overview/core-concepts/task_workflows.rst index 4f3c96e0f4..7aeb78e8dc 100644 --- a/docs/source/overview/core-concepts/task_workflows.rst +++ b/docs/source/overview/core-concepts/task_workflows.rst @@ -43,6 +43,16 @@ For example, the observation manager is responsible for computing the observatio computing the rewards, and the termination manager is responsible for computing the termination signal. This approach is known as the manager-based environment design in the framework. +.. image:: ../../_static/task-workflows/manager-based-light.svg + :class: only-light + :align: center + :alt: Manager-based Task Workflow + +.. image:: ../../_static/task-workflows/manager-based-dark.svg + :class: only-dark + :align: center + :alt: Manager-based Task Workflow + Manager-based environments promote modular implementations of tasks by decomposing the task into individual components that are managed by separate classes. Each component of the task, such as rewards, observations, termination can all be specified as individual configuration classes that are then passed to the corresponding @@ -94,6 +104,16 @@ of the environment. This approach does not require the manager classes. Instead, to implement their task through the APIs from the base classes :class:`envs.DirectRLEnv` or :class:`envs.DirectMARLEnv`. For users migrating from the `IsaacGymEnvs`_ and `OmniIsaacGymEnvs`_ framework, this workflow may be more familiar. +.. image:: ../../_static/task-workflows/direct-based-light.svg + :class: only-light + :align: center + :alt: Direct-based Task Workflow + +.. image:: ../../_static/task-workflows/direct-based-dark.svg + :class: only-dark + :align: center + :alt: Direct-based Task Workflow + When defining an environment with the direct-style implementation, we expect the user define a single class that implements the entire environment. The task class should inherit from the base classes :class:`envs.DirectRLEnv` or :class:`envs.DirectMARLEnv` and should have its corresponding configuration class that inherits from diff --git a/docs/source/overview/developer-guide/template.rst b/docs/source/overview/developer-guide/template.rst index defaa2b78c..e6c55b9b61 100644 --- a/docs/source/overview/developer-guide/template.rst +++ b/docs/source/overview/developer-guide/template.rst @@ -22,4 +22,36 @@ This template serves three distinct use cases: features and improvements in Isaac Lab. -To get started, please follow the instructions in the `extension template repository `_. +Installation +------------ + +Install Isaac Lab by following the `installation guide <../../setup/installation/index.html>`_. We recommend using the conda installation as it simplifies calling Python scripts from the terminal. + +Clone the extension template repository separately from the Isaac Lab installation (i.e. outside the IsaacLab directory): + +.. code:: bash + + # Option 1: HTTPS + git clone https://github.com/isaac-sim/IsaacLabExtensionTemplate.git + + # Option 2: SSH + git clone git@github.com:isaac-sim/IsaacLabExtensionTemplate.git + +Throughout the repository, the name ``ext_template`` only serves as an example and we provide a script to rename all the references to it automatically: + +.. code:: bash + + # Enter the repository + cd IsaacLabExtensionTemplate + + # Rename all occurrences of ext_template (in files/directories) to your_fancy_extension_name + python scripts/rename_template.py your_fancy_extension_name + +Using a python interpreter that has Isaac Lab installed, install the library: + +.. code:: bash + + python -m pip install -e exts/ext_template + + +For more details, please follow the instructions in the `extension template repository `_. diff --git a/docs/source/overview/environments.rst b/docs/source/overview/environments.rst index f42f3c34a5..00400c721d 100644 --- a/docs/source/overview/environments.rst +++ b/docs/source/overview/environments.rst @@ -61,6 +61,10 @@ Classic environments that are based on IsaacGymEnvs implementation of MuJoCo-sty | | | | | | |cartpole-depth-direct-link|| | +------------------+-----------------------------+-------------------------------------------------------------------------+ + | |cartpole| | |cartpole-resnet-link| | Move the cart to keep the pole upwards in the classic cartpole control | + | | | based off of features extracted from perceptive inputs with pre-trained | + | | |cartpole-theia-link| | frozen vision encoders | + +------------------+-----------------------------+-------------------------------------------------------------------------+ .. |humanoid| image:: ../_static/tasks/classic/humanoid.jpg .. |ant| image:: ../_static/tasks/classic/ant.jpg @@ -69,8 +73,11 @@ Classic environments that are based on IsaacGymEnvs implementation of MuJoCo-sty .. |humanoid-link| replace:: `Isaac-Humanoid-v0 `__ .. |ant-link| replace:: `Isaac-Ant-v0 `__ .. |cartpole-link| replace:: `Isaac-Cartpole-v0 `__ -.. |cartpole-rgb-link| replace:: `Isaac-Cartpole-RGB-Camera-v0 `__ -.. |cartpole-depth-link| replace:: `Isaac-Cartpole-Depth-Camera-v0 `__ +.. |cartpole-rgb-link| replace:: `Isaac-Cartpole-RGB-v0 `__ +.. |cartpole-depth-link| replace:: `Isaac-Cartpole-Depth-v0 `__ +.. |cartpole-resnet-link| replace:: `Isaac-Cartpole-RGB-ResNet18-v0 `__ +.. |cartpole-theia-link| replace:: `Isaac-Cartpole-RGB-TheiaTiny-v0 `__ + .. |humanoid-direct-link| replace:: `Isaac-Humanoid-Direct-v0 `__ .. |ant-direct-link| replace:: `Isaac-Ant-Direct-v0 `__ @@ -84,7 +91,7 @@ Manipulation Environments based on fixed-arm manipulation tasks. For many of these tasks, we include configurations with different arm action spaces. For example, -for the reach environment: +for the lift-cube environment: * |lift-cube-link|: Franka arm with joint position control * |lift-cube-ik-abs-link|: Franka arm with absolute IK control @@ -102,6 +109,8 @@ for the reach environment: +--------------------+-------------------------+-----------------------------------------------------------------------------+ | |lift-cube| | |lift-cube-link| | Pick a cube and bring it to a sampled target position with the Franka robot | +--------------------+-------------------------+-----------------------------------------------------------------------------+ + | |stack-cube| | |stack-cube-link| | Stack three cubes (bottom to top: blue, red, green) with the Franka robot | + +--------------------+-------------------------+-----------------------------------------------------------------------------+ | |cabi-franka| | |cabi-franka-link| | Grasp the handle of a cabinet's drawer and open it with the Franka robot | | | | | | | |franka-direct-link| | | @@ -125,6 +134,7 @@ for the reach environment: .. |cabi-franka| image:: ../_static/tasks/manipulation/franka_open_drawer.jpg .. |cube-allegro| image:: ../_static/tasks/manipulation/allegro_cube.jpg .. |cube-shadow| image:: ../_static/tasks/manipulation/shadow_cube.jpg +.. |stack-cube| image:: ../_static/tasks/manipulation/franka_stack.jpg .. |reach-franka-link| replace:: `Isaac-Reach-Franka-v0 `__ .. |reach-ur10-link| replace:: `Isaac-Reach-UR10-v0 `__ @@ -135,12 +145,46 @@ for the reach environment: .. |franka-direct-link| replace:: `Isaac-Franka-Cabinet-Direct-v0 `__ .. |cube-allegro-link| replace:: `Isaac-Repose-Cube-Allegro-v0 `__ .. |allegro-direct-link| replace:: `Isaac-Repose-Cube-Allegro-Direct-v0 `__ +.. |stack-cube-link| replace:: `Isaac-Stack-Cube-Franka-v0 `__ .. |cube-shadow-link| replace:: `Isaac-Repose-Cube-Shadow-Direct-v0 `__ .. |cube-shadow-ff-link| replace:: `Isaac-Repose-Cube-Shadow-OpenAI-FF-Direct-v0 `__ .. |cube-shadow-lstm-link| replace:: `Isaac-Repose-Cube-Shadow-OpenAI-LSTM-Direct-v0 `__ .. |cube-shadow-vis-link| replace:: `Isaac-Repose-Cube-Shadow-Vision-Direct-v0 `__ +Contact-rich Manipulation +~~~~~~~~~~~~ + +Environments based on contact-rich manipulation tasks such as peg insertion, gear meshing and nut-bolt fastening. + +These tasks share the same task configurations and control options. You can switch between them by specifying the task name. +For example: + +* |factory-peg-link|: Peg insertion with the Franka arm +* |factory-gear-link|: Gear meshing with the Franka arm +* |factory-nut-link|: Nut-Bolt fastening with the Franka arm + +.. table:: + :widths: 33 37 30 + + +--------------------+-------------------------+-----------------------------------------------------------------------------+ + | World | Environment ID | Description | + +====================+=========================+=============================================================================+ + | |factory-peg| | |factory-peg-link| | Insert peg into the socket with the Franka robot | + +--------------------+-------------------------+-----------------------------------------------------------------------------+ + | |factory-gear| | |factory-gear-link| | Insert and mesh gear into the base with other gears, using the Franka robot | + +--------------------+-------------------------+-----------------------------------------------------------------------------+ + | |factory-nut| | |factory-nut-link| | Thread the nut onto the first 2 threads of the bolt, using the Franka robot | + +--------------------+-------------------------+-----------------------------------------------------------------------------+ + +.. |factory-peg| image:: ../_static/tasks/factory/peg_insert.jpg +.. |factory-gear| image:: ../_static/tasks/factory/gear_mesh.jpg +.. |factory-nut| image:: ../_static/tasks/factory/nut_thread.jpg + +.. |factory-peg-link| replace:: `Isaac-Factory-PegInsert-Direct-v0 `__ +.. |factory-gear-link| replace:: `Isaac-Factory-GearMesh-Direct-v0 `__ +.. |factory-nut-link| replace:: `Isaac-Factory-NutThread-Direct-v0 `__ + Locomotion ~~~~~~~~~~ @@ -358,6 +402,18 @@ Comprehensive List of Environments - - Manager Based - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO) + * - Isaac-Factory-GearMesh-Direct-v0 + - + - Direct + - **rl_games** (PPO) + * - Isaac-Factory-NutThread-Direct-v0 + - + - Direct + - **rl_games** (PPO) + * - Isaac-Factory-PegInsert-Direct-v0 + - + - Direct + - **rl_games** (PPO) * - Isaac-Franka-Cabinet-Direct-v0 - - Direct @@ -410,6 +466,10 @@ Comprehensive List of Environments - - Manager Based - + * - Isaac-Reach-Franka-OSC-v0 + - Isaac-Reach-Franka-OSC-Play-v0 + - Manager Based + - **rsl_rl** (PPO) * - Isaac-Reach-Franka-v0 - Isaac-Reach-Franka-Play-v0 - Manager Based diff --git a/docs/source/overview/reinforcement-learning/rl_existing_scripts.rst b/docs/source/overview/reinforcement-learning/rl_existing_scripts.rst index 6cc88137f2..6ee865fd36 100644 --- a/docs/source/overview/reinforcement-learning/rl_existing_scripts.rst +++ b/docs/source/overview/reinforcement-learning/rl_existing_scripts.rst @@ -4,12 +4,12 @@ Reinforcement Learning Wrappers We provide wrappers to different reinforcement libraries. These wrappers convert the data from the environments into the respective libraries function argument and return types. -Stable-Baselines3 ------------------ + +RL-Games +-------- - Training an agent with - `Stable-Baselines3 `__ - on ``Isaac-Cartpole-v0``: + `RL-Games `__ on ``Isaac-Ant-v0``: .. tab-set:: :sync-group: os @@ -19,30 +19,65 @@ Stable-Baselines3 .. code:: bash - # install python module (for stable-baselines3) - ./isaaclab.sh -i sb3 + # install python module (for rl-games) + ./isaaclab.sh -i rl_games # run script for training - # note: we set the device to cpu since SB3 doesn't optimize for GPU anyway - ./isaaclab.sh -p source/standalone/workflows/sb3/train.py --task Isaac-Cartpole-v0 --headless --device cpu + ./isaaclab.sh -p source/standalone/workflows/rl_games/train.py --task Isaac-Ant-v0 --headless # run script for playing with 32 environments - ./isaaclab.sh -p source/standalone/workflows/sb3/play.py --task Isaac-Cartpole-v0 --num_envs 32 --checkpoint /PATH/TO/model.zip + ./isaaclab.sh -p source/standalone/workflows/rl_games/play.py --task Isaac-Ant-v0 --num_envs 32 --checkpoint /PATH/TO/model.pth # run script for recording video of a trained agent (requires installing `ffmpeg`) - ./isaaclab.sh -p source/standalone/workflows/sb3/play.py --task Isaac-Cartpole-v0 --headless --video --video_length 200 + ./isaaclab.sh -p source/standalone/workflows/rl_games/play.py --task Isaac-Ant-v0 --headless --video --video_length 200 .. tab-item:: :icon:`fa-brands fa-windows` Windows :sync: windows .. code:: batch - :: install python module (for stable-baselines3) - isaaclab.bat -i sb3 + :: install python module (for rl-games) + isaaclab.bat -i rl_games :: run script for training - :: note: we set the device to cpu since SB3 doesn't optimize for GPU anyway - isaaclab.bat -p source\standalone\workflows\sb3\train.py --task Isaac-Cartpole-v0 --headless --device cpu + isaaclab.bat -p source\standalone\workflows\rl_games\train.py --task Isaac-Ant-v0 --headless :: run script for playing with 32 environments - isaaclab.bat -p source\standalone\workflows\sb3\play.py --task Isaac-Cartpole-v0 --num_envs 32 --checkpoint /PATH/TO/model.zip + isaaclab.bat -p source\standalone\workflows\rl_games\play.py --task Isaac-Ant-v0 --num_envs 32 --checkpoint /PATH/TO/model.pth :: run script for recording video of a trained agent (requires installing `ffmpeg`) - isaaclab.bat -p source\standalone\workflows\sb3\play.py --task Isaac-Cartpole-v0 --headless --video --video_length 200 + isaaclab.bat -p source\standalone\workflows\rl_games\play.py --task Isaac-Ant-v0 --headless --video --video_length 200 + +RSL-RL +------ + +- Training an agent with + `RSL-RL `__ on ``Isaac-Reach-Franka-v0``: + + .. tab-set:: + :sync-group: os + + .. tab-item:: :icon:`fa-brands fa-linux` Linux + :sync: linux + + .. code:: bash + + # install python module (for rsl-rl) + ./isaaclab.sh -i rsl_rl + # run script for training + ./isaaclab.sh -p source/standalone/workflows/rsl_rl/train.py --task Isaac-Reach-Franka-v0 --headless + # run script for playing with 32 environments + ./isaaclab.sh -p source/standalone/workflows/rsl_rl/play.py --task Isaac-Reach-Franka-v0 --num_envs 32 --load_run run_folder_name --checkpoint model.pt + # run script for recording video of a trained agent (requires installing `ffmpeg`) + ./isaaclab.sh -p source/standalone/workflows/rsl_rl/play.py --task Isaac-Reach-Franka-v0 --headless --video --video_length 200 + + .. tab-item:: :icon:`fa-brands fa-windows` Windows + :sync: windows + + .. code:: batch + + :: install python module (for rsl-rl) + isaaclab.bat -i rsl_rl + :: run script for training + isaaclab.bat -p source\standalone\workflows\rsl_rl\train.py --task Isaac-Reach-Franka-v0 --headless + :: run script for playing with 32 environments + isaaclab.bat -p source\standalone\workflows\rsl_rl\play.py --task Isaac-Reach-Franka-v0 --num_envs 32 --load_run run_folder_name --checkpoint model.pt + :: run script for recording video of a trained agent (requires installing `ffmpeg`) + isaaclab.bat -p source\standalone\workflows\rsl_rl\play.py --task Isaac-Reach-Franka-v0 --headless --video --video_length 200 SKRL ---- @@ -129,48 +164,12 @@ SKRL :: run script for playing with 32 environments with the MAPPO algorithm (IPPO is also supported) isaaclab.bat -p source\standalone\workflows\skrl\play.py --task Isaac-Shadow-Hand-Over-Direct-v0 --num_envs 32 --algorithm MAPPO --checkpoint /PATH/TO/model.pt -RL-Games --------- - -- Training an agent with - `RL-Games `__ on ``Isaac-Ant-v0``: - - .. tab-set:: - :sync-group: os - - .. tab-item:: :icon:`fa-brands fa-linux` Linux - :sync: linux - - .. code:: bash - - # install python module (for rl-games) - ./isaaclab.sh -i rl_games - # run script for training - ./isaaclab.sh -p source/standalone/workflows/rl_games/train.py --task Isaac-Ant-v0 --headless - # run script for playing with 32 environments - ./isaaclab.sh -p source/standalone/workflows/rl_games/play.py --task Isaac-Ant-v0 --num_envs 32 --checkpoint /PATH/TO/model.pth - # run script for recording video of a trained agent (requires installing `ffmpeg`) - ./isaaclab.sh -p source/standalone/workflows/rl_games/play.py --task Isaac-Ant-v0 --headless --video --video_length 200 - - .. tab-item:: :icon:`fa-brands fa-windows` Windows - :sync: windows - - .. code:: batch - - :: install python module (for rl-games) - isaaclab.bat -i rl_games - :: run script for training - isaaclab.bat -p source\standalone\workflows\rl_games\train.py --task Isaac-Ant-v0 --headless - :: run script for playing with 32 environments - isaaclab.bat -p source\standalone\workflows\rl_games\play.py --task Isaac-Ant-v0 --num_envs 32 --checkpoint /PATH/TO/model.pth - :: run script for recording video of a trained agent (requires installing `ffmpeg`) - isaaclab.bat -p source\standalone\workflows\rl_games\play.py --task Isaac-Ant-v0 --headless --video --video_length 200 - -RSL-RL ------- +Stable-Baselines3 +----------------- - Training an agent with - `RSL-RL `__ on ``Isaac-Reach-Franka-v0``: + `Stable-Baselines3 `__ + on ``Isaac-Cartpole-v0``: .. tab-set:: :sync-group: os @@ -180,28 +179,30 @@ RSL-RL .. code:: bash - # install python module (for rsl-rl) - ./isaaclab.sh -i rsl_rl + # install python module (for stable-baselines3) + ./isaaclab.sh -i sb3 # run script for training - ./isaaclab.sh -p source/standalone/workflows/rsl_rl/train.py --task Isaac-Reach-Franka-v0 --headless + # note: we set the device to cpu since SB3 doesn't optimize for GPU anyway + ./isaaclab.sh -p source/standalone/workflows/sb3/train.py --task Isaac-Cartpole-v0 --headless --device cpu # run script for playing with 32 environments - ./isaaclab.sh -p source/standalone/workflows/rsl_rl/play.py --task Isaac-Reach-Franka-v0 --num_envs 32 --load_run run_folder_name --checkpoint model.pt + ./isaaclab.sh -p source/standalone/workflows/sb3/play.py --task Isaac-Cartpole-v0 --num_envs 32 --checkpoint /PATH/TO/model.zip # run script for recording video of a trained agent (requires installing `ffmpeg`) - ./isaaclab.sh -p source/standalone/workflows/rsl_rl/play.py --task Isaac-Reach-Franka-v0 --headless --video --video_length 200 + ./isaaclab.sh -p source/standalone/workflows/sb3/play.py --task Isaac-Cartpole-v0 --headless --video --video_length 200 .. tab-item:: :icon:`fa-brands fa-windows` Windows :sync: windows .. code:: batch - :: install python module (for rsl-rl) - isaaclab.bat -i rsl_rl + :: install python module (for stable-baselines3) + isaaclab.bat -i sb3 :: run script for training - isaaclab.bat -p source\standalone\workflows\rsl_rl\train.py --task Isaac-Reach-Franka-v0 --headless + :: note: we set the device to cpu since SB3 doesn't optimize for GPU anyway + isaaclab.bat -p source\standalone\workflows\sb3\train.py --task Isaac-Cartpole-v0 --headless --device cpu :: run script for playing with 32 environments - isaaclab.bat -p source\standalone\workflows\rsl_rl\play.py --task Isaac-Reach-Franka-v0 --num_envs 32 --load_run run_folder_name --checkpoint model.pt + isaaclab.bat -p source\standalone\workflows\sb3\play.py --task Isaac-Cartpole-v0 --num_envs 32 --checkpoint /PATH/TO/model.zip :: run script for recording video of a trained agent (requires installing `ffmpeg`) - isaaclab.bat -p source\standalone\workflows\rsl_rl\play.py --task Isaac-Reach-Franka-v0 --headless --video --video_length 200 + isaaclab.bat -p source\standalone\workflows\sb3\play.py --task Isaac-Cartpole-v0 --headless --video --video_length 200 All the scripts above log the training progress to `Tensorboard`_ in the ``logs`` directory in the root of the repository. The logs directory follows the pattern ``logs///``, where ```` diff --git a/docs/source/features/tiled_rendering.rst b/docs/source/overview/sensors/camera.rst similarity index 52% rename from docs/source/features/tiled_rendering.rst rename to docs/source/overview/sensors/camera.rst index 55d22561a4..0ef72196c5 100644 --- a/docs/source/features/tiled_rendering.rst +++ b/docs/source/overview/sensors/camera.rst @@ -1,28 +1,26 @@ -Tiled-Camera Rendering -====================== +.. _overview_sensors_camera: -.. currentmodule:: omni.isaac.lab + +Camera +======== + +Camera sensors are uniquely defined by the use of the ``render_product``, a structure for managing data generated by the rendering pipeline (images). Isaac Lab provides the ability to fully control how these renderings are created through camera parameters like focal length, pose, type, etc... and what kind of data you want to render through the use of Annotators, allowing you to record not only RGB, but also Instance segmentation, object pose, object ID, etc... + +Rendered images are unique among the supported data types in Isaac Lab due to the inherently large bandwidth requirements for moving those data. A single 800 x 600 image with 32-bit color (a single float per pixel) clocks in at just under 2 MB. If we render at 60 fps and record every frame, that camera needs to move 120 MB/s. Multiply this by the number of cameras in an environment and environments in a simulation, and you can quickly see how scaling a naive vectorization of camera data could lead to bandwidth challenges. NVIDIA's Isaac Lab leverages our expertise in GPU hardware to provide an API that specifically addresses these scaling challenges in the rendering pipeline. + +Tiled Rendering +~~~~~~~~~~~~~~~~~ .. note:: This feature is only available from Isaac Sim version 4.2.0 onwards. Tiled rendering in combination with image processing networks require heavy memory resources, especially - at larger resolutions. We recommend running at 512 cameras in the scene on RTX 4090 GPUs or similar. - + at larger resolutions. We recommend running 512 cameras in the scene on RTX 4090 GPUs or similar. -Tiled rendering APIs provide a vectorized interface for collecting data from camera sensors. -This is useful for reinforcement learning environments requiring vision in the loop. -Tiled rendering works by concatenating camera outputs from multiple cameras and rendering -one single large image instead of multiple smaller images that would have been produced -by each individual camera. This reduces the amount of time required for rendering and -provides a more efficient API for working with vision data. +The Tiled Rendering APIs provide a vectorized interface for collecting data from camera sensors. This is useful for reinforcement learning environments where parallelization can be exploited to accelerate data collection and thus the training loop. Tiled rendering works by using a single ``render_product`` for **all** clones of a single camera in the scene. The desired dimensions of a single image and the number of environments are used to compute a much larger ``render_product``, consisting of the tiled individual renders from the separate clones of the camera. When all cameras have populated their buffers the render product is "completed" and can be moved around as a single, large image, dramatically reducing the overhead for moving the data from the host to the device, for example. Only a single call is used to synchronize the device data, instead of one call per camera, and this is a big part of what makes the Tiled Rendering API more efficient for working with vision data. -Isaac Lab provides tiled rendering APIs for RGB, depth, along with other annotators through the :class:`~sensors.TiledCamera` -class. Configurations for the tiled rendering APIs can be defined through the :class:`~sensors.TiledCameraCfg` -class, specifying parameters such as the regex expression for all camera paths, the transform -for the cameras, the desired data type, the type of cameras to add to the scene, and the camera -resolution. +Isaac Lab provides tiled rendering APIs for RGB, depth, along with other annotators through the :class:`~sensors.TiledCamera` class. Configurations for the tiled rendering APIs can be defined through the :class:`~sensors.TiledCameraCfg` class, specifying parameters such as the regex expression for all camera paths, the transform for the cameras, the desired data type, the type of cameras to add to the scene, and the camera resolution. .. code-block:: python @@ -37,8 +35,7 @@ resolution. height=80, ) -To access the tiled rendering interface, a :class:`~sensors.TiledCamera` object can be created and used -to retrieve data from the cameras. +To access the tiled rendering interface, a :class:`~sensors.TiledCamera` object can be created and used to retrieve data from the cameras. .. code-block:: python @@ -46,19 +43,17 @@ to retrieve data from the cameras. data_type = "rgb" data = tiled_camera.data.output[data_type] -The returned data will be transformed into the shape (num_cameras, height, width, num_channels), which -can be used directly as observation for reinforcement learning. +The returned data will be transformed into the shape (num_cameras, height, width, num_channels), which can be used directly as observation for reinforcement learning. -When working with rendering, make sure to add the ``--enable_cameras`` argument when launching the -environment. For example: +When working with rendering, make sure to add the ``--enable_cameras`` argument when launching the environment. For example: .. code-block:: shell python source/standalone/workflows/rl_games/train.py --task=Isaac-Cartpole-RGB-Camera-Direct-v0 --headless --enable_cameras -Annotators and Data Types -------------------------- +Annotators +~~~~~~~~~~~~~~~~~ Both :class:`~sensors.TiledCamera` and :class:`~sensors.Camera` classes provide APIs for retrieving various types annotator data from replicator: @@ -76,6 +71,11 @@ Both :class:`~sensors.TiledCamera` and :class:`~sensors.Camera` classes provide RGB and RGBA ~~~~~~~~~~~~ +.. figure:: ../../_static/overview/overview_sensors_rgb.png + :align: center + :figwidth: 100% + :alt: A scene captured in RGB + ``rgb`` data type returns a 3-channel RGB colored image of type ``torch.uint8``, with dimension (B, H, W, 3). ``rgba`` data type returns a 4-channel RGBA colored image of type ``torch.uint8``, with dimension (B, H, W, 4). @@ -85,6 +85,11 @@ To convert the ``torch.uint8`` data to ``torch.float32``, divide the buffer by 2 Depth and Distances ~~~~~~~~~~~~~~~~~~~ +.. figure:: ../../_static/overview/overview_sensors_depth.png + :align: center + :figwidth: 100% + :alt: A scene captured in RGB + ``distance_to_camera`` returns a single-channel depth image with distance to the camera optical center. The dimension for this annotator is (B, H, W, 1) and has type ``torch.float32``. ``distance_to_image_plane`` returns a single-channel depth image with distances of 3D points from the camera plane along the camera's Z-axis. The dimension for this annotator is (B, H, W, 1) and has type ``torch.float32``. @@ -94,6 +99,11 @@ Depth and Distances Normals ~~~~~~~ +.. figure:: ../../_static/overview/overview_sensors_normals.png + :align: center + :figwidth: 100% + :alt: A scene captured in RGB + ``normals`` returns an image containing the local surface normal vectors at each pixel. The buffer has dimension (B, H, W, 3), containing the (x, y, z) information for each vector, and has data type ``torch.float32``. Motion Vectors @@ -104,6 +114,11 @@ Motion Vectors Semantic Segmentation ~~~~~~~~~~~~~~~~~~~~~ +.. figure:: ../../_static/overview/overview_sensors_semantic.png + :align: center + :figwidth: 100% + :alt: A scene captured in RGB + ``semantic_segmentation`` outputs semantic segmentation of each entity in the camera’s viewport that has semantic labels. In addition to the image buffer, an ``info`` dictionary can be retrieved with ``tiled_camera.data.info['semantic_segmentation']`` containing ID to labels information. - If ``colorize_semantic_segmentation=True`` in the camera config, a 4-channel RGBA image will be returned with dimension (B, H, W, 4) and type ``torch.uint8``. The info ``idToLabels`` dictionary will be the mapping from color to semantic labels. @@ -113,6 +128,11 @@ Semantic Segmentation Instance ID Segmentation ~~~~~~~~~~~~~~~~~~~~~~~~ +.. figure:: ../../_static/overview/overview_sensors_instanceID.png + :align: center + :figwidth: 100% + :alt: A scene captured in RGB + ``instance_id_segmentation_fast`` outputs instance ID segmentation of each entity in the camera’s viewport. The instance ID is unique for each prim in the scene with different paths. In addition to the image buffer, an ``info`` dictionary can be retrieved with ``tiled_camera.data.info['instance_id_segmentation_fast']`` containing ID to labels information. The main difference between ``instance_id_segmentation_fast`` and ``instance_segmentation_fast`` are that instance segmentation annotator goes down the hierarchy to the lowest level prim which has semantic labels, where instance ID segmentation always goes down to the leaf prim. @@ -124,8 +144,51 @@ The main difference between ``instance_id_segmentation_fast`` and ``instance_seg Instance Segmentation """"""""""""""""""""" +.. figure:: ../../_static/overview/overview_sensors_instance.png + :align: center + :figwidth: 100% + :alt: A scene captured in RGB + ``instance_segmentation_fast`` outputs instance segmentation of each entity in the camera’s viewport. In addition to the image buffer, an ``info`` dictionary can be retrieved with ``tiled_camera.data.info['instance_segmentation_fast']`` containing ID to labels and ID to semantic information. -- If ``colorize_instance_segmentation=True`` in the camera config, a 4-channel RGBA image will be returned with dimension (B, H, W, 4) and type ``torch.uint8``. The info ``idToLabels`` dictionary will be the mapping from color to USD prim path of that semantic entity. The info ``idToSemantics`` dictionary will be the mapping from color to semantic labels of that semantic entity. +- If ``colorize_instance_segmentation=True`` in the camera config, a 4-channel RGBA image will be returned with dimension (B, H, W, 4) and type ``torch.uint8``. + +- If ``colorize_instance_segmentation=False``, a buffer of dimension (B, H, W, 1) of type ``torch.int32`` will be returned, containing the instance ID of each pixel. + +The info ``idToLabels`` dictionary will be the mapping from color to USD prim path of that semantic entity. The info ``idToSemantics`` dictionary will be the mapping from color to semantic labels of that semantic entity. + + +Current Limitations +------------------- -- If ``colorize_instance_segmentation=False``, a buffer of dimension (B, H, W, 1) of type ``torch.int32`` will be returned, containing the instance ID of each pixel. The info ``idToLabels`` dictionary will be the mapping from instance ID to USD prim path of that semantic entity. The info ``idToSemantics`` dictionary will be the mapping from instance ID to semantic labels of that semantic entity. +Due to current limitations in the renderer, we can have only **one** :class:`~sensors.TiledCamera` instance in the scene. +For use cases that require a setup with more than one camera, we can imitate the multi-camera behavior by moving the location +of the camera in between render calls in a step. + +For example, in a stereo vision setup, the below snippet can be implemented: + +.. code-block:: python + + # render image from "first" camera + camera_data_1 = self._tiled_camera.data.output["rgb"].clone() / 255.0 + # update camera transform to the "second" camera location + self._tiled_camera.set_world_poses( + positions=pos, + orientations=rot, + convention="world" + ) + # step the renderer + self.sim.render() + self._tiled_camera.update(0, force_recompute=True) + # render image from "second" camera + camera_data_2 = self._tiled_camera.data.output["rgb"].clone() / 255.0 + +Note that this approach still limits the rendering resolution to be identical for all cameras. Currently, there is no workaround +to achieve different resolution images using :class:`~sensors.TiledCamera`. The best approach is to use the largest resolution out of all of the +desired resolutions and add additional scaling or cropping operations to the rendered output as a post-processing step. + +In addition, there may be visible quality differences when comparing render outputs of different numbers of environments. +Currently, any combined resolution that has a width less than 265 pixels or height less than 265 will automatically switch +to the DLAA anti-aliasing mode, which does not perform up-sampling during anti-aliasing. For resolutions larger than 265 in both +width and height dimensions, we default to using the "performance" DLSS mode for anti-aliasing for performance benefits. +Anti-aliasing modes and other rendering parameters can be specified in the :class:`~sim.RenderCfg`. diff --git a/docs/source/overview/sensors/contact_sensor.rst b/docs/source/overview/sensors/contact_sensor.rst new file mode 100644 index 0000000000..72239fd287 --- /dev/null +++ b/docs/source/overview/sensors/contact_sensor.rst @@ -0,0 +1,145 @@ +.. _overview_sensors_contact: + +Contact Sensor +================ + +.. figure:: ../../_static/overview/overview_sensors_contact_diagram.png + :align: center + :figwidth: 100% + :alt: A contact sensor with filtering + +The contact sensor is designed to return the net contact force acting on a given ridgid body. The sensor is written to behave as a physical object, and so the "scope" of the contact sensor is limited to the body (or bodies) that defines it. There are multiple ways to define this scope, depending on your need to filter the forces coming from the contact. + +By default, the reported force is the total contact force, but your application may only care about contact forces due to specific objects. Retrieving contact forces from specific objects requires filtering, and this can only be done in a "many-to-one" way. A multi-legged robot that needs filterable contact information for its feet would require one sensor per foot to be defined in the environment, but a robotic hand with contact sensors on the tips of each finger can be defined with a single sensor. + +Consider a simple environment with an Anymal Quadruped and a block + +.. code-block:: python + + @configclass + class ContactSensorsSceneCfg(InteractiveSceneCfg): + """Design the scene with sensors on the robot.""" + + # ground plane + ground = AssetBaseCfg(prim_path="/World/defaultGroundPlane", spawn=sim_utils.GroundPlaneCfg()) + + # lights + dome_light = AssetBaseCfg( + prim_path="/World/Light", spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75)) + ) + + # robot + robot = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + # Rigid Object + cube = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube", + spawn=sim_utils.CuboidCfg( + size=(0.5,0.5,0.1), + rigid_props=sim_utils.RigidBodyPropertiesCfg(), + mass_props=sim_utils.MassPropertiesCfg(mass=100.0), + collision_props=sim_utils.CollisionPropertiesCfg(), + physics_material=sim_utils.RigidBodyMaterialCfg(static_friction=1.0), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 1.0, 0.0), metallic=0.2), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.5, 0.5, 0.05)), + ) + + contact_forces_LF = ContactSensorCfg( + prim_path="{ENV_REGEX_NS}/Robot/LF_FOOT", + update_period=0.0, + history_length=6, + debug_vis=True, + filter_prim_paths_expr=["{ENV_REGEX_NS}/Cube"], + ) + + contact_forces_RF = ContactSensorCfg( + prim_path="{ENV_REGEX_NS}/Robot/RF_FOOT", + update_period=0.0, + history_length=6, + debug_vis=True, + filter_prim_paths_expr=["{ENV_REGEX_NS}/Cube"], + ) + + contact_forces_H = ContactSensorCfg( + prim_path="{ENV_REGEX_NS}/Robot/.*H_FOOT", + update_period=0.0, + history_length=6, + debug_vis=True, + ) + +We define the sensors on the feet of the robot in two different ways. The front feet are independent sensors (one sensor body per foot) and the "Cube" is placed under the left foot. The hind feet are defined as a single sensor with multiple bodies. + +We can then run the scene and print the data from the sensors + +.. code-block:: python + + def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): + . + . + . + # Simulate physics + while simulation_app.is_running(): + . + . + . + # print information from the sensors + print("-------------------------------") + print(scene["contact_forces_LF"]) + print("Received force matrix of: ", scene["contact_forces_LF"].data.force_matrix_w) + print("Received contact force of: ", scene["contact_forces_LF"].data.net_forces_w) + print("-------------------------------") + print(scene["contact_forces_RF"]) + print("Received force matrix of: ", scene["contact_forces_RF"].data.force_matrix_w) + print("Received contact force of: ", scene["contact_forces_RF"].data.net_forces_w) + print("-------------------------------") + print(scene["contact_forces_H"]) + print("Received force matrix of: ", scene["contact_forces_H"].data.force_matrix_w) + print("Received contact force of: ", scene["contact_forces_H"].data.net_forces_w) + +Here, we print both the net contact force and the filtered force matrix for each contact sensor defined in the scene. The front left and front right feet report the following + +.. code-block:: bash + + ------------------------------- + Contact sensor @ '/World/envs/env_.*/Robot/LF_FOOT': + view type : + update period (s) : 0.0 + number of bodies : 1 + body names : ['LF_FOOT'] + + Received force matrix of: tensor([[[[-1.3923e-05, 1.5727e-04, 1.1032e+02]]]], device='cuda:0') + Received contact force of: tensor([[[-1.3923e-05, 1.5727e-04, 1.1032e+02]]], device='cuda:0') + ------------------------------- + Contact sensor @ '/World/envs/env_.*/Robot/RF_FOOT': + view type : + update period (s) : 0.0 + number of bodies : 1 + body names : ['RF_FOOT'] + + Received force matrix of: tensor([[[[0., 0., 0.]]]], device='cuda:0') + Received contact force of: tensor([[[1.3529e-05, 0.0000e+00, 1.0069e+02]]], device='cuda:0') + +Notice that even with filtering, both sensors report the net contact force acting on the foot. However only the left foot has a non zero "force matrix", because the right foot isn't standing on the filtered body, ``/World/envs/env_.*/Cube``. Now, checkout the data coming from the hind feet! + +.. code-block:: bash + + ------------------------------- + Contact sensor @ '/World/envs/env_.*/Robot/.*H_FOOT': + view type : + update period (s) : 0.0 + number of bodies : 2 + body names : ['LH_FOOT', 'RH_FOOT'] + + Received force matrix of: None + Received contact force of: tensor([[[9.7227e-06, 0.0000e+00, 7.2364e+01], + [2.4322e-05, 0.0000e+00, 1.8102e+02]]], device='cuda:0') + +In this case, the contact sensor has two bodies: the left and right hind feet. When the force matrix is queried, the result is ``None`` because this is a many body sensor, and presently Isaac Lab only supports "many to one" contact force filtering. Unlike the single body contact sensor, the reported force tensor has multiple entries, with each "row" corresponding to the contact force on a single body of the sensor (matching the ordering at construction). + +.. dropdown:: Code for contact_sensor.py + :icon: code + + .. literalinclude:: ../../../../source/standalone/demos/sensors/contact_sensor.py + :language: python + :linenos: diff --git a/docs/source/overview/sensors/frame_transformer.rst b/docs/source/overview/sensors/frame_transformer.rst new file mode 100644 index 0000000000..d60dd9304c --- /dev/null +++ b/docs/source/overview/sensors/frame_transformer.rst @@ -0,0 +1,158 @@ +.. _overview_sensors_frame_transformer: + +Frame Transformer +==================== + +.. figure:: ../../_static/overview/overview_sensors_frame_transformer.png + :align: center + :figwidth: 100% + :alt: A diagram outlining the basic geometry of frame transformations + +.. + Do YOU want to know where things are relative to other things at a glance? Then the frame transformer is the sensor for you!* + +One of the most common operations that needs to be performed within a physics simulation is the frame transformation: rewriting a vector or quaternion in the basis of an arbitrary euclidean coordinate system. There are many ways to accomplish this within Isaac and USD, but these methods can be cumbersome to implement within Isaac Lab's GPU based simulation and cloned environments. To mitigate this problem, we have designed the Frame Transformer Sensor, that tracks and calculate the relative frame transformations for rigid bodies of interest to the scene. + +The sensory is minimally defined by a source frame and a list of target frames. These definitions take the form of a prim path (for the source) and list of regex capable prim paths the rigid bodies to be tracked (for the targets). + +.. code-block:: python + + @configclass + class FrameTransformerSensorSceneCfg(InteractiveSceneCfg): + """Design the scene with sensors on the robot.""" + + # ground plane + ground = AssetBaseCfg(prim_path="/World/defaultGroundPlane", spawn=sim_utils.GroundPlaneCfg()) + + # lights + dome_light = AssetBaseCfg( + prim_path="/World/Light", spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75)) + ) + + # robot + robot = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + # Rigid Object + cube = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube", + spawn=sim_utils.CuboidCfg( + size=(1,1,1), + rigid_props=sim_utils.RigidBodyPropertiesCfg(), + mass_props=sim_utils.MassPropertiesCfg(mass=100.0), + collision_props=sim_utils.CollisionPropertiesCfg(), + physics_material=sim_utils.RigidBodyMaterialCfg(static_friction=1.0), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 1.0, 0.0), metallic=0.2), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(5, 0, 0.5)), + ) + + specific_transforms = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/base", + target_frames=[ + FrameTransformerCfg.FrameCfg(prim_path="{ENV_REGEX_NS}/Robot/LF_FOOT"), + FrameTransformerCfg.FrameCfg(prim_path="{ENV_REGEX_NS}/Robot/RF_FOOT"), + ], + debug_vis=True, + ) + + cube_transform = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/base", + target_frames=[ + FrameTransformerCfg.FrameCfg(prim_path="{ENV_REGEX_NS}/Cube") + ], + debug_vis=False, + ) + + robot_transforms = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/base", + target_frames=[ + FrameTransformerCfg.FrameCfg(prim_path="{ENV_REGEX_NS}/Robot/.*") + ], + debug_vis=False, + ) + + +We can now run the scene and query the sensor for data + +.. code-block:: python + + def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): + . + . + . + # Simulate physics + while simulation_app.is_running(): + . + . + . + + # print information from the sensors + print("-------------------------------") + print(scene["specific_transforms"]) + print("relative transforms:", scene["specific_transforms"].data.target_pos_source) + print("relative orientations:", scene["specific_transforms"].data.target_quat_source) + print("-------------------------------") + print(scene["cube_transform"]) + print("relative transform:", scene["cube_transform"].data.target_pos_source) + print("-------------------------------") + print(scene["robot_transforms"]) + print("relative transforms:", scene["robot_transforms"].data.target_pos_source) + +Let's take a look at the result for tracking specific objects. First, we can take a look at the data coming from the sensors on the feet + +.. code-block:: bash + + ------------------------------- + FrameTransformer @ '/World/envs/env_.*/Robot/base': + tracked body frames: ['base', 'LF_FOOT', 'RF_FOOT'] + number of envs: 1 + source body frame: base + target frames (count: ['LF_FOOT', 'RF_FOOT']): 2 + + relative transforms: tensor([[[ 0.4658, 0.3085, -0.4840], + [ 0.4487, -0.2959, -0.4828]]], device='cuda:0') + relative orientations: tensor([[[ 0.9623, 0.0072, -0.2717, -0.0020], + [ 0.9639, 0.0052, -0.2663, -0.0014]]], device='cuda:0') + +.. figure:: ../../_static/overview/overview_sensors_ft_visualizer.png + :align: center + :figwidth: 100% + :alt: The frame transformer visualizer + +By activating the visualizer, we can see that the frames of the feet are rotated "upward" slightly. We can also see the explicit relative positions and rotations by querying the sensor for data, which returns these values as a list with the same order as the tracked frames. This becomes even more apparent if we examine the transforms specified by regex. + +.. code-block:: bash + + ------------------------------- + FrameTransformer @ '/World/envs/env_.*/Robot/base': + tracked body frames: ['base', 'LF_FOOT', 'LF_HIP', 'LF_SHANK', 'LF_THIGH', 'LH_FOOT', 'LH_HIP', 'LH_SHANK', 'LH_THIGH', 'RF_FOOT', 'RF_HIP', 'RF_SHANK', 'RF_THIGH', 'RH_FOOT', 'RH_HIP', 'RH_SHANK', 'RH_THIGH', 'base'] + number of envs: 1 + source body frame: base + target frames (count: ['LF_FOOT', 'LF_HIP', 'LF_SHANK', 'LF_THIGH', 'LH_FOOT', 'LH_HIP', 'LH_SHANK', 'LH_THIGH', 'RF_FOOT', 'RF_HIP', 'RF_SHANK', 'RF_THIGH', 'RH_FOOT', 'RH_HIP', 'RH_SHANK', 'RH_THIGH', 'base']): 17 + + relative transforms: tensor([[[ 4.6581e-01, 3.0846e-01, -4.8398e-01], + [ 2.9990e-01, 1.0400e-01, -1.7062e-09], + [ 2.1409e-01, 2.9177e-01, -2.4214e-01], + [ 3.5980e-01, 1.8780e-01, 1.2608e-03], + [-4.8813e-01, 3.0973e-01, -4.5927e-01], + [-2.9990e-01, 1.0400e-01, 2.7044e-09], + [-2.1495e-01, 2.9264e-01, -2.4198e-01], + [-3.5980e-01, 1.8780e-01, 1.5582e-03], + [ 4.4871e-01, -2.9593e-01, -4.8277e-01], + [ 2.9990e-01, -1.0400e-01, -2.7057e-09], + [ 1.9971e-01, -2.8554e-01, -2.3778e-01], + [ 3.5980e-01, -1.8781e-01, -9.1049e-04], + [-5.0090e-01, -2.9095e-01, -4.5746e-01], + [-2.9990e-01, -1.0400e-01, 6.3592e-09], + [-2.1860e-01, -2.8251e-01, -2.5163e-01], + [-3.5980e-01, -1.8779e-01, -1.8792e-03], + [ 0.0000e+00, 0.0000e+00, 0.0000e+00]]], device='cuda:0') + +Here, the sensor is tracking all rigid body children of ``Robot/base``, but this expression is **inclusive**, meaning that the source body itself is also a target. This can be seen both by examining the source and target list, where ``base`` appears twice, and also in the returned data, where the sensor returns the relative transform to itself, (0, 0, 0). + +.. dropdown:: Code for frame_transformer_sensor.py + :icon: code + + .. literalinclude:: ../../../../source/standalone/demos/sensors/frame_transformer_sensor.py + :language: python + :linenos: diff --git a/docs/source/overview/sensors/index.rst b/docs/source/overview/sensors/index.rst new file mode 100644 index 0000000000..496b4464a1 --- /dev/null +++ b/docs/source/overview/sensors/index.rst @@ -0,0 +1,20 @@ +.. _overview_sensors: + +Sensors +========= + +In this section, we will overview the various sensor APIs provided by Isaac Lab. + +Every sensor in Isaac Lab inherits from the ``SensorBase`` abstract class that provides the core functionality inherent to all sensors, which is to provide access to "measurements" of the scene. These measurements can take many forms such as ray-casting results, camera rendered images, or even simply ground truth data queried directly from the simulation (such as poses). Whatever the data may be, we can think of the sensor as having a buffer that is periodically updated with measurements by querying the scene. This ``update_period`` is defined in "simulated" seconds, meaning that even if the flow of time in the simulation is dilated relative to the real world, the sensor will update at the appropriate rate. The ``SensorBase`` is also designed with vectorizability in mind, holding the buffers for all copies of the sensor across cloned environments. + +Updating the buffers is done by overriding the ``_update_buffers_impl`` abstract method of the ``SensorBase`` class. On every time-step of the simulation, ``dt``, all sensors are queried for an update. During this query, the total time since the last update is incremented by ``dt`` for every buffer managed by that particular sensor. If the total time is greater than or equal to the ``update_period`` for a buffer, then that buffer is flagged to be updated on the next query. + +The following pages describe the available sensors in more detail: + +.. toctree:: + :maxdepth: 1 + + camera + contact_sensor + frame_transformer + ray_caster diff --git a/docs/source/overview/sensors/ray_caster.rst b/docs/source/overview/sensors/ray_caster.rst new file mode 100644 index 0000000000..8c45d00c44 --- /dev/null +++ b/docs/source/overview/sensors/ray_caster.rst @@ -0,0 +1,109 @@ +.. _overview_sensors_ray_caster: + +Ray Caster +============= + +.. figure:: ../../_static/overview/overview_sensors_rc_patterns.png + :align: center + :figwidth: 100% + :alt: A diagram outlining the basic geometry of frame transformations + +The Ray Caster sensor (and the ray caster camera) are similar to RTX based rendering in that they both involve casting rays. The difference here is that the rays cast by the Ray Caster sensor return strictly collision information along the cast, and the direction of each individual ray can be specified. They do not bounce, nor are they affected by things like materials or opacity. For each ray specified by the sensor, a line is traced along the path of the ray and the location of first collision with the specified mesh is returned. This is the method used by some of our quadruped examples to measure the local height field. + +To keep the sensor performant when there are many cloned environments, the line tracing is done directly in `Warp `_. This is the reason why specific meshes need to be identified to cast against: that mesh data is loaded onto the device by warp when the sensor is initialized. As a consequence, the current iteration of this sensor only works for literally static meshes (meshes that *are not changed from the defaults specified in their USD file*). This constraint will be removed in future releases. + +Using a ray caster sensor requires a **pattern** and a parent xform to be attached to. The pattern defines how the rays are cast, while the prim properties defines the orientation and position of the sensor (additional offsets can be specified for more exact placement). Isaac Lab supports a number of ray casting pattern configurations, including a generic LIDAR and grid pattern. + +.. code-block:: python + + @configclass + class RaycasterSensorSceneCfg(InteractiveSceneCfg): + """Design the scene with sensors on the robot.""" + + # ground plane with texture for interesting casting results + ground = AssetBaseCfg( + prim_path="/World/Ground", + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Environments/Terrains/rough_plane.usd", + scale = (1,1,1), + ) + ) + + # lights + dome_light = AssetBaseCfg( + prim_path="/World/Light", spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75)) + ) + + # robot + robot = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + ray_caster = RayCasterCfg( + prim_path="{ENV_REGEX_NS}/Robot/base/lidar_cage", + update_period = 1/60, + offset=RayCasterCfg.OffsetCfg(pos=(0,0,0.5)), + mesh_prim_paths=["/World/Ground"], + attach_yaw_only=True, + pattern_cfg=patterns.LidarPatternCfg( + channels = 100, + vertical_fov_range=[-90, 90], + horizontal_fov_range = [-90,90], + horizontal_res=1.0), + debug_vis=True, + ) + + +Notice that the units on the pattern config is in degrees! Also, we enable visualization here to explicitly show the pattern in the rendering, but this is not required and should be disabled for performance tuning. + +.. figure:: ../../_static/overview/overview_sensors_rc_visualizer.png + :align: center + :figwidth: 100% + :alt: Lidar Pattern visualized + +Querying the sensor for data can be done at simulation run time like any other sensor. + +.. code-block:: python + + def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): + . + . + . + # Simulate physics + while simulation_app.is_running(): + . + . + . + # print information from the sensors + print("-------------------------------") + print(scene["ray_caster"]) + print("Ray cast hit results: ", scene["ray_caster"].data.ray_hits_w) + + +.. code-block:: bash + + ------------------------------- + Ray-caster @ '/World/envs/env_.*/Robot/base/lidar_cage': + view type : + update period (s) : 0.016666666666666666 + number of meshes : 1 + number of sensors : 1 + number of rays/sensor: 18000 + total number of rays : 18000 + Ray cast hit results: tensor([[[-0.3698, 0.0357, 0.0000], + [-0.3698, 0.0357, 0.0000], + [-0.3698, 0.0357, 0.0000], + ..., + [ inf, inf, inf], + [ inf, inf, inf], + [ inf, inf, inf]]], device='cuda:0') + ------------------------------- + +Here we can see the data returned by the sensor itself. Notice first that there are 3 closed brackets at the beginning and the end: this is because the data returned is batched by the number of sensors. The ray cast pattern itself has also been flattened, and so the dimensions of the array are ``[N, B, 3]`` where ``N`` is the number of sensors, ``B`` is the number of cast rays in the pattern, and 3 is the dimension of the casting space. Finally, notice that the first several values in this casting pattern are the same: this is because the lidar pattern is spherical and we have specified our FOV to be hemispherical, which includes the poles. In this configuration, the "flattening pattern" becomes apparent: the first 180 entries will be the same because it's the bottom pole of this hemisphere, and there will be 180 of them because our horizontal FOV is 180 degrees with a resolution of 1 degree. + +You can use this script to experiment with pattern configurations and build an intuition about how the data is stored by altering the ``triggered`` variable on line 99. + +.. dropdown:: Code for raycaster_sensor.py + :icon: code + + .. literalinclude:: ../../../../source/standalone/demos/sensors/raycaster_sensor.py + :language: python + :linenos: diff --git a/docs/source/refs/reference_architecture/index.rst b/docs/source/refs/reference_architecture/index.rst new file mode 100644 index 0000000000..338b8a4415 --- /dev/null +++ b/docs/source/refs/reference_architecture/index.rst @@ -0,0 +1,375 @@ +Reference Architecture +====================== + +This document presents an overview of the end-to-end robot learning process with +Isaac Lab and Isaac Sim. This is demonstrated using a reference architecture that highlights +the major building blocks for training and deployment workflows. It provides a comprehensive, +user-friendly guide on the entire process of developing applications from training to deploying +the trained model in the real world, including links to demos, working examples, and documentation. + +Who is this document for? +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +This document is designed to assist robotics developers and researchers working with NVIDIA Isaac Lab +in the robot learning field, including those at research labs, Original Equipment Manufacturers (OEM), +Solutions Providers, Solutions Integrators (SI), and independent software vendors (ISV). It offers +guidance on utilizing Isaac Lab’s robot training framework and workflows as a foundational starting +point for environment configuration, task design, and policy training and testing. + + + +.. image:: ../../_static/reference-architecture/isaac-lab-ra-light.svg + :class: only-light + :align: center + :alt: Isaac Lab Reference Architecture + +.. image:: ../../_static/reference-architecture/isaac-lab-ra-dark.svg + :class: only-dark + :align: center + :alt: Isaac Lab Reference Architecture + + +| + +The reference architecture for Isaac Lab comprises the following components: + +1. :ref:`Asset Input` +2. :ref:`Configuration - Assets & Scene` +3. :ref:`Robot Learning Task Design` +4. :ref:`Register with Gymnasium` +5. :ref:`Environment Wrapping` +6. :ref:`Run Training` +7. :ref:`Run Testing` + + + + +Components +~~~~~~~~~~~ +In this section, we will briefly discuss the individual blocks for creating a +sample reference application in Isaac Lab. + + +.. _ra-asset-input: + +Component 1 - Asset Input +--------------------------- +Isaac Lab accepts URDF, MJCF XML or USD files for the assets. The first step to training using Isaac Lab is to +have the USD file of your asset and the USD or URDF file of your robot. This can be achieved in +the following ways: + + +1. Design your assets or robot in Isaac Sim and export the USD file. + +2. Design your assets or robot in any software of your choice and export it to USD using Isaac Sim converters. Isaac Sim supports the different converters/importers to USD such as the `CAD Converter`_, `URDF Importer`_, `MJCF Importer`_, `Onshape Importer`_, etc. More details are found in the `Importing Assets section`_ in the `Isaac Sim Reference Architecture`_. + +3. If you already have the URDF or MJCF file of your robot, you do not need to convert to USD as Isaac Lab takes URDF and MJCF XML. + + +.. _ra-configuration: + +Component 2 - Configuration (Assets and Scene) +------------------------------------------------------ + +Asset Configuration +^^^^^^^^^^^^^^^^^^^^^^^^ + +Given that you have the asset file for your robot and other assets such as environment objects based on the task, the next step is to import them into Isaac Lab. Isaac Lab uses asset configuration classes to spawn various objects (or prims) into the scene using Python. The first step is to write a configuration class to define the properties for the assets needed to complete the task. For example, a simple go-to-goal task for a mobile robot will include the robot asset, an object like cubes to signify the goal pose visually, lights, ground plane, etc. Isaac Lab understands these assets using the configuration classes. Isaac Lab provides various sim-ready assets such as physically accurate +3D objects that encompass accurate physical properties and behavior. It also provides connected data streams to represent the real world in simulated digital worlds such as `robots `__ +like ANYbotics Anymal, Unitree H1 Humanoid, etc. as well as `sensors `__. We provide these assets configuration classes. Users can also define their own assets using the configuration classes. + +Follow the tutorial on `how to write an Articulation and ArticulationCfg class `__. + +Scene Configuration +^^^^^^^^^^^^^^^^^^^^^^^^ + +Given the individual asset configurations, the next step is to put all the assets together into a +scene. The scene configuration is a simple config class that initializes all the assets in the +scene that are needed for the task and for visualization. This is an example for the +`Cartpole example scene configuration `__, +which includes the cartpole, ground plane, and dome light. + + +.. _ra-robot-learning-task-design: + +Component 3 - Robot Learning Task Design +------------------------------------------------------ +Now, we have the scene for the task, but we need to define the robot learning task. We will focus on +`reinforcement learning (RL) `__ algorithm here. We define the RL task +that the agent is going to do. RL tasks are defined as a Markov Decision Process (MDP), +which is a stochastic decision-making process where optional decisions are made for the agents +considering their current state and environment they interact with. The environment provides the +agents’ current state or observations, and executes the actions provided by the agent. +The environment responds to the agents by providing the next states, reward of taking the +action, done flag and information about the current episode. Therefore, different components +of the MDP formulation (the environment) – states, actions, rewards, reset, done, etc. — must +be defined by the user for the agent to perform the given task. + +In Isaac Lab, we provide two different workflows for designing environments. + +Manager-based +^^^^^^^^^^^^^^^^^ +.. image:: ../../_static/task-workflows/manager-based-light.svg + :class: only-light + :align: center + :alt: Manager-based Task Workflow + +.. image:: ../../_static/task-workflows/manager-based-dark.svg + :class: only-dark + :align: center + :alt: Manager-based Task Workflow + +This workflow is modular, and the environment is decomposed into individual components (or managers) +that handle the different aspects of the environment, such as computing observations, +applying actions, and applying randomization. As a user, you define different configuration classes +for each component. + +- An RL task should have the following configuration classes: + + - Observations Config: Defines the agents’ observations for the task. + - Actions Config: Defines the agent’s action type, i.e. how the output of the agent are mapped to + the robot's control inputs. + - Rewards Config: Defines the reward function for the task + - Terminations Config: Defines the conditions for termination of an episode or when the task + is completed. + +- You can add other optional configuration classes such as Event Config which defines the set of randomizations and noisification for the agent and environment, Curriculum Config for tasks that require `curriculum learning`_ and Commands Config for tasks where the input is from a controller/setpoint controls e.g. a gamepad controller. + +.. tip:: + + To learn more on how you can design your own manager-based environment, see :ref:`tutorial-create-manager-rl-env`. + + + +Direct +^^^^^^^^ +.. image:: ../../_static/task-workflows/direct-based-light.svg + :class: only-light + :align: center + :alt: Direct-based Task Workflow + +.. image:: ../../_static/task-workflows/direct-based-dark.svg + :class: only-dark + :align: center + :alt: Direct-based Task Workflow + +In this workflow, you implement a single class that is responsible for computing observations, applying actions, and computing rewards. This workflow allows for direct control of the environment logic. + +.. tip:: + To learn more on how you can design your own direct environment, see :ref:`tutorial-create-direct-rl-env`. + +Users can choose from Isaac Lab’s large suite of pre-configured environments or users can define +their own environments. For more technical information about the two workflows, please see the +`documentation `__. + + +In addition to designing the RL task, you will need to design your agent’s model, the neural +network policy and value function. To train the RL agent to solve the task, you need to define +the hyperparameters such as number of epochs, learning rate, etc. for training and the +policy/value model architecture. This is defined in the training configuration file specific +to the RL library you want to use. Examples are created under the agent's folder in each task directory. +See an example of `RSL-RL `__ for Anymal-B. + + +.. _ra-register-gym: + +Component 4 - Register with Gymnasium +------------------------------------------------------ + +The next step is to register the environments with the gymnasium registry to allow you to create the environment using the unique environment name. +Registration is a way to make the environment accessible and reusable across different +RL algorithms and experiments. This is common in the RL community. Follow the tutorial on +`Registering an Environment `__ to learn more about how to register in your own environment. + +.. _ra-env-wrap: + +Component 5 - Environment Wrapping +------------------------------------------------------ +In running your RL task, you might want to change the behavior of your environment without +changing the environment itself. For example, you might want to create functions to modify +observations or rewards, record videos, or enforce time limits. Isaac Lab utilizes the API +available in the `gymnasium.Wrapper `__ class to create interfaces to the simulated environments. + +Some wrappers include: + +* `Video Wrappers `__ +* `RL Libraries Wrappers `__ + +Most RL libraries expect their own variation of an environment interface. This means the +data types needed by each library differs. Isaac Lab provides its own wrappers to convert +the environment into the expected interface by the RL library a user wants to use. These are +specified in the `Isaac Lab utils wrapper module `__. + +See the `full list `__ of other wrappers APIs. For more information on how these wrappers work, +please refer to the `Wrapping environments `__ documentation. + +Adding your own wrappers +^^^^^^^^^^^^^^^^^^^^^^^^ + +You can define your own wrappers by adding them to the Isaac Lab utils wrapper module. More information is available `on the GitHub page for wrapping environments `__. + +.. _ra-run-training: + +Component 6 - Run Training +--------------------------- + +Finally, the last step is to run the training of the RL agent. Isaac Lab provides scripts which utilizes four popular RL libraries for training the models (GPU-based training): + +* `StableBaselines3 `__ +* `RSL-RL `__ +* `RL-Games `__ +* `SKRL `__ + + +.. note:: + + Isaac Lab does not provide the implementation of these RL libraries. They are already implemented by different authors. We provide the environments and framework wrappers for the RL libraries. + + + +If you want to integrate a different version of the provided algorithms or your learning library, you can follow +`these instructions `__. + + + +Single GPU Training +^^^^^^^^^^^^^^^^^^^^^^^^ +.. image:: ../../_static/reference-architecture/single-gpu-training-light.svg + :class: only-light + :align: center + :alt: Single GPU Training Data Flow + +.. image:: ../../_static/reference-architecture/single-gpu-training-dark.svg + :class: only-dark + :align: center + :alt: Single GPU Training Data Flow + +Isaac Lab supports training massively parallel environments to speed up RL training and provides rich data for the model to train. +For single GPU training, the following steps show how training works in Isaac Sim and Isaac Lab: + +1. **In Isaac Sim** + +* Isaac Sim provides the asset states such as robot and sensor states, including the observations defined in the task observation config class. + +2. **In Isaac Lab** + +* Randomizations are added to the states defined in the event configuration class to obtain the observation for the task. Randomizations are however optional. If not defined, the states are the observations. +* The observations are computed as PyTorch tensors, and it can optionally include the action provided by the trained model based on the task. + +3. **In the RL library** + +* The observation is passed to the policy. +* The policy is trained to output the right actions for the robot using RL library algorithms such as PPO, TRPO, etc. +* The actions can serve either as a setpoint for a controller that generates the action to the robot or used directly as the action to the robot based on the task. +* Action types such as joint position for a quadruped is an input to a joint controller, velocity of 1 or 0 is used to control the cart directly in the cartpole task, etc. +* In addition, based on how the task is defined, the previous action can be part of the next set of observations that is sent. + +4. **In Isaac Sim** + +* The actions from the policy are sent back to Isaac Sim to control the agent that is learning i.e. the robot. This is the physics simulation (sim) step. This generates the next states in Isaac Sim and the rewards are calculated in Isaac Lab. + +5. **Rendering** + +* The scene can be rendered to produce the cameras' images. + + +The next state is then passed in the flow till the training reaches the specified training steps or epochs. The final product is the trained model/agent. + + + +Multi-GPU and Multi-Node Training +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +.. image:: ../../_static/reference-architecture/multi-gpu-training-light.svg + :class: only-light + :align: center + :alt: Multi GPU Training Data Flow + +.. image:: ../../_static/reference-architecture/multi-gpu-training-dark.svg + :class: only-dark + :align: center + :alt: Multi GPU Training Data Flow + + +Isaac Lab supports scaling up training by taking advantage of multi-GPU and multi-node training on Linux. Follow the tutorial on `Multi-GPU training `__ and `Multi-Node training `__ to get started. + + +Cloud-Based Training +^^^^^^^^^^^^^^^^^^^^^^^^ +Isaac Lab can be deployed alongside Isaac Sim onto the public clouds with `Isaac Automator `__. AWS, GCP, Azure, and Alibaba Cloud are currently supported. Follow the tutorial on `how to run Isaac Lab in the cloud `__. + +.. note:: + + Both multi-GPU and multi-node jobs can be easily scaled across heterogeneous environments with `OSMO `__, a cloud-native, orchestration platform for scheduling complex multi-stage and multi-container heterogeneous computing workflows. Isaac Lab also provides the tools to run your RL task in Docker. See more details on `container deployment `__. + +.. _ra-run-testing: + +Component 7: Run Testing +----------------------------- +Isaac Lab provides scripts for `testing/playing the trained policy `__ on the environment and functions for converting the trained model from .pt to +.jit and .onnx for deployment. + + +Deployment on Physical Robots +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. image:: ../../_static/reference-architecture/deployment-light.svg + :class: only-light + :align: center + :alt: Isaac Lab Trained Policy Deployment + +.. image:: ../../_static/reference-architecture/deployment-dark.svg + :class: only-dark + :align: center + :alt: Isaac Lab Trained Policy Deployment + + +To deploy your trained model on a real robot, you would need what is shown in the flow diagram. Note, this is a sample reference architecture, hence it can be tweaked for a different application. +First, you need a robot with the required sensors and processing computer such as `NVIDIA Jetson `__ to deploy on. Next, you need a state estimator for your robot. The state estimator should be able to deliver the list of observations used for training. + +Once the observations are extracted, they are passed into the model which delivers the action using the model inferencing runtime. The commanded action from the model serves as setpoints for the action controller. The action controller outputs scaled actions which are then used to control the robot to get to the next state, and this continues till the task is done. + +NVIDIA Isaac platform provides some tools for state estimation, including visual slam and inferencing engines such as `TensorRT `__. Other inferencing runtime includes `OnnxRuntime `__, direct inferencing on the PyTorch model, etc. + + + + +Summary +~~~~~~~~~~~ + +This document presents a reference architecture for Isaac Lab that has undergone SQA testing. We have provided a user-friendly guide to end-to-end robot learning with Isaac Lab and Isaac Sim from training to real-world deployment, including demos, examples, and documentation links. + + +How to Get Started +~~~~~~~~~~~~~~~~~~~~~~ +Check out our resources on using Isaac Lab with your robots. + +Review Our Documentation & Samples Resources + +* `Isaac Lab Tutorials`_ +* `Fast-Track Robot Learning in Simulation Using NVIDIA Isaac Lab`_ +* `Supercharge Robotics Workflows with AI and Simulation Using NVIDIA Isaac Sim 4.0 and NVIDIA Isaac Lab`_ +* `Closing the Sim-to-Real Gap: Training Spot Quadruped Locomotion with NVIDIA Isaac Lab `__ +* `Additional Resources`_ + +Learn More About Featured NVIDIA Solutions + +* `Scale AI-Enabled Robotics Development Workloads with NVIDIA OSMO`_ +* `Parkour and More: How Simulation-Based RL Helps to Push the Boundaries in Legged Locomotion (GTC session) `__ +* `Isaac Perceptor`_ +* `Isaac Manipulator`_ + +.. _curriculum learning: https://arxiv.org/abs/2109.11978 +.. _CAD Converter: https://docs.omniverse.nvidia.com/extensions/latest/ext_cad-converter.html +.. _URDF Importer: https://docs.omniverse.nvidia.com/isaacsim/latest/advanced_tutorials/tutorial_advanced_import_urdf.html +.. _MJCF Importer: https://docs.omniverse.nvidia.com/isaacsim/latest/advanced_tutorials/tutorial_advanced_import_mjcf.html#import-mjcf +.. _Onshape Importer: https://docs.omniverse.nvidia.com/extensions/latest/ext_onshape.html +.. _Isaac Sim Reference Architecture: https://docs.omniverse.nvidia.com/isaacsim/latest/isaac_sim_reference_architecture.html +.. _Importing Assets section: https://docs.omniverse.nvidia.com/isaacsim/latest/isaac_sim_reference_architecture.html#importing-assets + +.. _Scale AI-Enabled Robotics Development Workloads with NVIDIA OSMO: https://developer.nvidia.com/blog/scale-ai-enabled-robotics-development-workloads-with-nvidia-osmo/ +.. _Isaac Perceptor: https://developer.nvidia.com/isaac/perceptor +.. _Isaac Manipulator: https://developer.nvidia.com/isaac/manipulator +.. _Additional Resources: https://isaac-sim.github.io/IsaacLab/main/source/refs/additional_resources.html +.. _Isaac Lab Tutorials: file:///home/oomotuyi/isaac/IsaacLab/docs/_build/current/source/tutorials/index.html +.. _Fast-Track Robot Learning in Simulation Using NVIDIA Isaac Lab: https://developer.nvidia.com/blog/fast-track-robot-learning-in-simulation-using-nvidia-isaac-lab/ +.. _Supercharge Robotics Workflows with AI and Simulation Using NVIDIA Isaac Sim 4.0 and NVIDIA Isaac Lab: https://developer.nvidia.com/blog/supercharge-robotics-workflows-with-ai-and-simulation-using-nvidia-isaac-sim-4-0-and-nvidia-isaac-lab/ diff --git a/docs/source/refs/snippets/tutorial_modify_direct_rl_env.py b/docs/source/refs/snippets/tutorial_modify_direct_rl_env.py index 8073886840..215ef3bd4f 100644 --- a/docs/source/refs/snippets/tutorial_modify_direct_rl_env.py +++ b/docs/source/refs/snippets/tutorial_modify_direct_rl_env.py @@ -28,8 +28,8 @@ # [end-h1_env-import] # [start-h1_env-spaces] -num_actions = 19 -num_observations = 69 +action_space = 19 +observation_space = 69 # [end-h1_env-spaces] # [start-h1_env-robot] diff --git a/docs/source/setup/ecosystem.rst b/docs/source/setup/ecosystem.rst new file mode 100644 index 0000000000..107c95843d --- /dev/null +++ b/docs/source/setup/ecosystem.rst @@ -0,0 +1,32 @@ +Isaac Lab Ecosystem +=================== + +Isaac Lab is built on top of Isaac Sim to provide a unified and flexible framework +for robot learning that exploits latest simulation technologies. It is designed to be modular and extensible, +and aims to simplify common workflows in robotics research (such as RL, learning from demonstrations, and +motion planning). While it includes some pre-built environments, sensors, and tasks, its main goal is to +provide an open-sourced, unified, and easy-to-use interface for developing and testing custom environments +and robot learning algorithms. + +Working with Isaac Lab requires the installation of Isaac Sim, which is packaged with core robotics tools +that Isaac Lab depends on, including URDF and MJCF importers, simulation managers, and ROS features. Isaac +Sim also builds on top of the Nvidia Omniverse platform, leveraging advanced physics simulation from PhysX, +photorealistic rendering technologies, and Universal Scene Description (USD) for scene creation. + +Isaac Lab not only inherits the capabilities of Isaac Sim, but also adds a number +of new features that pertain to robot learning research. For example, including actuator dynamics in the +simulation, procedural terrain generation, and support to collect data from human demonstrations. + +.. image:: ../_static/setup/ecosystem-light.jpg + :class: only-light + :align: center + :alt: The Isaac Lab, Isaac Sim, and Nvidia Omniverse ecosystem + +.. image:: ../_static/setup/ecosystem-dark.jpg + :class: only-dark + :align: center + :alt: The Isaac Lab, Isaac Sim, and Nvidia Omniverse ecosystem + + +For a detailed explanation of Nvidia's development journey of robot learning frameworks, please visit +the `FAQ page `_. diff --git a/docs/source/setup/installation/binaries_installation.rst b/docs/source/setup/installation/binaries_installation.rst index e9c4469bca..73138ebd17 100644 --- a/docs/source/setup/installation/binaries_installation.rst +++ b/docs/source/setup/installation/binaries_installation.rst @@ -38,7 +38,7 @@ To check the minimum system requirements,refer to the documentation :sync: windows On Windows systems, by default,Isaac Sim is installed in the directory - ``C:\Users\%USERPROFILE%\AppData\Local\ov\pkg\isaac_sim-*``, with ``*`` corresponding to the Isaac Sim version. + ``%USERPROFILE%\AppData\Local\ov\pkg\isaac_sim-*``, with ``*`` corresponding to the Isaac Sim version. Verifying the Isaac Sim installation ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -66,7 +66,7 @@ variables to your terminal for the remaining of the installation instructions: .. code:: batch :: Isaac Sim root directory - set ISAACSIM_PATH="C:\Users\%USERPROFILE%\AppData\Local\ov\pkg\isaac-sim-4.2.0" + set ISAACSIM_PATH="%USERPROFILE%\AppData\Local\ov\pkg\isaac-sim-4.2.0" :: Isaac Sim python executable set ISAACSIM_PYTHON_EXE="%ISAACSIM_PATH:"=%\python.bat" @@ -407,8 +407,14 @@ top of the repository: The above command should launch the simulator and display a window with a black -ground plane. You can exit the script by pressing ``Ctrl+C`` on your terminal. +viewport. You can exit the script by pressing ``Ctrl+C`` on your terminal. On Windows machines, please terminate the process from Command Prompt using ``Ctrl+Break`` or ``Ctrl+fn+B``. +.. figure:: ../../_static/setup/verify_install.jpg + :align: center + :figwidth: 100% + :alt: Simulator with a black window. + + If you see this, then the installation was successful! |:tada:| diff --git a/docs/source/setup/installation/index.rst b/docs/source/setup/installation/index.rst index cae19383b2..90c8937cfc 100644 --- a/docs/source/setup/installation/index.rst +++ b/docs/source/setup/installation/index.rst @@ -1,5 +1,5 @@ -Installation Guide -=================== +Local Installation +================== .. image:: https://img.shields.io/badge/IsaacSim-4.2.0-silver.svg :target: https://developer.nvidia.com/isaac-sim @@ -31,14 +31,9 @@ Installation Guide For the full list of system requirements for Isaac Sim, please refer to the `Isaac Sim system requirements `_. -As an experimental feature since Isaac Sim 4.0 release, Isaac Sim can also be installed through pip. -This simplifies the installation -process by avoiding the need to download the Omniverse Launcher and installing Isaac Sim through -the launcher. Therefore, there are two ways to install Isaac Lab: .. toctree:: :maxdepth: 2 - Option 1: Installation using Isaac Sim pip - Option 2: Installation using Isaac Sim binaries - cloud_installation + Pip installation (recommended for Ubuntu 22.04 and Windows) + Binary installation (recommended for Ubuntu 20.04) diff --git a/docs/source/setup/installation/pip_installation.rst b/docs/source/setup/installation/pip_installation.rst index cc014fd887..83a74194ce 100644 --- a/docs/source/setup/installation/pip_installation.rst +++ b/docs/source/setup/installation/pip_installation.rst @@ -88,7 +88,14 @@ compatibility issues with some Linux distributions. If you encounter any issues, pip install --upgrade pip -- Then, install the Isaac Sim packages necessary for running Isaac Lab: +- Then, install the Isaac Sim packages + + .. code-block:: bash + + pip install isaacsim==4.2.0.2 isaacsim-extscache-physics==4.2.0.2 isaacsim-extscache-kit==4.2.0.2 isaacsim-extscache-kit-sdk==4.2.0.2 --extra-index-url https://pypi.nvidia.com + + +- To install a minimal set of packages for running Isaac Lab only, the following command can be used. Note that you cannot run ``isaacsim`` with this. .. code-block:: bash @@ -310,8 +317,15 @@ top of the repository: The above command should launch the simulator and display a window with a black -ground plane. You can exit the script by pressing ``Ctrl+C`` on your terminal. +viewport as shown below. You can exit the script by pressing ``Ctrl+C`` on your terminal. On Windows machines, please terminate the process from Command Prompt using ``Ctrl+Break`` or ``Ctrl+fn+B``. + +.. figure:: ../../_static/setup/verify_install.jpg + :align: center + :figwidth: 100% + :alt: Simulator with a black window. + + If you see this, then the installation was successful! |:tada:| diff --git a/docs/source/tutorials/01_assets/run_articulation.rst b/docs/source/tutorials/01_assets/run_articulation.rst index dc8d14c669..dd456d6c6d 100644 --- a/docs/source/tutorials/01_assets/run_articulation.rst +++ b/docs/source/tutorials/01_assets/run_articulation.rst @@ -66,8 +66,8 @@ Similar to a rigid object, an articulation also has a root state. This state cor articulation tree. On top of the root state, an articulation also has joint states. These states correspond to the joint positions and velocities. -To reset the articulation, we first set the root state by calling the :meth:`Articulation.write_root_state_to_sim` -method. Similarly, we set the joint states by calling the :meth:`Articulation.write_joint_state_to_sim` method. +To reset the articulation, we first set the root state by calling the :meth:`Articulation.write_root_link_pose_to_sim` and :meth:`Articulation.write_root_com_velocity_to_sim` +methods. Similarly, we set the joint states by calling the :meth:`Articulation.write_joint_state_to_sim` method. Finally, we call the :meth:`Articulation.reset` method to reset any internal buffers and caches. .. literalinclude:: ../../../../source/standalone/tutorials/01_assets/run_articulation.py diff --git a/docs/source/tutorials/01_assets/run_deformable_object.rst b/docs/source/tutorials/01_assets/run_deformable_object.rst index 4041d7e7ed..8edb2fdf4a 100644 --- a/docs/source/tutorials/01_assets/run_deformable_object.rst +++ b/docs/source/tutorials/01_assets/run_deformable_object.rst @@ -149,7 +149,7 @@ the average position of all the nodes in the mesh. .. literalinclude:: ../../../../source/standalone/tutorials/01_assets/run_deformable_object.py :language: python :start-at: # update buffers - :end-at: print(f"Root position (in world): {cube_object.data.root_pos_w[:, :3]}") + :end-at: print(f"Root position (in world): {cube_object.data.root_link_pos_w[:, :3]}") The Code Execution diff --git a/docs/source/tutorials/01_assets/run_rigid_object.rst b/docs/source/tutorials/01_assets/run_rigid_object.rst index f9fe141467..e4249c45e0 100644 --- a/docs/source/tutorials/01_assets/run_rigid_object.rst +++ b/docs/source/tutorials/01_assets/run_rigid_object.rst @@ -96,7 +96,7 @@ desired state of the rigid object prim into the world frame before setting it. We use the :attr:`assets.RigidObject.data.default_root_state` attribute to get the default root state of the spawned rigid object prims. This default state can be configured from the :attr:`assets.RigidObjectCfg.init_state` attribute, which we left as identity in this tutorial. We then randomize the translation of the root state and -set the desired state of the rigid object prim using the :meth:`assets.RigidObject.write_root_state_to_sim` method. +set the desired state of the rigid object prim using the :meth:`assets.RigidObject.write_root_link_pose_to_sim` and :meth:`assets.RigidObject.write_root_com_velocity_to_sim` methods. As the name suggests, this method writes the root state of the rigid object prim into the simulation buffer. .. literalinclude:: ../../../../source/standalone/tutorials/01_assets/run_rigid_object.py diff --git a/docs/source/tutorials/03_envs/create_direct_rl_env.rst b/docs/source/tutorials/03_envs/create_direct_rl_env.rst index ab5136106a..a4b945be9d 100644 --- a/docs/source/tutorials/03_envs/create_direct_rl_env.rst +++ b/docs/source/tutorials/03_envs/create_direct_rl_env.rst @@ -48,9 +48,9 @@ config should define the number of actions and observations for the environment. @configclass class CartpoleEnvCfg(DirectRLEnvCfg): ... - num_actions = 1 - num_observations = 4 - num_states = 0 + action_space = 1 + observation_space = 4 + state_space = 0 The config class can also be used to define task-specific attributes, such as scaling for reward terms and thresholds for reset conditions. diff --git a/docs/source/tutorials/03_envs/create_manager_rl_env.rst b/docs/source/tutorials/03_envs/create_manager_rl_env.rst index 1ff7c71990..63f710965b 100644 --- a/docs/source/tutorials/03_envs/create_manager_rl_env.rst +++ b/docs/source/tutorials/03_envs/create_manager_rl_env.rst @@ -36,7 +36,7 @@ For this tutorial, we use the cartpole environment defined in ``omni.isaac.lab_t .. literalinclude:: ../../../../source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/cartpole_env_cfg.py :language: python - :emphasize-lines: 63-68, 124-149, 152-162, 165-169, 187-192 + :emphasize-lines: 117-141, 144-154, 172-174 :linenos: The script for running the environment ``run_cartpole_rl_env.py`` is present in the @@ -117,13 +117,8 @@ For various goal-conditioned tasks, it is useful to specify the goals or command handled through the :class:`managers.CommandManager`. The command manager handles resampling and updating the commands at each step. It can also be used to provide the commands as an observation to the agent. -For this simple task, we do not use any commands. This is specified by using a command term with the -:class:`envs.mdp.NullCommandCfg` configuration. However, you can see an example of command definitions in the -locomotion or manipulation tasks. - -.. literalinclude:: ../../../../source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/cartpole_env_cfg.py - :language: python - :pyobject: CommandsCfg +For this simple task, we do not use any commands. Hence, we leave this attribute as its default value, which is None. +You can see an example of how to define a command manager in the other locomotion or manipulation tasks. Defining curriculum ------------------- @@ -134,11 +129,6 @@ we provide a :class:`managers.CurriculumManager` class that can be used to defin In this tutorial we don't implement a curriculum for simplicity, but you can see an example of a curriculum definition in the other locomotion or manipulation tasks. -We use a simple pass-through curriculum to define a curriculum manager that does not modify the environment. - -.. literalinclude:: ../../../../source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/cartpole_env_cfg.py - :language: python - :pyobject: CurriculumCfg Tying it all together --------------------- diff --git a/docs/source/tutorials/05_controllers/run_diff_ik.rst b/docs/source/tutorials/05_controllers/run_diff_ik.rst index 99600b0591..ab451ab8ca 100644 --- a/docs/source/tutorials/05_controllers/run_diff_ik.rst +++ b/docs/source/tutorials/05_controllers/run_diff_ik.rst @@ -79,7 +79,7 @@ joint positions, current end-effector pose, and the Jacobian matrix. While the attribute :attr:`assets.ArticulationData.joint_pos` provides the joint positions, we only want the joint positions of the robot's arm, and not the gripper. Similarly, while -the attribute :attr:`assets.ArticulationData.body_state_w` provides the state of all the +the attribute :attr:`assets.ArticulationData.body_link_state_w` provides the state of all the robot's bodies, we only want the state of the robot's end-effector. Thus, we need to index into these arrays to obtain the desired quantities. diff --git a/docs/source/tutorials/05_controllers/run_osc.rst b/docs/source/tutorials/05_controllers/run_osc.rst new file mode 100644 index 0000000000..661056b1a8 --- /dev/null +++ b/docs/source/tutorials/05_controllers/run_osc.rst @@ -0,0 +1,178 @@ +Using an operational space controller +===================================== + +.. currentmodule:: omni.isaac.lab + +Sometimes, controlling the end-effector pose of the robot using a differential IK controller is not sufficient. +For example, we might want to enforce a very specific pose tracking error dynamics in the task space, actuate the robot +with joint effort/torque commands, or apply a contact force at a specific direction while controlling the motion of +the other directions (e.g., washing the surface of the table with a cloth). In such tasks, we can use an +operational space controller (OSC). + +.. rubric:: References for the operational space control: + +1. O Khatib. A unified approach for motion and force control of robot manipulators: + The operational space formulation. IEEE Journal of Robotics and Automation, 3(1):43–53, 1987. URL http://dx.doi.org/10.1109/JRA.1987.1087068. + +2. Robot Dynamics Lecture Notes by Marco Hutter (ETH Zurich). URL https://ethz.ch/content/dam/ethz/special-interest/mavt/robotics-n-intelligent-systems/rsl-dam/documents/RobotDynamics2017/RD_HS2017script.pdf + +In this tutorial, we will learn how to use an OSC to control the robot. +We will use the :class:`controllers.OperationalSpaceController` class to apply a constant force perpendicular to a +tilted wall surface while tracking a desired end-effector pose in all the other directions. + +The Code +~~~~~~~~ + +The tutorial corresponds to the ``run_osc.py`` script in the +``source/standalone/tutorials/05_controllers`` directory. + + +.. dropdown:: Code for run_osc.py + :icon: code + + .. literalinclude:: ../../../../source/standalone/tutorials/05_controllers/run_osc.py + :language: python + :linenos: + + +Creating an Operational Space Controller +---------------------------------------- + +The :class:`~controllers.OperationalSpaceController` class computes the joint +efforts/torques for a robot to do simultaneous motion and force control in task space. + +The reference frame of this task space could be an arbitrary coordinate frame in Euclidean space. By default, +it is the robot's base frame. However, in certain cases, it could be easier to define target coordinates w.r.t. a +different frame. In such cases, the pose of this task reference frame, w.r.t. to the robot's base frame, should be +provided in the ``set_command`` method's ``current_task_frame_pose_b`` argument. For example, in this tutorial, it +makes sense to define the target commands w.r.t. a frame that is parallel to the wall surface, as the force control +direction would be then only nonzero in the z-axis of this frame. The target pose, which is set to have the same +orientation as the wall surface, is such a candidate and is used as the task frame in this tutorial. Therefore, all +the arguments to the :class:`~controllers.OperationalSpaceControllerCfg` should be set with this task reference frame +in mind. + +For the motion control, the task space targets could be given as absolute (i.e., defined w.r.t. the robot base, +``target_types: "pose_abs"``) or relative the the end-effector's current pose (i.e., ``target_types: "pose_rel"``). +For the force control, the task space targets could be given as absolute (i.e., defined w.r.t. the robot base, +``target_types: "force_abs"``). If it is desired to apply pose and force control simultaneously, the ``target_types`` +should be a list such as ``["pose_abs", "wrench_abs"]`` or ``["pose_rel", "wrench_abs"]``. + +The axes that the motion and force control will be applied can be specified using the ``motion_control_axes_task`` and +``force_control_axes_task`` arguments, respectively. These lists should consist of 0/1 for all six axes (position and +rotation) and be complementary to each other (e.g., for the x-axis, if the ``motion_control_axes_task`` is ``0``, the +``force_control_axes_task`` should be ``1``). + +For the motion control axes, desired stiffness, and damping ratio values can be specified using the +``motion_control_stiffness`` and ``motion_damping_ratio_task`` arguments, which can be a scalar (same value for all +axes) or a list of six scalars, one value corresponding to each axis. If desired, the stiffness and damping ratio +values could be a command parameter (e.g., to learn the values using RL or change them on the go). For this, +``impedance_mode`` should be either ``"variable_kp"`` to include the stiffness values within the command or +``"variable"`` to include both the stiffness and damping ratio values. In these cases, ``motion_stiffness_limits_task`` +and ``motion_damping_limits_task`` should be set as well, which puts bounds on the stiffness and damping ratio values. + +For contact force control, it is possible to apply an open-loop force control by not setting the +``contact_wrench_stiffness_task``, or apply a closed-loop force control (with the feed-forward term) by setting +the desired stiffness values using the ``contact_wrench_stiffness_task`` argument, which can be a scalar or a list +of six scalars. Please note that, currently, only the linear part of the contact wrench (hence the first three +elements of the ``contact_wrench_stiffness_task``) is considered in the closed-loop control, as the rotational part +cannot be measured with the contact sensors. + +For the motion control, ``inertial_dynamics_decoupling`` should be set to ``True`` to use the robot's inertia matrix +to decouple the desired accelerations in the task space. This is important for the motion control to be accurate, +especially for rapid movements. This inertial decoupling accounts for the coupling between all the six motion axes. +If desired, the inertial coupling between the translational and rotational axes could be ignored by setting the +``partial_inertial_dynamics_decoupling`` to ``True``. + +If it is desired to include the gravity compensation in the operational space command, the ``gravity_compensation`` +should be set to ``True``. + +The included OSC implementation performs the computation in a batched format and uses PyTorch operations. + +In this tutorial, we will use ``"pose_abs"`` for controlling the motion in all axes except the z-axis and +``"wrench_abs"`` for controlling the force in the z-axis. Moreover, we will include the full inertia decoupling in +the motion control and not include the gravity compensation, as the gravity is disabled from the robot configuration. +Finally, we set the impedance mode to ``"variable_kp"`` to dynamically change the stiffness values +(``motion_damping_ratio_task`` is set to ``1``: the kd values adapt according to kp values to maintain a critically +damped response). + +.. literalinclude:: ../../../../source/standalone/tutorials/05_controllers/run_osc.py + :language: python + :start-at: # Create the OSC + :end-at: osc = OperationalSpaceController(osc_cfg, num_envs=scene.num_envs, device=sim.device) + +Updating the states of the robot +-------------------------------------------- + +The OSC implementation is a computation-only class. Thus, it expects the user to provide the necessary information +about the robot. This includes the robot's Jacobian matrix, mass/inertia matrix, end-effector pose, velocity, and +contact force, all in the root frame. Moreover, the user should provide gravity compensation vector, if desired. + +.. literalinclude:: ../../../../source/standalone/tutorials/05_controllers/run_osc.py + :language: python + :start-at: # Update robot states + :end-before: return jacobian_b, mass_matrix, gravity, ee_pose_b, ee_vel_b, root_pose_w, ee_pose_w, ee_force_b + + +Computing robot command +----------------------- + +The OSC separates the operation of setting the desired command and computing the desired joint positions. +To set the desired command, the user should provide command vector, which includes the target commands +(i.e., in the order they appear in the ``target_types`` argument of the OSC configuration), +and the desired stiffness and damping ratio values if the impedance_mode is set to ``"variable_kp"`` or ``"variable"``. +They should be all in the same coordinate frame as the task frame (e.g., indicated with ``_task`` subscript) and +concatanated together. + +In this tutorial, the desired wrench is already defined w.r.t. the task frame, and the desired pose is transformed +to the task frame as the following: + +.. literalinclude:: ../../../../source/standalone/tutorials/05_controllers/run_osc.py + :language: python + :start-at: # Convert the target commands to the task frame + :end-at: return command, task_frame_pose_b + +The OSC command is set with the command vector in the task frame, the end-effector pose in the base frame, and the +task (reference) frame pose in the base frame as the following. This information is needed, as the internal +computations are done in the base frame. + +.. literalinclude:: ../../../../source/standalone/tutorials/05_controllers/run_osc.py + :language: python + :start-at: # set the osc command + :end-at: osc.set_command(command=command, current_ee_pose_b=ee_pose_b, current_task_frame_pose_b=task_frame_pose_b) + +The joint effort/torque values are computed using the provided robot states and the desired command as the following: + +.. literalinclude:: ../../../../source/standalone/tutorials/05_controllers/run_osc.py + :language: python + :start-at: # compute the joint commands + :end-at: ) + + +The computed joint effort/torque targets can then be applied on the robot. + +.. literalinclude:: ../../../../source/standalone/tutorials/05_controllers/run_osc.py + :language: python + :start-at: # apply actions + :end-at: robot.write_data_to_sim() + + +The Code Execution +~~~~~~~~~~~~~~~~~~ + +You can now run the script and see the result: + +.. code-block:: bash + + ./isaaclab.sh -p source/standalone/tutorials/05_controllers/run_osc.py --num_envs 128 + +The script will start a simulation with 128 robots. The robots will be controlled using the OSC. +The current and desired end-effector poses should be displayed using frame markers in addition to the red tilted wall. +You should see that the robot reaches the desired pose while applying a constant force perpendicular to the wall +surface. + +.. figure:: ../../_static/tutorials/tutorial_operational_space_controller.jpg + :align: center + :figwidth: 100% + :alt: result of run_osc.py + +To stop the simulation, you can either close the window or press ``Ctrl+C`` in the terminal. diff --git a/docs/source/tutorials/index.rst b/docs/source/tutorials/index.rst index d3d582a5c4..3e6a09a1e6 100644 --- a/docs/source/tutorials/index.rst +++ b/docs/source/tutorials/index.rst @@ -101,3 +101,4 @@ tutorials show you how to use motion generators to control the robots at the tas :titlesonly: 05_controllers/run_diff_ik + 05_controllers/run_osc diff --git a/isaaclab.bat b/isaaclab.bat index b415ef1a13..09c6818e37 100644 --- a/isaaclab.bat +++ b/isaaclab.bat @@ -519,9 +519,9 @@ if "%arg%"=="-i" ( call :extract_python_exe pushd %ISAACLAB_PATH%\docs call !python_exe! -m pip install -r requirements.txt >nul - call !python_exe! -m sphinx -b html -d _build\doctrees . _build\html + call !python_exe! -m sphinx -b html -d _build\doctrees . _build\current echo [INFO] To open documentation on default browser, run: - echo xdg-open "%ISAACLAB_PATH%\docs\_build\html\index.html" + echo xdg-open "%ISAACLAB_PATH%\docs\_build\current\index.html" popd >nul shift goto :end diff --git a/isaaclab.sh b/isaaclab.sh index a604706e70..eedf6636ed 100755 --- a/isaaclab.sh +++ b/isaaclab.sh @@ -328,7 +328,7 @@ while [[ $# -gt 0 ]]; do fi # run the formatter over the repository # check if pre-commit is installed - if [ ! command -v pre-commit &>/dev/null ]; then + if ! command -v pre-commit &>/dev/null; then echo "[INFO] Installing pre-commit..." pip install pre-commit fi @@ -350,7 +350,7 @@ while [[ $# -gt 0 ]]; do python_exe=$(extract_python_exe) echo "[INFO] Using python from: ${python_exe}" shift # past argument - ${python_exe} $@ + ${python_exe} "$@" # exit neatly break ;; @@ -396,10 +396,10 @@ while [[ $# -gt 0 ]]; do cd ${ISAACLAB_PATH}/docs ${python_exe} -m pip install -r requirements.txt > /dev/null # build the documentation - ${python_exe} -m sphinx -b html -d _build/doctrees . _build/html + ${python_exe} -m sphinx -b html -d _build/doctrees . _build/current # open the documentation echo -e "[INFO] To open documentation on default browser, run:" - echo -e "\n\t\txdg-open $(pwd)/_build/html/index.html\n" + echo -e "\n\t\txdg-open $(pwd)/_build/current/index.html\n" # exit neatly cd - > /dev/null shift # past argument diff --git a/pyproject.toml b/pyproject.toml index 51d4375907..63ec9afd2a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,9 @@ extra_standard_library = [ "toml", "trimesh", "tqdm", + "torchvision", + "transformers", + "einops" # Needed for transformers, doesn't always auto-install ] # Imports from Isaac Sim and Omniverse known_third_party = [ diff --git a/source/apps/isaaclab.python.headless.kit b/source/apps/isaaclab.python.headless.kit index 3435606b7b..9d13464ecf 100644 --- a/source/apps/isaaclab.python.headless.kit +++ b/source/apps/isaaclab.python.headless.kit @@ -5,7 +5,7 @@ [package] title = "Isaac Lab Python Headless" description = "An app for running Isaac Lab headlessly" -version = "1.2.0" +version = "1.3.0" # That makes it browsable in UI with "experience" filter keywords = ["experience", "app", "isaaclab", "python", "headless"] diff --git a/source/apps/isaaclab.python.headless.rendering.kit b/source/apps/isaaclab.python.headless.rendering.kit index 5d14c551dc..cd39d0a797 100644 --- a/source/apps/isaaclab.python.headless.rendering.kit +++ b/source/apps/isaaclab.python.headless.rendering.kit @@ -9,7 +9,7 @@ [package] title = "Isaac Lab Python Headless Camera" description = "An app for running Isaac Lab headlessly with rendering enabled" -version = "1.2.0" +version = "1.3.0" # That makes it browsable in UI with "experience" filter keywords = ["experience", "app", "isaaclab", "python", "camera", "minimal"] @@ -35,6 +35,10 @@ app.folder = "${exe-path}/" app.name = "Isaac-Sim" app.version = "4.2.0" +# Disable print outs on extension startup information +# this only disables the app print_and_log function +app.enableStdoutOutput = false + # set the default ros bridge to disable on startup isaac.startup.ros_bridge_extension = "" diff --git a/source/apps/isaaclab.python.kit b/source/apps/isaaclab.python.kit index fdf924a6ad..7ea6fc1a05 100644 --- a/source/apps/isaaclab.python.kit +++ b/source/apps/isaaclab.python.kit @@ -5,7 +5,7 @@ [package] title = "Isaac Lab Python" description = "An app for running Isaac Lab" -version = "1.2.0" +version = "1.3.0" # That makes it browsable in UI with "experience" filter keywords = ["experience", "app", "usd"] diff --git a/source/apps/isaaclab.python.rendering.kit b/source/apps/isaaclab.python.rendering.kit index 539b109a19..38a48b7460 100644 --- a/source/apps/isaaclab.python.rendering.kit +++ b/source/apps/isaaclab.python.rendering.kit @@ -9,7 +9,7 @@ [package] title = "Isaac Lab Python Camera" description = "An app for running Isaac Lab with rendering enabled" -version = "1.2.0" +version = "1.3.0" # That makes it browsable in UI with "experience" filter keywords = ["experience", "app", "isaaclab", "python", "camera", "minimal"] @@ -35,6 +35,10 @@ app.folder = "${exe-path}/" app.name = "Isaac-Sim" app.version = "4.2.0" +# Disable print outs on extension startup information +# this only disables the app print_and_log function +app.enableStdoutOutput = false + # set the default ros bridge to disable on startup isaac.startup.ros_bridge_extension = "" diff --git a/source/extensions/omni.isaac.lab/config/extension.toml b/source/extensions/omni.isaac.lab/config/extension.toml index 517ac00610..05f5be440a 100644 --- a/source/extensions/omni.isaac.lab/config/extension.toml +++ b/source/extensions/omni.isaac.lab/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "0.25.1" +version = "0.30.0" # Description title = "Isaac Lab framework for Robot Learning" @@ -13,14 +13,12 @@ keywords = ["kit", "robotics", "learning", "ai"] [dependencies] "omni.isaac.core" = {} -"omni.isaac.ml_archive" = {} "omni.replicator.core" = {} [python.pipapi] requirements = [ "numpy", "prettytable==3.3.0", - "tensordict", "toml", "hidapi", "gymnasium==0.29.0", @@ -30,7 +28,6 @@ requirements = [ modules = [ "numpy", "prettytable", - "tensordict", "toml", "hid", "gymnasium", diff --git a/source/extensions/omni.isaac.lab/docs/CHANGELOG.rst b/source/extensions/omni.isaac.lab/docs/CHANGELOG.rst index e78abdb068..fe823cb371 100644 --- a/source/extensions/omni.isaac.lab/docs/CHANGELOG.rst +++ b/source/extensions/omni.isaac.lab/docs/CHANGELOG.rst @@ -1,6 +1,420 @@ Changelog --------- +0.30.0 (2024-12-16) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Previously, physx returns the rigid bodies and articulations velocities in the com of bodies rather than the link frame, while poses are in link frames. We now explicitly provide :attr:`body_link_state` and :attr:`body_com_state` APIs replacing the previous :attr:`body_state` API. Previous APIs are now marked as deprecated. Please update any code using the previous pose and velocity APIs to use the new ``*_link_*`` or ``*_com_*`` APIs in :attr:`omni.isaac_lab.assets.RigidBody`, :attr:`omni.isaac_lab.assets.RigidBodyCollection`, and :attr:`omni.isaac_lab.assets.Articulation`. + + +0.29.3 (2024-12-16) +~~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Fixed ordering of logging and resamping in the command manager, where we were logging the metrics after resampling the commands. This leads to incorrect logging of metrics when inside the resample call, the metrics tensors get reset. + + +0.29.2 (2024-12-16) +~~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Fixed errors within the calculations of :class:`omni.isaac.lab.controllers.OperationalSpaceController`. + +Added +^^^^^ + +* Added :class:`omni.isaac.lab.controllers.OperationalSpaceController` to API documentation. +* Added test cases for :class:`omni.isaac.lab.controllers.OperationalSpaceController`. +* Added a tutorial for :class:`omni.isaac.lab.controllers.OperationalSpaceController`. +* Added the implementation of :class:`omni.isaac.lab.envs.mdp.actions.OperationalSpaceControllerAction` class. + + +0.29.1 (2024-12-15) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Added call to update articulation kinematics after reset to ensure states are updated for non-rendering sensors. Previously, some changes in reset such as modifying joint states would not be reflected in the rigid body states immediately after reset. + + +0.29.0 (2024-12-15) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added UI interface to the Managers in the ManagerBasedEnv and MangerBasedRLEnv classes. +* Added UI widgets for :class:`LiveLinePlot` and :class:`ImagePlot`. +* Added ``ManagerLiveVisualizer/Cfg``: Given a ManagerBase (i.e. action_manager, observation_manager, etc) and a config file this class creates the the interface between managers and the UI. +* Added :class:`EnvLiveVisualizer`: A 'manager' of ManagerLiveVisualizer. This is added to the ManagerBasedEnv but is only called during the initialization of the managers in load_managers +* Added ``get_active_iterable_terms`` implementation methods to ActionManager, ObservationManager, CommandsManager, CurriculumManager, RewardManager, and TerminationManager. This method exports the active term data and labels for each manager and is called by ManagerLiveVisualizer. +* Additions to :class:`BaseEnvWindow` and :class:`RLEnvWindow` to register ManagerLiveVisualizer UI interfaces for the chosen managers. + + +0.28.0 (2024-12-15) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added observation history computation to :class:`omni.isaac.lab.manager.observation_manager.ObservationManager`. +* Added ``history_length`` and ``flatten_history_dim`` configuration parameters to :class:`omni.isaac.lab.manager.manager_term_cfg.ObservationTermCfg` +* Added ``history_length`` and ``flatten_history_dim`` configuration parameters to :class:`omni.isaac.lab.manager.manager_term_cfg.ObservationGroupCfg` +* Added full buffer property to :class:`omni.isaac.lab.utils.buffers.circular_buffer.CircularBuffer` + + +0.27.29 (2024-12-15) +~~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added action clip to all :class:`omni.isaac.lab.envs.mdp.actions`. + + +0.27.28 (2024-12-14) +~~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Added check for error below threshold in state machines to ensure the state has been reached. + + +0.27.27 (2024-12-13) +~~~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Fixed the shape of ``quat_w`` in the ``apply_actions`` method of :attr:`~omni.isaac.lab.env.mdp.NonHolonomicAction` (previously (N,B,4), now (N,4) since the number of root bodies B is required to be 1). Previously ``apply_actions`` errored because ``euler_xyz_from_quat`` requires inputs of shape (N,4). + + +0.27.26 (2024-12-11) +~~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Introduced an optional ``sensor_cfg`` parameter to the :meth:`~omni.isaac.lab.envs.mdp.rewards.base_height_l2` function, enabling the use of + :class:`~omni.isaac.lab.sensors.RayCaster` for height adjustments. For flat terrains, the function retains its previous behavior. +* Improved documentation to clarify the usage of the :meth:`~omni.isaac.lab.envs.mdp.rewards.base_height_l2` function in both flat and rough terrain settings. + + +0.27.25 (2024-12-11) +~~~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Modified :class:`omni.isaac.lab.envs.mdp.actions.DifferentialInverseKinematicsAction` class to use the geometric + Jacobian computed w.r.t. to the root frame of the robot. This helps ensure that root pose does not affect the tracking. + + +0.27.24 (2024-12-09) +~~~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Fixed the initial state recorder term in :class:`omni.isaac.lab.envs.mdp.recorders.InitialStateRecorder` to + return only the states of the specified environment IDs. + + +0.27.23 (2024-12-06) +~~~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Fixed the enforcement of :attr:`~omni.isaac.lab.actuators.ActuatorBaseCfg.velocity_limits` at the + :attr:`~omni.isaac.lab.assets.Articulation.root_physx_view` level. + + +0.27.22 (2024-12-06) +~~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* If a USD that contains an articulation root is loaded using a + :attr:`omni.isaac_lab.assets.RigidBody` we now fail unless the articulation root is explicitly + disabled. Using an articulation root for rigid bodies is not needed and decreases overall performance. + + +0.27.21 (2024-12-06) +~~~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Corrected the projection types of fisheye camera in :class:`omni.isaac.lab.sim.spawners.sensors.sensors_cfg.FisheyeCameraCfg`. + Earlier, the projection names used snakecase instead of camelcase. + + +0.27.20 (2024-12-06) +~~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added option to define the clipping behavior for depth images generated by + :class:`~omni.isaac.lab.sensors.RayCasterCamera`, :class:`~omni.isaac.lab.sensors.Camera`, and :class:`~omni.isaac.lab.sensors.TiledCamera` + +Changed +^^^^^^^ + +* Unified the clipping behavior for the depth images of all camera implementations. Per default, all values exceeding + the range are clipped to zero for both ``distance_to_image_plane`` and ``distance_to_camera`` depth images. Prev. + :class:`~omni.isaac.lab.sensors.RayCasterCamera` clipped the values to the maximum value of the depth image, + :class:`~omni.isaac.lab.sensors.Camera` did not clip them and had a different behavior for both types. + + +0.27.19 (2024-12-05) +~~~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Fixed the condition in ``isaaclab.sh`` that checks whether ``pre-commit`` is installed before attempting installation. + + +0.27.18 (2024-12-04) +~~~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Fixed the order of the incoming parameters in :class:`omni.isaac.lab.envs.DirectMARLEnv` to correctly use ``NoiseModel`` in marl-envs. + + +0.27.17 (2024-12-02) +~~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added :class:`~omni.isaac.lab.managers.RecorderManager` and its utility classes to record data from the simulation. +* Added :class:`~omni.isaac.lab.utils.datasets.EpisodeData` to store data for an episode. +* Added :class:`~omni.isaac.lab.utils.datasets.DatasetFileHandlerBase` as a base class for handling dataset files. +* Added :class:`~omni.isaac.lab.utils.datasets.HDF5DatasetFileHandler` as a dataset file handler implementation to + export and load episodes from HDF5 files. +* Added ``record_demos.py`` script to record human-teleoperated demos for a specified task and export to an HDF5 file. +* Added ``replay_demos.py`` script to replay demos loaded from an HDF5 file. + + +0.27.16 (2024-11-21) +~~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Changed :class:`omni.isaac.lab.envs.DirectMARLEnv` to inherit from ``Gymnasium.Env`` due to requirement from Gymnasium v1.0.0 requiring all environments to be a subclass of ``Gymnasium.Env`` when using the ``make`` interface. + + +0.27.15 (2024-11-09) +~~~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Fixed indexing in :meth:`omni.isaac.lab.assets.Articulation.write_joint_limits_to_sim` to correctly process non-None ``env_ids`` and ``joint_ids``. + + +0.27.14 (2024-10-23) +~~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added the class :class:`~omni.isaac.lab.assets.RigidObjectCollection` which allows to spawn + multiple objects in each environment and access/modify the quantities with a unified (env_ids, object_ids) API. + + +0.27.13 (2024-10-30) +~~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added the attributes :attr:`~omni.isaac.lab.sim.converters.MeshConverterCfg.translation`, :attr:`~omni.isaac.lab.sim.converters.MeshConverterCfg.rotation`, + :attr:`~omni.isaac.lab.sim.converters.MeshConverterCfg.scale` to translate, rotate, and scale meshes + when importing them with :class:`~omni.isaac.lab.sim.converters.MeshConverter`. + + +0.27.12 (2024-01-04) +~~~~~~~~~~~~~~~~~~~~ + +Removed +^^^^^^^ + +* Removed TensorDict usage in favor of Python dictionary in sensors + + +0.27.11 (2024-10-31) +~~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added support to define tuple of floats to scale observation terms by expanding the + :attr:`omni.isaac.lab.managers.manager_term_cfg.ObservationManagerCfg.scale` attribute. + + +0.27.10 (2024-11-01) +~~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Cached the PhysX view's joint paths before looping over them when processing fixed joint tendons + inside the :class:`Articulation` class. This helps improve the processing time for the tendons. + + +0.27.9 (2024-11-01) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added the :class:`omni.isaac.lab.utils.types.ArticulationActions` class to store the joint actions + for an articulation. Earlier, the class from Isaac Sim was being used. However, it used a different + type for the joint actions which was not compatible with the Isaac Lab framework. + + +0.27.8 (2024-11-01) +~~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Added sanity check if the term is a valid type inside the command manager. +* Corrected the iteration over ``group_cfg_items`` inside the observation manager. + + +0.27.7 (2024-10-28) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added frozen encoder feature extraction observation space with ResNet and Theia + + +0.27.6 (2024-10-25) +~~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Fixed usage of ``meshes`` property in :class:`omni.isaac.lab.sensors.RayCasterCamera` to use ``self.meshes`` instead of the undefined ``RayCaster.meshes``. +* Fixed issue in :class:`omni.isaac.lab.envs.ui.BaseEnvWindow` where undefined configs were being accessed when creating debug visualization elements in UI. + + +0.27.5 (2024-10-25) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added utilities for serializing/deserializing Gymnasium spaces. + + +0.27.4 (2024-10-18) +~~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Updated installation path instructions for Windows in the Isaac Lab documentation to remove redundancy in the use of %USERPROFILE% for path definitions. + + +0.27.3 (2024-10-22) +~~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Fixed the issue with using list or tuples of ``configclass`` within a ``configclass``. Earlier, the list of + configclass objects were not converted to dictionary properly when ``to_dict`` function was called. + + +0.27.2 (2024-10-21) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added ``--kit_args`` to :class:`~omni.isaac.lab.app.AppLauncher` to allow passing command line arguments directly to Omniverse Kit SDK. + + +0.27.1 (2024-10-20) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added :class:`~omni.isaac.lab.sim.RenderCfg` and the attribute :attr:`~omni.isaac.lab.sim.SimulationCfg.render` for + specifying render related settings. + + +0.27.0 (2024-10-14) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added a method to :class:`~omni.isaac.lab.utils.configclass` to check for attributes with values of + type ``MISSING``. This is useful when the user wants to check if a certain attribute has been set or not. +* Added the configuration validation check inside the constructor of all the core classes + (such as sensor base, asset base, scene and environment base classes). +* Added support for environments without commands by leaving the attribute + :attr:`omni.isaac.lab.envs.ManagerBasedRLEnvCfg.commands` as None. Before, this had to be done using + the class :class:`omni.isaac.lab.command_generators.NullCommandGenerator`. +* Moved the ``meshes`` attribute in the :class:`omni.isaac.lab.sensors.RayCaster` class from class variable to instance variable. + This prevents the meshes to overwrite each other. + + +0.26.0 (2024-10-16) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added Imu sensor implementation that directly accesses the physx view :class:`omni.isaac.lab.sensors.Imu`. The + sensor comes with a configuration class :class:`omni.isaac.lab.sensors.ImuCfg` and data class + :class:`omni.isaac.lab.sensors.ImuData`. +* Moved and renamed :meth:`omni.isaac.lab.sensors.camera.utils.convert_orientation_convention` to :meth:`omni.isaac.lab.utils.math.convert_camera_frame_orientation_convention` +* Moved :meth:`omni.isaac.lab.sensors.camera.utils.create_rotation_matrix_from_view` to :meth:`omni.isaac.lab.utils.math.create_rotation_matrix_from_view` + + +0.25.2 (2024-10-16) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added support for different Gymnasium spaces (``Box``, ``Discrete``, ``MultiDiscrete``, ``Tuple`` and ``Dict``) + to define observation, action and state spaces in the direct workflow. +* Added :meth:`sample_space` to environment utils to sample supported spaces where data containers are torch tensors. + +Changed +^^^^^^^ + +* Mark the :attr:`num_observations`, :attr:`num_actions` and :attr:`num_states` in :class:`DirectRLEnvCfg` as deprecated + in favor of :attr:`observation_space`, :attr:`action_space` and :attr:`state_space` respectively. +* Mark the :attr:`num_observations`, :attr:`num_actions` and :attr:`num_states` in :class:`DirectMARLEnvCfg` as deprecated + in favor of :attr:`observation_spaces`, :attr:`action_spaces` and :attr:`state_space` respectively. + + 0.25.1 (2024-10-10) ~~~~~~~~~~~~~~~~~~~ diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/actuators/actuator_base.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/actuators/actuator_base.py index 67926fce31..fba943bc15 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/actuators/actuator_base.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/actuators/actuator_base.py @@ -10,9 +10,8 @@ from collections.abc import Sequence from typing import TYPE_CHECKING -from omni.isaac.core.utils.types import ArticulationActions - import omni.isaac.lab.utils.string as string_utils +from omni.isaac.lab.utils.types import ArticulationActions if TYPE_CHECKING: from .actuator_cfg import ActuatorBaseCfg @@ -207,6 +206,7 @@ def _parse_joint_parameter( TypeError: If the parameter value is not of the expected type. TypeError: If the default value is not of the expected type. ValueError: If the parameter value is None and no default value is provided. + ValueError: If the default value tensor is the wrong shape. """ # create parameter buffer param = torch.zeros(self._num_envs, self.num_joints, device=self._device) @@ -231,7 +231,14 @@ def _parse_joint_parameter( param[:] = float(default_value) elif isinstance(default_value, torch.Tensor): # if tensor, then use the same tensor for all joints - param[:] = default_value.float() + if default_value.shape == (self._num_envs, self.num_joints): + param = default_value.float() + else: + raise ValueError( + "Invalid default value tensor shape.\n" + f"Got: {default_value.shape}\n" + f"Expected: {(self._num_envs, self.num_joints)}" + ) else: raise TypeError( f"Invalid type for default value: {type(default_value)} for " diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/actuators/actuator_net.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/actuators/actuator_net.py index 5cf034679d..c0d440533a 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/actuators/actuator_net.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/actuators/actuator_net.py @@ -18,9 +18,8 @@ from collections.abc import Sequence from typing import TYPE_CHECKING -from omni.isaac.core.utils.types import ArticulationActions - from omni.isaac.lab.utils.assets import read_file +from omni.isaac.lab.utils.types import ArticulationActions from .actuator_pd import DCMotor diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/actuators/actuator_pd.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/actuators/actuator_pd.py index 165567536f..563b29ba47 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/actuators/actuator_pd.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/actuators/actuator_pd.py @@ -9,9 +9,8 @@ from collections.abc import Sequence from typing import TYPE_CHECKING -from omni.isaac.core.utils.types import ArticulationActions - from omni.isaac.lab.utils import DelayBuffer, LinearInterpolation +from omni.isaac.lab.utils.types import ArticulationActions from .actuator_base import ActuatorBase @@ -63,7 +62,22 @@ def reset(self, *args, **kwargs): def compute( self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor ) -> ArticulationActions: - """Compute the aproximmate torques for the actuated joint (physX does not compute this explicitly).""" + """Process the actuator group actions and compute the articulation actions. + + In case of implicit actuator, the control action is directly returned as the computed action. + This function is a no-op and does not perform any computation on the input control action. + However, it computes the approximate torques for the actuated joint since PhysX does not compute + this quantity explicitly. + + Args: + control_action: The joint action instance comprising of the desired joint positions, joint velocities + and (feed-forward) joint efforts. + joint_pos: The current joint positions of the joints in the group. Shape is (num_envs, num_joints). + joint_vel: The current joint velocities of the joints in the group. Shape is (num_envs, num_joints). + + Returns: + The computed desired joint positions, joint velocities and joint efforts. + """ # store approximate torques for reward computation error_pos = control_action.joint_positions - joint_pos error_vel = control_action.joint_velocities - joint_vel diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/app/app_launcher.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/app/app_launcher.py index 0453cc4cc5..c6dc4b84dc 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/app/app_launcher.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/app/app_launcher.py @@ -185,6 +185,10 @@ def add_app_launcher_args(parser: argparse.ArgumentParser) -> None: * If headless and enable_cameras are False, the experience file is set to ``isaaclab.python.kit``. * If headless is True and enable_cameras is False, the experience file is set to ``isaaclab.python.headless.kit``. + * ``kit_args`` (str): Optional command line arguments to be passed to Omniverse Kit directly. + Arguments should be combined into a single string separated by space. + Example usage: --kit_args "--ext-folder=/path/to/ext1 --ext-folder=/path/to/ext2" + Args: parser: An argument parser instance to be extended with the AppLauncher specific options. """ @@ -271,6 +275,15 @@ def add_app_launcher_args(parser: argparse.ArgumentParser) -> None: " it is resolved relative to the `apps` folder in Isaac Sim and Isaac Lab (in that order)." ), ) + arg_group.add_argument( + "--kit_args", + type=str, + default="", + help=( + "Command line arguments for Omniverse Kit as a string separated by a space delimiter." + ' Example usage: --kit_args "--ext-folder=/path/to/ext1 --ext-folder=/path/to/ext2"' + ), + ) # Corresponding to the beginning of the function, # if we have removed -h/--help handling, we add it back. @@ -557,6 +570,12 @@ def _config_resolution(self, launcher_args: dict): " The file does not exist." ) + # Resolve additional arguments passed to Kit + self._kit_args = [] + if "kit_args" in launcher_args: + self._kit_args = [arg for arg in launcher_args["kit_args"].split()] + sys.argv += self._kit_args + # Resolve the absolute path of the experience file self._sim_experience_file = os.path.abspath(self._sim_experience_file) print(f"[INFO][AppLauncher]: Loading experience file: {self._sim_experience_file}") @@ -595,6 +614,9 @@ def _create_app(self): # remove the threadCount argument from sys.argv if it was added for distributed training pattern = r"--/plugins/carb\.tasking\.plugin/threadCount=\d+" sys.argv = [arg for arg in sys.argv if not re.match(pattern, arg)] + # remove additional OV args from sys.argv + if len(self._kit_args) > 0: + sys.argv = [arg for arg in sys.argv if arg not in self._kit_args] def _rendering_enabled(self) -> bool: """Check if rendering is required by the app.""" diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/__init__.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/__init__.py index f0d29abcae..f91e4db7e6 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/__init__.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/__init__.py @@ -43,3 +43,4 @@ from .asset_base_cfg import AssetBaseCfg from .deformable_object import DeformableObject, DeformableObjectCfg, DeformableObjectData from .rigid_object import RigidObject, RigidObjectCfg, RigidObjectData +from .rigid_object_collection import RigidObjectCollection, RigidObjectCollectionCfg, RigidObjectCollectionData diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation.py index a604254caa..602a9efaa1 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation.py @@ -16,13 +16,13 @@ import omni.isaac.core.utils.stage as stage_utils import omni.log import omni.physics.tensors.impl.api as physx -from omni.isaac.core.utils.types import ArticulationActions from pxr import PhysxSchema, UsdPhysics import omni.isaac.lab.sim as sim_utils import omni.isaac.lab.utils.math as math_utils import omni.isaac.lab.utils.string as string_utils from omni.isaac.lab.actuators import ActuatorBase, ActuatorBaseCfg, ImplicitActuator +from omni.isaac.lab.utils.types import ArticulationActions from ..asset_base import AssetBase from .articulation_data import ArticulationData @@ -98,6 +98,10 @@ def __init__(self, cfg: ArticulationCfg): """ super().__init__(cfg) + self._root_state_dep_warn = False + self._root_pose_dep_warn = False + self._root_vel_dep_warn = False + """ Properties """ @@ -280,15 +284,71 @@ def write_root_state_to_sim(self, root_state: torch.Tensor, env_ids: Sequence[in root_state: Root state in simulation frame. Shape is (len(env_ids), 13). env_ids: Environment indices. If None, then all indices are used. """ + + # deprecation warning + if not self._root_state_dep_warn: + omni.log.warn( + "DeprecationWarning: Articluation.write_root_state_to_sim will be removed in a future release. Please" + " use write_root_link_state_to_sim or write_root_com_state_to_sim instead." + ) + self._root_state_dep_warn = True + # set into simulation self.write_root_pose_to_sim(root_state[:, :7], env_ids=env_ids) self.write_root_velocity_to_sim(root_state[:, 7:], env_ids=env_ids) + def write_root_com_state_to_sim(self, root_state: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root center of mass state over selected environment indices into the simulation. + + The root state comprises of the cartesian position, quaternion orientation in (w, x, y, z), and linear + and angular velocity. All the quantities are in the simulation frame. + + Args: + root_state: Root state in simulation frame. Shape is (len(env_ids), 13). + env_ids: Environment indices. If None, then all indices are used. + """ + # set into simulation + self.write_root_com_pose_to_sim(root_state[:, :7], env_ids=env_ids) + self.write_root_com_velocity_to_sim(root_state[:, 7:], env_ids=env_ids) + + def write_root_link_state_to_sim(self, root_state: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root link state over selected environment indices into the simulation. + + The root state comprises of the cartesian position, quaternion orientation in (w, x, y, z), and linear + and angular velocity. All the quantities are in the simulation frame. + + Args: + root_state: Root state in simulation frame. Shape is (len(env_ids), 13). + env_ids: Environment indices. If None, then all indices are used. + """ + # set into simulation + self.write_root_link_pose_to_sim(root_state[:, :7], env_ids=env_ids) + self.write_root_link_velocity_to_sim(root_state[:, 7:], env_ids=env_ids) + def write_root_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] | None = None): """Set the root pose over selected environment indices into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + Args: + root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7). + env_ids: Environment indices. If None, then all indices are used. + """ + # deprecation warning + if not self._root_pose_dep_warn: + omni.log.warn( + "DeprecationWarning: Articluation.write_root_pos_to_sim will be removed in a future release. Please use" + " write_root_link_pose_to_sim or write_root_com_pose_to_sim instead." + ) + self._root_pose_dep_warn = True + + self.write_root_link_pose_to_sim(root_pose, env_ids) + + def write_root_link_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root link pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + Args: root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7). env_ids: Environment indices. If None, then all indices are used. @@ -300,22 +360,78 @@ def write_root_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] physx_env_ids = self._ALL_INDICES # note: we need to do this here since tensors are not set into simulation until step. # set into internal buffers - self._data.root_state_w[env_ids, :7] = root_pose.clone() + self._data.root_link_state_w[env_ids, :7] = root_pose.clone() + self._data._ignore_dep_warn = True + self._data.root_state_w[env_ids, :7] = self._data.root_link_state_w[env_ids, :7] + self._data._ignore_dep_warn = False # convert root quaternion from wxyz to xyzw - root_poses_xyzw = self._data.root_state_w[:, :7].clone() + root_poses_xyzw = self._data.root_link_state_w[:, :7].clone() root_poses_xyzw[:, 3:] = math_utils.convert_quat(root_poses_xyzw[:, 3:], to="xyzw") # Need to invalidate the buffer to trigger the update with the new root pose. self._data._body_state_w.timestamp = -1.0 + self._data._body_link_state_w.timestamp = -1.0 + self._data._body_com_state_w.timestamp = -1.0 # set into simulation self.root_physx_view.set_root_transforms(root_poses_xyzw, indices=physx_env_ids) + def write_root_com_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root center of mass pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + The orientation is the orientation of the principle axes of inertia. + + Args: + root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7). + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + if env_ids is None: + local_env_ids = slice(None) + + com_pos = self.data.com_pos_b[local_env_ids, 0, :] + com_quat = self.data.com_quat_b[local_env_ids, 0, :] + + root_link_pos, root_link_quat = math_utils.combine_frame_transforms( + root_pose[..., :3], + root_pose[..., 3:7], + math_utils.quat_rotate(math_utils.quat_inv(com_quat), -com_pos), + math_utils.quat_inv(com_quat), + ) + + root_link_pose = torch.cat((root_link_pos, root_link_quat), dim=-1) + self.write_root_link_pose_to_sim(root_pose=root_link_pose, env_ids=env_ids) + def write_root_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Sequence[int] | None = None): - """Set the root velocity over selected environment indices into the simulation. + """Set the root center of mass velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + NOTE: This sets the velocity of the root's center of mass rather than the roots frame. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6). + env_ids: Environment indices. If None, then all indices are used. + """ + # deprecation warning + if not self._root_vel_dep_warn: + omni.log.warn( + "DeprecationWarning: Articluation.write_root_velocity_to_sim will be removed in a future release." + " Please use write_root_link_velocity_to_sim or write_root_com_velocity_to_sim instead." + ) + self._root_vel_dep_warn = True + + self.write_root_com_velocity_to_sim(root_velocity=root_velocity, env_ids=env_ids) + + def write_root_com_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root center of mass velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + NOTE: This sets the velocity of the root's center of mass rather than the roots frame. Args: - root_velocity: Root velocities in simulation frame. Shape is (len(env_ids), 6). + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6). env_ids: Environment indices. If None, then all indices are used. """ + # resolve all indices physx_env_ids = env_ids if env_ids is None: @@ -323,10 +439,37 @@ def write_root_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Seque physx_env_ids = self._ALL_INDICES # note: we need to do this here since tensors are not set into simulation until step. # set into internal buffers - self._data.root_state_w[env_ids, 7:] = root_velocity.clone() + self._data.root_com_state_w[env_ids, 7:] = root_velocity.clone() + self._data._ignore_dep_warn = True + self._data.root_state_w[env_ids, 7:] = self._data.root_com_state_w[env_ids, 7:] + self._data._ignore_dep_warn = False self._data.body_acc_w[env_ids] = 0.0 # set into simulation - self.root_physx_view.set_root_velocities(self._data.root_state_w[:, 7:], indices=physx_env_ids) + self.root_physx_view.set_root_velocities(self._data.root_com_state_w[:, 7:], indices=physx_env_ids) + + def write_root_link_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root link velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + NOTE: This sets the velocity of the root's frame rather than the roots center of mass. + + Args: + root_velocity: Root frame velocities in simulation world frame. Shape is (len(env_ids), 6). + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + if env_ids is None: + local_env_ids = slice(None) + + root_com_velocity = root_velocity.clone() + quat = self.data.root_link_state_w[local_env_ids, 3:7] + com_pos_b = self.data.com_pos_b[local_env_ids, 0, :] + # transform given velocity to center of mass + root_com_velocity[:, :3] += torch.linalg.cross( + root_com_velocity[:, 3:], math_utils.quat_rotate(quat, com_pos_b), dim=-1 + ) + # write center of mass velocity to sim + self.write_root_com_velocity_to_sim(root_velocity=root_com_velocity, env_ids=env_ids) def write_joint_state_to_sim( self, @@ -360,6 +503,8 @@ def write_joint_state_to_sim( self._data.joint_acc[env_ids, joint_ids] = 0.0 # Need to invalidate the buffer to trigger the update with the new root pose. self._data._body_state_w.timestamp = -1.0 + self._data._body_link_state_w.timestamp = -1.0 + self._data._body_com_state_w.timestamp = -1.0 # set into simulation self.root_physx_view.set_dof_positions(self._data.joint_pos, indices=physx_env_ids) self.root_physx_view.set_dof_velocities(self._data.joint_vel, indices=physx_env_ids) @@ -424,6 +569,39 @@ def write_joint_damping_to_sim( # set into simulation self.root_physx_view.set_dof_dampings(self._data.joint_damping.cpu(), indices=physx_env_ids.cpu()) + def write_joint_velocity_limit_to_sim( + self, + limits: torch.Tensor | float, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ): + """Write joint max velocity to the simulation. + + Args: + limits: Joint max velocity. Shape is (len(env_ids), len(joint_ids)). + joint_ids: The joint indices to set the max velocity for. Defaults to None (all joints). + env_ids: The environment indices to set the max velocity for. Defaults to None (all environments). + """ + # resolve indices + physx_env_ids = env_ids + if env_ids is None: + env_ids = slice(None) + physx_env_ids = self._ALL_INDICES + if joint_ids is None: + joint_ids = slice(None) + # broadcast env_ids if needed to allow double indexing + if env_ids != slice(None) and joint_ids != slice(None): + env_ids = env_ids[:, None] + # move tensor to cpu if needed + if isinstance(limits, torch.Tensor): + limits = limits.to(self.device) + + # set into internal buffers + self._data.joint_velocity_limits = self.root_physx_view.get_dof_max_velocities().to(self.device) + self._data.joint_velocity_limits[env_ids, joint_ids] = limits + # set into simulation + self.root_physx_view.set_dof_max_velocities(self._data.joint_velocity_limits.cpu(), indices=physx_env_ids.cpu()) + def write_joint_effort_limit_to_sim( self, limits: torch.Tensor | float, @@ -540,8 +718,13 @@ def write_joint_limits_to_sim( # set into internal buffers self._data.joint_limits[env_ids, joint_ids] = limits # update default joint pos to stay within the new limits - if torch.any((self._data.default_joint_pos < limits[..., 0]) | (self._data.default_joint_pos > limits[..., 1])): - self._data.default_joint_pos = torch.clamp(self._data.default_joint_pos, limits[..., 0], limits[..., 1]) + if torch.any( + (self._data.default_joint_pos[env_ids, joint_ids] < limits[..., 0]) + | (self._data.default_joint_pos[env_ids, joint_ids] > limits[..., 1]) + ): + self._data.default_joint_pos[env_ids, joint_ids] = torch.clamp( + self._data.default_joint_pos[env_ids, joint_ids], limits[..., 0], limits[..., 1] + ) omni.log.warn( "Some default joint positions are outside of the range of the new joint limits. Default joint positions" " will be clamped to be within the new joint limits." @@ -1105,12 +1288,10 @@ def _process_actuators_cfg(self): self._has_implicit_actuators = False # cache the values coming from the usd - usd_stiffness = self.root_physx_view.get_dof_stiffnesses().clone() - usd_damping = self.root_physx_view.get_dof_dampings().clone() - usd_armature = self.root_physx_view.get_dof_armatures().clone() - usd_friction = self.root_physx_view.get_dof_friction_coefficients().clone() - usd_effort_limit = self.root_physx_view.get_dof_max_forces().clone() - usd_velocity_limit = self.root_physx_view.get_dof_max_velocities().clone() + self._data.default_joint_stiffness = self.root_physx_view.get_dof_stiffnesses().to(self.device).clone() + self._data.default_joint_damping = self.root_physx_view.get_dof_dampings().to(self.device).clone() + self._data.default_joint_armature = self.root_physx_view.get_dof_armatures().to(self.device).clone() + self._data.default_joint_friction = self.root_physx_view.get_dof_friction_coefficients().to(self.device).clone() # iterate over all actuator configurations for actuator_name, actuator_cfg in self.cfg.actuators.items(): @@ -1134,12 +1315,12 @@ def _process_actuators_cfg(self): ), num_envs=self.num_instances, device=self.device, - stiffness=usd_stiffness[:, joint_ids], - damping=usd_damping[:, joint_ids], - armature=usd_armature[:, joint_ids], - friction=usd_friction[:, joint_ids], - effort_limit=usd_effort_limit[:, joint_ids], - velocity_limit=usd_velocity_limit[:, joint_ids], + stiffness=self._data.default_joint_stiffness[:, joint_ids], + damping=self._data.default_joint_damping[:, joint_ids], + armature=self._data.default_joint_armature[:, joint_ids], + friction=self._data.default_joint_friction[:, joint_ids], + effort_limit=self.root_physx_view.get_dof_max_forces().to(self.device).clone()[:, joint_ids], + velocity_limit=self.root_physx_view.get_dof_max_velocities().to(self.device).clone()[:, joint_ids], ) # log information on actuator groups omni.log.info( @@ -1155,6 +1336,7 @@ def _process_actuators_cfg(self): self.write_joint_stiffness_to_sim(actuator.stiffness, joint_ids=actuator.joint_indices) self.write_joint_damping_to_sim(actuator.damping, joint_ids=actuator.joint_indices) self.write_joint_effort_limit_to_sim(actuator.effort_limit, joint_ids=actuator.joint_indices) + self.write_joint_velocity_limit_to_sim(actuator.velocity_limit, joint_ids=actuator.joint_indices) self.write_joint_armature_to_sim(actuator.armature, joint_ids=actuator.joint_indices) self.write_joint_friction_to_sim(actuator.friction, joint_ids=actuator.joint_indices) else: @@ -1163,16 +1345,12 @@ def _process_actuators_cfg(self): self.write_joint_stiffness_to_sim(0.0, joint_ids=actuator.joint_indices) self.write_joint_damping_to_sim(0.0, joint_ids=actuator.joint_indices) self.write_joint_effort_limit_to_sim(1.0e9, joint_ids=actuator.joint_indices) + self.write_joint_velocity_limit_to_sim(actuator.velocity_limit, joint_ids=actuator.joint_indices) self.write_joint_armature_to_sim(actuator.armature, joint_ids=actuator.joint_indices) self.write_joint_friction_to_sim(actuator.friction, joint_ids=actuator.joint_indices) - - # set the default joint parameters based on the changes from the actuators - self._data.default_joint_stiffness = self.root_physx_view.get_dof_stiffnesses().to(device=self.device).clone() - self._data.default_joint_damping = self.root_physx_view.get_dof_dampings().to(device=self.device).clone() - self._data.default_joint_armature = self.root_physx_view.get_dof_armatures().to(device=self.device).clone() - self._data.default_joint_friction = ( - self.root_physx_view.get_dof_friction_coefficients().to(device=self.device).clone() - ) + # Store the actual default stiffness and damping values for explicit actuators (not written the sim) + self._data.default_joint_stiffness[:, actuator.joint_indices] = actuator.stiffness + self._data.default_joint_damping[:, actuator.joint_indices] = actuator.damping # perform some sanity checks to ensure actuators are prepared correctly total_act_joints = sum(actuator.num_joints for actuator in self.actuators.values()) @@ -1190,10 +1368,11 @@ def _process_fixed_tendons(self): # parse fixed tendons properties if they exist if self.num_fixed_tendons > 0: stage = stage_utils.get_current_stage() + joint_paths = self.root_physx_view.dof_paths[0] # iterate over all joints to find tendons attached to them for j in range(self.num_joints): - usd_joint_path = self.root_physx_view.dof_paths[0][j] + usd_joint_path = joint_paths[j] # check whether joint has tendons - tendon name follows the joint name it is attached to joint = UsdPhysics.Joint.Get(stage, usd_joint_path) if joint.GetPrim().HasAPI(PhysxSchema.PhysxTendonAxisRootAPI): diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation_data.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation_data.py index 4b84f96ba3..a53b88ad6d 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation_data.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation_data.py @@ -6,6 +6,7 @@ import torch import weakref +import omni.log import omni.physics.tensors.impl.api as physx import omni.isaac.lab.utils.math as math_utils @@ -64,12 +65,21 @@ def __init__(self, root_physx_view: physx.ArticulationView, device: str): # Initialize the lazy buffers. self._root_state_w = TimestampedBuffer() + self._root_link_state_w = TimestampedBuffer() + self._root_com_state_w = TimestampedBuffer() self._body_state_w = TimestampedBuffer() + self._body_link_state_w = TimestampedBuffer() + self._body_com_state_w = TimestampedBuffer() self._body_acc_w = TimestampedBuffer() self._joint_pos = TimestampedBuffer() self._joint_acc = TimestampedBuffer() self._joint_vel = TimestampedBuffer() + # deprecation warning check + self._root_state_dep_warn = False + self._body_state_dep_warn = False + self._ignore_dep_warn = False + def update(self, dt: float): # update the simulation timestamp self._sim_timestamp += dt @@ -211,15 +221,12 @@ def update(self, dt: float): joint_damping: torch.Tensor = None """Joint damping provided to simulation. Shape is (num_instances, num_joints).""" - joint_armature: torch.Tensor = None - """Joint armature provided to simulation. Shape is (num_instances, num_joints).""" - - joint_friction: torch.Tensor = None - """Joint friction provided to simulation. Shape is (num_instances, num_joints).""" - joint_limits: torch.Tensor = None """Joint limits provided to simulation. Shape is (num_instances, num_joints, 2).""" + joint_velocity_limits: torch.Tensor = None + """Joint maximum velocity provided to simulation. Shape is (num_instances, num_joints).""" + ## # Fixed tendon properties. ## @@ -263,9 +270,17 @@ def update(self, dt: float): def root_state_w(self): """Root state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 13). - The position and quaternion are of the articulation root's actor frame. Meanwhile, the linear and angular - velocities are of the articulation root's center of mass frame. + The position and quaternion are of the articulation root's actor frame relative to the world. Meanwhile, + the linear and angular velocities are of the articulation root's center of mass frame. """ + + if not self._root_state_dep_warn and not self._ignore_dep_warn: + omni.log.warn( + "DeprecationWarning: root_state_w and it's derived properties will be deprecated in a future release." + " Please use root_link_state_w or root_com_state_w." + ) + self._root_state_dep_warn = True + if self._root_state_w.timestamp < self._sim_timestamp: # read data from simulation pose = self._root_physx_view.get_root_transforms().clone() @@ -276,6 +291,53 @@ def root_state_w(self): self._root_state_w.timestamp = self._sim_timestamp return self._root_state_w.data + @property + def root_link_state_w(self): + """Root state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 13). + + The position, quaternion, and linear/angular velocity are of the articulation root's actor frame relative to the + world. + """ + if self._root_link_state_w.timestamp < self._sim_timestamp: + # read data from simulation + pose = self._root_physx_view.get_root_transforms().clone() + pose[:, 3:7] = math_utils.convert_quat(pose[:, 3:7], to="wxyz") + velocity = self._root_physx_view.get_root_velocities().clone() + + # adjust linear velocity to link from center of mass + velocity[:, :3] += torch.linalg.cross( + velocity[:, 3:], math_utils.quat_rotate(pose[:, 3:7], -self.com_pos_b[:, 0, :]), dim=-1 + ) + # set the buffer data and timestamp + self._root_link_state_w.data = torch.cat((pose, velocity), dim=-1) + self._root_link_state_w.timestamp = self._sim_timestamp + + return self._root_link_state_w.data + + @property + def root_com_state_w(self): + """Root center of mass state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 13). + + The position, quaternion, and linear/angular velocity are of the articulation root link's center of mass frame + relative to the world. Center of mass frame is assumed to be the same orientation as the link rather than the + orientation of the principle inertia. + """ + if self._root_com_state_w.timestamp < self._sim_timestamp: + # read data from simulation (pose is of link) + pose = self._root_physx_view.get_root_transforms().clone() + pose[:, 3:7] = math_utils.convert_quat(pose[:, 3:7], to="wxyz") + velocity = self._root_physx_view.get_root_velocities() + + # adjust pose to center of mass + pos, quat = math_utils.combine_frame_transforms( + pose[:, :3], pose[:, 3:7], self.com_pos_b[:, 0, :], self.com_quat_b[:, 0, :] + ) + pose = torch.cat((pos, quat), dim=-1) + # set the buffer data and timestamp + self._root_com_state_w.data = torch.cat((pose, velocity), dim=-1) + self._root_com_state_w.timestamp = self._sim_timestamp + return self._root_com_state_w.data + @property def body_state_w(self): """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame. @@ -284,6 +346,14 @@ def body_state_w(self): The position and quaternion are of all the articulation links's actor frame. Meanwhile, the linear and angular velocities are of the articulation links's center of mass frame. """ + + if not self._body_state_dep_warn and not self._ignore_dep_warn: + omni.log.warn( + "DeprecationWarning: body_state_w and it's derived properties will be deprecated in a future release." + " Please use body_link_state_w or body_com_state_w." + ) + self._body_state_dep_warn = True + if self._body_state_w.timestamp < self._sim_timestamp: self._physics_sim_view.update_articulations_kinematic() # read data from simulation @@ -295,22 +365,73 @@ def body_state_w(self): self._body_state_w.timestamp = self._sim_timestamp return self._body_state_w.data + @property + def body_link_state_w(self): + """State of all bodies' link frame`[pos, quat, lin_vel, ang_vel]` in simulation world frame. + Shape is (num_instances, num_bodies, 13). + + The position, quaternion, and linear/angular velocity are of the body's link frame relative to the world. + """ + if self._body_link_state_w.timestamp < self._sim_timestamp: + self._physics_sim_view.update_articulations_kinematic() + # read data from simulation + pose = self._root_physx_view.get_link_transforms().clone() + pose[..., 3:7] = math_utils.convert_quat(pose[..., 3:7], to="wxyz") + velocity = self._root_physx_view.get_link_velocities() + + # adjust linear velocity to link from center of mass + velocity[..., :3] += torch.linalg.cross( + velocity[..., 3:], math_utils.quat_rotate(pose[..., 3:7], -self.com_pos_b), dim=-1 + ) + # set the buffer data and timestamp + self._body_link_state_w.data = torch.cat((pose, velocity), dim=-1) + self._body_link_state_w.timestamp = self._sim_timestamp + + return self._body_link_state_w.data + + @property + def body_com_state_w(self): + """State of all bodies center of mass `[pos, quat, lin_vel, ang_vel]` in simulation world frame. + Shape is (num_instances, num_bodies, 13). + + The position, quaternion, and linear/angular velocity are of the body's center of mass frame relative to the + world. Center of mass frame is assumed to be the same orientation as the link rather than the orientation of the + principle inertia. + """ + if self._body_com_state_w.timestamp < self._sim_timestamp: + self._physics_sim_view.update_articulations_kinematic() + # read data from simulation (pose is of link) + pose = self._root_physx_view.get_link_transforms().clone() + pose[..., 3:7] = math_utils.convert_quat(pose[..., 3:7], to="wxyz") + velocity = self._root_physx_view.get_link_velocities() + + # adjust pose to center of mass + pos, quat = math_utils.combine_frame_transforms( + pose[..., :3], pose[..., 3:7], self.com_pos_b, self.com_quat_b + ) + pose = torch.cat((pos, quat), dim=-1) + # set the buffer data and timestamp + self._body_com_state_w.data = torch.cat((pose, velocity), dim=-1) + self._body_com_state_w.timestamp = self._sim_timestamp + return self._body_com_state_w.data + @property def body_acc_w(self): - """Acceleration of all bodies. Shape is (num_instances, num_bodies, 6). + """Acceleration of all bodies (center of mass). Shape is (num_instances, num_bodies, 6). - This quantity is the acceleration of the articulation links' center of mass frame. + All values are relative to the world. """ if self._body_acc_w.timestamp < self._sim_timestamp: # read data from simulation and set the buffer data and timestamp self._body_acc_w.data = self._root_physx_view.get_link_accelerations() + self._body_acc_w.timestamp = self._sim_timestamp return self._body_acc_w.data @property def projected_gravity_b(self): """Projection of the gravity direction on base frame. Shape is (num_instances, 3).""" - return math_utils.quat_rotate_inverse(self.root_quat_w, self.GRAVITY_VEC_W) + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.GRAVITY_VEC_W) @property def heading_w(self): @@ -320,7 +441,7 @@ def heading_w(self): This quantity is computed by assuming that the forward-direction of the base frame is along x-direction, i.e. :math:`(1, 0, 0)`. """ - forward_w = math_utils.quat_apply(self.root_quat_w, self.FORWARD_VEC_B) + forward_w = math_utils.quat_apply(self.root_link_quat_w, self.FORWARD_VEC_B) return torch.atan2(forward_w[:, 1], forward_w[:, 0]) @property @@ -361,7 +482,7 @@ def joint_acc(self): def root_pos_w(self) -> torch.Tensor: """Root position in simulation world frame. Shape is (num_instances, 3). - This quantity is the position of the actor frame of the articulation root. + This quantity is the position of the actor frame of the articulation root relative to the world. """ return self.root_state_w[:, :3] @@ -369,7 +490,7 @@ def root_pos_w(self) -> torch.Tensor: def root_quat_w(self) -> torch.Tensor: """Root orientation (w, x, y, z) in simulation world frame. Shape is (num_instances, 4). - This quantity is the orientation of the actor frame of the articulation root. + This quantity is the orientation of the actor frame of the articulation root relative to the world. """ return self.root_state_w[:, 3:7] @@ -377,8 +498,8 @@ def root_quat_w(self) -> torch.Tensor: def root_vel_w(self) -> torch.Tensor: """Root velocity in simulation world frame. Shape is (num_instances, 6). - This quantity contains the linear and angular velocities of the articulation root's center of - mass frame. + This quantity contains the linear and angular velocities of the articulation root's center of mass frame + relative to the world. """ return self.root_state_w[:, 7:13] @@ -386,7 +507,7 @@ def root_vel_w(self) -> torch.Tensor: def root_lin_vel_w(self) -> torch.Tensor: """Root linear velocity in simulation world frame. Shape is (num_instances, 3). - This quantity is the linear velocity of the articulation root's center of mass frame. + This quantity is the linear velocity of the articulation root's center of mass frame relative to the world. """ return self.root_state_w[:, 7:10] @@ -394,7 +515,7 @@ def root_lin_vel_w(self) -> torch.Tensor: def root_ang_vel_w(self) -> torch.Tensor: """Root angular velocity in simulation world frame. Shape is (num_instances, 3). - This quantity is the angular velocity of the articulation root's center of mass frame. + This quantity is the angular velocity of the articulation root's center of mass frame relative to the world. """ return self.root_state_w[:, 10:13] @@ -402,8 +523,8 @@ def root_ang_vel_w(self) -> torch.Tensor: def root_lin_vel_b(self) -> torch.Tensor: """Root linear velocity in base frame. Shape is (num_instances, 3). - This quantity is the linear velocity of the articulation root's center of mass frame with - respect to the articulation root's actor frame. + This quantity is the linear velocity of the articulation root's center of mass frame relative to the world + with respect to the articulation root's actor frame. """ return math_utils.quat_rotate_inverse(self.root_quat_w, self.root_lin_vel_w) @@ -411,16 +532,163 @@ def root_lin_vel_b(self) -> torch.Tensor: def root_ang_vel_b(self) -> torch.Tensor: """Root angular velocity in base world frame. Shape is (num_instances, 3). - This quantity is the angular velocity of the articulation root's center of mass frame with respect to the - articulation root's actor frame. + This quantity is the angular velocity of the articulation root's center of mass frame relative to the world with + respect to the articulation root's actor frame. """ return math_utils.quat_rotate_inverse(self.root_quat_w, self.root_ang_vel_w) + # + # Derived Root Link Frame Properties + # + + @property + def root_link_pos_w(self) -> torch.Tensor: + """Root link position in simulation world frame. Shape is (num_instances, 3). + + This quantity is the position of the actor frame of the root rigid body relative to the world. + """ + if self._body_com_state_w.timestamp < self._sim_timestamp: + # read data from simulation (pose is of link) + pose = self._root_physx_view.get_root_transforms() + return pose[:, :3] + return self.root_link_state_w[:, :3] + + @property + def root_link_quat_w(self) -> torch.Tensor: + """Root link orientation (w, x, y, z) in simulation world frame. Shape is (num_instances, 4). + + This quantity is the orientation of the actor frame of the root rigid body. + """ + if self._body_com_state_w.timestamp < self._sim_timestamp: + # read data from simulation (pose is of link) + pose = self._root_physx_view.get_root_transforms().clone() + pose[:, 3:7] = math_utils.convert_quat(pose[:, 3:7], to="wxyz") + return pose[:, 3:7] + return self.root_link_state_w[:, 3:7] + + @property + def root_link_vel_w(self) -> torch.Tensor: + """Root link velocity in simulation world frame. Shape is (num_instances, 6). + + This quantity contains the linear and angular velocities of the actor frame of the root + rigid body relative to the world. + """ + return self.root_link_state_w[:, 7:13] + + @property + def root_link_lin_vel_w(self) -> torch.Tensor: + """Root linear velocity in simulation world frame. Shape is (num_instances, 3). + + This quantity is the linear velocity of the root rigid body's actor frame relative to the world. + """ + return self.root_link_state_w[:, 7:10] + + @property + def root_link_ang_vel_w(self) -> torch.Tensor: + """Root link angular velocity in simulation world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. + """ + return self.root_link_state_w[:, 10:13] + + @property + def root_link_lin_vel_b(self) -> torch.Tensor: + """Root link linear velocity in base frame. Shape is (num_instances, 3). + + This quantity is the linear velocity of the actor frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.root_link_lin_vel_w) + + @property + def root_link_ang_vel_b(self) -> torch.Tensor: + """Root link angular velocity in base world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the actor frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.root_link_ang_vel_w) + + # + # Root Center of Mass state properties + # + + @property + def root_com_pos_w(self) -> torch.Tensor: + """Root center of mass position in simulation world frame. Shape is (num_instances, 3). + + This quantity is the position of the actor frame of the root rigid body relative to the world. + """ + return self.root_com_state_w[:, :3] + + @property + def root_com_quat_w(self) -> torch.Tensor: + """Root center of mass orientation (w, x, y, z) in simulation world frame. Shape is (num_instances, 4). + + This quantity is the orientation of the actor frame of the root rigid body relative to the world. + """ + return self.root_com_state_w[:, 3:7] + + @property + def root_com_vel_w(self) -> torch.Tensor: + """Root center of mass velocity in simulation world frame. Shape is (num_instances, 6). + + This quantity contains the linear and angular velocities of the root rigid body's center of mass frame relative to the world. + """ + if self._root_com_state_w.timestamp < self._sim_timestamp: + # read data from simulation (pose is of link) + velocity = self._root_physx_view.get_root_velocities() + return velocity + return self.root_com_state_w[:, 7:13] + + @property + def root_com_lin_vel_w(self) -> torch.Tensor: + """Root center of mass linear velocity in simulation world frame. Shape is (num_instances, 3). + + This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. + """ + if self._root_com_state_w.timestamp < self._sim_timestamp: + # read data from simulation (pose is of link) + velocity = self._root_physx_view.get_root_velocities() + return velocity[:, 0:3] + return self.root_com_state_w[:, 7:10] + + @property + def root_com_ang_vel_w(self) -> torch.Tensor: + """Root center of mass angular velocity in simulation world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. + """ + if self._root_com_state_w.timestamp < self._sim_timestamp: + self._physics_sim_view.update_articulations_kinematic() + # read data from simulation (pose is of link) + velocity = self._root_physx_view.get_root_velocities() + return velocity[:, 3:6] + return self.root_com_state_w[:, 10:13] + + @property + def root_com_lin_vel_b(self) -> torch.Tensor: + """Root center of mass linear velocity in base frame. Shape is (num_instances, 3). + + This quantity is the linear velocity of the root rigid body's center of mass frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.root_com_lin_vel_w) + + @property + def root_com_ang_vel_b(self) -> torch.Tensor: + """Root center of mass angular velocity in base world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the root rigid body's center of mass frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.root_com_ang_vel_w) + @property def body_pos_w(self) -> torch.Tensor: """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). - This quantity is the position of the rigid bodies' actor frame. + This quantity is the position of the rigid bodies' actor frame relative to the world. """ return self.body_state_w[..., :3] @@ -428,7 +696,7 @@ def body_pos_w(self) -> torch.Tensor: def body_quat_w(self) -> torch.Tensor: """Orientation (w, x, y, z) of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 4). - This quantity is the orientation of the rigid bodies' actor frame. + This quantity is the orientation of the rigid bodies' actor frame relative to the world. """ return self.body_state_w[..., 3:7] @@ -436,7 +704,8 @@ def body_quat_w(self) -> torch.Tensor: def body_vel_w(self) -> torch.Tensor: """Velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 6). - This quantity contains the linear and angular velocities of the rigid bodies' center of mass frame. + This quantity contains the linear and angular velocities of the rigid bodies' center of mass frame relative + to the world. """ return self.body_state_w[..., 7:13] @@ -444,7 +713,7 @@ def body_vel_w(self) -> torch.Tensor: def body_lin_vel_w(self) -> torch.Tensor: """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). - This quantity is the linear velocity of the rigid bodies' center of mass frame. + This quantity is the linear velocity of the rigid bodies' center of mass frame relative to the world. """ return self.body_state_w[..., 7:10] @@ -452,7 +721,7 @@ def body_lin_vel_w(self) -> torch.Tensor: def body_ang_vel_w(self) -> torch.Tensor: """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). - This quantity is the angular velocity of the rigid bodies' center of mass frame. + This quantity is the angular velocity of the rigid bodies' center of mass frame relative to the world. """ return self.body_state_w[..., 10:13] @@ -460,7 +729,7 @@ def body_ang_vel_w(self) -> torch.Tensor: def body_lin_acc_w(self) -> torch.Tensor: """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). - This quantity is the linear acceleration of the rigid bodies' center of mass frame. + This quantity is the linear acceleration of the rigid bodies' center of mass frame relative to the world. """ return self.body_acc_w[..., 0:3] @@ -468,6 +737,138 @@ def body_lin_acc_w(self) -> torch.Tensor: def body_ang_acc_w(self) -> torch.Tensor: """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). - This quantity is the angular acceleration of the rigid bodies' center of mass frame. + This quantity is the angular acceleration of the rigid bodies' center of mass frame relative to the world. """ return self.body_acc_w[..., 3:6] + + # + # Link body properties + # + + @property + def body_link_pos_w(self) -> torch.Tensor: + """Positions of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the position of the rigid bodies' actor frame relative to the world. + """ + if self._body_link_state_w.timestamp < self._sim_timestamp: + self._physics_sim_view.update_articulations_kinematic() + # read data from simulation + pose = self._root_physx_view.get_link_transforms() + return pose[..., :3] + return self._body_link_state_w.data[..., :3] + + @property + def body_link_quat_w(self) -> torch.Tensor: + """Orientation (w, x, y, z) of all bodies in simulation world frame. Shape is (num_instances, 1, 4). + + This quantity is the orientation of the rigid bodies' actor frame relative to the world. + """ + if self._body_link_state_w.timestamp < self._sim_timestamp: + self._physics_sim_view.update_articulations_kinematic() + # read data from simulation + pose = self._root_physx_view.get_link_transforms().clone() + pose[..., 3:7] = math_utils.convert_quat(pose[..., 3:7], to="wxyz") + return pose[..., 3:7] + return self.body_link_state_w[..., 3:7] + + @property + def body_link_vel_w(self) -> torch.Tensor: + """Velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 6). + + This quantity contains the linear and angular velocities of the rigid bodies' center of mass frame + relative to the world. + """ + return self.body_link_state_w[..., 7:13] + + @property + def body_link_lin_vel_w(self) -> torch.Tensor: + """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the linear velocity of the rigid bodies' center of mass frame relative to the world. + """ + return self.body_link_state_w[..., 7:10] + + @property + def body_link_ang_vel_w(self) -> torch.Tensor: + """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the angular velocity of the rigid bodies' center of mass frame relative to the world. + """ + return self.body_link_state_w[..., 10:13] + + # + # Center of mass body properties + # + + @property + def body_com_pos_w(self) -> torch.Tensor: + """Positions of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the position of the rigid bodies' actor frame. + """ + return self.body_com_state_w[..., :3] + + @property + def body_com_quat_w(self) -> torch.Tensor: + """Orientation (w, x, y, z) of the prinicple axies of inertia of all bodies in simulation world frame. + + Shape is (num_instances, 1, 4). This quantity is the orientation of the rigid bodies' actor frame. + """ + return self.body_com_state_w[..., 3:7] + + @property + def body_com_vel_w(self) -> torch.Tensor: + """Velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 6). + + This quantity contains the linear and angular velocities of the rigid bodies' center of mass frame. + """ + if self._body_com_state_w.timestamp < self._sim_timestamp: + self._physics_sim_view.update_articulations_kinematic() + # read data from simulation (velocity is of com) + velocity = self._root_physx_view.get_link_velocities() + return velocity + return self.body_com_state_w[..., 7:13] + + @property + def body_com_lin_vel_w(self) -> torch.Tensor: + """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the linear velocity of the rigid bodies' center of mass frame. + """ + if self._body_com_state_w.timestamp < self._sim_timestamp: + self._physics_sim_view.update_articulations_kinematic() + # read data from simulation (velocity is of com) + velocity = self._root_physx_view.get_link_velocities() + return velocity[..., 0:3] + return self.body_com_state_w[..., 7:10] + + @property + def body_com_ang_vel_w(self) -> torch.Tensor: + """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the angular velocity of the rigid bodies' center of mass frame. + """ + if self._body_com_state_w.timestamp < self._sim_timestamp: + self._physics_sim_view.update_articulations_kinematic() + # read data from simulation (velocity is of com) + velocity = self._root_physx_view.get_link_velocities() + return velocity[..., 3:6] + return self.body_com_state_w[..., 10:13] + + @property + def com_pos_b(self) -> torch.Tensor: + """Center of mass of all of the bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + + This quantity is the center of mass location relative to its body frame. + """ + return self._root_physx_view.get_coms().to(self.device)[..., :3] + + @property + def com_quat_b(self) -> torch.Tensor: + """Orientation (w,x,y,z) of the prinicple axies of inertia of all of the bodies in simulation world frame. Shape is (num_instances, num_bodies, 4). + + This quantity is the orientation of the principles axes of inertia relative to its body frame. + """ + quat = self._root_physx_view.get_coms().to(self.device)[..., 3:7] + return math_utils.convert_quat(quat, to="wxyz") diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/asset_base.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/asset_base.py index 9a55a87ef0..8c66bb626a 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/asset_base.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/asset_base.py @@ -59,6 +59,8 @@ def __init__(self, cfg: AssetBaseCfg): Raises: RuntimeError: If no prims found at input prim path or prim path expression. """ + # check that the config is valid + cfg.validate() # store inputs self.cfg = cfg # flag for whether the asset is initialized diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/asset_base_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/asset_base_cfg.py index 6bea572dcb..62d047fb5c 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/asset_base_cfg.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/asset_base_cfg.py @@ -39,8 +39,9 @@ class InitialStateCfg: Defaults to (1.0, 0.0, 0.0, 0.0). """ - class_type: type[AssetBase] = MISSING - """The associated asset class. + class_type: type[AssetBase] = None + """The associated asset class. Defaults to None, which means that the asset will be spawned + but cannot be interacted with via the asset class. The class should inherit from :class:`omni.isaac.lab.assets.asset_base.AssetBase`. """ diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object/rigid_object.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object/rigid_object.py index e44ee7eddf..426a06fbf2 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object/rigid_object.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object/rigid_object.py @@ -55,6 +55,9 @@ def __init__(self, cfg: RigidObjectCfg): cfg: A configuration instance. """ super().__init__(cfg) + self._root_state_dep_warn = False + self._root_pose_dep_warn = False + self._root_vel_dep_warn = False """ Properties @@ -156,15 +159,70 @@ def write_root_state_to_sim(self, root_state: torch.Tensor, env_ids: Sequence[in root_state: Root state in simulation frame. Shape is (len(env_ids), 13). env_ids: Environment indices. If None, then all indices are used. """ + # deprecation warning + if not self._root_state_dep_warn: + omni.log.warn( + "DeprecationWarning: RigidObject.write_root_state_to_sim will be removed in a future release. Please" + " use write_root_link_state_to_sim or write_root_com_state_to_sim instead." + ) + self._root_state_dep_warn = True + # set into simulation self.write_root_pose_to_sim(root_state[:, :7], env_ids=env_ids) self.write_root_velocity_to_sim(root_state[:, 7:], env_ids=env_ids) + def write_root_com_state_to_sim(self, root_state: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root center of mass state over selected environment indices into the simulation. + + The root state comprises of the cartesian position, quaternion orientation in (w, x, y, z), and linear + and angular velocity. All the quantities are in the simulation frame. + + Args: + root_state: Root state in simulation frame. Shape is (len(env_ids), 13). + env_ids: Environment indices. If None, then all indices are used. + """ + # set into simulation + self.write_root_com_pose_to_sim(root_state[:, :7], env_ids=env_ids) + self.write_root_com_velocity_to_sim(root_state[:, 7:], env_ids=env_ids) + + def write_root_link_state_to_sim(self, root_state: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root link state over selected environment indices into the simulation. + + The root state comprises of the cartesian position, quaternion orientation in (w, x, y, z), and linear + and angular velocity. All the quantities are in the simulation frame. + + Args: + root_state: Root state in simulation frame. Shape is (len(env_ids), 13). + env_ids: Environment indices. If None, then all indices are used. + """ + # set into simulation + self.write_root_link_pose_to_sim(root_state[:, :7], env_ids=env_ids) + self.write_root_link_velocity_to_sim(root_state[:, 7:], env_ids=env_ids) + def write_root_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] | None = None): """Set the root pose over selected environment indices into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + Args: + root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7). + env_ids: Environment indices. If None, then all indices are used. + """ + # deprecation warning + if not self._root_pose_dep_warn: + omni.log.warn( + "DeprecationWarning: RigidObject.write_root_pose_to_sim will be removed in a future release. Please use" + " write_root_link_pose_to_sim or write_root_com_pose_to_sim instead." + ) + self._root_pose_dep_warn = True + + self.write_root_link_pose_to_sim(root_pose, env_ids) + + def write_root_link_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root link pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + Args: root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7). env_ids: Environment indices. If None, then all indices are used. @@ -176,20 +234,75 @@ def write_root_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] physx_env_ids = self._ALL_INDICES # note: we need to do this here since tensors are not set into simulation until step. # set into internal buffers - self._data.root_state_w[env_ids, :7] = root_pose.clone() + self._data.root_link_state_w[env_ids, :7] = root_pose.clone() + self._data._ignore_dep_warn = True + self._data.root_state_w[env_ids, :7] = self._data.root_link_state_w[env_ids, :7] + self._data._ignore_dep_warn = False # convert root quaternion from wxyz to xyzw - root_poses_xyzw = self._data.root_state_w[:, :7].clone() + root_poses_xyzw = self._data.root_link_state_w[:, :7].clone() root_poses_xyzw[:, 3:] = math_utils.convert_quat(root_poses_xyzw[:, 3:], to="xyzw") # set into simulation self.root_physx_view.set_transforms(root_poses_xyzw, indices=physx_env_ids) + def write_root_com_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root center of mass pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + The orientation is the orientation of the principle axes of inertia. + + Args: + root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7). + env_ids: Environment indices. If None, then all indices are used. + """ + + # resolve all indices + if env_ids is None: + local_env_ids = slice(None) + + com_pos = self.data.com_pos_b[local_env_ids, 0, :] + com_quat = self.data.com_quat_b[local_env_ids, 0, :] + + root_link_pos, root_link_quat = math_utils.combine_frame_transforms( + root_pose[..., :3], + root_pose[..., 3:7], + math_utils.quat_rotate(math_utils.quat_inv(com_quat), -com_pos), + math_utils.quat_inv(com_quat), + ) + + root_link_pose = torch.cat((root_link_pos, root_link_quat), dim=-1) + self.write_root_link_pose_to_sim(root_pose=root_link_pose, env_ids=env_ids) + def write_root_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Sequence[int] | None = None): - """Set the root velocity over selected environment indices into the simulation. + """Set the root center of mass velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + NOTE: This sets the velocity of the root's center of mass rather than the roots frame. Args: - root_velocity: Root velocities in simulation frame. Shape is (len(env_ids), 6). + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6). env_ids: Environment indices. If None, then all indices are used. """ + # deprecation warning + if not self._root_vel_dep_warn: + omni.log.warn( + "DeprecationWarning: RigidObject.write_root_velocity_to_sim will be removed in a future release. Please" + " use write_root_link_velocity_to_sim or write_root_com_velocity_to_sim instead." + ) + self._root_vel_dep_warn = True + + self.write_root_com_velocity_to_sim(root_velocity=root_velocity, env_ids=env_ids) + + def write_root_com_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root center of mass velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + NOTE: This sets the velocity of the root's center of mass rather than the roots frame. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6). + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices physx_env_ids = env_ids if env_ids is None: @@ -197,10 +310,37 @@ def write_root_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Seque physx_env_ids = self._ALL_INDICES # note: we need to do this here since tensors are not set into simulation until step. # set into internal buffers - self._data.root_state_w[env_ids, 7:] = root_velocity.clone() + self._data.root_com_state_w[env_ids, 7:] = root_velocity.clone() + self._data._ignore_dep_warn = True + self._data.root_state_w[env_ids, 7:] = self._data.root_com_state_w[env_ids, 7:] + self._data._ignore_dep_warn = False self._data.body_acc_w[env_ids] = 0.0 # set into simulation - self.root_physx_view.set_velocities(self._data.root_state_w[:, 7:], indices=physx_env_ids) + self.root_physx_view.set_velocities(self._data.root_com_state_w[:, 7:], indices=physx_env_ids) + + def write_root_link_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root link velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + NOTE: This sets the velocity of the root's frame rather than the roots center of mass. + + Args: + root_velocity: Root frame velocities in simulation world frame. Shape is (len(env_ids), 6). + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + if env_ids is None: + local_env_ids = slice(None) + + root_com_velocity = root_velocity.clone() + quat = self.data.root_link_state_w[local_env_ids, 3:7] + com_pos_b = self.data.com_pos_b[local_env_ids, 0, :] + # transform given velocity to center of mass + root_com_velocity[:, :3] += torch.linalg.cross( + root_com_velocity[:, 3:], math_utils.quat_rotate(quat, com_pos_b), dim=-1 + ) + # write center of mass velocity to sim + self.write_root_com_velocity_to_sim(root_velocity=root_com_velocity, env_ids=env_ids) """ Operations - Setters. @@ -257,6 +397,9 @@ def set_external_force_and_torque( self._external_torque_b[env_ids, body_ids] = torques else: self.has_external_wrench = False + # reset external wrench + self._external_force_b[env_ids] = 0.0 + self._external_torque_b[env_ids] = 0.0 """ Internal helper. @@ -288,6 +431,18 @@ def _initialize_impl(self): " Please ensure that there is only one rigid body in the prim path tree." ) + articulation_prims = sim_utils.get_all_matching_child_prims( + template_prim_path, predicate=lambda prim: prim.HasAPI(UsdPhysics.ArticulationRootAPI) + ) + if len(articulation_prims) != 0: + if articulation_prims[0].GetAttribute("physxArticulation:articulationEnabled").Get(): + raise RuntimeError( + f"Found an articulation root when resolving '{self.cfg.prim_path}' for rigid objects. These are" + f" located at: '{articulation_prims}' under '{template_prim_path}'. Please disable the articulation" + " root in the USD or from code by setting the parameter" + " 'ArticulationRootPropertiesCfg.articulation_enabled' to False in the spawn configuration." + ) + # resolve root prim back into regex expression root_prim_path = root_prims[0].GetPath().pathString root_prim_path_expr = self.cfg.prim_path + root_prim_path[len(template_prim_path) :] diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object/rigid_object_data.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object/rigid_object_data.py index 63d7be943d..dfb343064b 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object/rigid_object_data.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object/rigid_object_data.py @@ -6,6 +6,7 @@ import torch import weakref +import omni.log import omni.physics.tensors.impl.api as physx import omni.isaac.lab.utils.math as math_utils @@ -64,8 +65,14 @@ def __init__(self, root_physx_view: physx.RigidBodyView, device: str): # Initialize the lazy buffers. self._root_state_w = TimestampedBuffer() + self._root_link_state_w = TimestampedBuffer() + self._root_com_state_w = TimestampedBuffer() self._body_acc_w = TimestampedBuffer() + # deprecation warning check + self._root_state_dep_warn = False + self._ignore_dep_warn = False + def update(self, dt: float): """Updates the data for the rigid object. @@ -114,6 +121,13 @@ def root_state_w(self): The position and orientation are of the rigid body's actor frame. Meanwhile, the linear and angular velocities are of the rigid body's center of mass frame. """ + if not self._root_state_dep_warn and not self._ignore_dep_warn: + omni.log.warn( + "DeprecationWarning: root_state_w and it's derived properties will be deprecated in a future release." + " Please use root_link_state_w or root_com_state_w." + ) + self._root_state_dep_warn = True + if self._root_state_w.timestamp < self._sim_timestamp: # read data from simulation pose = self._root_physx_view.get_transforms().clone() @@ -124,6 +138,52 @@ def root_state_w(self): self._root_state_w.timestamp = self._sim_timestamp return self._root_state_w.data + @property + def root_link_state_w(self): + """Root state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 13). + + The position, quaternion, and linear/angular velocity are of the rigid body root frame relative to the + world. + """ + if self._root_link_state_w.timestamp < self._sim_timestamp: + # read data from simulation + pose = self._root_physx_view.get_transforms().clone() + pose[:, 3:7] = math_utils.convert_quat(pose[:, 3:7], to="wxyz") + velocity = self._root_physx_view.get_velocities().clone() + + # adjust linear velocity to link from center of mass + velocity[:, :3] += torch.linalg.cross( + velocity[:, 3:], math_utils.quat_rotate(pose[:, 3:7], -self.com_pos_b[:, 0, :]), dim=-1 + ) + # set the buffer data and timestamp + self._root_link_state_w.data = torch.cat((pose, velocity), dim=-1) + self._root_link_state_w.timestamp = self._sim_timestamp + + return self._root_link_state_w.data + + @property + def root_com_state_w(self): + """Root center of mass state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 13). + + The position, quaternion, and linear/angular velocity are of the rigid body's center of mass frame + relative to the world. Center of mass frame is the orientation principle axes of inertia. + """ + if self._root_com_state_w.timestamp < self._sim_timestamp: + # read data from simulation (pose is of link) + pose = self._root_physx_view.get_transforms().clone() + pose[:, 3:7] = math_utils.convert_quat(pose[:, 3:7], to="wxyz") + velocity = self._root_physx_view.get_velocities() + + # adjust pose to center of mass + pos, quat = math_utils.combine_frame_transforms( + pose[:, :3], pose[:, 3:7], self.com_pos_b[:, 0, :], self.com_quat_b[:, 0, :] + ) + pose = torch.cat((pos, quat), dim=-1) + # set the buffer data and timestamp + self._root_com_state_w.data = torch.cat((pos, quat, velocity), dim=-1) + self._root_com_state_w.timestamp = self._sim_timestamp + return self._root_com_state_w.data + @property def body_state_w(self): """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame. Shape is (num_instances, 1, 13). @@ -131,8 +191,34 @@ def body_state_w(self): The position and orientation are of the rigid bodies' actor frame. Meanwhile, the linear and angular velocities are of the rigid bodies' center of mass frame. """ + + omni.log.warn( + "DeprecationWarning: body_state_w and it's derived properties will be deprecated in a future release." + " Please use body_link_state_w or body_com_state_w." + ) + return self.root_state_w.view(-1, 1, 13) + @property + def body_link_state_w(self): + """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame. + Shape is (num_instances,1, 13). + + The position, quaternion, and linear/angular velocity are of the body's link frame relative to the world. + """ + return self.root_link_state_w.view(-1, 1, 13) + + @property + def body_com_state_w(self): + """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame. + Shape is (num_instances, num_bodies, 13). + + The position, quaternion, and linear/angular velocity are of the body's center of mass frame relative to the + world. Center of mass frame is assumed to be the same orientation as the link rather than the orientation of the + principle inertia. + """ + return self.root_com_state_w.view(-1, 1, 13) + @property def body_acc_w(self): """Acceleration of all bodies. Shape is (num_instances, 1, 6). @@ -148,7 +234,7 @@ def body_acc_w(self): @property def projected_gravity_b(self): """Projection of the gravity direction on base frame. Shape is (num_instances, 3).""" - return math_utils.quat_rotate_inverse(self.root_quat_w, self.GRAVITY_VEC_W) + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.GRAVITY_VEC_W) @property def heading_w(self): @@ -158,7 +244,7 @@ def heading_w(self): This quantity is computed by assuming that the forward-direction of the base frame is along x-direction, i.e. :math:`(1, 0, 0)`. """ - forward_w = math_utils.quat_apply(self.root_quat_w, self.FORWARD_VEC_B) + forward_w = math_utils.quat_apply(self.root_link_quat_w, self.FORWARD_VEC_B) return torch.atan2(forward_w[:, 1], forward_w[:, 0]) ## @@ -193,7 +279,7 @@ def root_vel_w(self) -> torch.Tensor: def root_lin_vel_w(self) -> torch.Tensor: """Root linear velocity in simulation world frame. Shape is (num_instances, 3). - This quantity is the linear velocity of the root rigid body's center of mass frame. + This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. """ return self.root_state_w[:, 7:10] @@ -212,7 +298,7 @@ def root_lin_vel_b(self) -> torch.Tensor: This quantity is the linear velocity of the root rigid body's center of mass frame with respect to the rigid body's actor frame. """ - return math_utils.quat_rotate_inverse(self.root_quat_w, self.root_lin_vel_w) + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.root_lin_vel_w) @property def root_ang_vel_b(self) -> torch.Tensor: @@ -221,7 +307,145 @@ def root_ang_vel_b(self) -> torch.Tensor: This quantity is the angular velocity of the root rigid body's center of mass frame with respect to the rigid body's actor frame. """ - return math_utils.quat_rotate_inverse(self.root_quat_w, self.root_ang_vel_w) + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.root_ang_vel_w) + + @property + def root_link_pos_w(self) -> torch.Tensor: + """Root link position in simulation world frame. Shape is (num_instances, 3). + + This quantity is the position of the actor frame of the root rigid body relative to the world. + """ + if self._root_link_state_w.timestamp < self._sim_timestamp: + # read data from simulation + pose = self._root_physx_view.get_transforms() + return pose[:, :3] + return self.root_link_state_w[:, :3] + + @property + def root_link_quat_w(self) -> torch.Tensor: + """Root link orientation (w, x, y, z) in simulation world frame. Shape is (num_instances, 4). + + This quantity is the orientation of the actor frame of the root rigid body. + """ + if self._root_link_state_w.timestamp < self._sim_timestamp: + # read data from simulation + pose = self._root_physx_view.get_transforms().clone() + pose[:, 3:7] = math_utils.convert_quat(pose[:, 3:7], to="wxyz") + return pose[:, 3:7] + return self.root_link_state_w[:, 3:7] + + @property + def root_link_vel_w(self) -> torch.Tensor: + """Root link velocity in simulation world frame. Shape is (num_instances, 6). + + This quantity contains the linear and angular velocities of the actor frame of the root + rigid body relative to the world. + """ + return self.root_link_state_w[:, 7:13] + + @property + def root_link_lin_vel_w(self) -> torch.Tensor: + """Root linear velocity in simulation world frame. Shape is (num_instances, 3). + + This quantity is the linear velocity of the root rigid body's actor frame relative to the world. + """ + return self.root_link_state_w[:, 7:10] + + @property + def root_link_ang_vel_w(self) -> torch.Tensor: + """Root link angular velocity in simulation world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. + """ + return self.root_link_state_w[:, 10:13] + + @property + def root_link_lin_vel_b(self) -> torch.Tensor: + """Root link linear velocity in base frame. Shape is (num_instances, 3). + + This quantity is the linear velocity of the actor frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.root_link_lin_vel_w) + + @property + def root_link_ang_vel_b(self) -> torch.Tensor: + """Root link angular velocity in base world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the actor frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.root_link_ang_vel_w) + + @property + def root_com_pos_w(self) -> torch.Tensor: + """Root center of mass position in simulation world frame. Shape is (num_instances, 3). + + This quantity is the position of the actor frame of the root rigid body relative to the world. + """ + return self.root_com_state_w[:, :3] + + @property + def root_com_quat_w(self) -> torch.Tensor: + """Root center of mass orientation (w, x, y, z) in simulation world frame. Shape is (num_instances, 4). + + This quantity is the orientation of the actor frame of the root rigid body relative to the world. + """ + return self.root_com_state_w[:, 3:7] + + @property + def root_com_vel_w(self) -> torch.Tensor: + """Root center of mass velocity in simulation world frame. Shape is (num_instances, 6). + + This quantity contains the linear and angular velocities of the root rigid body's center of mass frame relative to the world. + """ + if self._root_link_state_w.timestamp < self._sim_timestamp: + # read data from simulation + velocity = self._root_physx_view.get_velocities() + return velocity + return self.root_com_state_w[:, 7:13] + + @property + def root_com_lin_vel_w(self) -> torch.Tensor: + """Root center of mass linear velocity in simulation world frame. Shape is (num_instances, 3). + + This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. + """ + if self._root_link_state_w.timestamp < self._sim_timestamp: + # read data from simulation + velocity = self._root_physx_view.get_velocities() + return velocity[:, 0:3] + return self.root_com_state_w[:, 7:10] + + @property + def root_com_ang_vel_w(self) -> torch.Tensor: + """Root center of mass angular velocity in simulation world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. + """ + if self._root_link_state_w.timestamp < self._sim_timestamp: + # read data from simulation + velocity = self._root_physx_view.get_velocities() + return velocity[:, 3:6] + return self.root_com_state_w[:, 10:13] + + @property + def root_com_lin_vel_b(self) -> torch.Tensor: + """Root center of mass linear velocity in base frame. Shape is (num_instances, 3). + + This quantity is the linear velocity of the root rigid body's center of mass frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.root_com_lin_vel_w) + + @property + def root_com_ang_vel_b(self) -> torch.Tensor: + """Root center of mass angular velocity in base world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the root rigid body's center of mass frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.root_com_ang_vel_w) @property def body_pos_w(self) -> torch.Tensor: @@ -278,3 +502,109 @@ def body_ang_acc_w(self) -> torch.Tensor: This quantity is the angular acceleration of the rigid bodies' center of mass frame. """ return self.body_acc_w[..., 3:6] + + # + # Link body properties + # + + @property + def body_link_pos_w(self) -> torch.Tensor: + """Positions of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the position of the rigid bodies' actor frame relative to the world. + """ + return self.body_link_state_w[..., :3] + + @property + def body_link_quat_w(self) -> torch.Tensor: + """Orientation (w, x, y, z) of all bodies in simulation world frame. Shape is (num_instances, 1, 4). + + This quantity is the orientation of the rigid bodies' actor frame relative to the world. + """ + return self.body_link_state_w[..., 3:7] + + @property + def body_link_vel_w(self) -> torch.Tensor: + """Velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 6). + + This quantity contains the linear and angular velocities of the rigid bodies' center of mass frame + relative to the world. + """ + return self.body_link_state_w[..., 7:13] + + @property + def body_link_lin_vel_w(self) -> torch.Tensor: + """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the linear velocity of the rigid bodies' center of mass frame relative to the world. + """ + return self.body_link_state_w[..., 7:10] + + @property + def body_link_ang_vel_w(self) -> torch.Tensor: + """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the angular velocity of the rigid bodies' center of mass frame relative to the world. + """ + return self.body_link_state_w[..., 10:13] + + # + # Center of mass body properties + # + + @property + def body_com_pos_w(self) -> torch.Tensor: + """Positions of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the position of the rigid bodies' actor frame. + """ + return self.body_com_state_w[..., :3] + + @property + def body_com_quat_w(self) -> torch.Tensor: + """Orientation (w, x, y, z) of the prinicple axies of inertia of all bodies in simulation world frame. + + Shape is (num_instances, 1, 4). This quantity is the orientation of the rigid bodies' actor frame. + """ + return self.body_com_state_w[..., 3:7] + + @property + def body_com_vel_w(self) -> torch.Tensor: + """Velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 6). + + This quantity contains the linear and angular velocities of the rigid bodies' center of mass frame. + """ + return self.body_com_state_w[..., 7:13] + + @property + def body_com_lin_vel_w(self) -> torch.Tensor: + """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the linear velocity of the rigid bodies' center of mass frame. + """ + return self.body_com_state_w[..., 7:10] + + @property + def body_com_ang_vel_w(self) -> torch.Tensor: + """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the angular velocity of the rigid bodies' center of mass frame. + """ + return self.body_com_state_w[..., 10:13] + + @property + def com_pos_b(self) -> torch.Tensor: + """Center of mass of all of the bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the center of mass location relative to its body frame. + """ + return self._root_physx_view.get_coms().to(self.device)[..., :3].view(-1, 1, 3) + + @property + def com_quat_b(self) -> torch.Tensor: + """Orientation (w,x,y,z) of the prinicple axies of inertia of all of the bodies in simulation world frame. Shape is (num_instances, 1, 4). + + This quantity is the orientation of the principles axes of inertia relative to its body frame. + """ + quat = self._root_physx_view.get_coms().to(self.device)[..., 3:7] + return math_utils.convert_quat(quat, to="wxyz").view(-1, 1, 4) diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/__init__.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/__init__.py new file mode 100644 index 0000000000..c0f4b5a46c --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-module for rigid object collection.""" + +from .rigid_object_collection import RigidObjectCollection +from .rigid_object_collection_cfg import RigidObjectCollectionCfg +from .rigid_object_collection_data import RigidObjectCollectionData diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/rigid_object_collection.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/rigid_object_collection.py new file mode 100644 index 0000000000..acea81ff06 --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/rigid_object_collection.py @@ -0,0 +1,686 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import re +import torch +import weakref +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import omni.kit.app +import omni.log +import omni.physics.tensors.impl.api as physx +import omni.timeline +from pxr import UsdPhysics + +import omni.isaac.lab.sim as sim_utils +import omni.isaac.lab.utils.math as math_utils +import omni.isaac.lab.utils.string as string_utils + +from ..asset_base import AssetBase +from .rigid_object_collection_data import RigidObjectCollectionData + +if TYPE_CHECKING: + from .rigid_object_collection_cfg import RigidObjectCollectionCfg + + +class RigidObjectCollection(AssetBase): + """A rigid object collection class. + + This class represents a collection of rigid objects in the simulation, where the state of the rigid objects can be + accessed and modified using a batched ``(env_ids, object_ids)`` API. + + For each rigid body in the collection, the root prim of the asset must have the `USD RigidBodyAPI`_ + applied to it. This API is used to define the simulation properties of the rigid bodies. On playing the + simulation, the physics engine will automatically register the rigid bodies and create a corresponding + rigid body handle. This handle can be accessed using the :attr:`root_physx_view` attribute. + + .. note:: + Rigid objects in the collection are uniquely identified via the key of the dictionary + :attr:`~omni.isaac.lab.assets.RigidObjectCollectionCfg.rigid_objects` in :class:`~omni.isaac.lab.assets.RigidObjectCollectionCfg`. + This differs from the class :class:`~omni.isaac.lab.assets.RigidObject`, where a rigid object is identified by + the name of the Xform where the `USD RigidBodyAPI`_ is applied. This would not be possible for the rigid object + collection since the :attr:`~omni.isaac.lab.assets.RigidObjectCollectionCfg.rigid_objects` dictionary could + contain the same rigid object multiple times, leading to ambiguity. + + .. _`USD RigidBodyAPI`: https://openusd.org/dev/api/class_usd_physics_rigid_body_a_p_i.html + """ + + cfg: RigidObjectCollectionCfg + """Configuration instance for the rigid object collection.""" + + def __init__(self, cfg: RigidObjectCollectionCfg): + """Initialize the rigid object collection. + + Args: + cfg: A configuration instance. + """ + # check that the config is valid + cfg.validate() + # store inputs + self.cfg = cfg + # flag for whether the asset is initialized + self._is_initialized = False + for rigid_object_cfg in self.cfg.rigid_objects.values(): + # check if the rigid object path is valid + # note: currently the spawner does not work if there is a regex pattern in the leaf + # For example, if the prim path is "/World/Object_[1,2]" since the spawner will not + # know which prim to spawn. This is a limitation of the spawner and not the asset. + asset_path = rigid_object_cfg.prim_path.split("/")[-1] + asset_path_is_regex = re.match(r"^[a-zA-Z0-9/_]+$", asset_path) is None + # spawn the asset + if rigid_object_cfg.spawn is not None and not asset_path_is_regex: + rigid_object_cfg.spawn.func( + rigid_object_cfg.prim_path, + rigid_object_cfg.spawn, + translation=rigid_object_cfg.init_state.pos, + orientation=rigid_object_cfg.init_state.rot, + ) + # check that spawn was successful + matching_prims = sim_utils.find_matching_prims(rigid_object_cfg.prim_path) + if len(matching_prims) == 0: + raise RuntimeError(f"Could not find prim with path {rigid_object_cfg.prim_path}.") + + # stores object names + self._object_names_list = [] + + # note: Use weakref on all callbacks to ensure that this object can be deleted when its destructor is called. + # add callbacks for stage play/stop + # The order is set to 10 which is arbitrary but should be lower priority than the default order of 0 + timeline_event_stream = omni.timeline.get_timeline_interface().get_timeline_event_stream() + self._initialize_handle = timeline_event_stream.create_subscription_to_pop_by_type( + int(omni.timeline.TimelineEventType.PLAY), + lambda event, obj=weakref.proxy(self): obj._initialize_callback(event), + order=10, + ) + self._invalidate_initialize_handle = timeline_event_stream.create_subscription_to_pop_by_type( + int(omni.timeline.TimelineEventType.STOP), + lambda event, obj=weakref.proxy(self): obj._invalidate_initialize_callback(event), + order=10, + ) + + self._debug_vis_handle = None + + self._root_state_dep_warn = False + self._root_pose_dep_warn = False + self._root_vel_dep_warn = False + + """ + Properties + """ + + @property + def data(self) -> RigidObjectCollectionData: + return self._data + + @property + def num_instances(self) -> int: + """Number of instances of the collection.""" + return self.root_physx_view.count // self.num_objects + + @property + def num_objects(self) -> int: + """Number of objects in the collection. + + This corresponds to the distinct number of rigid bodies in the collection. + """ + return len(self.object_names) + + @property + def object_names(self) -> list[str]: + """Ordered names of objects in the rigid object collection.""" + return self._object_names_list + + @property + def root_physx_view(self) -> physx.RigidBodyView: + """Rigid body view for the rigid body collection (PhysX). + + Note: + Use this view with caution. It requires handling of tensors in a specific way. + """ + return self._root_physx_view + + """ + Operations. + """ + + def reset(self, env_ids: torch.Tensor | None = None, object_ids: slice | torch.Tensor | None = None): + """Resets all internal buffers of selected environments and objects. + + Args: + env_ids: The indices of the object to reset. Defaults to None (all instances). + object_ids: The indices of the object to reset. Defaults to None (all objects). + """ + # resolve all indices + if env_ids is None: + env_ids = self._ALL_ENV_INDICES + if object_ids is None: + object_ids = self._ALL_OBJ_INDICES + # reset external wrench + self._external_force_b[env_ids[:, None], object_ids] = 0.0 + self._external_torque_b[env_ids[:, None], object_ids] = 0.0 + + def write_data_to_sim(self): + """Write external wrench to the simulation. + + Note: + We write external wrench to the simulation here since this function is called before the simulation step. + This ensures that the external wrench is applied at every simulation step. + """ + # write external wrench + if self.has_external_wrench: + self.root_physx_view.apply_forces_and_torques_at_position( + force_data=self.reshape_data_to_view(self._external_force_b), + torque_data=self.reshape_data_to_view(self._external_torque_b), + position_data=None, + indices=self._env_obj_ids_to_view_ids(self._ALL_ENV_INDICES, self._ALL_OBJ_INDICES), + is_global=False, + ) + + def update(self, dt: float): + self._data.update(dt) + + """ + Operations - Finders. + """ + + def find_objects( + self, name_keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[torch.Tensor, list[str]]: + """Find objects in the collection based on the name keys. + + Please check the :meth:`omni.isaac.lab.utils.string_utils.resolve_matching_names` function for more + information on the name matching. + + Args: + name_keys: A regular expression or a list of regular expressions to match the object names. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple containing the object indices and names. + """ + obj_ids, obj_names = string_utils.resolve_matching_names(name_keys, self.object_names, preserve_order) + return torch.tensor(obj_ids, device=self.device), obj_names + + """ + Operations - Write to simulation. + """ + + def write_object_state_to_sim( + self, + object_state: torch.Tensor, + env_ids: torch.Tensor | None = None, + object_ids: slice | torch.Tensor | None = None, + ): + """Set the object state over selected environment and object indices into the simulation. + + The object state comprises of the cartesian position, quaternion orientation in (w, x, y, z), and linear + and angular velocity. All the quantities are in the simulation frame. + + Args: + object_state: Object state in simulation frame. Shape is (len(env_ids), len(object_ids), 13). + env_ids: Environment indices. If None, then all indices are used. + object_ids: Object indices. If None, then all indices are used. + """ + # deprecation warning + if not self._root_state_dep_warn: + omni.log.warn( + "DeprecationWarning: RigidObjectCollection.write_object_state_to_sim will be removed in a future" + " release. Please use write_object_link_state_to_sim or write_object_com_state_to_sim instead." + ) + self._root_state_dep_warn = True + + # set into simulation + self.write_object_pose_to_sim(object_state[..., :7], env_ids=env_ids, object_ids=object_ids) + self.write_object_velocity_to_sim(object_state[..., 7:], env_ids=env_ids, object_ids=object_ids) + + def write_object_com_state_to_sim( + self, + object_state: torch.Tensor, + env_ids: torch.Tensor | None = None, + object_ids: slice | torch.Tensor | None = None, + ): + """Set the object center of mass state over selected environment indices into the simulation. + The object state comprises of the cartesian position, quaternion orientation in (w, x, y, z), and linear + and angular velocity. All the quantities are in the simulation frame. + + Args: + object_state: Object state in simulation frame. Shape is (len(env_ids), len(object_ids), 13). + env_ids: Environment indices. If None, then all indices are used. + object_ids: Object indices. If None, then all indices are used. + """ + # set into simulation + self.write_object_com_pose_to_sim(object_state[..., :7], env_ids=env_ids, object_ids=object_ids) + self.write_object_com_velocity_to_sim(object_state[..., 7:], env_ids=env_ids, object_ids=object_ids) + + def write_object_link_state_to_sim( + self, + object_state: torch.Tensor, + env_ids: torch.Tensor | None = None, + object_ids: slice | torch.Tensor | None = None, + ): + """Set the object link state over selected environment indices into the simulation. + The object state comprises of the cartesian position, quaternion orientation in (w, x, y, z), and linear + and angular velocity. All the quantities are in the simulation frame. + + Args: + object_state: Object state in simulation frame. Shape is (len(env_ids), len(object_ids), 13). + env_ids: Environment indices. If None, then all indices are used. + object_ids: Object indices. If None, then all indices are used. + """ + # set into simulation + self.write_object_link_pose_to_sim(object_state[..., :7], env_ids=env_ids, object_ids=object_ids) + self.write_object_link_velocity_to_sim(object_state[..., 7:], env_ids=env_ids, object_ids=object_ids) + + def write_object_pose_to_sim( + self, + object_pose: torch.Tensor, + env_ids: torch.Tensor | None = None, + object_ids: slice | torch.Tensor | None = None, + ): + """Set the object pose over selected environment and object indices into the simulation. + + The object pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + + Args: + object_pose: Object poses in simulation frame. Shape is (len(env_ids), len(object_ids), 7). + env_ids: Environment indices. If None, then all indices are used. + object_ids: Object indices. If None, then all indices are used. + """ + # deprecation warning + if not self._root_pose_dep_warn: + omni.log.warn( + "DeprecationWarning: RigidObjectCollection.write_object_pose_to_sim will be removed in a future" + " release. Please use write_object_link_pose_to_sim or write_object_com_pose_to_sim instead." + ) + self._root_pose_dep_warn = True + + self.write_object_link_pose_to_sim(object_pose, env_ids, object_ids) + + def write_object_link_pose_to_sim( + self, + object_pose: torch.Tensor, + env_ids: torch.Tensor | None = None, + object_ids: slice | torch.Tensor | None = None, + ): + """Set the object pose over selected environment and object indices into the simulation. + + The object pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + + Args: + object_pose: Object poses in simulation frame. Shape is (len(env_ids), len(object_ids), 7). + env_ids: Environment indices. If None, then all indices are used. + object_ids: Object indices. If None, then all indices are used. + """ + + # resolve all indices + # -- env_ids + if env_ids is None: + env_ids = self._ALL_ENV_INDICES + # -- object_ids + if object_ids is None: + object_ids = self._ALL_OBJ_INDICES + # note: we need to do this here since tensors are not set into simulation until step. + # set into internal buffers + self._data.object_link_state_w[env_ids[:, None], object_ids, :7] = object_pose.clone() + self._data._ignore_dep_warn = True + self._data.object_state_w[env_ids[:, None], object_ids, :7] = object_pose.clone() + self._data._ignore_dep_warn = False + # convert the quaternion from wxyz to xyzw + poses_xyzw = self._data.object_link_state_w[..., :7].clone() + poses_xyzw[..., 3:] = math_utils.convert_quat(poses_xyzw[..., 3:], to="xyzw") + # set into simulation + view_ids = self._env_obj_ids_to_view_ids(env_ids, object_ids) + self.root_physx_view.set_transforms(self.reshape_data_to_view(poses_xyzw), indices=view_ids) + + def write_object_com_pose_to_sim( + self, + object_pose: torch.Tensor, + env_ids: torch.Tensor | None = None, + object_ids: slice | torch.Tensor | None = None, + ): + """Set the object center of mass pose over selected environment indices into the simulation. + The object pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + The orientation is the orientation of the principle axes of inertia. + Args: + object_pose: Object poses in simulation frame. Shape is (len(env_ids), len(object_ids), 7). + env_ids: Environment indices. If None, then all indices are used. + object_ids: Object indices. If None, then all indices are used. + """ + + # resolve all indices + if env_ids is None: + local_env_ids = slice(env_ids) + else: + local_env_ids = env_ids + if object_ids is None: + local_object_ids = slice(object_ids) + else: + local_object_ids = object_ids + + com_pos = self.data.com_pos_b[local_env_ids, local_object_ids, :] + com_quat = self.data.com_quat_b[local_env_ids, local_object_ids, :] + + object_link_pos, object_link_quat = math_utils.combine_frame_transforms( + object_pose[..., :3], + object_pose[..., 3:7], + math_utils.quat_rotate(math_utils.quat_inv(com_quat), -com_pos), + math_utils.quat_inv(com_quat), + ) + + object_link_pose = torch.cat((object_link_pos, object_link_quat), dim=-1) + self.write_object_link_pose_to_sim(object_pose=object_link_pose, env_ids=env_ids, object_ids=object_ids) + + def write_object_velocity_to_sim( + self, + object_velocity: torch.Tensor, + env_ids: torch.Tensor | None = None, + object_ids: slice | torch.Tensor | None = None, + ): + """Set the object velocity over selected environment and object indices into the simulation. + + Args: + object_velocity: Object velocities in simulation frame. Shape is (len(env_ids), len(object_ids), 6). + env_ids: Environment indices. If None, then all indices are used. + object_ids: Object indices. If None, then all indices are used. + """ + # deprecation warning + if not self._root_vel_dep_warn: + omni.log.warn( + "DeprecationWarning: RigidObjectCollection.write_object_velocity_to_sim will be removed in a future" + " release. Please use write_object_link_velocity_to_sim or write_object_com_velocity_to_sim instead." + ) + self._root_vel_dep_warn = True + + self.write_object_com_velocity_to_sim(object_velocity=object_velocity, env_ids=env_ids, object_ids=object_ids) + + def write_object_com_velocity_to_sim( + self, + object_velocity: torch.Tensor, + env_ids: torch.Tensor | None = None, + object_ids: slice | torch.Tensor | None = None, + ): + """Set the object center of mass velocity over selected environment and object indices into the simulation. + + Args: + object_velocity: Object velocities in simulation frame. Shape is (len(env_ids), len(object_ids), 6). + env_ids: Environment indices. If None, then all indices are used. + object_ids: Object indices. If None, then all indices are used. + """ + # resolve all indices + # -- env_ids + if env_ids is None: + env_ids = self._ALL_ENV_INDICES + # -- object_ids + if object_ids is None: + object_ids = self._ALL_OBJ_INDICES + + self._data.object_com_state_w[env_ids[:, None], object_ids, 7:] = object_velocity.clone() + self._data._ignore_dep_warn = True + self._data.object_state_w[env_ids[:, None], object_ids, 7:] = object_velocity.clone() + self._data._ignore_dep_warn = False + self._data.object_acc_w[env_ids[:, None], object_ids] = 0.0 + + # set into simulation + view_ids = self._env_obj_ids_to_view_ids(env_ids, object_ids) + self.root_physx_view.set_velocities( + self.reshape_data_to_view(self._data.object_com_state_w[..., 7:]), indices=view_ids + ) + + def write_object_link_velocity_to_sim( + self, + object_velocity: torch.Tensor, + env_ids: torch.Tensor | None = None, + object_ids: slice | torch.Tensor | None = None, + ): + """Set the object link velocity over selected environment indices into the simulation. + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + NOTE: This sets the velocity of the object's frame rather than the objects center of mass. + Args: + object_velocity: Object velocities in simulation frame. Shape is (len(env_ids), len(object_ids), 6). + env_ids: Environment indices. If None, then all indices are used. + object_ids: Object indices. If None, then all indices are used. + """ + # resolve all indices + if env_ids is None: + local_env_ids = slice(env_ids) + else: + local_env_ids = env_ids + if object_ids is None: + local_object_ids = slice(object_ids) + else: + local_object_ids = object_ids + + object_com_velocity = object_velocity.clone() + quat = self.data.object_link_state_w[local_env_ids, local_object_ids, 3:7] + com_pos_b = self.data.com_pos_b[local_env_ids, local_object_ids, :] + # transform given velocity to center of mass + object_com_velocity[..., :3] += torch.linalg.cross( + object_com_velocity[..., 3:], math_utils.quat_rotate(quat, com_pos_b), dim=-1 + ) + # write center of mass velocity to sim + self.write_object_com_velocity_to_sim( + object_velocity=object_com_velocity, env_ids=env_ids, object_ids=object_ids + ) + + """ + Operations - Setters. + """ + + def set_external_force_and_torque( + self, + forces: torch.Tensor, + torques: torch.Tensor, + object_ids: slice | torch.Tensor | None = None, + env_ids: torch.Tensor | None = None, + ): + """Set external force and torque to apply on the objects' bodies in their local frame. + + For many applications, we want to keep the applied external force on rigid bodies constant over a period of + time (for instance, during the policy control). This function allows us to store the external force and torque + into buffers which are then applied to the simulation at every step. + + .. caution:: + If the function is called with empty forces and torques, then this function disables the application + of external wrench to the simulation. + + .. code-block:: python + + # example of disabling external wrench + asset.set_external_force_and_torque(forces=torch.zeros(0, 0, 3), torques=torch.zeros(0, 0, 3)) + + .. note:: + This function does not apply the external wrench to the simulation. It only fills the buffers with + the desired values. To apply the external wrench, call the :meth:`write_data_to_sim` function + right before the simulation step. + + Args: + forces: External forces in bodies' local frame. Shape is (len(env_ids), len(object_ids), 3). + torques: External torques in bodies' local frame. Shape is (len(env_ids), len(object_ids), 3). + object_ids: Object indices to apply external wrench to. Defaults to None (all objects). + env_ids: Environment indices to apply external wrench to. Defaults to None (all instances). + """ + if forces.any() or torques.any(): + self.has_external_wrench = True + # resolve all indices + # -- env_ids + if env_ids is None: + env_ids = self._ALL_ENV_INDICES + # -- object_ids + if object_ids is None: + object_ids = self._ALL_OBJ_INDICES + # set into internal buffers + self._external_force_b[env_ids[:, None], object_ids] = forces + self._external_torque_b[env_ids[:, None], object_ids] = torques + else: + self.has_external_wrench = False + + """ + Internal helper. + """ + + def _initialize_impl(self): + # create simulation view + self._physics_sim_view = physx.create_simulation_view(self._backend) + self._physics_sim_view.set_subspace_roots("/") + root_prim_path_exprs = [] + for name, rigid_object_cfg in self.cfg.rigid_objects.items(): + # obtain the first prim in the regex expression (all others are assumed to be a copy of this) + template_prim = sim_utils.find_first_matching_prim(rigid_object_cfg.prim_path) + if template_prim is None: + raise RuntimeError(f"Failed to find prim for expression: '{rigid_object_cfg.prim_path}'.") + template_prim_path = template_prim.GetPath().pathString + + # find rigid root prims + root_prims = sim_utils.get_all_matching_child_prims( + template_prim_path, predicate=lambda prim: prim.HasAPI(UsdPhysics.RigidBodyAPI) + ) + if len(root_prims) == 0: + raise RuntimeError( + f"Failed to find a rigid body when resolving '{rigid_object_cfg.prim_path}'." + " Please ensure that the prim has 'USD RigidBodyAPI' applied." + ) + if len(root_prims) > 1: + raise RuntimeError( + f"Failed to find a single rigid body when resolving '{rigid_object_cfg.prim_path}'." + f" Found multiple '{root_prims}' under '{template_prim_path}'." + " Please ensure that there is only one rigid body in the prim path tree." + ) + + # check that no rigid object has an articulation root API, which decreases simulation performance + articulation_prims = sim_utils.get_all_matching_child_prims( + template_prim_path, predicate=lambda prim: prim.HasAPI(UsdPhysics.ArticulationRootAPI) + ) + if len(articulation_prims) != 0: + if articulation_prims[0].GetAttribute("physxArticulation:articulationEnabled").Get(): + raise RuntimeError( + f"Found an articulation root when resolving '{rigid_object_cfg.prim_path}' in the rigid object" + f" collection. These are located at: '{articulation_prims}' under '{template_prim_path}'." + " Please disable the articulation root in the USD or from code by setting the parameter" + " 'ArticulationRootPropertiesCfg.articulation_enabled' to False in the spawn configuration." + ) + + # resolve root prim back into regex expression + root_prim_path = root_prims[0].GetPath().pathString + root_prim_path_expr = rigid_object_cfg.prim_path + root_prim_path[len(template_prim_path) :] + root_prim_path_exprs.append(root_prim_path_expr.replace(".*", "*")) + + self._object_names_list.append(name) + + # -- object view + self._root_physx_view = self._physics_sim_view.create_rigid_body_view(root_prim_path_exprs) + + # check if the rigid body was created + if self._root_physx_view._backend is None: + raise RuntimeError("Failed to create rigid body collection. Please check PhysX logs.") + + # log information about the rigid body + omni.log.info(f"Number of instances: {self.num_instances}") + omni.log.info(f"Number of distinct objects: {self.num_objects}") + omni.log.info(f"Object names: {self.object_names}") + + # container for data access + self._data = RigidObjectCollectionData(self.root_physx_view, self.num_objects, self.device) + + # create buffers + self._create_buffers() + # process configuration + self._process_cfg() + # update the rigid body data + self.update(0.0) + + def _create_buffers(self): + """Create buffers for storing data.""" + # constants + self._ALL_ENV_INDICES = torch.arange(self.num_instances, dtype=torch.long, device=self.device) + self._ALL_OBJ_INDICES = torch.arange(self.num_objects, dtype=torch.long, device=self.device) + + # external forces and torques + self.has_external_wrench = False + self._external_force_b = torch.zeros((self.num_instances, self.num_objects, 3), device=self.device) + self._external_torque_b = torch.zeros_like(self._external_force_b) + + # set information about rigid body into data + self._data.object_names = self.object_names + self._data.default_mass = self.reshape_view_to_data(self.root_physx_view.get_masses().clone()) + self._data.default_inertia = self.reshape_view_to_data(self.root_physx_view.get_inertias().clone()) + + def _process_cfg(self): + """Post processing of configuration parameters.""" + # default state + # -- object state + default_object_states = [] + for rigid_object_cfg in self.cfg.rigid_objects.values(): + default_object_state = ( + tuple(rigid_object_cfg.init_state.pos) + + tuple(rigid_object_cfg.init_state.rot) + + tuple(rigid_object_cfg.init_state.lin_vel) + + tuple(rigid_object_cfg.init_state.ang_vel) + ) + default_object_state = ( + torch.tensor(default_object_state, dtype=torch.float, device=self.device) + .repeat(self.num_instances, 1) + .unsqueeze(1) + ) + default_object_states.append(default_object_state) + # concatenate the default state for each object + default_object_states = torch.cat(default_object_states, dim=1) + self._data.default_object_state = default_object_states + + def reshape_view_to_data(self, data: torch.Tensor) -> torch.Tensor: + """Reshapes and arranges the data coming from the :attr:`root_physx_view` to (num_instances, num_objects, data_size). + + Args: + data: The data coming from the :attr:`root_physx_view`. Shape is (num_instances*num_objects, data_size). + + Returns: + The reshaped data. Shape is (num_instances, num_objects, data_size). + """ + return torch.einsum("ijk -> jik", data.reshape(self.num_objects, self.num_instances, -1)) + + def reshape_data_to_view(self, data: torch.Tensor) -> torch.Tensor: + """Reshapes and arranges the data to the be consistent with data from the :attr:`root_physx_view`. + + Args: + data: The data to be reshaped. Shape is (num_instances, num_objects, data_size). + + Returns: + The reshaped data. Shape is (num_instances*num_objects, data_size). + """ + return torch.einsum("ijk -> jik", data).reshape(self.num_objects * self.num_instances, *data.shape[2:]) + + def _env_obj_ids_to_view_ids( + self, env_ids: torch.Tensor, object_ids: Sequence[int] | slice | torch.Tensor + ) -> torch.Tensor: + """Converts environment and object indices to indices consistent with data from :attr:`root_physx_view`. + + Args: + env_ids: Environment indices. + object_ids: Object indices. + + Returns: + The view indices. + """ + # the order is env_0/object_0, env_0/object_1, env_0/object_..., env_1/object_0, env_1/object_1, ... + # return a flat tensor of indices + if isinstance(object_ids, slice): + object_ids = self._ALL_OBJ_INDICES + elif isinstance(object_ids, Sequence): + object_ids = torch.tensor(object_ids, device=self.device) + return (object_ids.unsqueeze(1) * self.num_instances + env_ids).flatten() + + """ + Internal simulation callbacks. + """ + + def _invalidate_initialize_callback(self, event): + """Invalidates the scene elements.""" + # call parent + super()._invalidate_initialize_callback(event) + # set all existing views to None to invalidate them + self._physics_sim_view = None + self._root_physx_view = None diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/rigid_object_collection_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/rigid_object_collection_cfg.py new file mode 100644 index 0000000000..bf21fd6091 --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/rigid_object_collection_cfg.py @@ -0,0 +1,28 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import MISSING + +from omni.isaac.lab.assets.rigid_object import RigidObjectCfg +from omni.isaac.lab.utils import configclass + +from .rigid_object_collection import RigidObjectCollection + + +@configclass +class RigidObjectCollectionCfg: + """Configuration parameters for a rigid object collection.""" + + class_type: type = RigidObjectCollection + """The associated asset class. + + The class should inherit from :class:`omni.isaac.lab.assets.asset_base.AssetBase`. + """ + + rigid_objects: dict[str, RigidObjectCfg] = MISSING + """Dictionary of rigid object configurations to spawn. + + The keys are the names for the objects, which are used as unique identifiers throughout the code. + """ diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/rigid_object_collection_data.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/rigid_object_collection_data.py new file mode 100644 index 0000000000..accae3befc --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/rigid_object_collection_data.py @@ -0,0 +1,469 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch +import weakref + +import omni.log +import omni.physics.tensors.impl.api as physx + +import omni.isaac.lab.utils.math as math_utils +from omni.isaac.lab.utils.buffers import TimestampedBuffer + + +class RigidObjectCollectionData: + """Data container for a rigid object collection. + + This class contains the data for a rigid object collection in the simulation. The data includes the state of + all the bodies in the collection. The data is stored in the simulation world frame unless otherwise specified. + The data is in the order ``(num_instances, num_objects, data_size)``, where data_size is the size of the data. + + For a rigid body, there are two frames of reference that are used: + + - Actor frame: The frame of reference of the rigid body prim. This typically corresponds to the Xform prim + with the rigid body schema. + - Center of mass frame: The frame of reference of the center of mass of the rigid body. + + Depending on the settings of the simulation, the actor frame and the center of mass frame may be the same. + This needs to be taken into account when interpreting the data. + + The data is lazily updated, meaning that the data is only updated when it is accessed. This is useful + when the data is expensive to compute or retrieve. The data is updated when the timestamp of the buffer + is older than the current simulation timestamp. The timestamp is updated whenever the data is updated. + """ + + def __init__(self, root_physx_view: physx.RigidBodyView, num_objects: int, device: str): + """Initializes the data. + + Args: + root_physx_view: The root rigid body view. + num_objects: The number of objects in the collection. + device: The device used for processing. + """ + # Set the parameters + self.device = device + self.num_objects = num_objects + # Set the root rigid body view + # note: this is stored as a weak reference to avoid circular references between the asset class + # and the data container. This is important to avoid memory leaks. + self._root_physx_view: physx.RigidBodyView = weakref.proxy(root_physx_view) + self.num_instances = self._root_physx_view.count // self.num_objects + + # Set initial time stamp + self._sim_timestamp = 0.0 + + # Obtain global physics sim view + physics_sim_view = physx.create_simulation_view("torch") + physics_sim_view.set_subspace_roots("/") + gravity = physics_sim_view.get_gravity() + # Convert to direction vector + gravity_dir = torch.tensor((gravity[0], gravity[1], gravity[2]), device=self.device) + gravity_dir = math_utils.normalize(gravity_dir.unsqueeze(0)).squeeze(0) + + # Initialize constants + self.GRAVITY_VEC_W = gravity_dir.repeat(self.num_instances, self.num_objects, 1) + self.FORWARD_VEC_B = torch.tensor((1.0, 0.0, 0.0), device=self.device).repeat( + self.num_instances, self.num_objects, 1 + ) + + # Initialize the lazy buffers. + self._object_state_w = TimestampedBuffer() + self._object_link_state_w = TimestampedBuffer() + self._object_com_state_w = TimestampedBuffer() + self._object_acc_w = TimestampedBuffer() + + # deprecation warning check + self._root_state_dep_warn = False + self._ignore_dep_warn = False + + def update(self, dt: float): + """Updates the data for the rigid object collection. + + Args: + dt: The time step for the update. This must be a positive value. + """ + # update the simulation timestamp + self._sim_timestamp += dt + + ## + # Names. + ## + + object_names: list[str] = None + """Object names in the order parsed by the simulation view.""" + + ## + # Defaults. + ## + + default_object_state: torch.Tensor = None + """Default object state ``[pos, quat, lin_vel, ang_vel]`` in local environment frame. Shape is (num_instances, num_objects, 13). + + The position and quaternion are of each object's rigid body's actor frame. Meanwhile, the linear and angular velocities are + of the center of mass frame. + """ + + default_mass: torch.Tensor = None + """Default object mass read from the simulation. Shape is (num_instances, num_objects, 1).""" + + default_inertia: torch.Tensor = None + """Default object inertia tensor read from the simulation. Shape is (num_instances, num_objects, 9). + + The inertia is the inertia tensor relative to the center of mass frame. The values are stored in + the order :math:`[I_{xx}, I_{xy}, I_{xz}, I_{yx}, I_{yy}, I_{yz}, I_{zx}, I_{zy}, I_{zz}]`. + """ + + ## + # Properties. + ## + + @property + def object_state_w(self): + """Object state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_objects, 13). + + The position and orientation are of the rigid body's actor frame. Meanwhile, the linear and angular + velocities are of the rigid body's center of mass frame. + """ + + if not self._root_state_dep_warn and not self._ignore_dep_warn: + omni.log.warn( + "DeprecationWarning: object_state_w and it's derived properties will be deprecated in a future release." + " Please use object_link_state_w or object_com_state_w." + ) + self._root_state_dep_warn = True + + if self._object_state_w.timestamp < self._sim_timestamp: + # read data from simulation + pose = self._reshape_view_to_data(self._root_physx_view.get_transforms().clone()) + pose[..., 3:7] = math_utils.convert_quat(pose[..., 3:7], to="wxyz") + velocity = self._reshape_view_to_data(self._root_physx_view.get_velocities()) + # set the buffer data and timestamp + self._object_state_w.data = torch.cat((pose, velocity), dim=-1) + self._object_state_w.timestamp = self._sim_timestamp + return self._object_state_w.data + + @property + def object_link_state_w(self): + """Object center of mass state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_objects, 13). + + The position, quaternion, and linear/angular velocity are of the rigid body root frame relative to the + world. + """ + if self._object_link_state_w.timestamp < self._sim_timestamp: + # read data from simulation + pose = self._reshape_view_to_data(self._root_physx_view.get_transforms().clone()) + pose[..., 3:7] = math_utils.convert_quat(pose[..., 3:7], to="wxyz") + velocity = self._reshape_view_to_data(self._root_physx_view.get_velocities()) + + # adjust linear velocity to link from center of mass + velocity[..., :3] += torch.linalg.cross( + velocity[..., 3:], math_utils.quat_rotate(pose[..., 3:7], -self.com_pos_b[..., :]), dim=-1 + ) + + # set the buffer data and timestamp + self._object_link_state_w.data = torch.cat((pose, velocity), dim=-1) + self._object_link_state_w.timestamp = self._sim_timestamp + return self._object_link_state_w.data + + @property + def object_com_state_w(self): + """Object state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_objects, 13). + + The position, quaternion, and linear/angular velocity are of the rigid body's center of mass frame + relative to the world. Center of mass frame is the orientation principle axes of inertia. + """ + + if self._object_com_state_w.timestamp < self._sim_timestamp: + # read data from simulation + pose = self._reshape_view_to_data(self._root_physx_view.get_transforms().clone()) + pose[..., 3:7] = math_utils.convert_quat(pose[..., 3:7], to="wxyz") + velocity = self._reshape_view_to_data(self._root_physx_view.get_velocities()) + + # adjust pose to center of mass + pos, quat = math_utils.combine_frame_transforms( + pose[..., :3], pose[..., 3:7], self.com_pos_b[..., :], self.com_quat_b[..., :] + ) + + # set the buffer data and timestamp + self._object_com_state_w.data = torch.cat((pos, quat, velocity), dim=-1) + self._object_com_state_w.timestamp = self._sim_timestamp + return self._object_com_state_w.data + + @property + def object_acc_w(self): + """Acceleration of all objects. Shape is (num_instances, num_objects, 6). + + This quantity is the acceleration of the rigid bodies' center of mass frame. + """ + if self._object_acc_w.timestamp < self._sim_timestamp: + # note: we use finite differencing to compute acceleration + self._object_acc_w.data = self._reshape_view_to_data(self._root_physx_view.get_accelerations().clone()) + self._object_acc_w.timestamp = self._sim_timestamp + return self._object_acc_w.data + + @property + def projected_gravity_b(self): + """Projection of the gravity direction on base frame. Shape is (num_instances, num_objects, 3).""" + return math_utils.quat_rotate_inverse(self.object_link_quat_w, self.GRAVITY_VEC_W) + + @property + def heading_w(self): + """Yaw heading of the base frame (in radians). Shape is (num_instances, num_objects,). + + Note: + This quantity is computed by assuming that the forward-direction of the base + frame is along x-direction, i.e. :math:`(1, 0, 0)`. + """ + forward_w = math_utils.quat_apply(self.object_link_quat_w, self.FORWARD_VEC_B) + return torch.atan2(forward_w[..., 1], forward_w[..., 0]) + + ## + # Derived properties. + ## + + @property + def object_pos_w(self) -> torch.Tensor: + """Object position in simulation world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the position of the actor frame of the rigid bodies. + """ + return self.object_state_w[..., :3] + + @property + def object_quat_w(self) -> torch.Tensor: + """Object orientation (w, x, y, z) in simulation world frame. Shape is (num_instances, num_objects, 4). + + This quantity is the orientation of the actor frame of the rigid bodies. + """ + return self.object_state_w[..., 3:7] + + @property + def object_vel_w(self) -> torch.Tensor: + """Object velocity in simulation world frame. Shape is (num_instances, num_objects, 6). + + This quantity contains the linear and angular velocities of the rigid bodies' center of mass frame. + """ + return self.object_state_w[..., 7:13] + + @property + def object_lin_vel_w(self) -> torch.Tensor: + """Object linear velocity in simulation world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the linear velocity of the rigid bodies' center of mass frame. + """ + return self.object_state_w[..., 7:10] + + @property + def object_ang_vel_w(self) -> torch.Tensor: + """Object angular velocity in simulation world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the angular velocity of the rigid bodies' center of mass frame. + """ + return self.object_state_w[..., 10:13] + + @property + def object_lin_vel_b(self) -> torch.Tensor: + """Object linear velocity in base frame. Shape is (num_instances, num_objects, 3). + + This quantity is the linear velocity of the rigid bodies' center of mass frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.object_quat_w, self.object_lin_vel_w) + + @property + def object_ang_vel_b(self) -> torch.Tensor: + """Object angular velocity in base world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the angular velocity of the rigid bodies' center of mass frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.object_quat_w, self.object_ang_vel_w) + + @property + def object_lin_acc_w(self) -> torch.Tensor: + """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the linear acceleration of the rigid bodies' center of mass frame. + """ + return self.object_acc_w[..., 0:3] + + @property + def object_ang_acc_w(self) -> torch.Tensor: + """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the angular acceleration of the rigid bodies' center of mass frame. + """ + return self.object_acc_w[..., 3:6] + + @property + def object_link_pos_w(self) -> torch.Tensor: + """Object link position in simulation world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the position of the actor frame of the rigid bodies. + """ + if self._object_state_w.timestamp < self._sim_timestamp: + # read data from simulation + pose = self._reshape_view_to_data(self._root_physx_view.get_transforms().clone()) + return pose[..., :3] + return self.object_link_state_w[..., :3] + + @property + def object_link_quat_w(self) -> torch.Tensor: + """Object link orientation (w, x, y, z) in simulation world frame. Shape is (num_instances, num_objects, 4). + + This quantity is the orientation of the actor frame of the rigid bodies. + """ + if self._object_state_w.timestamp < self._sim_timestamp: + # read data from simulation + pose = self._reshape_view_to_data(self._root_physx_view.get_transforms().clone()) + pose[..., 3:7] = math_utils.convert_quat(pose[..., 3:7], to="wxyz") + return pose[..., 3:7] + return self.object_link_state_w[..., 3:7] + + @property + def object_link_vel_w(self) -> torch.Tensor: + """Object link velocity in simulation world frame. Shape is (num_instances, num_objects, 6). + + This quantity contains the linear and angular velocities of the rigid bodies' actor frame. + """ + return self.object_link_state_w[..., 7:13] + + @property + def object_link_lin_vel_w(self) -> torch.Tensor: + """Object link linear velocity in simulation world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the linear velocity of the rigid bodies' actor frame. + """ + return self.object_link_state_w[..., 7:10] + + @property + def object_link_ang_vel_w(self) -> torch.Tensor: + """Object link angular velocity in simulation world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the angular velocity of the rigid bodies' actor frame. + """ + return self.object_link_state_w[..., 10:13] + + @property + def object_link_lin_vel_b(self) -> torch.Tensor: + """Object link linear velocity in base frame. Shape is (num_instances, num_objects, 3). + + This quantity is the linear velocity of the actor frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.object_link_quat_w, self.object_link_lin_vel_w) + + @property + def object_link_ang_vel_b(self) -> torch.Tensor: + """Object link angular velocity in base world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the angular velocity of the actor frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.object_link_quat_w, self.object_link_ang_vel_w) + + @property + def object_com_pos_w(self) -> torch.Tensor: + """Object center of mass position in simulation world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the position of the center of mass frame of the rigid bodies. + """ + return self.object_com_state_w[..., :3] + + @property + def object_com_quat_w(self) -> torch.Tensor: + """Object center of mass orientation (w, x, y, z) in simulation world frame. Shape is (num_instances, num_objects, 4). + + This quantity is the orientation of the center of mass frame of the rigid bodies. + """ + return self.object_com_state_w[..., 3:7] + + @property + def object_com_vel_w(self) -> torch.Tensor: + """Object center of mass velocity in simulation world frame. Shape is (num_instances, num_objects, 6). + + This quantity contains the linear and angular velocities of the rigid bodies' center of mass frame. + """ + if self._object_state_w.timestamp < self._sim_timestamp: + # read data from simulation + velocity = self._reshape_view_to_data(self._root_physx_view.get_velocities()) + return velocity + return self.object_com_state_w[..., 7:13] + + @property + def object_com_lin_vel_w(self) -> torch.Tensor: + """Object center of mass linear velocity in simulation world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the linear velocity of the rigid bodies' center of mass frame. + """ + if self._object_state_w.timestamp < self._sim_timestamp: + # read data from simulation + velocity = self._reshape_view_to_data(self._root_physx_view.get_velocities()) + return velocity[..., 0:3] + return self.object_com_state_w[..., 7:10] + + @property + def object_com_ang_vel_w(self) -> torch.Tensor: + """Object center of mass angular velocity in simulation world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the angular velocity of the rigid bodies' center of mass frame. + """ + if self._object_state_w.timestamp < self._sim_timestamp: + # read data from simulation + velocity = self._reshape_view_to_data(self._root_physx_view.get_velocities()) + return velocity[..., 3:6] + return self.object_com_state_w[..., 10:13] + + @property + def object_com_lin_vel_b(self) -> torch.Tensor: + """Object center of mass linear velocity in base frame. Shape is (num_instances, num_objects, 3). + + This quantity is the linear velocity of the center of mass frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.object_link_quat_w, self.object_com_lin_vel_w) + + @property + def object_com_ang_vel_b(self) -> torch.Tensor: + """Object center of mass angular velocity in base world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the angular velocity of the center of mass frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.object_link_quat_w, self.object_com_ang_vel_w) + + @property + def com_pos_b(self) -> torch.Tensor: + """Center of mass of all of the bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the center of mass location relative to its body frame. + """ + pos = self._root_physx_view.get_coms().to(self.device)[..., :3] + return self._reshape_view_to_data(pos) + + @property + def com_quat_b(self) -> torch.Tensor: + """Orientation (w,x,y,z) of the prinicple axies of inertia of all of the bodies in simulation world frame. Shape is (num_instances, 1, 4). + + This quantity is the orientation of the principles axes of inertia relative to its body frame. + """ + quat = self._root_physx_view.get_coms().to(self.device)[..., 3:7].view(self.num_instances, self.num_objects, 4) + quat_wxyz = math_utils.convert_quat(quat, to="wxyz") + return self._reshape_view_to_data(quat_wxyz) + + ## + # Helpers. + ## + + def _reshape_view_to_data(self, data: torch.Tensor) -> torch.Tensor: + """Reshapes and arranges the data from the physics view to (num_instances, num_objects, data_size). + + Args: + data: The data from the physics view. Shape is (num_instances*num_objects, data_size). + + Returns: + The reshaped data. Shape is (num_objects, num_instances, data_size). + """ + return torch.einsum("ijk -> jik", data.reshape(self.num_objects, self.num_instances, -1)) diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/controllers/__init__.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/controllers/__init__.py index e97b0ac0ab..5c509e2067 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/controllers/__init__.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/controllers/__init__.py @@ -13,3 +13,5 @@ from .differential_ik import DifferentialIKController from .differential_ik_cfg import DifferentialIKControllerCfg +from .operational_space import OperationalSpaceController +from .operational_space_cfg import OperationalSpaceControllerCfg diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/controllers/operational_space.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/controllers/operational_space.py index 5b57b14c5a..9f278248e2 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/controllers/operational_space.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/controllers/operational_space.py @@ -3,92 +3,40 @@ # # SPDX-License-Identifier: BSD-3-Clause -import torch -from collections.abc import Sequence -from dataclasses import MISSING - -from omni.isaac.lab.utils import configclass -from omni.isaac.lab.utils.math import apply_delta_pose, compute_pose_error - - -@configclass -class OperationSpaceControllerCfg: - """Configuration for operation-space controller.""" - - command_types: Sequence[str] = MISSING - """Type of command. - - It has two sub-strings joined by underscore: - - type of command mode: "position", "pose", "force" - - type of command resolving: "abs" (absolute), "rel" (relative) - """ - - impedance_mode: str = MISSING - """Type of gains for motion control: "fixed", "variable", "variable_kp".""" - - uncouple_motion_wrench: bool = False - """Whether to decouple the wrench computation from task-space pose (motion) error.""" - - motion_control_axes: Sequence[int] = (1, 1, 1, 1, 1, 1) - """Motion direction to control. Mark as 0/1 for each axis.""" - force_control_axes: Sequence[int] = (0, 0, 0, 0, 0, 0) - """Force direction to control. Mark as 0/1 for each axis.""" - - inertial_compensation: bool = False - """Whether to perform inertial compensation for motion control (inverse dynamics).""" - - gravity_compensation: bool = False - """Whether to perform gravity compensation.""" - - stiffness: float | Sequence[float] = MISSING - """The positional gain for determining wrenches based on task-space pose error.""" - - damping_ratio: float | Sequence[float] | None = None - """The damping ratio is used in-conjunction with positional gain to compute wrenches - based on task-space velocity error. - - The following math operation is performed for computing velocity gains: - :math:`d_gains = 2 * sqrt(p_gains) * damping_ratio`. - """ - - stiffness_limits: tuple[float, float] = (0, 300) - """Minimum and maximum values for positional gains. - - Note: Used only when :obj:`impedance_mode` is "variable" or "variable_kp". - """ - - damping_ratio_limits: tuple[float, float] = (0, 100) - """Minimum and maximum values for damping ratios used to compute velocity gains. +from __future__ import annotations - Note: Used only when :obj:`impedance_mode` is "variable". - """ +import torch +from typing import TYPE_CHECKING - force_stiffness: float | Sequence[float] = None - """The positional gain for determining wrenches for closed-loop force control. +from omni.isaac.lab.utils.math import ( + apply_delta_pose, + combine_frame_transforms, + compute_pose_error, + matrix_from_quat, + subtract_frame_transforms, +) - If obj:`None`, then open-loop control of desired forces is performed. - """ +if TYPE_CHECKING: + from .operational_space_cfg import OperationalSpaceControllerCfg - position_command_scale: tuple[float, float, float] = (1.0, 1.0, 1.0) - """Scaling of the position command received. Used only in relative mode.""" - rotation_command_scale: tuple[float, float, float] = (1.0, 1.0, 1.0) - """Scaling of the rotation command received. Used only in relative mode.""" - -class OperationSpaceController: - """Operation-space controller. +class OperationalSpaceController: + """Operational-space controller. Reference: - [1] https://ethz.ch/content/dam/ethz/special-interest/mavt/robotics-n-intelligent-systems/rsl-dam/documents/RobotDynamics2017/RD_HS2017script.pdf + + 1. `A unified approach for motion and force control of robot manipulators: The operational space formulation `_ + by Oussama Khatib (Stanford University) + 2. `Robot Dynamics Lecture Notes `_ + by Marco Hutter (ETH Zurich) """ - def __init__(self, cfg: OperationSpaceControllerCfg, num_robots: int, num_dof: int, device: str): - """Initialize operation-space controller. + def __init__(self, cfg: OperationalSpaceControllerCfg, num_envs: int, device: str): + """Initialize operational-space controller. Args: - cfg: The configuration for operation-space controller. - num_robots: The number of robots to control. - num_dof: The number of degrees of freedom of the robot. + cfg: The configuration for operational-space controller. + num_envs: The number of environments. device: The device to use for computations. Raises: @@ -96,79 +44,105 @@ def __init__(self, cfg: OperationSpaceControllerCfg, num_robots: int, num_dof: i """ # store inputs self.cfg = cfg - self.num_robots = num_robots - self.num_dof = num_dof + self.num_envs = num_envs self._device = device # resolve tasks-pace target dimensions self.target_list = list() - for command_type in self.cfg.command_types: - if "position" in command_type: - self.target_list.append(3) - elif command_type == "pose_rel": + for command_type in self.cfg.target_types: + if command_type == "pose_rel": self.target_list.append(6) elif command_type == "pose_abs": self.target_list.append(7) - elif command_type == "force_abs": + elif command_type == "wrench_abs": self.target_list.append(6) else: raise ValueError(f"Invalid control command: {command_type}.") self.target_dim = sum(self.target_list) # create buffers - # -- selection matrices - self._selection_matrix_motion = torch.diag( - torch.tensor(self.cfg.motion_control_axes, dtype=torch.float, device=self._device) + # -- selection matrices, which might be defined in the task reference frame different from the root frame + self._selection_matrix_motion_task = torch.diag_embed( + torch.tensor(self.cfg.motion_control_axes_task, dtype=torch.float, device=self._device) + .unsqueeze(0) + .repeat(self.num_envs, 1) ) - self._selection_matrix_force = torch.diag( - torch.tensor(self.cfg.force_control_axes, dtype=torch.float, device=self._device) + self._selection_matrix_force_task = torch.diag_embed( + torch.tensor(self.cfg.contact_wrench_control_axes_task, dtype=torch.float, device=self._device) + .unsqueeze(0) + .repeat(self.num_envs, 1) ) + # -- selection matrices in root frame + self._selection_matrix_motion_b = torch.zeros_like(self._selection_matrix_motion_task) + self._selection_matrix_force_b = torch.zeros_like(self._selection_matrix_force_task) # -- commands - self._task_space_target = torch.zeros(self.num_robots, self.target_dim, device=self._device) - # -- scaling of command - self._position_command_scale = torch.diag(torch.tensor(self.cfg.position_command_scale, device=self._device)) - self._rotation_command_scale = torch.diag(torch.tensor(self.cfg.rotation_command_scale, device=self._device)) + self._task_space_target_task = torch.zeros(self.num_envs, self.target_dim, device=self._device) + # -- buffers for motion/force control + self.desired_ee_pose_task = None + self.desired_ee_pose_b = None + self.desired_ee_wrench_task = None + self.desired_ee_wrench_b = None # -- motion control gains - self._p_gains = torch.zeros(self.num_robots, 6, device=self._device) - self._p_gains[:] = torch.tensor(self.cfg.stiffness, device=self._device) - self._d_gains = 2 * torch.sqrt(self._p_gains) * torch.tensor(self.cfg.damping_ratio, device=self._device) + self._motion_p_gains_task = torch.diag_embed( + torch.ones(self.num_envs, 6, device=self._device) + * torch.tensor(self.cfg.motion_stiffness_task, dtype=torch.float, device=self._device) + ) + # -- -- zero out the axes that are not motion controlled, as keeping them non-zero will cause other axes + # -- -- to act due to coupling + self._motion_p_gains_task[:] = self._selection_matrix_motion_task @ self._motion_p_gains_task[:] + self._motion_d_gains_task = torch.diag_embed( + 2 + * torch.diagonal(self._motion_p_gains_task, dim1=-2, dim2=-1).sqrt() + * torch.as_tensor(self.cfg.motion_damping_ratio_task, dtype=torch.float, device=self._device).reshape(1, -1) + ) + # -- -- motion control gains in root frame + self._motion_p_gains_b = torch.zeros_like(self._motion_p_gains_task) + self._motion_d_gains_b = torch.zeros_like(self._motion_d_gains_task) # -- force control gains - if self.cfg.force_stiffness is not None: - self._p_wrench_gains = torch.zeros(self.num_robots, 6, device=self._device) - self._p_wrench_gains[:] = torch.tensor(self.cfg.force_stiffness, device=self._device) + if self.cfg.contact_wrench_stiffness_task is not None: + self._contact_wrench_p_gains_task = torch.diag_embed( + torch.ones(self.num_envs, 6, device=self._device) + * torch.tensor(self.cfg.contact_wrench_stiffness_task, dtype=torch.float, device=self._device) + ) + self._contact_wrench_p_gains_task[:] = ( + self._selection_matrix_force_task @ self._contact_wrench_p_gains_task[:] + ) + # -- -- force control gains in root frame + self._contact_wrench_p_gains_b = torch.zeros_like(self._contact_wrench_p_gains_task) else: - self._p_wrench_gains = None + self._contact_wrench_p_gains_task = None + self._contact_wrench_p_gains_b = None # -- position gain limits - self._p_gains_limits = torch.zeros(self.num_robots, 6, device=self._device) - self._p_gains_limits[..., 0], self._p_gains_limits[..., 1] = ( - self.cfg.stiffness_limits[0], - self.cfg.stiffness_limits[1], + self._motion_p_gains_limits = torch.zeros(self.num_envs, 6, 2, device=self._device) + self._motion_p_gains_limits[..., 0], self._motion_p_gains_limits[..., 1] = ( + self.cfg.motion_stiffness_limits_task[0], + self.cfg.motion_stiffness_limits_task[1], ) # -- damping ratio limits - self._damping_ratio_limits = torch.zeros_like(self._p_gains_limits) - self._damping_ratio_limits[..., 0], self._damping_ratio_limits[..., 1] = ( - self.cfg.damping_ratio_limits[0], - self.cfg.damping_ratio_limits[1], + self._motion_damping_ratio_limits = torch.zeros_like(self._motion_p_gains_limits) + self._motion_damping_ratio_limits[..., 0], self._motion_damping_ratio_limits[..., 1] = ( + self.cfg.motion_damping_ratio_limits_task[0], + self.cfg.motion_damping_ratio_limits_task[1], ) - # -- storing outputs - self._desired_torques = torch.zeros(self.num_robots, self.num_dof, self.num_dof, device=self._device) + # -- end-effector contact wrench + self._ee_contact_wrench_b = torch.zeros(self.num_envs, 6, device=self._device) """ Properties. """ @property - def num_actions(self) -> int: + def action_dim(self) -> int: """Dimension of the action space of controller.""" # impedance mode if self.cfg.impedance_mode == "fixed": - # task-space pose + # task-space targets return self.target_dim elif self.cfg.impedance_mode == "variable_kp": - # task-space pose + stiffness + # task-space targets + stiffness return self.target_dim + 6 elif self.cfg.impedance_mode == "variable": - # task-space pose + stiffness + damping + # task-space targets + stiffness + damping return self.target_dim + 6 + 6 else: raise ValueError(f"Invalid impedance mode: {self.cfg.impedance_mode}.") @@ -177,184 +151,301 @@ def num_actions(self) -> int: Operations. """ - def initialize(self): - """Initialize the internals.""" - pass - - def reset_idx(self, robot_ids: torch.Tensor = None): + def reset(self): """Reset the internals.""" - pass + self.desired_ee_pose_b = None + self.desired_ee_pose_task = None + self.desired_ee_wrench_b = None + self.desired_ee_wrench_task = None - def set_command(self, command: torch.Tensor): - """Set target end-effector pose or force command. + def set_command( + self, + command: torch.Tensor, + current_ee_pose_b: torch.Tensor | None = None, + current_task_frame_pose_b: torch.Tensor | None = None, + ): + """Set the task-space targets and impedance parameters. Args: - command: The target end-effector pose or force command. + command (torch.Tensor): A concatenated tensor of shape (``num_envs``, ``action_dim``) containing task-space + targets (i.e., pose/wrench) and impedance parameters. + current_ee_pose_b (torch.Tensor, optional): Current end-effector pose, in root frame, of shape + (``num_envs``, 7), containing position and quaternion ``(w, x, y, z)``. Required for relative + commands. Defaults to None. + current_task_frame_pose_b: Current pose of the task frame, in root frame, in which the targets and the + (motion/wrench) control axes are defined. It is a tensor of shape (``num_envs``, 7), + containing position and the quaternion ``(w, x, y, z)``. Defaults to None. + + Format: + Task-space targets, ordered according to 'command_types': + + Absolute pose: shape (``num_envs``, 7), containing position and quaternion ``(w, x, y, z)``. + Relative pose: shape (``num_envs``, 6), containing delta position and rotation in axis-angle form. + Absolute wrench: shape (``num_envs``, 6), containing force and torque. + + Impedance parameters: stiffness for ``variable_kp``, or stiffness, followed by damping ratio for + ``variable``: + + Stiffness: shape (``num_envs``, 6) + Damping ratio: shape (``num_envs``, 6) + + Raises: + ValueError: When the command dimensions are invalid. + ValueError: When an invalid impedance mode is provided. + ValueError: When the current end-effector pose is not provided for the ``pose_rel`` command. + ValueError: When an invalid control command is provided. """ - # check input size - if command.shape != (self.num_robots, self.num_actions): + # Check the input dimensions + if command.shape != (self.num_envs, self.action_dim): raise ValueError( - f"Invalid command shape '{command.shape}'. Expected: '{(self.num_robots, self.num_actions)}'." + f"Invalid command shape '{command.shape}'. Expected: '{(self.num_envs, self.action_dim)}'." ) - # impedance mode + + # Resolve the impedance parameters if self.cfg.impedance_mode == "fixed": - # joint positions - self._task_space_target[:] = command + # task space targets (i.e., pose/wrench) + self._task_space_target_task[:] = command elif self.cfg.impedance_mode == "variable_kp": # split input command - task_space_command, stiffness = torch.tensor_split(command, (self.target_dim, 6), dim=-1) + task_space_command, stiffness = torch.split(command, [self.target_dim, 6], dim=-1) # format command - stiffness = stiffness.clip_(min=self._p_gains_limits[0], max=self._p_gains_limits[1]) - # joint positions + stiffness - self._task_space_target[:] = task_space_command.squeeze(dim=-1) - self._p_gains[:] = stiffness - self._d_gains[:] = 2 * torch.sqrt(self._p_gains) # critically damped + stiffness = stiffness.clip_( + min=self._motion_p_gains_limits[..., 0], max=self._motion_p_gains_limits[..., 1] + ) + # task space targets + stiffness + self._task_space_target_task[:] = task_space_command.squeeze(dim=-1) + self._motion_p_gains_task[:] = torch.diag_embed(stiffness) + self._motion_p_gains_task[:] = self._selection_matrix_motion_task @ self._motion_p_gains_task[:] + self._motion_d_gains_task = torch.diag_embed( + 2 + * torch.diagonal(self._motion_p_gains_task, dim1=-2, dim2=-1).sqrt() + * torch.as_tensor(self.cfg.motion_damping_ratio_task, dtype=torch.float, device=self._device).reshape( + 1, -1 + ) + ) elif self.cfg.impedance_mode == "variable": # split input command - task_space_command, stiffness, damping_ratio = torch.tensor_split(command, 3, dim=-1) + task_space_command, stiffness, damping_ratio = torch.split(command, [self.target_dim, 6, 6], dim=-1) # format command - stiffness = stiffness.clip_(min=self._p_gains_limits[0], max=self._p_gains_limits[1]) - damping_ratio = damping_ratio.clip_(min=self._damping_ratio_limits[0], max=self._damping_ratio_limits[1]) - # joint positions + stiffness + damping - self._task_space_target[:] = task_space_command - self._p_gains[:] = stiffness - self._d_gains[:] = 2 * torch.sqrt(self._p_gains) * damping_ratio + stiffness = stiffness.clip_( + min=self._motion_p_gains_limits[..., 0], max=self._motion_p_gains_limits[..., 1] + ) + damping_ratio = damping_ratio.clip_( + min=self._motion_damping_ratio_limits[..., 0], max=self._motion_damping_ratio_limits[..., 1] + ) + # task space targets + stiffness + damping + self._task_space_target_task[:] = task_space_command + self._motion_p_gains_task[:] = torch.diag_embed(stiffness) + self._motion_p_gains_task[:] = self._selection_matrix_motion_task @ self._motion_p_gains_task[:] + self._motion_d_gains_task[:] = torch.diag_embed( + 2 * torch.diagonal(self._motion_p_gains_task, dim1=-2, dim2=-1).sqrt() * damping_ratio + ) else: raise ValueError(f"Invalid impedance mode: {self.cfg.impedance_mode}.") + if current_task_frame_pose_b is None: + current_task_frame_pose_b = torch.tensor( + [[0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]] * self.num_envs, device=self._device + ) + + # Resolve the target commands + target_groups = torch.split(self._task_space_target_task, self.target_list, dim=1) + for command_type, target in zip(self.cfg.target_types, target_groups): + if command_type == "pose_rel": + # check input is provided + if current_ee_pose_b is None: + raise ValueError("Current pose is required for 'pose_rel' command.") + # Transform the current pose from base/root frame to task frame + current_ee_pos_task, current_ee_rot_task = subtract_frame_transforms( + current_task_frame_pose_b[:, :3], + current_task_frame_pose_b[:, 3:], + current_ee_pose_b[:, :3], + current_ee_pose_b[:, 3:], + ) + # compute targets in task frame + desired_ee_pos_task, desired_ee_rot_task = apply_delta_pose( + current_ee_pos_task, current_ee_rot_task, target + ) + self.desired_ee_pose_task = torch.cat([desired_ee_pos_task, desired_ee_rot_task], dim=-1) + elif command_type == "pose_abs": + # compute targets + self.desired_ee_pose_task = target.clone() + elif command_type == "wrench_abs": + # compute targets + self.desired_ee_wrench_task = target.clone() + else: + raise ValueError(f"Invalid control command: {command_type}.") + + # Rotation of task frame wrt root frame, converts a coordinate from task frame to root frame. + R_task_b = matrix_from_quat(current_task_frame_pose_b[:, 3:]) + # Rotation of root frame wrt task frame, converts a coordinate from root frame to task frame. + R_b_task = R_task_b.mT + + # Transform motion control stiffness gains from task frame to root frame + self._motion_p_gains_b[:, 0:3, 0:3] = R_task_b @ self._motion_p_gains_task[:, 0:3, 0:3] @ R_b_task + self._motion_p_gains_b[:, 3:6, 3:6] = R_task_b @ self._motion_p_gains_task[:, 3:6, 3:6] @ R_b_task + + # Transform motion control damping gains from task frame to root frame + self._motion_d_gains_b[:, 0:3, 0:3] = R_task_b @ self._motion_d_gains_task[:, 0:3, 0:3] @ R_b_task + self._motion_d_gains_b[:, 3:6, 3:6] = R_task_b @ self._motion_d_gains_task[:, 3:6, 3:6] @ R_b_task + + # Transform contact wrench gains from task frame to root frame (if applicable) + if self._contact_wrench_p_gains_task is not None and self._contact_wrench_p_gains_b is not None: + self._contact_wrench_p_gains_b[:, 0:3, 0:3] = ( + R_task_b @ self._contact_wrench_p_gains_task[:, 0:3, 0:3] @ R_b_task + ) + self._contact_wrench_p_gains_b[:, 3:6, 3:6] = ( + R_task_b @ self._contact_wrench_p_gains_task[:, 3:6, 3:6] @ R_b_task + ) + + # Transform selection matrices from target frame to base frame + self._selection_matrix_motion_b[:, 0:3, 0:3] = ( + R_task_b @ self._selection_matrix_motion_task[:, 0:3, 0:3] @ R_b_task + ) + self._selection_matrix_motion_b[:, 3:6, 3:6] = ( + R_task_b @ self._selection_matrix_motion_task[:, 3:6, 3:6] @ R_b_task + ) + self._selection_matrix_force_b[:, 0:3, 0:3] = ( + R_task_b @ self._selection_matrix_force_task[:, 0:3, 0:3] @ R_b_task + ) + self._selection_matrix_force_b[:, 3:6, 3:6] = ( + R_task_b @ self._selection_matrix_force_task[:, 3:6, 3:6] @ R_b_task + ) + + # Transform desired pose from task frame to root frame + if self.desired_ee_pose_task is not None: + self.desired_ee_pose_b = torch.zeros_like(self.desired_ee_pose_task) + self.desired_ee_pose_b[:, :3], self.desired_ee_pose_b[:, 3:] = combine_frame_transforms( + current_task_frame_pose_b[:, :3], + current_task_frame_pose_b[:, 3:], + self.desired_ee_pose_task[:, :3], + self.desired_ee_pose_task[:, 3:], + ) + + # Transform desired wrenches to root frame + if self.desired_ee_wrench_task is not None: + self.desired_ee_wrench_b = torch.zeros_like(self.desired_ee_wrench_task) + self.desired_ee_wrench_b[:, :3] = (R_task_b @ self.desired_ee_wrench_task[:, :3].unsqueeze(-1)).squeeze(-1) + self.desired_ee_wrench_b[:, 3:] = (R_task_b @ self.desired_ee_wrench_task[:, 3:].unsqueeze(-1)).squeeze( + -1 + ) + torch.cross(current_task_frame_pose_b[:, :3], self.desired_ee_wrench_b[:, :3], dim=-1) + def compute( self, - jacobian: torch.Tensor, - ee_pose: torch.Tensor | None = None, - ee_vel: torch.Tensor | None = None, - ee_force: torch.Tensor | None = None, + jacobian_b: torch.Tensor, + current_ee_pose_b: torch.Tensor | None = None, + current_ee_vel_b: torch.Tensor | None = None, + current_ee_force_b: torch.Tensor | None = None, mass_matrix: torch.Tensor | None = None, gravity: torch.Tensor | None = None, ) -> torch.Tensor: """Performs inference with the controller. Args: - jacobian: The Jacobian matrix of the end-effector. - ee_pose: The current end-effector pose. It is a tensor of shape - (num_robots, 7), which contains the position and quaternion (w, x, y, z). Defaults to None. - ee_vel: The current end-effector velocity. It is a tensor of shape - (num_robots, 6), which contains the linear and angular velocities. Defaults to None. - ee_force: The current external force on the end-effector. - It is a tensor of shape (num_robots, 3), which contains the linear force. Defaults to None. - mass_matrix: The joint-space inertial matrix. Defaults to None. - gravity: The joint-space gravity vector. Defaults to None. + jacobian_b: The Jacobian matrix of the end-effector in root frame. It is a tensor of shape + (``num_envs``, 6, ``num_DoF``). + current_ee_pose_b: The current end-effector pose in root frame. It is a tensor of shape + (``num_envs``, 7), which contains the position and quaternion ``(w, x, y, z)``. Defaults to ``None``. + current_ee_vel_b: The current end-effector velocity in root frame. It is a tensor of shape + (``num_envs``, 6), which contains the linear and angular velocities. Defaults to None. + current_ee_force_b: The current external force on the end-effector in root frame. + It is a tensor of shape (``num_envs``, 3), which contains the linear force. Defaults to ``None``. + mass_matrix: The joint-space mass/inertia matrix. It is a tensor of shape (``num_envs``, ``num_DoF``, + ``num_DoF``). Defaults to ``None``. + gravity: The joint-space gravity vector. It is a tensor of shape (``num_envs``, ``num_DoF``). Defaults + to ``None``. Raises: - ValueError: When the end-effector pose is not provided for the 'position_rel' command. - ValueError: When the end-effector pose is not provided for the 'position_abs' command. - ValueError: When the end-effector pose is not provided for the 'pose_rel' command. - ValueError: When an invalid command type is provided. - ValueError: When motion-control is enabled but the end-effector pose or velocity is not provided. - ValueError: When force-control is enabled but the end-effector force is not provided. - ValueError: When inertial compensation is enabled but the mass matrix is not provided. + ValueError: When motion-control is enabled but the current end-effector pose or velocity is not provided. + ValueError: When inertial dynamics decoupling is enabled but the mass matrix is not provided. + ValueError: When the current end-effector pose is not provided for the ``pose_rel`` command. + ValueError: When closed-loop force control is enabled but the current end-effector force is not provided. ValueError: When gravity compensation is enabled but the gravity vector is not provided. Returns: - The target joint torques commands. + Tensor: The joint efforts computed by the controller. It is a tensor of shape (``num_envs``, ``num_DoF``). """ - # buffers for motion/force control - desired_ee_pos = None - desired_ee_rot = None - desired_ee_force = None - # resolve the commands - target_groups = torch.split(self._task_space_target, self.target_list) - for command_type, target in zip(self.cfg.command_types, target_groups): - if command_type == "position_rel": - # check input is provided - if ee_pose is None: - raise ValueError("End-effector pose is required for 'position_rel' command.") - # scale command - target @= self._position_command_scale - # compute targets - desired_ee_pos = ee_pose[:, :3] + target - desired_ee_rot = ee_pose[:, 3:] - elif command_type == "position_abs": - # check input is provided - if ee_pose is None: - raise ValueError("End-effector pose is required for 'position_abs' command.") - # compute targets - desired_ee_pos = target - desired_ee_rot = ee_pose[:, 3:] - elif command_type == "pose_rel": - # check input is provided - if ee_pose is None: - raise ValueError("End-effector pose is required for 'pose_rel' command.") - # scale command - target[:, 0:3] @= self._position_command_scale - target[:, 3:6] @= self._rotation_command_scale - # compute targets - desired_ee_pos, desired_ee_rot = apply_delta_pose(ee_pose[:, :3], ee_pose[:, 3:], target) - elif command_type == "pose_abs": - # compute targets - desired_ee_pos = target[:, 0:3] - desired_ee_rot = target[:, 3:7] - elif command_type == "force_abs": - # compute targets - desired_ee_force = target - else: - raise ValueError(f"Invalid control command: {self.cfg.command_type}.") - # reset desired joint torques - self._desired_torques[:] = 0 - # compute for motion-control - if desired_ee_pos is not None: + # deduce number of DoF + num_DoF = jacobian_b.shape[2] + # create joint effort vector + joint_efforts = torch.zeros(self.num_envs, num_DoF, device=self._device) + + # compute joint efforts for motion-control + if self.desired_ee_pose_b is not None: # check input is provided - if ee_pose is None or ee_vel is None: - raise ValueError("End-effector pose and velocity are required for motion control.") + if current_ee_pose_b is None or current_ee_vel_b is None: + raise ValueError("Current end-effector pose and velocity are required for motion control.") # -- end-effector tracking error - pose_error = compute_pose_error( - ee_pose[:, :3], ee_pose[:, 3:], desired_ee_pos, desired_ee_rot, rot_error_type="axis_angle" + pose_error_b = torch.cat( + compute_pose_error( + current_ee_pose_b[:, :3], + current_ee_pose_b[:, 3:], + self.desired_ee_pose_b[:, :3], + self.desired_ee_pose_b[:, 3:], + rot_error_type="axis_angle", + ), + dim=-1, ) - velocity_error = -ee_vel # zero target velocity - # -- desired end-effector acceleration (spring damped system) - des_ee_acc = self._p_gains * pose_error + self._d_gains * velocity_error - # -- inertial compensation - if self.cfg.inertial_compensation: + velocity_error_b = -current_ee_vel_b # zero target velocity. The target is assumed to be stationary. + # -- desired end-effector acceleration (spring-damper system) + des_ee_acc_b = self._motion_p_gains_b @ pose_error_b.unsqueeze( + -1 + ) + self._motion_d_gains_b @ velocity_error_b.unsqueeze(-1) + # -- Inertial dynamics decoupling + if self.cfg.inertial_dynamics_decoupling: # check input is provided if mass_matrix is None: - raise ValueError("Mass matrix is required for inertial compensation.") - # compute task-space dynamics quantities - # wrench = (J M^(-1) J^T)^(-1) * \ddot(x_des) + raise ValueError("Mass matrix is required for inertial decoupling.") + # Compute operational space mass matrix mass_matrix_inv = torch.inverse(mass_matrix) - if self.cfg.uncouple_motion_wrench: - # decoupled-mass matrices - lambda_pos = torch.inverse(jacobian[:, 0:3] @ mass_matrix_inv * jacobian[:, 0:3].T) - lambda_ori = torch.inverse(jacobian[:, 3:6] @ mass_matrix_inv * jacobian[:, 3:6].T) - # desired end-effector wrench (from pseudo-dynamics) - decoupled_force = lambda_pos @ des_ee_acc[:, 0:3] - decoupled_torque = lambda_ori @ des_ee_acc[:, 3:6] - des_motion_wrench = torch.cat(decoupled_force, decoupled_torque) + if self.cfg.partial_inertial_dynamics_decoupling: + # Create a zero tensor representing the mass matrix, to fill in the non-zero elements later + os_mass_matrix = torch.zeros(self.num_envs, 6, 6, device=self._device) + # Fill in the translational and rotational parts of the inertia separately, ignoring their coupling + os_mass_matrix[:, 0:3, 0:3] = torch.inverse( + jacobian_b[:, 0:3] @ mass_matrix_inv @ jacobian_b[:, 0:3].mT + ) + os_mass_matrix[:, 3:6, 3:6] = torch.inverse( + jacobian_b[:, 3:6] @ mass_matrix_inv @ jacobian_b[:, 3:6].mT + ) else: - # coupled dynamics - lambda_full = torch.inverse(jacobian @ mass_matrix_inv * jacobian.T) - # desired end-effector wrench (from pseudo-dynamics) - des_motion_wrench = lambda_full @ des_ee_acc + # Calculate the operational space mass matrix fully accounting for the couplings + os_mass_matrix = torch.inverse(jacobian_b @ mass_matrix_inv @ jacobian_b.mT) + # (Generalized) operational space command forces + # F = (J M^(-1) J^T)^(-1) * \ddot(x_des) = M_task * \ddot(x_des) + os_command_forces_b = os_mass_matrix @ des_ee_acc_b else: - # task-space impedance control - # wrench = \ddot(x_des) - des_motion_wrench = des_ee_acc - # -- joint-space wrench - self._desired_torques += jacobian.T @ self._selection_matrix_motion @ des_motion_wrench - - # compute for force control - if desired_ee_force is not None: - # -- task-space wrench - if self.cfg.stiffness is not None: + # Task-space impedance control: command forces = \ddot(x_des). + # Please note that the definition of task-space impedance control varies in literature. + # This implementation ignores the inertial term. For inertial decoupling, + # use inertial_dynamics_decoupling=True. + os_command_forces_b = des_ee_acc_b + # -- joint-space commands + joint_efforts += (jacobian_b.mT @ self._selection_matrix_motion_b @ os_command_forces_b).squeeze(-1) + + # compute joint efforts for contact wrench/force control + if self.desired_ee_wrench_b is not None: + # -- task-space contact wrench + if self.cfg.contact_wrench_stiffness_task is not None: # check input is provided - if ee_force is None: - raise ValueError("End-effector force is required for closed-loop force control.") - # closed-loop control - des_force_wrench = desired_ee_force + self._p_wrench_gains * (desired_ee_force - ee_force) + if current_ee_force_b is None: + raise ValueError("Current end-effector force is required for closed-loop force control.") + # We can only measure the force component at the contact, so only apply the feedback for only the force + # component, keep the control of moment components open loop + self._ee_contact_wrench_b[:, 0:3] = current_ee_force_b + self._ee_contact_wrench_b[:, 3:6] = self.desired_ee_wrench_b[:, 3:6] + # closed-loop control with feedforward term + os_contact_wrench_command_b = self.desired_ee_wrench_b.unsqueeze( + -1 + ) + self._contact_wrench_p_gains_b @ (self.desired_ee_wrench_b - self._ee_contact_wrench_b).unsqueeze( + -1 + ) else: # open-loop control - des_force_wrench = desired_ee_force - # -- joint-space wrench - self._desired_torques += jacobian.T @ self._selection_matrix_force @ des_force_wrench + os_contact_wrench_command_b = self.desired_ee_wrench_b.unsqueeze(-1) + # -- joint-space commands + joint_efforts += (jacobian_b.mT @ self._selection_matrix_force_b @ os_contact_wrench_command_b).squeeze(-1) # add gravity compensation (bias correction) if self.cfg.gravity_compensation: @@ -362,6 +453,6 @@ def compute( if gravity is None: raise ValueError("Gravity vector is required for gravity compensation.") # add gravity compensation - self._desired_torques += gravity + joint_efforts += gravity - return self._desired_torques + return joint_efforts diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/controllers/operational_space_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/controllers/operational_space_cfg.py new file mode 100644 index 0000000000..8befaee6aa --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/controllers/operational_space_cfg.py @@ -0,0 +1,77 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from collections.abc import Sequence +from dataclasses import MISSING + +from omni.isaac.lab.utils import configclass + +from .operational_space import OperationalSpaceController + + +@configclass +class OperationalSpaceControllerCfg: + """Configuration for operational-space controller.""" + + class_type: type = OperationalSpaceController + """The associated controller class.""" + + target_types: Sequence[str] = MISSING + """Type of task-space targets. + + It has two sub-strings joined by underscore: + - type of task-space target: ``"pose"``, ``"wrench"`` + - reference for the task-space targets: ``"abs"`` (absolute), ``"rel"`` (relative, only for pose) + """ + + motion_control_axes_task: Sequence[int] = (1, 1, 1, 1, 1, 1) + """Motion direction to control in task reference frame. Mark as ``0/1`` for each axis.""" + + contact_wrench_control_axes_task: Sequence[int] = (0, 0, 0, 0, 0, 0) + """Contact wrench direction to control in task reference frame. Mark as 0/1 for each axis.""" + + inertial_dynamics_decoupling: bool = False + """Whether to perform inertial dynamics decoupling for motion control (inverse dynamics).""" + + partial_inertial_dynamics_decoupling: bool = False + """Whether to ignore the inertial coupling between the translational & rotational motions.""" + + gravity_compensation: bool = False + """Whether to perform gravity compensation.""" + + impedance_mode: str = "fixed" + """Type of gains for motion control: ``"fixed"``, ``"variable"``, ``"variable_kp"``.""" + + motion_stiffness_task: float | Sequence[float] = (100.0, 100.0, 100.0, 100.0, 100.0, 100.0) + """The positional gain for determining operational space command forces based on task-space pose error.""" + + motion_damping_ratio_task: float | Sequence[float] = (1.0, 1.0, 1.0, 1.0, 1.0, 1.0) + """The damping ratio is used in-conjunction with positional gain to compute operational space command forces + based on task-space velocity error. + + The following math operation is performed for computing velocity gains: + :math:`d_gains = 2 * sqrt(p_gains) * damping_ratio`. + """ + + motion_stiffness_limits_task: tuple[float, float] = (0, 1000) + """Minimum and maximum values for positional gains. + + Note: Used only when :obj:`impedance_mode` is ``"variable"`` or ``"variable_kp"``. + """ + + motion_damping_ratio_limits_task: tuple[float, float] = (0, 100) + """Minimum and maximum values for damping ratios used to compute velocity gains. + + Note: Used only when :obj:`impedance_mode` is ``"variable"``. + """ + + contact_wrench_stiffness_task: float | Sequence[float] | None = None + """The proportional gain for determining operational space command forces for closed-loop contact force control. + + If ``None``, then open-loop control of desired contact wrench is performed. + + Note: since only the linear forces could be measured at the moment, + only the first three elements are used for the feedback loop. + """ diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/__init__.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/__init__.py index 97701e50cc..667c0b5049 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/__init__.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/__init__.py @@ -52,4 +52,4 @@ from .manager_based_env_cfg import ManagerBasedEnvCfg from .manager_based_rl_env import ManagerBasedRLEnv from .manager_based_rl_env_cfg import ManagerBasedRLEnvCfg -from .utils import multi_agent_to_single_agent, multi_agent_with_one_agent +from .utils.marl import multi_agent_to_single_agent, multi_agent_with_one_agent diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/common.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/common.py index d0c99f8ad8..7d6b02d309 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/common.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/common.py @@ -5,6 +5,7 @@ from __future__ import annotations +import gymnasium as gym import torch from typing import Dict, Literal, TypeVar @@ -62,6 +63,9 @@ class ViewerCfg: # Types. ## +SpaceType = TypeVar("SpaceType", gym.spaces.Space, int, set, tuple, list, dict) +"""A sentinel object to indicate a valid space type to specify states, observations and actions.""" + VecEnvObs = Dict[str, torch.Tensor | Dict[str, torch.Tensor]] """Observation returned by the environment. diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/direct_marl_env.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/direct_marl_env.py index 2fad4a0243..55b8a9501d 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/direct_marl_env.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/direct_marl_env.py @@ -14,6 +14,7 @@ import weakref from abc import abstractmethod from collections.abc import Sequence +from dataclasses import MISSING from typing import Any, ClassVar import omni.isaac.core.utils.torch as torch_utils @@ -30,9 +31,10 @@ from .common import ActionType, AgentID, EnvStepReturn, ObsType, StateType from .direct_marl_env_cfg import DirectMARLEnvCfg from .ui import ViewportCameraController +from .utils.spaces import sample_space, spec_to_gym_space -class DirectMARLEnv: +class DirectMARLEnv(gym.Env): """The superclass for the direct workflow to design multi-agent environments. This class implements the core functionality for multi-agent reinforcement learning (MARL) @@ -72,6 +74,8 @@ def __init__(self, cfg: DirectMARLEnvCfg, render_mode: str | None = None, **kwar RuntimeError: If a simulation context already exists. The environment must always create one since it configures the simulation context and controls the simulation. """ + # check that the config is valid + cfg.validate() # store inputs to class self.cfg = cfg # store the render mode @@ -102,7 +106,7 @@ def __init__(self, cfg: DirectMARLEnvCfg, render_mode: str | None = None, **kwar if self.cfg.sim.render_interval < self.cfg.decimation: msg = ( f"The render interval ({self.cfg.sim.render_interval}) is smaller than the decimation " - f"({self.cfg.decimation}). Multiple multiple render calls will happen for each environment step." + f"({self.cfg.decimation}). Multiple render calls will happen for each environment step." "If this is not intended, set the render interval to be equal to the decimation." ) omni.log.warn(msg) @@ -164,10 +168,6 @@ def __init__(self, cfg: DirectMARLEnvCfg, render_mode: str | None = None, **kwar # -- init buffers self.episode_length_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.long) self.reset_buf = torch.zeros(self.num_envs, dtype=torch.bool, device=self.sim.device) - self.actions = { - agent: torch.zeros(self.num_envs, self.cfg.num_actions[agent], device=self.sim.device) - for agent in self.cfg.possible_agents - } # setup the observation, state and action spaces self._configure_env_spaces() @@ -175,13 +175,13 @@ def __init__(self, cfg: DirectMARLEnvCfg, render_mode: str | None = None, **kwar # setup noise cfg for adding action and observation noise if self.cfg.action_noise_model: self._action_noise_model: dict[AgentID, NoiseModel] = { - agent: noise_model.class_type(self.num_envs, noise_model, self.device) + agent: noise_model.class_type(noise_model, num_envs=self.num_envs, device=self.device) for agent, noise_model in self.cfg.action_noise_model.items() if noise_model is not None } if self.cfg.observation_noise_model: self._observation_noise_model: dict[AgentID, NoiseModel] = { - agent: noise_model.class_type(self.num_envs, noise_model, self.device) + agent: noise_model.class_type(noise_model, num_envs=self.num_envs, device=self.device) for agent, noise_model in self.cfg.observation_noise_model.items() if noise_model is not None } @@ -406,16 +406,19 @@ def state(self) -> StateType | None: """Returns the state for the environment. The state-space is used for centralized training or asymmetric actor-critic architectures. It is configured - using the :attr:`DirectMARLEnvCfg.num_states` parameter. + using the :attr:`DirectMARLEnvCfg.state_space` parameter. Returns: - The states for the environment, or None if :attr:`DirectMARLEnvCfg.num_states` parameter is zero. + The states for the environment, or None if :attr:`DirectMARLEnvCfg.state_space` parameter is zero. """ - if not self.cfg.num_states: + if not self.cfg.state_space: return None # concatenate and return the observations as state - if self.cfg.num_states < 0: - self.state_buf = torch.cat([self.obs_dict[agent] for agent in self.cfg.possible_agents], dim=-1) + # FIXME: This implementation assumes the spaces are fundamental ones. Fix it to support composite spaces + if isinstance(self.cfg.state_space, int) and self.cfg.state_space < 0: + self.state_buf = torch.cat( + [self.obs_dict[agent].reshape(self.num_envs, -1) for agent in self.cfg.possible_agents], dim=-1 + ) # compute and return custom environment state else: self.state_buf = self._get_states() @@ -568,25 +571,45 @@ def _configure_env_spaces(self): self.agents = self.cfg.possible_agents self.possible_agents = self.cfg.possible_agents + # show deprecation message and overwrite configuration + if self.cfg.num_actions is not None: + omni.log.warn("DirectMARLEnvCfg.num_actions is deprecated. Use DirectMARLEnvCfg.action_spaces instead.") + if isinstance(self.cfg.action_spaces, type(MISSING)): + self.cfg.action_spaces = self.cfg.num_actions + if self.cfg.num_observations is not None: + omni.log.warn( + "DirectMARLEnvCfg.num_observations is deprecated. Use DirectMARLEnvCfg.observation_spaces instead." + ) + if isinstance(self.cfg.observation_spaces, type(MISSING)): + self.cfg.observation_spaces = self.cfg.num_observations + if self.cfg.num_states is not None: + omni.log.warn("DirectMARLEnvCfg.num_states is deprecated. Use DirectMARLEnvCfg.state_space instead.") + if isinstance(self.cfg.state_space, type(MISSING)): + self.cfg.state_space = self.cfg.num_states + # set up observation and action spaces self.observation_spaces = { - agent: gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.cfg.num_observations[agent],)) - for agent in self.cfg.possible_agents + agent: spec_to_gym_space(self.cfg.observation_spaces[agent]) for agent in self.cfg.possible_agents } self.action_spaces = { - agent: gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.cfg.num_actions[agent],)) - for agent in self.cfg.possible_agents + agent: spec_to_gym_space(self.cfg.action_spaces[agent]) for agent in self.cfg.possible_agents } # set up state space - if not self.cfg.num_states: + if not self.cfg.state_space: self.state_space = None - if self.cfg.num_states < 0: - self.state_space = gym.spaces.Box( - low=-np.inf, high=np.inf, shape=(sum(self.cfg.num_observations.values()),) + if isinstance(self.cfg.state_space, int) and self.cfg.state_space < 0: + self.state_space = gym.spaces.flatten_space( + gym.spaces.Tuple([self.observation_spaces[agent] for agent in self.cfg.possible_agents]) ) else: - self.state_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.cfg.num_states,)) + self.state_space = spec_to_gym_space(self.cfg.state_space) + + # instantiate actions (needed for tasks for which the observations computation is dependent on the actions) + self.actions = { + agent: sample_space(self.action_spaces[agent], self.sim.device, batch_size=self.num_envs, fill_value=0) + for agent in self.cfg.possible_agents + } def _reset_idx(self, env_ids: Sequence[int]): """Reset environments based on specified indices. @@ -664,8 +687,8 @@ def _get_observations(self) -> dict[AgentID, ObsType]: def _get_states(self) -> StateType: """Compute and return the states for the environment. - This method is only called (and therefore has to be implemented) when the :attr:`DirectMARLEnvCfg.num_states` - parameter is greater than zero. + This method is only called (and therefore has to be implemented) when the :attr:`DirectMARLEnvCfg.state_space` + parameter is not a number less than or equal to zero. Returns: The states for the environment. diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/direct_marl_env_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/direct_marl_env_cfg.py index 3dcf364f5c..40ecb64297 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/direct_marl_env_cfg.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/direct_marl_env_cfg.py @@ -10,7 +10,7 @@ from omni.isaac.lab.utils import configclass from omni.isaac.lab.utils.noise import NoiseModelCfg -from .common import AgentID, ViewerCfg +from .common import AgentID, SpaceType, ViewerCfg from .ui import BaseEnvWindow @@ -104,11 +104,39 @@ class DirectMARLEnvCfg: Please refer to the :class:`omni.isaac.lab.managers.EventManager` class for more details. """ - num_observations: dict[AgentID, int] = MISSING - """The dimension of the observation space from each agent.""" + observation_spaces: dict[AgentID, SpaceType] = MISSING + """Observation space definition for each agent. + + The space can be defined either using Gymnasium :py:mod:`~gymnasium.spaces` (when a more detailed + specification of the space is desired) or basic Python data types (for simplicity). + + .. list-table:: + :header-rows: 1 + + * - Gymnasium space + - Python data type + * - :class:`~gymnasium.spaces.Box` + - Integer or list of integers (e.g.: ``7``, ``[64, 64, 3]``) + * - :class:`~gymnasium.spaces.Discrete` + - Single-element set (e.g.: ``{2}``) + * - :class:`~gymnasium.spaces.MultiDiscrete` + - List of single-element sets (e.g.: ``[{2}, {5}]``) + * - :class:`~gymnasium.spaces.Dict` + - Dictionary (e.g.: ``{"joints": 7, "rgb": [64, 64, 3], "gripper": {2}}``) + * - :class:`~gymnasium.spaces.Tuple` + - Tuple (e.g.: ``(7, [64, 64, 3], {2})``) + """ - num_states: int = MISSING - """The dimension of the state space from each environment instance. + num_observations: dict[AgentID, int] | None = None + """The dimension of the observation space for each agent. + + .. warning:: + + This attribute is deprecated. Use :attr:`~omni.isaac.lab.envs.DirectMARLEnvCfg.observation_spaces` instead. + """ + + state_space: SpaceType = MISSING + """State space definition. The following values are supported: @@ -116,6 +144,33 @@ class DirectMARLEnvCfg: * 0: No state-space will be constructed (`state_space` is None). This is useful to save computational resources when the algorithm to be trained does not need it. * greater than 0: Custom state-space dimension to be provided by the task implementation. + + The space can be defined either using Gymnasium :py:mod:`~gymnasium.spaces` (when a more detailed + specification of the space is desired) or basic Python data types (for simplicity). + + .. list-table:: + :header-rows: 1 + + * - Gymnasium space + - Python data type + * - :class:`~gymnasium.spaces.Box` + - Integer or list of integers (e.g.: ``7``, ``[64, 64, 3]``) + * - :class:`~gymnasium.spaces.Discrete` + - Single-element set (e.g.: ``{2}``) + * - :class:`~gymnasium.spaces.MultiDiscrete` + - List of single-element sets (e.g.: ``[{2}, {5}]``) + * - :class:`~gymnasium.spaces.Dict` + - Dictionary (e.g.: ``{"joints": 7, "rgb": [64, 64, 3], "gripper": {2}}``) + * - :class:`~gymnasium.spaces.Tuple` + - Tuple (e.g.: ``(7, [64, 64, 3], {2})``) + """ + + num_states: int | None = None + """The dimension of the state space from each environment instance. + + .. warning:: + + This attribute is deprecated. Use :attr:`~omni.isaac.lab.envs.DirectMARLEnvCfg.state_space` instead. """ observation_noise_model: dict[AgentID, NoiseModelCfg | None] | None = None @@ -124,8 +179,36 @@ class DirectMARLEnvCfg: Please refer to the :class:`omni.isaac.lab.utils.noise.NoiseModel` class for more details. """ - num_actions: dict[AgentID, int] = MISSING - """The dimension of the action space for each agent.""" + action_spaces: dict[AgentID, SpaceType] = MISSING + """Action space definition for each agent. + + The space can be defined either using Gymnasium :py:mod:`~gymnasium.spaces` (when a more detailed + specification of the space is desired) or basic Python data types (for simplicity). + + .. list-table:: + :header-rows: 1 + + * - Gymnasium space + - Python data type + * - :class:`~gymnasium.spaces.Box` + - Integer or list of integers (e.g.: ``7``, ``[64, 64, 3]``) + * - :class:`~gymnasium.spaces.Discrete` + - Single-element set (e.g.: ``{2}``) + * - :class:`~gymnasium.spaces.MultiDiscrete` + - List of single-element sets (e.g.: ``[{2}, {5}]``) + * - :class:`~gymnasium.spaces.Dict` + - Dictionary (e.g.: ``{"joints": 7, "rgb": [64, 64, 3], "gripper": {2}}``) + * - :class:`~gymnasium.spaces.Tuple` + - Tuple (e.g.: ``(7, [64, 64, 3], {2})``) + """ + + num_actions: dict[AgentID, int] | None = None + """The dimension of the action space for each agent. + + .. warning:: + + This attribute is deprecated. Use :attr:`~omni.isaac.lab.envs.DirectMARLEnvCfg.action_spaces` instead. + """ action_noise_model: dict[AgentID, NoiseModelCfg | None] | None = None """The noise model applied to the actions provided to the environment. Default is None, which means no noise is added. diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/direct_rl_env.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/direct_rl_env.py index 8564b4a916..5de734c4eb 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/direct_rl_env.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/direct_rl_env.py @@ -14,6 +14,7 @@ import weakref from abc import abstractmethod from collections.abc import Sequence +from dataclasses import MISSING from typing import Any, ClassVar import omni.isaac.core.utils.torch as torch_utils @@ -30,6 +31,7 @@ from .common import VecEnvObs, VecEnvStepReturn from .direct_rl_env_cfg import DirectRLEnvCfg from .ui import ViewportCameraController +from .utils.spaces import sample_space, spec_to_gym_space class DirectRLEnv(gym.Env): @@ -77,6 +79,8 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs RuntimeError: If a simulation context already exists. The environment must always create one since it configures the simulation context and controls the simulation. """ + # check that the config is valid + cfg.validate() # store inputs to class self.cfg = cfg # store the render mode @@ -107,7 +111,7 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs if self.cfg.sim.render_interval < self.cfg.decimation: msg = ( f"The render interval ({self.cfg.sim.render_interval}) is smaller than the decimation " - f"({self.cfg.decimation}). Multiple multiple render calls will happen for each environment step." + f"({self.cfg.decimation}). Multiple render calls will happen for each environment step." "If this is not intended, set the render interval to be equal to the decimation." ) omni.log.warn(msg) @@ -171,7 +175,6 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs self.reset_terminated = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) self.reset_time_outs = torch.zeros_like(self.reset_terminated) self.reset_buf = torch.zeros(self.num_envs, dtype=torch.bool, device=self.sim.device) - self.actions = torch.zeros(self.num_envs, self.cfg.num_actions, device=self.sim.device) # setup the action and observation spaces for Gym self._configure_gym_env_spaces() @@ -270,6 +273,10 @@ def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) indices = torch.arange(self.num_envs, dtype=torch.int64, device=self.device) self._reset_idx(indices) + # update articulation kinematics + self.scene.write_data_to_sim() + self.sim.forward() + # if sensors are added to the scene, make sure we render to reflect changes in reset if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset: self.sim.render() @@ -343,6 +350,9 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1) if len(reset_env_ids) > 0: self._reset_idx(reset_env_ids) + # update articulation kinematics + self.scene.write_data_to_sim() + self.sim.forward() # if sensors are added to the scene, make sure we render to reflect changes in reset if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset: self.sim.render() @@ -507,27 +517,40 @@ def set_debug_vis(self, debug_vis: bool) -> bool: def _configure_gym_env_spaces(self): """Configure the action and observation spaces for the Gym environment.""" - # observation space (unbounded since we don't impose any limits) - self.num_actions = self.cfg.num_actions - self.num_observations = self.cfg.num_observations - self.num_states = self.cfg.num_states + # show deprecation message and overwrite configuration + if self.cfg.num_actions is not None: + omni.log.warn("DirectRLEnvCfg.num_actions is deprecated. Use DirectRLEnvCfg.action_space instead.") + if isinstance(self.cfg.action_space, type(MISSING)): + self.cfg.action_space = self.cfg.num_actions + if self.cfg.num_observations is not None: + omni.log.warn( + "DirectRLEnvCfg.num_observations is deprecated. Use DirectRLEnvCfg.observation_space instead." + ) + if isinstance(self.cfg.observation_space, type(MISSING)): + self.cfg.observation_space = self.cfg.num_observations + if self.cfg.num_states is not None: + omni.log.warn("DirectRLEnvCfg.num_states is deprecated. Use DirectRLEnvCfg.state_space instead.") + if isinstance(self.cfg.state_space, type(MISSING)): + self.cfg.state_space = self.cfg.num_states # set up spaces self.single_observation_space = gym.spaces.Dict() - self.single_observation_space["policy"] = gym.spaces.Box( - low=-np.inf, high=np.inf, shape=(self.num_observations,) - ) - self.single_action_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.num_actions,)) + self.single_observation_space["policy"] = spec_to_gym_space(self.cfg.observation_space) + self.single_action_space = spec_to_gym_space(self.cfg.action_space) # batch the spaces for vectorized environments self.observation_space = gym.vector.utils.batch_space(self.single_observation_space["policy"], self.num_envs) self.action_space = gym.vector.utils.batch_space(self.single_action_space, self.num_envs) # optional state space for asymmetric actor-critic architectures - if self.num_states > 0: - self.single_observation_space["critic"] = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.num_states,)) + self.state_space = None + if self.cfg.state_space: + self.single_observation_space["critic"] = spec_to_gym_space(self.cfg.state_space) self.state_space = gym.vector.utils.batch_space(self.single_observation_space["critic"], self.num_envs) + # instantiate actions (needed for tasks for which the observations computation is dependent on the actions) + self.actions = sample_space(self.single_action_space, self.sim.device, batch_size=self.num_envs, fill_value=0) + def _reset_idx(self, env_ids: Sequence[int]): """Reset environments based on specified indices. @@ -601,7 +624,7 @@ def _get_states(self) -> VecEnvObs | None: """Compute and return the states for the environment. The state-space is used for asymmetric actor-critic architectures. It is configured - using the :attr:`DirectRLEnvCfg.num_states` parameter. + using the :attr:`DirectRLEnvCfg.state_space` parameter. Returns: The states for the environment. If the environment does not have a state-space, the function diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/direct_rl_env_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/direct_rl_env_cfg.py index ad8c6c18c8..4e4f1725c8 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/direct_rl_env_cfg.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/direct_rl_env_cfg.py @@ -10,7 +10,7 @@ from omni.isaac.lab.utils import configclass from omni.isaac.lab.utils.noise import NoiseModelCfg -from .common import ViewerCfg +from .common import SpaceType, ViewerCfg from .ui import BaseEnvWindow @@ -98,19 +98,74 @@ class DirectRLEnvCfg: Please refer to the :class:`omni.isaac.lab.scene.InteractiveSceneCfg` class for more details. """ - events: object = None + events: object | None = None """Event settings. Defaults to None, in which case no events are applied through the event manager. Please refer to the :class:`omni.isaac.lab.managers.EventManager` class for more details. """ - num_observations: int = MISSING - """The dimension of the observation space from each environment instance.""" + observation_space: SpaceType = MISSING + """Observation space definition. + + The space can be defined either using Gymnasium :py:mod:`~gymnasium.spaces` (when a more detailed + specification of the space is desired) or basic Python data types (for simplicity). + + .. list-table:: + :header-rows: 1 + + * - Gymnasium space + - Python data type + * - :class:`~gymnasium.spaces.Box` + - Integer or list of integers (e.g.: ``7``, ``[64, 64, 3]``) + * - :class:`~gymnasium.spaces.Discrete` + - Single-element set (e.g.: ``{2}``) + * - :class:`~gymnasium.spaces.MultiDiscrete` + - List of single-element sets (e.g.: ``[{2}, {5}]``) + * - :class:`~gymnasium.spaces.Dict` + - Dictionary (e.g.: ``{"joints": 7, "rgb": [64, 64, 3], "gripper": {2}}``) + * - :class:`~gymnasium.spaces.Tuple` + - Tuple (e.g.: ``(7, [64, 64, 3], {2})``) + """ + + num_observations: int | None = None + """The dimension of the observation space from each environment instance. + + .. warning:: + + This attribute is deprecated. Use :attr:`~omni.isaac.lab.envs.DirectRLEnvCfg.observation_space` instead. + """ - num_states: int = 0 - """The dimension of the state-space from each environment instance. Default is 0, which means no state-space is defined. + state_space: SpaceType | None = None + """State space definition. This is useful for asymmetric actor-critic and defines the observation space for the critic. + + The space can be defined either using Gymnasium :py:mod:`~gymnasium.spaces` (when a more detailed + specification of the space is desired) or basic Python data types (for simplicity). + + .. list-table:: + :header-rows: 1 + + * - Gymnasium space + - Python data type + * - :class:`~gymnasium.spaces.Box` + - Integer or list of integers (e.g.: ``7``, ``[64, 64, 3]``) + * - :class:`~gymnasium.spaces.Discrete` + - Single-element set (e.g.: ``{2}``) + * - :class:`~gymnasium.spaces.MultiDiscrete` + - List of single-element sets (e.g.: ``[{2}, {5}]``) + * - :class:`~gymnasium.spaces.Dict` + - Dictionary (e.g.: ``{"joints": 7, "rgb": [64, 64, 3], "gripper": {2}}``) + * - :class:`~gymnasium.spaces.Tuple` + - Tuple (e.g.: ``(7, [64, 64, 3], {2})``) + """ + + num_states: int | None = None + """The dimension of the state-space from each environment instance. + + .. warning:: + + This attribute is deprecated. Use :attr:`~omni.isaac.lab.envs.DirectRLEnvCfg.state_space` instead. """ observation_noise_model: NoiseModelCfg | None = None @@ -119,8 +174,36 @@ class DirectRLEnvCfg: Please refer to the :class:`omni.isaac.lab.utils.noise.NoiseModel` class for more details. """ - num_actions: int = MISSING - """The dimension of the action space for each environment.""" + action_space: SpaceType = MISSING + """Action space definition. + + The space can be defined either using Gymnasium :py:mod:`~gymnasium.spaces` (when a more detailed + specification of the space is desired) or basic Python data types (for simplicity). + + .. list-table:: + :header-rows: 1 + + * - Gymnasium space + - Python data type + * - :class:`~gymnasium.spaces.Box` + - Integer or list of integers (e.g.: ``7``, ``[64, 64, 3]``) + * - :class:`~gymnasium.spaces.Discrete` + - Single-element set (e.g.: ``{2}``) + * - :class:`~gymnasium.spaces.MultiDiscrete` + - List of single-element sets (e.g.: ``[{2}, {5}]``) + * - :class:`~gymnasium.spaces.Dict` + - Dictionary (e.g.: ``{"joints": 7, "rgb": [64, 64, 3], "gripper": {2}}``) + * - :class:`~gymnasium.spaces.Tuple` + - Tuple (e.g.: ``(7, [64, 64, 3], {2})``) + """ + + num_actions: int | None = None + """The dimension of the action space for each environment. + + .. warning:: + + This attribute is deprecated. Use :attr:`~omni.isaac.lab.envs.DirectRLEnvCfg.action_space` instead. + """ action_noise_model: NoiseModelCfg | None = None """The noise model applied to the actions provided to the environment. Default is None, which means no noise is added. diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_env.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_env.py index 2e4ac663dc..c2a7b4116c 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_env.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_env.py @@ -11,9 +11,10 @@ import omni.isaac.core.utils.torch as torch_utils import omni.log -from omni.isaac.lab.managers import ActionManager, EventManager, ObservationManager +from omni.isaac.lab.managers import ActionManager, EventManager, ObservationManager, RecorderManager from omni.isaac.lab.scene import InteractiveScene from omni.isaac.lab.sim import SimulationContext +from omni.isaac.lab.ui.widgets import ManagerLiveVisualizer from omni.isaac.lab.utils.timer import Timer from .common import VecEnvObs @@ -45,6 +46,9 @@ class ManagerBasedEnv: This includes resetting the scene to a default state, applying random pushes to the robot at different intervals of time, or randomizing properties such as mass and friction coefficients. This is useful for training and evaluating the robot in a variety of scenarios. + * **Recorder Manager**: The recorder manager that handles recording data produced during different steps + in the simulation. This includes recording in the beginning and end of a reset and a step. The recorded data + is distinguished per episode, per environment and can be exported through a dataset file handler to a file. The environment provides a unified interface for interacting with the simulation. However, it does not include task-specific quantities such as the reward function, or the termination conditions. These @@ -69,6 +73,8 @@ def __init__(self, cfg: ManagerBasedEnvCfg): RuntimeError: If a simulation context already exists. The environment must always create one since it configures the simulation context and controls the simulation. """ + # check that the config is valid + cfg.validate() # store inputs to class self.cfg = cfg # initialize internal variables @@ -103,7 +109,7 @@ def __init__(self, cfg: ManagerBasedEnvCfg): if self.cfg.sim.render_interval < self.cfg.decimation: msg = ( f"The render interval ({self.cfg.sim.render_interval}) is smaller than the decimation " - f"({self.cfg.decimation}). Multiple multiple render calls will happen for each environment step. " + f"({self.cfg.decimation}). Multiple render calls will happen for each environment step. " "If this is not intended, set the render interval to be equal to the decimation." ) omni.log.warn(msg) @@ -143,6 +149,8 @@ def __init__(self, cfg: ManagerBasedEnvCfg): # we need to do this here after all the managers are initialized # this is because they dictate the sensors and commands right now if self.sim.has_gui() and self.cfg.ui_window_class_type is not None: + # setup live visualizers + self.setup_manager_visualizers() self._window = self.cfg.ui_window_class_type(self, window_name="IsaacLab") else: # if no window, then we don't need to store the window @@ -151,6 +159,9 @@ def __init__(self, cfg: ManagerBasedEnvCfg): # allocate dictionary to store metrics self.extras = {} + # initialize observation buffers + self.obs_buf = {} + def __del__(self): """Cleanup for the environment.""" self.close() @@ -206,6 +217,9 @@ def load_managers(self): """ # prepare the managers + # -- recorder manager + self.recorder_manager = RecorderManager(self.cfg.recorders, self) + print("[INFO] Recorder Manager: ", self.recorder_manager) # -- action manager self.action_manager = ActionManager(self.cfg.actions, self) print("[INFO] Action Manager: ", self.action_manager) @@ -222,19 +236,30 @@ def load_managers(self): if self.__class__ == ManagerBasedEnv and "startup" in self.event_manager.available_modes: self.event_manager.apply(mode="startup") + def setup_manager_visualizers(self): + """Creates live visualizers for manager terms.""" + + self.manager_visualizers = { + "action_manager": ManagerLiveVisualizer(manager=self.action_manager), + "observation_manager": ManagerLiveVisualizer(manager=self.observation_manager), + } + """ Operations - MDP. """ - def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) -> tuple[VecEnvObs, dict]: - """Resets all the environments and returns observations. + def reset( + self, seed: int | None = None, env_ids: Sequence[int] | None = None, options: dict[str, Any] | None = None + ) -> tuple[VecEnvObs, dict]: + """Resets the specified environments and returns observations. - This function calls the :meth:`_reset_idx` function to reset all the environments. + This function calls the :meth:`_reset_idx` function to reset the specified environments. However, certain operations, such as procedural terrain generation, that happened during initialization are not repeated. Args: seed: The seed to use for randomization. Defaults to None, in which case the seed is not set. + env_ids: The environment ids to reset. Defaults to None, in which case all environments are reset. options: Additional information to specify how the environment is reset. Defaults to None. Note: @@ -243,20 +268,83 @@ def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) Returns: A tuple containing the observations and extras. """ + if env_ids is None: + env_ids = torch.arange(self.num_envs, dtype=torch.int64, device=self.device) + + # trigger recorder terms for pre-reset calls + self.recorder_manager.record_pre_reset(env_ids) + # set the seed if seed is not None: self.seed(seed) # reset state of scene - indices = torch.arange(self.num_envs, dtype=torch.int64, device=self.device) - self._reset_idx(indices) + self._reset_idx(env_ids) + + # update articulation kinematics + self.scene.write_data_to_sim() + self.sim.forward() + # if sensors are added to the scene, make sure we render to reflect changes in reset + if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset: + self.sim.render() + + # trigger recorder terms for post-reset calls + self.recorder_manager.record_post_reset(env_ids) + + # compute observations + self.obs_buf = self.observation_manager.compute() + + # return observations + return self.obs_buf, self.extras + + def reset_to( + self, + state: dict[str, dict[str, dict[str, torch.Tensor]]], + env_ids: Sequence[int] | None, + seed: int | None = None, + is_relative: bool = False, + ) -> None: + """Resets specified environments to known states. + + Note that this is different from reset() function as it resets the environments to specific states + + Args: + state: The state to reset the specified environments to. + env_ids: The environment ids to reset. Defaults to None, in which case all environments are reset. + seed: The seed to use for randomization. Defaults to None, in which case the seed is not set. + is_relative: If set to True, the state is considered relative to the environment origins. Defaults to False. + """ + # reset all envs in the scene if env_ids is None + if env_ids is None: + env_ids = torch.arange(self.num_envs, dtype=torch.int64, device=self.device) + + # trigger recorder terms for pre-reset calls + self.recorder_manager.record_pre_reset(env_ids) + + # set the seed + if seed is not None: + self.seed(seed) + + self._reset_idx(env_ids) + + # set the state + self.scene.reset_to(state, env_ids, is_relative=is_relative) + + # update articulation kinematics + self.sim.forward() # if sensors are added to the scene, make sure we render to reflect changes in reset if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset: self.sim.render() + # trigger recorder terms for post-reset calls + self.recorder_manager.record_post_reset(env_ids) + + # compute observations + self.obs_buf = self.observation_manager.compute() + # return observations - return self.observation_manager.compute(), self.extras + return self.obs_buf, self.extras def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]: """Execute one time-step of the environment's dynamics. @@ -276,6 +364,8 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]: # process actions self.action_manager.process_action(action.to(self.device)) + self.recorder_manager.record_pre_step() + # check if we need to do rendering within the physics loop # note: checked here once to avoid multiple checks within the loop is_rendering = self.sim.has_gui() or self.sim.has_rtx_sensors() @@ -301,8 +391,12 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]: if "interval" in self.event_manager.available_modes: self.event_manager.apply(mode="interval", dt=self.step_dt) + # -- compute observations + self.obs_buf = self.observation_manager.compute() + self.recorder_manager.record_post_step() + # return observations and extras - return self.observation_manager.compute(), self.extras + return self.obs_buf, self.extras @staticmethod def seed(seed: int = -1) -> int: @@ -332,6 +426,7 @@ def close(self): del self.action_manager del self.observation_manager del self.event_manager + del self.recorder_manager del self.scene # clear callbacks and instance self.sim.clear_all_callbacks() @@ -373,3 +468,6 @@ def _reset_idx(self, env_ids: Sequence[int]): # -- event manager info = self.event_manager.reset(env_ids) self.extras["log"].update(info) + # -- recorder manager + info = self.recorder_manager.reset(env_ids) + self.extras["log"].update(info) diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_env_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_env_cfg.py index 4c6436c9dd..6c64927dd7 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_env_cfg.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_env_cfg.py @@ -13,6 +13,7 @@ import omni.isaac.lab.envs.mdp as mdp from omni.isaac.lab.managers import EventTermCfg as EventTerm +from omni.isaac.lab.managers import RecorderManagerBaseCfg as DefaultEmptyRecorderManagerCfg from omni.isaac.lab.scene import InteractiveSceneCfg from omni.isaac.lab.sim import SimulationCfg from omni.isaac.lab.utils import configclass @@ -78,6 +79,12 @@ class ManagerBasedEnvCfg: Please refer to the :class:`omni.isaac.lab.scene.InteractiveSceneCfg` class for more details. """ + recorders: object = DefaultEmptyRecorderManagerCfg() + """Recorder settings. Defaults to recording nothing. + + Please refer to the :class:`omni.isaac.lab.managers.RecorderManager` class for more details. + """ + observations: object = MISSING """Observation space settings. diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_rl_env.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_rl_env.py index 0b899c0c33..fda4dc11e2 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_rl_env.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_rl_env.py @@ -16,6 +16,7 @@ from omni.isaac.version import get_version from omni.isaac.lab.managers import CommandManager, CurriculumManager, RewardManager, TerminationManager +from omni.isaac.lab.ui.widgets import ManagerLiveVisualizer from .common import VecEnvStepReturn from .manager_based_env import ManagerBasedEnv @@ -132,6 +133,18 @@ def load_managers(self): if "startup" in self.event_manager.available_modes: self.event_manager.apply(mode="startup") + def setup_manager_visualizers(self): + """Creates live visualizers for manager terms.""" + + self.manager_visualizers = { + "action_manager": ManagerLiveVisualizer(manager=self.action_manager), + "observation_manager": ManagerLiveVisualizer(manager=self.observation_manager), + "command_manager": ManagerLiveVisualizer(manager=self.command_manager), + "termination_manager": ManagerLiveVisualizer(manager=self.termination_manager), + "reward_manager": ManagerLiveVisualizer(manager=self.reward_manager), + "curriculum_manager": ManagerLiveVisualizer(manager=self.curriculum_manager), + } + """ Operations - MDP """ @@ -158,6 +171,8 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: # process actions self.action_manager.process_action(action.to(self.device)) + self.recorder_manager.record_pre_step() + # check if we need to do rendering within the physics loop # note: checked here once to avoid multiple checks within the loop is_rendering = self.sim.has_gui() or self.sim.has_rtx_sensors() @@ -190,14 +205,29 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: # -- reward computation self.reward_buf = self.reward_manager.compute(dt=self.step_dt) + if len(self.recorder_manager.active_terms) > 0: + # update observations for recording if needed + self.obs_buf = self.observation_manager.compute() + self.recorder_manager.record_post_step() + # -- reset envs that terminated/timed-out and log the episode information reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1) if len(reset_env_ids) > 0: + # trigger recorder terms for pre-reset calls + self.recorder_manager.record_pre_reset(reset_env_ids) + self._reset_idx(reset_env_ids) + # update articulation kinematics + self.scene.write_data_to_sim() + self.sim.forward() + # if sensors are added to the scene, make sure we render to reflect changes in reset if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset: self.sim.render() + # trigger recorder terms for post-reset calls + self.recorder_manager.record_post_reset(reset_env_ids) + # -- update command self.command_manager.compute(dt=self.step_dt) # -- step interval events @@ -353,6 +383,9 @@ def _reset_idx(self, env_ids: Sequence[int]): # -- termination manager info = self.termination_manager.reset(env_ids) self.extras["log"].update(info) + # -- recorder manager + info = self.recorder_manager.reset(env_ids) + self.extras["log"].update(info) # reset the episode length buffer self.episode_length_buf[env_ids] = 0 diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_rl_env_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_rl_env_cfg.py index b0def63606..93195d4d55 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_rl_env_cfg.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_rl_env_cfg.py @@ -67,14 +67,14 @@ class ManagerBasedRLEnvCfg(ManagerBasedEnvCfg): Please refer to the :class:`omni.isaac.lab.managers.TerminationManager` class for more details. """ - curriculum: object = MISSING - """Curriculum settings. + curriculum: object | None = None + """Curriculum settings. Defaults to None, in which case no curriculum is applied. Please refer to the :class:`omni.isaac.lab.managers.CurriculumManager` class for more details. """ - commands: object = MISSING - """Command settings. + commands: object | None = None + """Command settings. Defaults to None, in which case no commands are generated. Please refer to the :class:`omni.isaac.lab.managers.CommandManager` class for more details. """ diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/__init__.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/__init__.py index f46bf684b6..d02079a0df 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/__init__.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/__init__.py @@ -20,5 +20,6 @@ from .curriculums import * # noqa: F401, F403 from .events import * # noqa: F401, F403 from .observations import * # noqa: F401, F403 +from .recorders import * # noqa: F401, F403 from .rewards import * # noqa: F401, F403 from .terminations import * # noqa: F401, F403 diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/actions_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/actions_cfg.py index 6d07ceb855..09ffd48522 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/actions_cfg.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/actions_cfg.py @@ -5,7 +5,7 @@ from dataclasses import MISSING -from omni.isaac.lab.controllers import DifferentialIKControllerCfg +from omni.isaac.lab.controllers import DifferentialIKControllerCfg, OperationalSpaceControllerCfg from omni.isaac.lab.managers.action_manager import ActionTerm, ActionTermCfg from omni.isaac.lab.utils import configclass @@ -248,3 +248,58 @@ class OffsetCfg: """Scale factor for the action. Defaults to 1.0.""" controller: DifferentialIKControllerCfg = MISSING """The configuration for the differential IK controller.""" + + +@configclass +class OperationalSpaceControllerActionCfg(ActionTermCfg): + """Configuration for operational space controller action term. + + See :class:`OperationalSpaceControllerAction` for more details. + """ + + @configclass + class OffsetCfg: + """The offset pose from parent frame to child frame. + + On many robots, end-effector frames are fictitious frames that do not have a corresponding + rigid body. In such cases, it is easier to define this transform w.r.t. their parent rigid body. + For instance, for the Franka Emika arm, the end-effector is defined at an offset to the the + "panda_hand" frame. + """ + + pos: tuple[float, float, float] = (0.0, 0.0, 0.0) + """Translation w.r.t. the parent frame. Defaults to (0.0, 0.0, 0.0).""" + rot: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0) + """Quaternion rotation ``(w, x, y, z)`` w.r.t. the parent frame. Defaults to (1.0, 0.0, 0.0, 0.0).""" + + class_type: type[ActionTerm] = task_space_actions.OperationalSpaceControllerAction + + joint_names: list[str] = MISSING + """List of joint names or regex expressions that the action will be mapped to.""" + + body_name: str = MISSING + """Name of the body or frame for which motion/force control is performed.""" + + body_offset: OffsetCfg | None = None + """Offset of target frame w.r.t. to the body frame. Defaults to None, in which case no offset is applied.""" + + task_frame_rel_path: str = None + """The path of a ``RigidObject``, relative to the sub-environment, representing task frame. Defaults to None.""" + + controller_cfg: OperationalSpaceControllerCfg = MISSING + """The configuration for the operational space controller.""" + + position_scale: float = 1.0 + """Scale factor for the position targets. Defaults to 1.0.""" + + orientation_scale: float = 1.0 + """Scale factor for the orientation (quad for ``pose_abs`` or axis-angle for ``pose_rel``). Defaults to 1.0.""" + + wrench_scale: float = 1.0 + """Scale factor for the wrench targets. Defaults to 1.0.""" + + stiffness_scale: float = 1.0 + """Scale factor for the stiffness commands. Defaults to 1.0.""" + + damping_ratio_scale: float = 1.0 + """Scale factor for the damping ratio commands. Defaults to 1.0.""" diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/binary_joint_actions.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/binary_joint_actions.py index 5c2ba3fa15..45a7f14f10 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/binary_joint_actions.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/binary_joint_actions.py @@ -40,9 +40,10 @@ class BinaryJointAction(ActionTerm): cfg: actions_cfg.BinaryJointActionCfg """The configuration of the action term.""" - _asset: Articulation """The articulation asset on which the action term is applied.""" + _clip: torch.Tensor + """The clip applied to the input action.""" def __init__(self, cfg: actions_cfg.BinaryJointActionCfg, env: ManagerBasedEnv) -> None: # initialize the action term @@ -83,6 +84,17 @@ def __init__(self, cfg: actions_cfg.BinaryJointActionCfg, env: ManagerBasedEnv) ) self._close_command[index_list] = torch.tensor(value_list, device=self.device) + # parse clip + if self.cfg.clip is not None: + if isinstance(cfg.clip, dict): + self._clip = torch.tensor([[-float("inf"), float("inf")]], device=self.device).repeat( + self.num_envs, self.action_dim, 1 + ) + index_list, _, value_list = string_utils.resolve_matching_names_values(self.cfg.clip, self._joint_names) + self._clip[:, index_list] = torch.tensor(value_list, device=self.device) + else: + raise ValueError(f"Unsupported clip type: {type(cfg.clip)}. Supported types are dict.") + """ Properties. """ @@ -115,6 +127,10 @@ def process_actions(self, actions: torch.Tensor): binary_mask = actions < 0 # compute the command self._processed_actions = torch.where(binary_mask, self._close_command, self._open_command) + if self.cfg.clip is not None: + self._processed_actions = torch.clamp( + self._processed_actions, min=self._clip[:, :, 0], max=self._clip[:, :, 1] + ) def reset(self, env_ids: Sequence[int] | None = None) -> None: self._raw_actions[env_ids] = 0.0 diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/joint_actions.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/joint_actions.py index ee5586b7f2..e2f95987aa 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/joint_actions.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/joint_actions.py @@ -50,6 +50,8 @@ class JointAction(ActionTerm): """The scaling factor applied to the input action.""" _offset: torch.Tensor | float """The offset applied to the input action.""" + _clip: torch.Tensor + """The clip applied to the input action.""" def __init__(self, cfg: actions_cfg.JointActionCfg, env: ManagerBasedEnv) -> None: # initialize the action term @@ -94,6 +96,16 @@ def __init__(self, cfg: actions_cfg.JointActionCfg, env: ManagerBasedEnv) -> Non self._offset[:, index_list] = torch.tensor(value_list, device=self.device) else: raise ValueError(f"Unsupported offset type: {type(cfg.offset)}. Supported types are float and dict.") + # parse clip + if self.cfg.clip is not None: + if isinstance(cfg.clip, dict): + self._clip = torch.tensor([[-float("inf"), float("inf")]], device=self.device).repeat( + self.num_envs, self.action_dim, 1 + ) + index_list, _, value_list = string_utils.resolve_matching_names_values(self.cfg.clip, self._joint_names) + self._clip[:, index_list] = torch.tensor(value_list, device=self.device) + else: + raise ValueError(f"Unsupported clip type: {type(cfg.clip)}. Supported types are dict.") """ Properties. @@ -120,6 +132,11 @@ def process_actions(self, actions: torch.Tensor): self._raw_actions[:] = actions # apply the affine transformations self._processed_actions = self._raw_actions * self._scale + self._offset + # clip actions + if self.cfg.clip is not None: + self._processed_actions = torch.clamp( + self._processed_actions, min=self._clip[:, :, 0], max=self._clip[:, :, 1] + ) def reset(self, env_ids: Sequence[int] | None = None) -> None: self._raw_actions[env_ids] = 0.0 diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/joint_actions_to_limits.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/joint_actions_to_limits.py index 4345b4dec6..b69ac4beda 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/joint_actions_to_limits.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/joint_actions_to_limits.py @@ -44,6 +44,8 @@ class JointPositionToLimitsAction(ActionTerm): """The articulation asset on which the action term is applied.""" _scale: torch.Tensor | float """The scaling factor applied to the input action.""" + _clip: torch.Tensor + """The clip applied to the input action.""" def __init__(self, cfg: actions_cfg.JointPositionToLimitsActionCfg, env: ManagerBasedEnv): # initialize the action term @@ -76,6 +78,16 @@ def __init__(self, cfg: actions_cfg.JointPositionToLimitsActionCfg, env: Manager self._scale[:, index_list] = torch.tensor(value_list, device=self.device) else: raise ValueError(f"Unsupported scale type: {type(cfg.scale)}. Supported types are float and dict.") + # parse clip + if self.cfg.clip is not None: + if isinstance(cfg.clip, dict): + self._clip = torch.tensor([[-float("inf"), float("inf")]], device=self.device).repeat( + self.num_envs, self.action_dim, 1 + ) + index_list, _, value_list = string_utils.resolve_matching_names_values(self.cfg.clip, self._joint_names) + self._clip[:, index_list] = torch.tensor(value_list, device=self.device) + else: + raise ValueError(f"Unsupported clip type: {type(cfg.clip)}. Supported types are dict.") """ Properties. @@ -102,6 +114,10 @@ def process_actions(self, actions: torch.Tensor): self._raw_actions[:] = actions # apply affine transformations self._processed_actions = self._raw_actions * self._scale + if self.cfg.clip is not None: + self._processed_actions = torch.clamp( + self._processed_actions, min=self._clip[:, :, 0], max=self._clip[:, :, 1] + ) # rescale the position targets if configured # this is useful when the input actions are in the range [-1, 1] if self.cfg.rescale_to_limits: @@ -183,6 +199,8 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: # check if specific environment ids are provided if env_ids is None: env_ids = slice(None) + else: + env_ids = env_ids[:, None] super().reset(env_ids) # reset history to current joint positions self._prev_applied_actions[env_ids, :] = self._asset.data.joint_pos[env_ids, self._joint_ids] diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/non_holonomic_actions.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/non_holonomic_actions.py index fc9ed89d6e..4d002810e5 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/non_holonomic_actions.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/non_holonomic_actions.py @@ -11,6 +11,7 @@ import omni.log +import omni.isaac.lab.utils.string as string_utils from omni.isaac.lab.assets.articulation import Articulation from omni.isaac.lab.managers.action_manager import ActionTerm from omni.isaac.lab.utils.math import euler_xyz_from_quat @@ -59,6 +60,8 @@ class NonHolonomicAction(ActionTerm): """The scaling factor applied to the input action. Shape is (1, 2).""" _offset: torch.Tensor """The offset applied to the input action. Shape is (1, 2).""" + _clip: torch.Tensor + """The clip applied to the input action.""" def __init__(self, cfg: actions_cfg.NonHolonomicActionCfg, env: ManagerBasedEnv): # initialize the action term @@ -104,6 +107,16 @@ def __init__(self, cfg: actions_cfg.NonHolonomicActionCfg, env: ManagerBasedEnv) # save the scale and offset as tensors self._scale = torch.tensor(self.cfg.scale, device=self.device).unsqueeze(0) self._offset = torch.tensor(self.cfg.offset, device=self.device).unsqueeze(0) + # parse clip + if self.cfg.clip is not None: + if isinstance(cfg.clip, dict): + self._clip = torch.tensor([[-float("inf"), float("inf")]], device=self.device).repeat( + self.num_envs, self.action_dim, 1 + ) + index_list, _, value_list = string_utils.resolve_matching_names_values(self.cfg.clip, self._joint_names) + self._clip[:, index_list] = torch.tensor(value_list, device=self.device) + else: + raise ValueError(f"Unsupported clip type: {type(cfg.clip)}. Supported types are dict.") """ Properties. @@ -129,10 +142,15 @@ def process_actions(self, actions): # store the raw actions self._raw_actions[:] = actions self._processed_actions = self.raw_actions * self._scale + self._offset + # clip actions + if self.cfg.clip is not None: + self._processed_actions = torch.clamp( + self._processed_actions, min=self._clip[:, :, 0], max=self._clip[:, :, 1] + ) def apply_actions(self): # obtain current heading - quat_w = self._asset.data.body_quat_w[:, self._body_idx] + quat_w = self._asset.data.body_link_quat_w[:, self._body_idx].view(self.num_envs, 4) yaw_w = euler_xyz_from_quat(quat_w)[2] # compute joint velocities targets self._joint_vel_command[:, 0] = torch.cos(yaw_w) * self.processed_actions[:, 0] # x diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/task_space_actions.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/task_space_actions.py index 96154116c8..8a96a2523c 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/task_space_actions.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/task_space_actions.py @@ -10,11 +10,16 @@ from typing import TYPE_CHECKING import omni.log +from pxr import UsdPhysics import omni.isaac.lab.utils.math as math_utils +import omni.isaac.lab.utils.string as string_utils from omni.isaac.lab.assets.articulation import Articulation from omni.isaac.lab.controllers.differential_ik import DifferentialIKController +from omni.isaac.lab.controllers.operational_space import OperationalSpaceController from omni.isaac.lab.managers.action_manager import ActionTerm +from omni.isaac.lab.sensors import ContactSensor, ContactSensorCfg, FrameTransformer, FrameTransformerCfg +from omni.isaac.lab.sim.utils import find_matching_prims if TYPE_CHECKING: from omni.isaac.lab.envs import ManagerBasedEnv @@ -42,6 +47,8 @@ class DifferentialInverseKinematicsAction(ActionTerm): """The articulation asset on which the action term is applied.""" _scale: torch.Tensor """The scaling factor applied to the input action. Shape is (1, action_dim).""" + _clip: torch.Tensor + """The clip applied to the input action.""" def __init__(self, cfg: actions_cfg.DifferentialInverseKinematicsActionCfg, env: ManagerBasedEnv): # initialize the action term @@ -101,6 +108,17 @@ def __init__(self, cfg: actions_cfg.DifferentialInverseKinematicsActionCfg, env: else: self._offset_pos, self._offset_rot = None, None + # parse clip + if self.cfg.clip is not None: + if isinstance(cfg.clip, dict): + self._clip = torch.tensor([[-float("inf"), float("inf")]], device=self.device).repeat( + self.num_envs, self.action_dim, 1 + ) + index_list, _, value_list = string_utils.resolve_matching_names_values(self.cfg.clip, self._joint_names) + self._clip[:, index_list] = torch.tensor(value_list, device=self.device) + else: + raise ValueError(f"Unsupported clip type: {type(cfg.clip)}. Supported types are dict.") + """ Properties. """ @@ -117,6 +135,19 @@ def raw_actions(self) -> torch.Tensor: def processed_actions(self) -> torch.Tensor: return self._processed_actions + @property + def jacobian_w(self) -> torch.Tensor: + return self._asset.root_physx_view.get_jacobians()[:, self._jacobi_body_idx, :, self._jacobi_joint_ids] + + @property + def jacobian_b(self) -> torch.Tensor: + jacobian = self.jacobian_w + base_rot = self._asset.data.root_link_quat_w + base_rot_matrix = math_utils.matrix_from_quat(math_utils.quat_inv(base_rot)) + jacobian[:, :3, :] = torch.bmm(base_rot_matrix, jacobian[:, :3, :]) + jacobian[:, 3:, :] = torch.bmm(base_rot_matrix, jacobian[:, 3:, :]) + return jacobian + """ Operations. """ @@ -125,6 +156,10 @@ def process_actions(self, actions: torch.Tensor): # store the raw actions self._raw_actions[:] = actions self._processed_actions[:] = self.raw_actions * self._scale + if self.cfg.clip is not None: + self._processed_actions = torch.clamp( + self._processed_actions, min=self._clip[:, :, 0], max=self._clip[:, :, 1] + ) # obtain quantities from simulation ee_pos_curr, ee_quat_curr = self._compute_frame_pose() # set command into controller @@ -157,12 +192,12 @@ def _compute_frame_pose(self) -> tuple[torch.Tensor, torch.Tensor]: A tuple of the body's position and orientation in the root frame. """ # obtain quantities from simulation - ee_pose_w = self._asset.data.body_state_w[:, self._body_idx, :7] - root_pose_w = self._asset.data.root_state_w[:, :7] + ee_pos_w = self._asset.data.body_link_pos_w[:, self._body_idx] + ee_quat_w = self._asset.data.body_link_quat_w[:, self._body_idx] + root_pos_w = self._asset.data.root_link_pos_w + root_quat_w = self._asset.data.root_link_quat_w # compute the pose of the body in the root frame - ee_pose_b, ee_quat_b = math_utils.subtract_frame_transforms( - root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_pose_w[:, 0:3], ee_pose_w[:, 3:7] - ) + ee_pose_b, ee_quat_b = math_utils.subtract_frame_transforms(root_pos_w, root_quat_w, ee_pos_w, ee_quat_w) # account for the offset if self.cfg.body_offset is not None: ee_pose_b, ee_quat_b = math_utils.combine_frame_transforms( @@ -178,7 +213,7 @@ def _compute_frame_jacobian(self): the right Jacobian from the parent body Jacobian. """ # read the parent jacobian - jacobian = self._asset.root_physx_view.get_jacobians()[:, self._jacobi_body_idx, :, self._jacobi_joint_ids] + jacobian = self.jacobian_b # account for the offset if self.cfg.body_offset is not None: # Modify the jacobian to account for the offset @@ -192,3 +227,431 @@ def _compute_frame_jacobian(self): jacobian[:, 3:, :] = torch.bmm(math_utils.matrix_from_quat(self._offset_rot), jacobian[:, 3:, :]) return jacobian + + +class OperationalSpaceControllerAction(ActionTerm): + r"""Operational space controller action term. + + This action term performs pre-processing of the raw actions for operational space control. + + """ + + cfg: actions_cfg.OperationalSpaceControllerActionCfg + """The configuration of the action term.""" + _asset: Articulation + """The articulation asset on which the action term is applied.""" + _contact_sensor: ContactSensor = None + """The contact sensor for the end-effector body.""" + _task_frame_transformer: FrameTransformer = None + """The frame transformer for the task frame.""" + + def __init__(self, cfg: actions_cfg.OperationalSpaceControllerActionCfg, env: ManagerBasedEnv): + # initialize the action term + super().__init__(cfg, env) + + self._sim_dt = env.sim.get_physics_dt() + + # resolve the joints over which the action term is applied + self._joint_ids, self._joint_names = self._asset.find_joints(self.cfg.joint_names) + self._num_DoF = len(self._joint_ids) + # parse the ee body index + body_ids, body_names = self._asset.find_bodies(self.cfg.body_name) + if len(body_ids) != 1: + raise ValueError( + f"Expected one match for the ee body name: {self.cfg.body_name}. Found {len(body_ids)}: {body_names}." + ) + # save only the first ee body index + self._ee_body_idx = body_ids[0] + self._ee_body_name = body_names[0] + # check if articulation is fixed-base + # if fixed-base then the jacobian for the base is not computed + # this means that number of bodies is one less than the articulation's number of bodies + if self._asset.is_fixed_base: + self._jacobi_ee_body_idx = self._ee_body_idx - 1 + self._jacobi_joint_idx = self._joint_ids + else: + self._jacobi_ee_body_idx = self._ee_body_idx + self._jacobi_joint_idx = [i + 6 for i in self._joint_ids] + + # log info for debugging + omni.log.info( + f"Resolved joint names for the action term {self.__class__.__name__}:" + f" {self._joint_names} [{self._joint_ids}]" + ) + omni.log.info( + f"Resolved ee body name for the action term {self.__class__.__name__}:" + f" {self._ee_body_name} [{self._ee_body_idx}]" + ) + # Avoid indexing across all joints for efficiency + if self._num_DoF == self._asset.num_joints: + self._joint_ids = slice(None) + + # convert the fixed offsets to torch tensors of batched shape + if self.cfg.body_offset is not None: + self._offset_pos = torch.tensor(self.cfg.body_offset.pos, device=self.device).repeat(self.num_envs, 1) + self._offset_rot = torch.tensor(self.cfg.body_offset.rot, device=self.device).repeat(self.num_envs, 1) + else: + self._offset_pos, self._offset_rot = None, None + + # create contact sensor if any of the command is wrench_abs, and if stiffness is provided + if ( + "wrench_abs" in self.cfg.controller_cfg.target_types + and self.cfg.controller_cfg.contact_wrench_stiffness_task is not None + ): + self._contact_sensor_cfg = ContactSensorCfg(prim_path=self._asset.cfg.prim_path + "/" + self._ee_body_name) + self._contact_sensor = ContactSensor(self._contact_sensor_cfg) + if not self._contact_sensor.is_initialized: + self._contact_sensor._initialize_impl() + self._contact_sensor._is_initialized = True + + # Initialize the task frame transformer if a relative path for the RigidObject, representing the task frame, + # is provided. + if self.cfg.task_frame_rel_path is not None: + # The source RigidObject can be any child of the articulation asset (we will not use it), + # hence, we will use the first RigidObject child. + root_rigidbody_path = self._first_RigidObject_child_path() + task_frame_transformer_path = "/World/envs/env_.*/" + self.cfg.task_frame_rel_path + task_frame_transformer_cfg = FrameTransformerCfg( + prim_path=root_rigidbody_path, + target_frames=[ + FrameTransformerCfg.FrameCfg( + name="task_frame", + prim_path=task_frame_transformer_path, + ), + ], + ) + self._task_frame_transformer = FrameTransformer(task_frame_transformer_cfg) + if not self._task_frame_transformer.is_initialized: + self._task_frame_transformer._initialize_impl() + self._task_frame_transformer._is_initialized = True + # create tensor for task frame pose wrt the root frame + self._task_frame_pose_b = torch.zeros(self.num_envs, 7, device=self.device) + else: + # create an empty reference for task frame pose + self._task_frame_pose_b = None + + # create the operational space controller + self._osc = OperationalSpaceController(cfg=self.cfg.controller_cfg, num_envs=self.num_envs, device=self.device) + + # create tensors for raw and processed actions + self._raw_actions = torch.zeros(self.num_envs, self.action_dim, device=self.device) + self._processed_actions = torch.zeros_like(self.raw_actions) + + # create tensors for the dynamic-related quantities + self._jacobian_b = torch.zeros(self.num_envs, 6, self._num_DoF, device=self.device) + self._mass_matrix = torch.zeros(self.num_envs, self._num_DoF, self._num_DoF, device=self.device) + self._gravity = torch.zeros(self.num_envs, self._num_DoF, device=self.device) + + # create tensors for the ee states + self._ee_pose_w = torch.zeros(self.num_envs, 7, device=self.device) + self._ee_pose_b = torch.zeros(self.num_envs, 7, device=self.device) + self._ee_pose_b_no_offset = torch.zeros(self.num_envs, 7, device=self.device) # The original ee without offset + self._ee_vel_w = torch.zeros(self.num_envs, 6, device=self.device) + self._ee_vel_b = torch.zeros(self.num_envs, 6, device=self.device) + self._ee_force_w = torch.zeros(self.num_envs, 3, device=self.device) # Only the forces are used for now + self._ee_force_b = torch.zeros(self.num_envs, 3, device=self.device) # Only the forces are used for now + + # create the joint effort tensor + self._joint_efforts = torch.zeros(self.num_envs, self._num_DoF, device=self.device) + + # save the scale as tensors + self._position_scale = torch.tensor(self.cfg.position_scale, device=self.device) + self._orientation_scale = torch.tensor(self.cfg.orientation_scale, device=self.device) + self._wrench_scale = torch.tensor(self.cfg.wrench_scale, device=self.device) + self._stiffness_scale = torch.tensor(self.cfg.stiffness_scale, device=self.device) + self._damping_ratio_scale = torch.tensor(self.cfg.damping_ratio_scale, device=self.device) + + # indexes for the various command elements (e.g., pose_rel, stifness, etc.) within the command tensor + self._pose_abs_idx = None + self._pose_rel_idx = None + self._wrench_abs_idx = None + self._stiffness_idx = None + self._damping_ratio_idx = None + self._resolve_command_indexes() + + """ + Properties. + """ + + @property + def action_dim(self) -> int: + """Dimension of the action space of operational space control.""" + return self._osc.action_dim + + @property + def raw_actions(self) -> torch.Tensor: + """Raw actions for operational space control.""" + return self._raw_actions + + @property + def processed_actions(self) -> torch.Tensor: + """Processed actions for operational space control.""" + return self._processed_actions + + @property + def jacobian_w(self) -> torch.Tensor: + return self._asset.root_physx_view.get_jacobians()[:, self._jacobi_ee_body_idx, :, self._jacobi_joint_idx] + + @property + def jacobian_b(self) -> torch.Tensor: + jacobian = self.jacobian_w + base_rot = self._asset.data.root_link_quat_w + base_rot_matrix = math_utils.matrix_from_quat(math_utils.quat_inv(base_rot)) + jacobian[:, :3, :] = torch.bmm(base_rot_matrix, jacobian[:, :3, :]) + jacobian[:, 3:, :] = torch.bmm(base_rot_matrix, jacobian[:, 3:, :]) + return jacobian + + """ + Operations. + """ + + def process_actions(self, actions: torch.Tensor): + """Pre-processes the raw actions and sets them as commands for for operational space control. + + Args: + actions (torch.Tensor): The raw actions for operational space control. It is a tensor of + shape (``num_envs``, ``action_dim``). + """ + + # Update ee pose, which would be used by relative targets (i.e., pose_rel) + self._compute_ee_pose() + + # Update task frame pose w.r.t. the root frame. + self._compute_task_frame_pose() + + # Pre-process the raw actions for operational space control. + self._preprocess_actions(actions) + + # set command into controller + self._osc.set_command( + command=self._processed_actions, + current_ee_pose_b=self._ee_pose_b, + current_task_frame_pose_b=self._task_frame_pose_b, + ) + + def apply_actions(self): + """Computes the joint efforts for operational space control and applies them to the articulation.""" + + # Update the relevant states and dynamical quantities + self._compute_dynamic_quantities() + self._compute_ee_jacobian() + self._compute_ee_pose() + self._compute_ee_velocity() + self._compute_ee_force() + # Calculate the joint efforts + self._joint_efforts[:] = self._osc.compute( + jacobian_b=self._jacobian_b, + current_ee_pose_b=self._ee_pose_b, + current_ee_vel_b=self._ee_vel_b, + current_ee_force_b=self._ee_force_b, + mass_matrix=self._mass_matrix, + gravity=self._gravity, + ) + self._asset.set_joint_effort_target(self._joint_efforts, joint_ids=self._joint_ids) + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Resets the raw actions and the sensors if available. + + Args: + env_ids (Sequence[int] | None): The environment indices to reset. If ``None``, all environments are reset. + """ + self._raw_actions[env_ids] = 0.0 + if self._contact_sensor is not None: + self._contact_sensor.reset(env_ids) + if self._task_frame_transformer is not None: + self._task_frame_transformer.reset(env_ids) + + """ + Helper functions. + + """ + + def _first_RigidObject_child_path(self): + """Finds the first ``RigidObject`` child under the articulation asset. + + Raises: + ValueError: If no child ``RigidObject`` is found under the articulation asset. + + Returns: + str: The path to the first ``RigidObject`` child under the articulation asset. + """ + child_prims = find_matching_prims(self._asset.cfg.prim_path + "/.*") + rigid_child_prim = None + # Loop through the list and stop at the first RigidObject found + for prim in child_prims: + if prim.HasAPI(UsdPhysics.RigidBodyAPI): + rigid_child_prim = prim + break + if rigid_child_prim is None: + raise ValueError("No child rigid body found under the expression: '{self._asset.cfg.prim_path}'/.") + rigid_child_prim_path = rigid_child_prim.GetPath().pathString + # Remove the specific env index from the path string + rigid_child_prim_path = self._asset.cfg.prim_path + "/" + rigid_child_prim_path.split("/")[-1] + return rigid_child_prim_path + + def _resolve_command_indexes(self): + """Resolves the indexes for the various command elements within the command tensor. + + Raises: + ValueError: If any command index is left unresolved. + """ + # First iterate over the target types to find the indexes of the different command elements + cmd_idx = 0 + for target_type in self.cfg.controller_cfg.target_types: + if target_type == "pose_abs": + self._pose_abs_idx = cmd_idx + cmd_idx += 7 + elif target_type == "pose_rel": + self._pose_rel_idx = cmd_idx + cmd_idx += 6 + elif target_type == "wrench_abs": + self._wrench_abs_idx = cmd_idx + cmd_idx += 6 + else: + raise ValueError("Undefined target_type for OSC within OperationalSpaceControllerAction.") + # Then iterate over the impedance parameters depending on the impedance mode + if ( + self.cfg.controller_cfg.impedance_mode == "variable_kp" + or self.cfg.controller_cfg.impedance_mode == "variable" + ): + self._stiffness_idx = cmd_idx + cmd_idx += 6 + if self.cfg.controller_cfg.impedance_mode == "variable": + self._damping_ratio_idx = cmd_idx + cmd_idx += 6 + + # Check if any command is left unresolved + if self.action_dim != cmd_idx: + raise ValueError("Not all command indexes have been resolved.") + + def _compute_dynamic_quantities(self): + """Computes the dynamic quantities for operational space control.""" + + self._mass_matrix[:] = self._asset.root_physx_view.get_mass_matrices()[:, self._joint_ids, :][ + :, :, self._joint_ids + ] + self._gravity[:] = self._asset.root_physx_view.get_generalized_gravity_forces()[:, self._joint_ids] + + def _compute_ee_jacobian(self): + """Computes the geometric Jacobian of the ee body frame in root frame. + + This function accounts for the target frame offset and applies the necessary transformations to obtain + the right Jacobian from the parent body Jacobian. + """ + # Get the Jacobian in root frame + self._jacobian_b[:] = self.jacobian_b + + # account for the offset + if self.cfg.body_offset is not None: + # Modify the jacobian to account for the offset + # -- translational part + # v_link = v_ee + w_ee x r_link_ee = v_J_ee * q + w_J_ee * q x r_link_ee + # = (v_J_ee + w_J_ee x r_link_ee ) * q + # = (v_J_ee - r_link_ee_[x] @ w_J_ee) * q + self._jacobian_b[:, 0:3, :] += torch.bmm(-math_utils.skew_symmetric_matrix(self._offset_pos), self._jacobian_b[:, 3:, :]) # type: ignore + # -- rotational part + # w_link = R_link_ee @ w_ee + self._jacobian_b[:, 3:, :] = torch.bmm(math_utils.matrix_from_quat(self._offset_rot), self._jacobian_b[:, 3:, :]) # type: ignore + + def _compute_ee_pose(self): + """Computes the pose of the ee frame in root frame.""" + # Obtain quantities from simulation + self._ee_pose_w[:, 0:3] = self._asset.data.body_link_pos_w[:, self._ee_body_idx] + self._ee_pose_w[:, 3:7] = self._asset.data.body_link_quat_w[:, self._ee_body_idx] + # Compute the pose of the ee body in the root frame + self._ee_pose_b_no_offset[:, 0:3], self._ee_pose_b_no_offset[:, 3:7] = math_utils.subtract_frame_transforms( + self._asset.data.root_link_pos_w, + self._asset.data.root_link_quat_w, + self._ee_pose_w[:, 0:3], + self._ee_pose_w[:, 3:7], + ) + # Account for the offset + if self.cfg.body_offset is not None: + self._ee_pose_b[:, 0:3], self._ee_pose_b[:, 3:7] = math_utils.combine_frame_transforms( + self._ee_pose_b_no_offset[:, 0:3], self._ee_pose_b_no_offset[:, 3:7], self._offset_pos, self._offset_rot + ) + else: + self._ee_pose_b[:] = self._ee_pose_b_no_offset + + def _compute_ee_velocity(self): + """Computes the velocity of the ee frame in root frame.""" + # Extract end-effector velocity in the world frame + self._ee_vel_w[:] = self._asset.data.body_com_vel_w[:, self._ee_body_idx, :] + # Compute the relative velocity in the world frame + relative_vel_w = self._ee_vel_w - self._asset.data.root_com_vel_w + + # Convert ee velocities from world to root frame + self._ee_vel_b[:, 0:3] = math_utils.quat_rotate_inverse( + self._asset.data.root_link_quat_w, relative_vel_w[:, 0:3] + ) + self._ee_vel_b[:, 3:6] = math_utils.quat_rotate_inverse( + self._asset.data.root_link_quat_w, relative_vel_w[:, 3:6] + ) + + # Account for the offset + if self.cfg.body_offset is not None: + # Compute offset vector in root frame + r_offset_b = math_utils.quat_rotate(self._ee_pose_b_no_offset[:, 3:7], self._offset_pos) + # Adjust the linear velocity to account for the offset + self._ee_vel_b[:, :3] += torch.cross(self._ee_vel_b[:, 3:], r_offset_b, dim=-1) + # Angular velocity is not affected by the offset + + def _compute_ee_force(self): + """Computes the contact forces on the ee frame in root frame.""" + # Obtain contact forces only if the contact sensor is available + if self._contact_sensor is not None: + self._contact_sensor.update(self._sim_dt) + self._ee_force_w[:] = self._contact_sensor.data.net_forces_w[:, 0, :] # type: ignore + # Rotate forces and torques into root frame + self._ee_force_b[:] = math_utils.quat_rotate_inverse(self._asset.data.root_link_quat_w, self._ee_force_w) + + def _compute_task_frame_pose(self): + """Computes the pose of the task frame in root frame.""" + # Update task frame pose if task frame rigidbody is provided + if self._task_frame_transformer is not None and self._task_frame_pose_b is not None: + self._task_frame_transformer.update(self._sim_dt) + # Calculate the pose of the task frame in the root frame + self._task_frame_pose_b[:, :3], self._task_frame_pose_b[:, 3:] = math_utils.subtract_frame_transforms( + self._asset.data.root_link_pos_w, + self._asset.data.root_link_quat_w, + self._task_frame_transformer.data.target_pos_w[:, 0, :], + self._task_frame_transformer.data.target_quat_w[:, 0, :], + ) + + def _preprocess_actions(self, actions: torch.Tensor): + """Pre-processes the raw actions for operational space control. + + Args: + actions (torch.Tensor): The raw actions for operational space control. It is a tensor of + shape (``num_envs``, ``action_dim``). + """ + # Store the raw actions. Please note that the actions contain task space targets + # (in the order of the target_types), and possibly the impedance parameters depending on impedance_mode. + self._raw_actions[:] = actions + # Initialize the processed actions with raw actions. + self._processed_actions[:] = self._raw_actions + # Go through the command types one by one, and apply the pre-processing if needed. + if self._pose_abs_idx is not None: + self._processed_actions[:, self._pose_abs_idx : self._pose_abs_idx + 3] *= self._position_scale + self._processed_actions[:, self._pose_abs_idx + 3 : self._pose_abs_idx + 7] *= self._orientation_scale + if self._pose_rel_idx is not None: + self._processed_actions[:, self._pose_rel_idx : self._pose_rel_idx + 3] *= self._position_scale + self._processed_actions[:, self._pose_rel_idx + 3 : self._pose_rel_idx + 6] *= self._orientation_scale + if self._wrench_abs_idx is not None: + self._processed_actions[:, self._wrench_abs_idx : self._wrench_abs_idx + 6] *= self._wrench_scale + if self._stiffness_idx is not None: + self._processed_actions[:, self._stiffness_idx : self._stiffness_idx + 6] *= self._stiffness_scale + self._processed_actions[:, self._stiffness_idx : self._stiffness_idx + 6] = torch.clamp( + self._processed_actions[:, self._stiffness_idx : self._stiffness_idx + 6], + min=self.cfg.controller_cfg.motion_stiffness_limits_task[0], + max=self.cfg.controller_cfg.motion_stiffness_limits_task[1], + ) + if self._damping_ratio_idx is not None: + self._processed_actions[ + :, self._damping_ratio_idx : self._damping_ratio_idx + 6 + ] *= self._damping_ratio_scale + self._processed_actions[:, self._damping_ratio_idx : self._damping_ratio_idx + 6] = torch.clamp( + self._processed_actions[:, self._damping_ratio_idx : self._damping_ratio_idx + 6], + min=self.cfg.controller_cfg.motion_damping_ratio_limits_task[0], + max=self.cfg.controller_cfg.motion_damping_ratio_limits_task[1], + ) diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/commands_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/commands_cfg.py index d548f554db..d19bea60e2 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/commands_cfg.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/commands_cfg.py @@ -37,29 +37,46 @@ class UniformVelocityCommandCfg(CommandTermCfg): asset_name: str = MISSING """Name of the asset in the environment for which the commands are generated.""" - heading_command: bool = MISSING - """Whether to use heading command or angular velocity command. + + heading_command: bool = False + """Whether to use heading command or angular velocity command. Defaults to False. If True, the angular velocity command is computed from the heading error, where the target heading is sampled uniformly from provided range. Otherwise, the angular velocity command is sampled uniformly from provided range. """ - heading_control_stiffness: float = MISSING - """Scale factor to convert the heading error to angular velocity command.""" - rel_standing_envs: float = MISSING - """Probability threshold for environments where the robots that are standing still.""" - rel_heading_envs: float = MISSING - """Probability threshold for environments where the robots follow the heading-based angular velocity command - (the others follow the sampled angular velocity command).""" + + heading_control_stiffness: float = 1.0 + """Scale factor to convert the heading error to angular velocity command. Defaults to 1.0.""" + + rel_standing_envs: float = 0.0 + """The sampled probability of environments that should be standing still. Defaults to 0.0.""" + + rel_heading_envs: float = 1.0 + """The sampled probability of environments where the robots follow the heading-based angular velocity command + (the others follow the sampled angular velocity command). Defaults to 1.0. + + This parameter is only used if :attr:`heading_command` is True. + """ @configclass class Ranges: """Uniform distribution ranges for the velocity commands.""" - lin_vel_x: tuple[float, float] = MISSING # min max [m/s] - lin_vel_y: tuple[float, float] = MISSING # min max [m/s] - ang_vel_z: tuple[float, float] = MISSING # min max [rad/s] - heading: tuple[float, float] = MISSING # min max [rad] + lin_vel_x: tuple[float, float] = MISSING + """Range for the linear-x velocity command (in m/s).""" + + lin_vel_y: tuple[float, float] = MISSING + """Range for the linear-y velocity command (in m/s).""" + + ang_vel_z: tuple[float, float] = MISSING + """Range for the angular-z velocity command (in rad/s).""" + + heading: tuple[float, float] | None = None + """Range for the heading command (in rad). Defaults to None. + + This parameter is only used if :attr:`~UniformVelocityCommandCfg.heading_command` is True. + """ ranges: Ranges = MISSING """Distribution ranges for the velocity commands.""" @@ -91,15 +108,17 @@ class Ranges: """Normal distribution ranges for the velocity commands.""" mean_vel: tuple[float, float, float] = MISSING - """Mean velocity for the normal distribution. + """Mean velocity for the normal distribution (in m/s). The tuple contains the mean linear-x, linear-y, and angular-z velocity. """ + std_vel: tuple[float, float, float] = MISSING - """Standard deviation for the normal distribution. + """Standard deviation for the normal distribution (in m/s). The tuple contains the standard deviation linear-x, linear-y, and angular-z velocity. """ + zero_prob: tuple[float, float, float] = MISSING """Probability of zero velocity for the normal distribution. @@ -118,6 +137,7 @@ class UniformPoseCommandCfg(CommandTermCfg): asset_name: str = MISSING """Name of the asset in the environment for which the commands are generated.""" + body_name: str = MISSING """Name of the body in the asset for which the commands are generated.""" @@ -131,12 +151,23 @@ class UniformPoseCommandCfg(CommandTermCfg): class Ranges: """Uniform distribution ranges for the pose commands.""" - pos_x: tuple[float, float] = MISSING # min max [m] - pos_y: tuple[float, float] = MISSING # min max [m] - pos_z: tuple[float, float] = MISSING # min max [m] - roll: tuple[float, float] = MISSING # min max [rad] - pitch: tuple[float, float] = MISSING # min max [rad] - yaw: tuple[float, float] = MISSING # min max [rad] + pos_x: tuple[float, float] = MISSING + """Range for the x position (in m).""" + + pos_y: tuple[float, float] = MISSING + """Range for the y position (in m).""" + + pos_z: tuple[float, float] = MISSING + """Range for the z position (in m).""" + + roll: tuple[float, float] = MISSING + """Range for the roll angle (in rad).""" + + pitch: tuple[float, float] = MISSING + """Range for the pitch angle (in rad).""" + + yaw: tuple[float, float] = MISSING + """Range for the yaw angle (in rad).""" ranges: Ranges = MISSING """Ranges for the commands.""" @@ -175,8 +206,10 @@ class Ranges: pos_x: tuple[float, float] = MISSING """Range for the x position (in m).""" + pos_y: tuple[float, float] = MISSING """Range for the y position (in m).""" + heading: tuple[float, float] = MISSING """Heading range for the position commands (in rad). diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/pose_2d_command.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/pose_2d_command.py index f933760e07..f38f5fbd6c 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/pose_2d_command.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/pose_2d_command.py @@ -81,7 +81,9 @@ def command(self) -> torch.Tensor: def _update_metrics(self): # logs data - self.metrics["error_pos_2d"] = torch.norm(self.pos_command_w[:, :2] - self.robot.data.root_pos_w[:, :2], dim=1) + self.metrics["error_pos_2d"] = torch.norm( + self.pos_command_w[:, :2] - self.robot.data.root_link_pos_w[:, :2], dim=1 + ) self.metrics["error_heading"] = torch.abs(wrap_to_pi(self.heading_command_w - self.robot.data.heading_w)) def _resample_command(self, env_ids: Sequence[int]): @@ -95,7 +97,7 @@ def _resample_command(self, env_ids: Sequence[int]): if self.cfg.simple_heading: # set heading command to point towards target - target_vec = self.pos_command_w[env_ids] - self.robot.data.root_pos_w[env_ids] + target_vec = self.pos_command_w[env_ids] - self.robot.data.root_link_pos_w[env_ids] target_direction = torch.atan2(target_vec[:, 1], target_vec[:, 0]) flipped_target_direction = wrap_to_pi(target_direction + torch.pi) @@ -116,8 +118,8 @@ def _resample_command(self, env_ids: Sequence[int]): def _update_command(self): """Re-target the position command to the current root state.""" - target_vec = self.pos_command_w - self.robot.data.root_pos_w[:, :3] - self.pos_command_b[:] = quat_rotate_inverse(yaw_quat(self.robot.data.root_quat_w), target_vec) + target_vec = self.pos_command_w - self.robot.data.root_link_pos_w[:, :3] + self.pos_command_b[:] = quat_rotate_inverse(yaw_quat(self.robot.data.root_link_quat_w), target_vec) self.heading_command_b[:] = wrap_to_pi(self.heading_command_w - self.robot.data.heading_w) def _set_debug_vis_impl(self, debug_vis: bool): @@ -182,7 +184,7 @@ def _resample_command(self, env_ids: Sequence[int]): if self.cfg.simple_heading: # set heading command to point towards target - target_vec = self.pos_command_w[env_ids] - self.robot.data.root_pos_w[env_ids] + target_vec = self.pos_command_w[env_ids] - self.robot.data.root_link_pos_w[env_ids] target_direction = torch.atan2(target_vec[:, 1], target_vec[:, 0]) flipped_target_direction = wrap_to_pi(target_direction + torch.pi) diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/pose_command.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/pose_command.py index 328834540e..4afb50d5ad 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/pose_command.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/pose_command.py @@ -92,8 +92,8 @@ def command(self) -> torch.Tensor: def _update_metrics(self): # transform command from base frame to simulation world frame self.pose_command_w[:, :3], self.pose_command_w[:, 3:] = combine_frame_transforms( - self.robot.data.root_pos_w, - self.robot.data.root_quat_w, + self.robot.data.root_link_pos_w, + self.robot.data.root_link_quat_w, self.pose_command_b[:, :3], self.pose_command_b[:, 3:], ) @@ -101,8 +101,8 @@ def _update_metrics(self): pos_error, rot_error = compute_pose_error( self.pose_command_w[:, :3], self.pose_command_w[:, 3:], - self.robot.data.body_state_w[:, self.body_idx, :3], - self.robot.data.body_state_w[:, self.body_idx, 3:7], + self.robot.data.body_link_state_w[:, self.body_idx, :3], + self.robot.data.body_link_state_w[:, self.body_idx, 3:7], ) self.metrics["position_error"] = torch.norm(pos_error, dim=-1) self.metrics["orientation_error"] = torch.norm(rot_error, dim=-1) @@ -151,5 +151,5 @@ def _debug_vis_callback(self, event): # -- goal pose self.goal_pose_visualizer.visualize(self.pose_command_w[:, :3], self.pose_command_w[:, 3:]) # -- current body pose - body_pose_w = self.robot.data.body_state_w[:, self.body_idx] - self.current_pose_visualizer.visualize(body_pose_w[:, :3], body_pose_w[:, 3:7]) + body_link_state_w = self.robot.data.body_link_state_w[:, self.body_idx] + self.current_pose_visualizer.visualize(body_link_state_w[:, :3], body_link_state_w[:, 3:7]) diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/velocity_command.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/velocity_command.py index 4a35adc5fd..93a4f51d10 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/velocity_command.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/velocity_command.py @@ -11,6 +11,8 @@ from collections.abc import Sequence from typing import TYPE_CHECKING +import omni.log + import omni.isaac.lab.utils.math as math_utils from omni.isaac.lab.assets import Articulation from omni.isaac.lab.managers import CommandTerm @@ -49,10 +51,25 @@ def __init__(self, cfg: UniformVelocityCommandCfg, env: ManagerBasedEnv): Args: cfg: The configuration of the command generator. env: The environment. + + Raises: + ValueError: If the heading command is active but the heading range is not provided. """ # initialize the base class super().__init__(cfg, env) + # check configuration + if self.cfg.heading_command and self.cfg.ranges.heading is None: + raise ValueError( + "The velocity command has heading commands active (heading_command=True) but the `ranges.heading`" + " parameter is set to None." + ) + if self.cfg.ranges.heading and not self.cfg.heading_command: + omni.log.warn( + f"The velocity command has the 'ranges.heading' attribute set to '{self.cfg.ranges.heading}'" + " but the heading command is not active. Consider setting the flag for the heading command to True." + ) + # obtain the robot asset # -- robot self.robot: Articulation = env.scene[cfg.asset_name] @@ -97,10 +114,10 @@ def _update_metrics(self): max_command_step = max_command_time / self._env.step_dt # logs data self.metrics["error_vel_xy"] += ( - torch.norm(self.vel_command_b[:, :2] - self.robot.data.root_lin_vel_b[:, :2], dim=-1) / max_command_step + torch.norm(self.vel_command_b[:, :2] - self.robot.data.root_com_lin_vel_b[:, :2], dim=-1) / max_command_step ) self.metrics["error_vel_yaw"] += ( - torch.abs(self.vel_command_b[:, 2] - self.robot.data.root_ang_vel_b[:, 2]) / max_command_step + torch.abs(self.vel_command_b[:, 2] - self.robot.data.root_com_ang_vel_b[:, 2]) / max_command_step ) def _resample_command(self, env_ids: Sequence[int]): @@ -167,11 +184,11 @@ def _debug_vis_callback(self, event): return # get marker location # -- base state - base_pos_w = self.robot.data.root_pos_w.clone() + base_pos_w = self.robot.data.root_link_pos_w.clone() base_pos_w[:, 2] += 0.5 # -- resolve the scales and quaternions vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2]) - vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(self.robot.data.root_lin_vel_b[:, :2]) + vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(self.robot.data.root_com_lin_vel_b[:, :2]) # display markers self.goal_vel_visualizer.visualize(base_pos_w, vel_des_arrow_quat, vel_des_arrow_scale) self.current_vel_visualizer.visualize(base_pos_w, vel_arrow_quat, vel_arrow_scale) @@ -192,7 +209,7 @@ def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torc zeros = torch.zeros_like(heading_angle) arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle) # convert everything back from base to world frame - base_quat_w = self.robot.data.root_quat_w + base_quat_w = self.robot.data.root_link_quat_w arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat) return arrow_scale, arrow_quat diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/events.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/events.py index 3eaeb650f9..44a58374e1 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/events.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/events.py @@ -300,54 +300,57 @@ def randomize_actuator_gains( .. tip:: For implicit actuators, this function uses CPU tensors to assign the actuator gains into the simulation. In such cases, it is recommended to use this function only during the initialization of the environment. - - Raises: - NotImplementedError: If the joint indices are in explicit motor mode. This operation is currently - not supported for explicit actuator models. """ - # extract the used quantities (to enable type-hinting) + # Extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] - # resolve environment ids + # Resolve environment ids if env_ids is None: env_ids = torch.arange(env.scene.num_envs, device=asset.device) - # resolve joint indices - if asset_cfg.joint_ids == slice(None): - joint_ids_list = range(asset.num_joints) - joint_ids = slice(None) # for optimization purposes - else: - joint_ids_list = asset_cfg.joint_ids - joint_ids = torch.tensor(asset_cfg.joint_ids, dtype=torch.int, device=asset.device) - - # check if none of the joint indices are in explicit motor mode - for joint_index in joint_ids_list: - for act_name, actuator in asset.actuators.items(): - # if joint indices are a slice (i.e., all joints are captured) or the joint index is in the actuator - if actuator.joint_indices == slice(None) or joint_index in actuator.joint_indices: - if not isinstance(actuator, ImplicitActuator): - raise NotImplementedError( - "Event term 'randomize_actuator_stiffness_and_damping' is performed on asset" - f" '{asset_cfg.name}' on the joint '{asset.joint_names[joint_index]}' ('{joint_index}') which" - f" uses an explicit actuator model '{act_name}<{actuator.__class__.__name__}>'. This operation" - " is currently not supported for explicit actuator models." - ) + def randomize(data: torch.Tensor, params: tuple[float, float]) -> torch.Tensor: + return _randomize_prop_by_op( + data, params, dim_0_ids=None, dim_1_ids=actuator_indices, operation=operation, distribution=distribution + ) - # sample joint properties from the given ranges and set into the physics simulation - # -- stiffness - if stiffness_distribution_params is not None: - stiffness = asset.data.default_joint_stiffness.to(asset.device).clone() - stiffness = _randomize_prop_by_op( - stiffness, stiffness_distribution_params, env_ids, joint_ids, operation=operation, distribution=distribution - )[env_ids][:, joint_ids] - asset.write_joint_stiffness_to_sim(stiffness, joint_ids=joint_ids, env_ids=env_ids) - # -- damping - if damping_distribution_params is not None: - damping = asset.data.default_joint_damping.to(asset.device).clone() - damping = _randomize_prop_by_op( - damping, damping_distribution_params, env_ids, joint_ids, operation=operation, distribution=distribution - )[env_ids][:, joint_ids] - asset.write_joint_damping_to_sim(damping, joint_ids=joint_ids, env_ids=env_ids) + # Loop through actuators and randomize gains + for actuator in asset.actuators.values(): + if isinstance(asset_cfg.joint_ids, slice): + # we take all the joints of the actuator + actuator_indices = slice(None) + if isinstance(actuator.joint_indices, slice): + global_indices = slice(None) + else: + global_indices = torch.tensor(actuator.joint_indices, device=asset.device) + elif isinstance(actuator.joint_indices, slice): + # we take the joints defined in the asset config + global_indices = actuator_indices = torch.tensor(asset_cfg.joint_ids, device=asset.device) + else: + # we take the intersection of the actuator joints and the asset config joints + actuator_joint_indices = torch.tensor(actuator.joint_indices, device=asset.device) + asset_joint_ids = torch.tensor(asset_cfg.joint_ids, device=asset.device) + # the indices of the joints in the actuator that have to be randomized + actuator_indices = torch.nonzero(torch.isin(actuator_joint_indices, asset_joint_ids)).view(-1) + if len(actuator_indices) == 0: + continue + # maps actuator indices that have to be randomized to global joint indices + global_indices = actuator_joint_indices[actuator_indices] + # Randomize stiffness + if stiffness_distribution_params is not None: + stiffness = actuator.stiffness[env_ids].clone() + stiffness[:, actuator_indices] = asset.data.default_joint_stiffness[env_ids][:, global_indices].clone() + randomize(stiffness, stiffness_distribution_params) + actuator.stiffness[env_ids] = stiffness + if isinstance(actuator, ImplicitActuator): + asset.write_joint_stiffness_to_sim(stiffness, joint_ids=actuator.joint_indices, env_ids=env_ids) + # Randomize damping + if damping_distribution_params is not None: + damping = actuator.damping[env_ids].clone() + damping[:, actuator_indices] = asset.data.default_joint_damping[env_ids][:, global_indices].clone() + randomize(damping, damping_distribution_params) + actuator.damping[env_ids] = damping + if isinstance(actuator, ImplicitActuator): + asset.write_joint_damping_to_sim(damping, joint_ids=actuator.joint_indices, env_ids=env_ids) def randomize_joint_parameters( @@ -621,13 +624,13 @@ def push_by_setting_velocity( asset: RigidObject | Articulation = env.scene[asset_cfg.name] # velocities - vel_w = asset.data.root_vel_w[env_ids] + vel_w = asset.data.root_com_vel_w[env_ids] # sample random velocities range_list = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] ranges = torch.tensor(range_list, device=asset.device) vel_w[:] = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], vel_w.shape, device=asset.device) # set the velocities into the physics simulation - asset.write_root_velocity_to_sim(vel_w, env_ids=env_ids) + asset.write_root_com_velocity_to_sim(vel_w, env_ids=env_ids) def reset_root_state_uniform( @@ -671,8 +674,8 @@ def reset_root_state_uniform( velocities = root_states[:, 7:13] + rand_samples # set into the physics simulation - asset.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids) - asset.write_root_velocity_to_sim(velocities, env_ids=env_ids) + asset.write_root_link_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids) + asset.write_root_com_velocity_to_sim(velocities, env_ids=env_ids) def reset_root_state_with_random_orientation( @@ -723,8 +726,8 @@ def reset_root_state_with_random_orientation( velocities = root_states[:, 7:13] + rand_samples # set into the physics simulation - asset.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids) - asset.write_root_velocity_to_sim(velocities, env_ids=env_ids) + asset.write_root_link_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids) + asset.write_root_com_velocity_to_sim(velocities, env_ids=env_ids) def reset_root_state_from_terrain( @@ -790,8 +793,8 @@ def reset_root_state_from_terrain( velocities = asset.data.default_root_state[:, 7:13] + rand_samples # set into the physics simulation - asset.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids) - asset.write_root_velocity_to_sim(velocities, env_ids=env_ids) + asset.write_root_link_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids) + asset.write_root_com_velocity_to_sim(velocities, env_ids=env_ids) def reset_joints_by_scale( @@ -911,14 +914,16 @@ def reset_scene_to_default(env: ManagerBasedEnv, env_ids: torch.Tensor): default_root_state = rigid_object.data.default_root_state[env_ids].clone() default_root_state[:, 0:3] += env.scene.env_origins[env_ids] # set into the physics simulation - rigid_object.write_root_state_to_sim(default_root_state, env_ids=env_ids) + rigid_object.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids=env_ids) + rigid_object.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids=env_ids) # articulations for articulation_asset in env.scene.articulations.values(): # obtain default and deal with the offset for env origins default_root_state = articulation_asset.data.default_root_state[env_ids].clone() default_root_state[:, 0:3] += env.scene.env_origins[env_ids] # set into the physics simulation - articulation_asset.write_root_state_to_sim(default_root_state, env_ids=env_ids) + articulation_asset.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids=env_ids) + articulation_asset.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids=env_ids) # obtain default joint positions default_joint_pos = articulation_asset.data.default_joint_pos[env_ids].clone() default_joint_vel = articulation_asset.data.default_joint_vel[env_ids].clone() @@ -967,7 +972,8 @@ def _randomize_prop_by_op( dim_0_ids = slice(None) else: n_dim_0 = len(dim_0_ids) - dim_0_ids = dim_0_ids[:, None] + 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] diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/observations.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/observations.py index c770915337..8df6dbbead 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/observations.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/observations.py @@ -17,11 +17,14 @@ import omni.isaac.lab.utils.math as math_utils from omni.isaac.lab.assets import Articulation, RigidObject from omni.isaac.lab.managers import SceneEntityCfg -from omni.isaac.lab.sensors import Camera, RayCaster, RayCasterCamera, TiledCamera +from omni.isaac.lab.managers.manager_base import ManagerTermBase +from omni.isaac.lab.managers.manager_term_cfg import ObservationTermCfg +from omni.isaac.lab.sensors import Camera, Imu, RayCaster, RayCasterCamera, TiledCamera if TYPE_CHECKING: from omni.isaac.lab.envs import ManagerBasedEnv, ManagerBasedRLEnv + """ Root state. """ @@ -31,21 +34,21 @@ def base_pos_z(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg( """Root height in the simulation world frame.""" # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] - return asset.data.root_pos_w[:, 2].unsqueeze(-1) + return asset.data.root_link_pos_w[:, 2].unsqueeze(-1) def base_lin_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Root linear velocity in the asset's root frame.""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return asset.data.root_lin_vel_b + return asset.data.root_com_lin_vel_b def base_ang_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Root angular velocity in the asset's root frame.""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return asset.data.root_ang_vel_b + return asset.data.root_com_ang_vel_b def projected_gravity(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: @@ -59,7 +62,7 @@ def root_pos_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg( """Asset root position in the environment frame.""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return asset.data.root_pos_w - env.scene.env_origins + return asset.data.root_link_pos_w - env.scene.env_origins def root_quat_w( @@ -74,7 +77,7 @@ def root_quat_w( # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - quat = asset.data.root_quat_w + quat = asset.data.root_link_quat_w # make the quaternion real-part positive if configured return math_utils.quat_unique(quat) if make_quat_unique else quat @@ -83,14 +86,14 @@ def root_lin_vel_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntity """Asset root linear velocity in the environment frame.""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return asset.data.root_lin_vel_w + return asset.data.root_com_lin_vel_w def root_ang_vel_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Asset root angular velocity in the environment frame.""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return asset.data.root_ang_vel_w + return asset.data.root_com_ang_vel_w """ @@ -182,6 +185,52 @@ def body_incoming_wrench(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg) -> tor return link_incoming_forces.view(env.num_envs, -1) +def imu_orientation(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("imu")) -> torch.Tensor: + """Imu sensor orientation in the simulation world frame. + + Args: + env: The environment. + asset_cfg: The SceneEntity associated with an IMU sensor. Defaults to SceneEntityCfg("imu"). + + Returns: + Orientation in the world frame in (w, x, y, z) quaternion form. Shape is (num_envs, 4). + """ + # extract the used quantities (to enable type-hinting) + asset: Imu = env.scene[asset_cfg.name] + # return the orientation quaternion + return asset.data.quat_w + + +def imu_ang_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("imu")) -> torch.Tensor: + """Imu sensor angular velocity w.r.t. environment origin expressed in the sensor frame. + + Args: + env: The environment. + asset_cfg: The SceneEntity associated with an IMU sensor. Defaults to SceneEntityCfg("imu"). + + Returns: + The angular velocity (rad/s) in the sensor frame. Shape is (num_envs, 3). + """ + # extract the used quantities (to enable type-hinting) + asset: Imu = env.scene[asset_cfg.name] + # return the angular velocity + return asset.data.ang_vel_b + + +def imu_lin_acc(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("imu")) -> torch.Tensor: + """Imu sensor linear acceleration w.r.t. the environment origin expressed in sensor frame. + + Args: + env: The environment. + asset_cfg: The SceneEntity associated with an IMU sensor. Defaults to SceneEntityCfg("imu"). + + Returns: + The linear acceleration (m/s^2) in the sensor frame. Shape is (num_envs, 3). + """ + asset: Imu = env.scene[asset_cfg.name] + return asset.data.lin_acc_b + + def image( env: ManagerBasedEnv, sensor_cfg: SceneEntityCfg = SceneEntityCfg("tiled_camera"), @@ -231,6 +280,230 @@ def image( return images.clone() +class image_features(ManagerTermBase): + """Extracted image features from a pre-trained frozen encoder. + + This term uses models from the model zoo in PyTorch and extracts features from the images. + + It calls the :func:`image` function to get the images and then processes them using the model zoo. + + A user can provide their own model zoo configuration to use different models for feature extraction. + The model zoo configuration should be a dictionary that maps different model names to a dictionary + that defines the model, preprocess and inference functions. The dictionary should have the following + entries: + + - "model": A callable that returns the model when invoked without arguments. + - "reset": A callable that resets the model. This is useful when the model has a state that needs to be reset. + - "inference": A callable that, when given the model and the images, returns the extracted features. + + If the model zoo configuration is not provided, the default model zoo configurations are used. The default + model zoo configurations include the models from Theia :cite:`shang2024theia` and ResNet :cite:`he2016deep`. + These models are loaded from `Hugging-Face transformers `_ and + `PyTorch torchvision `_ respectively. + + Args: + sensor_cfg: The sensor configuration to poll. Defaults to SceneEntityCfg("tiled_camera"). + data_type: The sensor data type. Defaults to "rgb". + convert_perspective_to_orthogonal: Whether to orthogonalize perspective depth images. + This is used only when the data type is "distance_to_camera". Defaults to False. + model_zoo_cfg: A user-defined dictionary that maps different model names to their respective configurations. + Defaults to None. If None, the default model zoo configurations are used. + model_name: The name of the model to use for inference. Defaults to "resnet18". + model_device: The device to store and infer the model on. This is useful when offloading the computation + from the environment simulation device. Defaults to the environment device. + inference_kwargs: Additional keyword arguments to pass to the inference function. Defaults to None, + which means no additional arguments are passed. + + Returns: + The extracted features tensor. Shape is (num_envs, feature_dim). + + Raises: + ValueError: When the model name is not found in the provided model zoo configuration. + ValueError: When the model name is not found in the default model zoo configuration. + """ + + def __init__(self, cfg: ObservationTermCfg, env: ManagerBasedEnv): + # initialize the base class + super().__init__(cfg, env) + + # extract parameters from the configuration + self.model_zoo_cfg: dict = cfg.params.get("model_zoo_cfg") # type: ignore + self.model_name: str = cfg.params.get("model_name", "resnet18") # type: ignore + self.model_device: str = cfg.params.get("model_device", env.device) # type: ignore + + # List of Theia models - These are configured through `_prepare_theia_transformer_model` function + default_theia_models = [ + "theia-tiny-patch16-224-cddsv", + "theia-tiny-patch16-224-cdiv", + "theia-small-patch16-224-cdiv", + "theia-base-patch16-224-cdiv", + "theia-small-patch16-224-cddsv", + "theia-base-patch16-224-cddsv", + ] + # List of ResNet models - These are configured through `_prepare_resnet_model` function + default_resnet_models = ["resnet18", "resnet34", "resnet50", "resnet101"] + + # Check if model name is specified in the model zoo configuration + if self.model_zoo_cfg is not None and self.model_name not in self.model_zoo_cfg: + raise ValueError( + f"Model name '{self.model_name}' not found in the provided model zoo configuration." + " Please add the model to the model zoo configuration or use a different model name." + f" Available models in the provided list: {list(self.model_zoo_cfg.keys())}." + "\nHint: If you want to use a default model, consider using one of the following models:" + f" {default_theia_models + default_resnet_models}. In this case, you can remove the" + " 'model_zoo_cfg' parameter from the observation term configuration." + ) + if self.model_zoo_cfg is None: + if self.model_name in default_theia_models: + model_config = self._prepare_theia_transformer_model(self.model_name, self.model_device) + elif self.model_name in default_resnet_models: + model_config = self._prepare_resnet_model(self.model_name, self.model_device) + else: + raise ValueError( + f"Model name '{self.model_name}' not found in the default model zoo configuration." + f" Available models: {default_theia_models + default_resnet_models}." + ) + else: + model_config = self.model_zoo_cfg[self.model_name] + + # Retrieve the model, preprocess and inference functions + self._model = model_config["model"]() + self._reset_fn = model_config.get("reset") + self._inference_fn = model_config["inference"] + + def reset(self, env_ids: torch.Tensor | None = None): + # reset the model if a reset function is provided + # this might be useful when the model has a state that needs to be reset + # for example: video transformers + if self._reset_fn is not None: + self._reset_fn(self._model, env_ids) + + def __call__( + self, + env: ManagerBasedEnv, + sensor_cfg: SceneEntityCfg = SceneEntityCfg("tiled_camera"), + data_type: str = "rgb", + convert_perspective_to_orthogonal: bool = False, + model_zoo_cfg: dict | None = None, + model_name: str = "resnet18", + model_device: str | None = None, + inference_kwargs: dict | None = None, + ) -> torch.Tensor: + # obtain the images from the sensor + image_data = image( + env=env, + sensor_cfg=sensor_cfg, + data_type=data_type, + convert_perspective_to_orthogonal=convert_perspective_to_orthogonal, + normalize=False, # we pre-process based on model + ) + # store the device of the image + image_device = image_data.device + # forward the images through the model + features = self._inference_fn(self._model, image_data, **(inference_kwargs or {})) + + # move the features back to the image device + return features.detach().to(image_device) + + """ + Helper functions. + """ + + def _prepare_theia_transformer_model(self, model_name: str, model_device: str) -> dict: + """Prepare the Theia transformer model for inference. + + Args: + model_name: The name of the Theia transformer model to prepare. + model_device: The device to store and infer the model on. + + Returns: + A dictionary containing the model and inference functions. + """ + from transformers import AutoModel + + def _load_model() -> torch.nn.Module: + """Load the Theia transformer model.""" + model = AutoModel.from_pretrained(f"theaiinstitute/{model_name}", trust_remote_code=True).eval() + return model.to(model_device) + + def _inference(model, images: torch.Tensor) -> torch.Tensor: + """Inference the Theia transformer model. + + Args: + model: The Theia transformer model. + images: The preprocessed image tensor. Shape is (num_envs, height, width, channel). + + Returns: + The extracted features tensor. Shape is (num_envs, feature_dim). + """ + # Move the image to the model device + image_proc = images.to(model_device) + # permute the image to (num_envs, channel, height, width) + image_proc = image_proc.permute(0, 3, 1, 2).float() / 255.0 + # Normalize the image + mean = torch.tensor([0.485, 0.456, 0.406], device=model_device).view(1, 3, 1, 1) + std = torch.tensor([0.229, 0.224, 0.225], device=model_device).view(1, 3, 1, 1) + image_proc = (image_proc - mean) / std + + # Taken from Transformers; inference converted to be GPU only + features = model.backbone.model(pixel_values=image_proc, interpolate_pos_encoding=True) + return features.last_hidden_state[:, 1:] + + # return the model, preprocess and inference functions + return {"model": _load_model, "inference": _inference} + + def _prepare_resnet_model(self, model_name: str, model_device: str) -> dict: + """Prepare the ResNet model for inference. + + Args: + model_name: The name of the ResNet model to prepare. + model_device: The device to store and infer the model on. + + Returns: + A dictionary containing the model and inference functions. + """ + from torchvision import models + + def _load_model() -> torch.nn.Module: + """Load the ResNet model.""" + # map the model name to the weights + resnet_weights = { + "resnet18": "ResNet18_Weights.IMAGENET1K_V1", + "resnet34": "ResNet34_Weights.IMAGENET1K_V1", + "resnet50": "ResNet50_Weights.IMAGENET1K_V1", + "resnet101": "ResNet101_Weights.IMAGENET1K_V1", + } + + # load the model + model = getattr(models, model_name)(weights=resnet_weights[model_name]).eval() + return model.to(model_device) + + def _inference(model, images: torch.Tensor) -> torch.Tensor: + """Inference the ResNet model. + + Args: + model: The ResNet model. + images: The preprocessed image tensor. Shape is (num_envs, channel, height, width). + + Returns: + The extracted features tensor. Shape is (num_envs, feature_dim). + """ + # move the image to the model device + image_proc = images.to(model_device) + # permute the image to (num_envs, channel, height, width) + image_proc = image_proc.permute(0, 3, 1, 2).float() / 255.0 + # normalize the image + mean = torch.tensor([0.485, 0.456, 0.406], device=model_device).view(1, 3, 1, 1) + std = torch.tensor([0.229, 0.224, 0.225], device=model_device).view(1, 3, 1, 1) + image_proc = (image_proc - mean) / std + + # forward the image through the model + return model(image_proc) + + # return the model, preprocess and inference functions + return {"model": _load_model, "inference": _inference} + + """ Actions. """ diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/recorders/__init__.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/recorders/__init__.py new file mode 100644 index 0000000000..207a732367 --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/recorders/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Various recorder terms that can be used in the environment.""" + +from .recorders import * +from .recorders_cfg import * diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/recorders/recorders.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/recorders/recorders.py new file mode 100644 index 0000000000..026c22f9a3 --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/recorders/recorders.py @@ -0,0 +1,44 @@ +# Copyright (c) 2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from collections.abc import Sequence + +from omni.isaac.lab.managers.recorder_manager import RecorderTerm + + +class InitialStateRecorder(RecorderTerm): + """Recorder term that records the initial state of the environment after reset.""" + + def record_post_reset(self, env_ids: Sequence[int] | None): + def extract_env_ids_values(value): + nonlocal env_ids + if isinstance(value, dict): + return {k: extract_env_ids_values(v) for k, v in value.items()} + return value[env_ids] + + return "initial_state", extract_env_ids_values(self._env.scene.get_state(is_relative=True)) + + +class PostStepStatesRecorder(RecorderTerm): + """Recorder term that records the state of the environment at the end of each step.""" + + def record_post_step(self): + return "states", self._env.scene.get_state(is_relative=True) + + +class PreStepActionsRecorder(RecorderTerm): + """Recorder term that records the actions in the beginning of each step.""" + + def record_pre_step(self): + return "actions", self._env.action_manager.action + + +class PreStepFlatPolicyObservationsRecorder(RecorderTerm): + """Recorder term that records the policy group observations in each step.""" + + def record_pre_step(self): + return "obs", self._env.obs_buf["policy"] diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/recorders/recorders_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/recorders/recorders_cfg.py new file mode 100644 index 0000000000..16c1eee34a --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/recorders/recorders_cfg.py @@ -0,0 +1,56 @@ +# Copyright (c) 2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from omni.isaac.lab.managers.recorder_manager import RecorderManagerBaseCfg, RecorderTerm, RecorderTermCfg +from omni.isaac.lab.utils import configclass + +from . import recorders + +## +# State recorders. +## + + +@configclass +class InitialStateRecorderCfg(RecorderTermCfg): + """Configuration for the initial state recorder term.""" + + class_type: type[RecorderTerm] = recorders.InitialStateRecorder + + +@configclass +class PostStepStatesRecorderCfg(RecorderTermCfg): + """Configuration for the step state recorder term.""" + + class_type: type[RecorderTerm] = recorders.PostStepStatesRecorder + + +@configclass +class PreStepActionsRecorderCfg(RecorderTermCfg): + """Configuration for the step action recorder term.""" + + class_type: type[RecorderTerm] = recorders.PreStepActionsRecorder + + +@configclass +class PreStepFlatPolicyObservationsRecorderCfg(RecorderTermCfg): + """Configuration for the step policy observation recorder term.""" + + class_type: type[RecorderTerm] = recorders.PreStepFlatPolicyObservationsRecorder + + +## +# Recorder manager configurations. +## + + +@configclass +class ActionStateRecorderManagerCfg(RecorderManagerBaseCfg): + """Recorder configurations for recording actions and states.""" + + record_initial_state = InitialStateRecorderCfg() + record_post_step_states = PostStepStatesRecorderCfg() + record_pre_step_actions = PreStepActionsRecorderCfg() + record_pre_step_flat_policy_observations = PreStepFlatPolicyObservationsRecorderCfg() diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/rewards.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/rewards.py index 7f03bf41f0..cd5564a08b 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/rewards.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/rewards.py @@ -18,7 +18,7 @@ from omni.isaac.lab.managers import SceneEntityCfg from omni.isaac.lab.managers.manager_base import ManagerTermBase from omni.isaac.lab.managers.manager_term_cfg import RewardTermCfg -from omni.isaac.lab.sensors import ContactSensor +from omni.isaac.lab.sensors import ContactSensor, RayCaster if TYPE_CHECKING: from omni.isaac.lab.envs import ManagerBasedRLEnv @@ -77,14 +77,14 @@ def lin_vel_z_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntity """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] - return torch.square(asset.data.root_lin_vel_b[:, 2]) + return torch.square(asset.data.root_com_lin_vel_b[:, 2]) 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] - return torch.sum(torch.square(asset.data.root_ang_vel_b[:, :2]), dim=1) + return torch.sum(torch.square(asset.data.root_com_ang_vel_b[:, :2]), dim=1) def flat_orientation_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: @@ -98,17 +98,28 @@ def flat_orientation_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = Scen def base_height_l2( - env: ManagerBasedRLEnv, target_height: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") + 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: - Currently, it assumes a flat terrain, i.e. the target height is in the world frame. + 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] - # TODO: Fix this for rough-terrain. - return torch.square(asset.data.root_pos_w[:, 2] - target_height) + if sensor_cfg is not None: + sensor: RayCaster = env.scene[sensor_cfg.name] + # Adjust the target height using the sensor data + adjusted_target_height = target_height + sensor.data.pos_w[:, 2] + else: + # Use the provided target height directly for flat terrain + adjusted_target_height = target_height + # Compute the L2 squared penalty + return torch.square(asset.data.root_link_pos_w[:, 2] - adjusted_target_height) def body_lin_acc_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: @@ -159,7 +170,7 @@ def joint_acc_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntity return torch.sum(torch.square(asset.data.joint_acc[:, asset_cfg.joint_ids]), dim=1) -def joint_deviation_l1(env, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: +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] @@ -281,7 +292,7 @@ def track_lin_vel_xy_exp( 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]), + torch.square(env.command_manager.get_command(command_name)[:, :2] - asset.data.root_com_lin_vel_b[:, :2]), dim=1, ) return torch.exp(-lin_vel_error / std**2) @@ -294,5 +305,7 @@ def track_ang_vel_z_exp( # 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]) + ang_vel_error = torch.square( + env.command_manager.get_command(command_name)[:, 2] - asset.data.root_com_ang_vel_b[:, 2] + ) return torch.exp(-ang_vel_error / std**2) diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/terminations.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/terminations.py index a91bb39581..8f4017597b 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/terminations.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/terminations.py @@ -69,7 +69,7 @@ def root_height_below_minimum( """ # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return asset.data.root_pos_w[:, 2] < minimum_height + return asset.data.root_link_pos_w[:, 2] < minimum_height """ diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/ui/base_env_window.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/ui/base_env_window.py index a37a8723ec..4ae47df8e3 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/ui/base_env_window.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/ui/base_env_window.py @@ -16,6 +16,8 @@ import omni.usd from pxr import PhysxSchema, Sdf, Usd, UsdGeom, UsdPhysics +from omni.isaac.lab.ui.widgets import ManagerLiveVisualizer + if TYPE_CHECKING: import omni.ui @@ -57,6 +59,9 @@ def __init__(self, env: ManagerBasedEnv, window_name: str = "IsaacLab"): *self.env.scene.articulations.keys(), ] + # Listeners for environment selection changes + self._ui_listeners: list[ManagerLiveVisualizer] = [] + print("Creating window for environment.") # create window for UI self.ui_window = omni.ui.Window( @@ -80,6 +85,10 @@ def __init__(self, env: ManagerBasedEnv, window_name: str = "IsaacLab"): self._build_viewer_frame() # create collapsable frame for debug visualization self._build_debug_vis_frame() + with self.ui_window_elements["debug_frame"]: + with self.ui_window_elements["debug_vstack"]: + self._visualize_manager(title="Actions", class_name="action_manager") + self._visualize_manager(title="Observations", class_name="observation_manager") def __del__(self): """Destructor for the window.""" @@ -200,9 +209,6 @@ def _build_debug_vis_frame(self): that has it implemented. If the element does not have a debug visualization implemented, a label is created instead. """ - # import omni.isaac.ui.ui_utils as ui_utils - # import omni.ui - # create collapsable frame for debug visualization self.ui_window_elements["debug_frame"] = omni.ui.CollapsableFrame( title="Scene Debug Visualization", @@ -234,6 +240,26 @@ def _build_debug_vis_frame(self): if elem is not None: self._create_debug_vis_ui_element(name, elem) + def _visualize_manager(self, title: str, class_name: str) -> None: + """Checks if the attribute with the name 'class_name' can be visualized. If yes, create vis interface. + + Args: + title: The title of the manager visualization frame. + class_name: The name of the manager to visualize. + """ + + if hasattr(self.env, class_name) and class_name in self.env.manager_visualizers: + manager = self.env.manager_visualizers[class_name] + if hasattr(manager, "has_debug_vis_implementation"): + self._create_debug_vis_ui_element(title, manager) + else: + print( + f"ManagerLiveVisualizer cannot be created for manager: {class_name}, has_debug_vis_implementation" + " does not exist" + ) + else: + print(f"ManagerLiveVisualizer cannot be created for manager: {class_name}, Manager does not exist") + """ Custom callbacks for UI elements. """ @@ -357,6 +383,9 @@ def _set_viewer_env_index_fn(self, model: omni.ui.SimpleIntModel): raise ValueError("Viewport camera controller is not initialized! Please check the rendering mode.") # store the desired env index, UI is 1-indexed vcc.set_view_env_index(model.as_int - 1) + # notify additional listeners + for listener in self._ui_listeners: + listener.set_env_selection(model.as_int - 1) """ Helper functions - UI building. @@ -379,14 +408,30 @@ def _create_debug_vis_ui_element(self, name: str, elem: object): alignment=omni.ui.Alignment.LEFT_CENTER, tooltip=text, ) + has_cfg = hasattr(elem, "cfg") and elem.cfg is not None + is_checked = False + if has_cfg: + is_checked = (hasattr(elem.cfg, "debug_vis") and elem.cfg.debug_vis) or ( + hasattr(elem, "debug_vis") and elem.debug_vis + ) self.ui_window_elements[f"{name}_cb"] = SimpleCheckBox( model=omni.ui.SimpleBoolModel(), enabled=elem.has_debug_vis_implementation, - checked=elem.cfg.debug_vis, + checked=is_checked, on_checked_fn=lambda value, e=weakref.proxy(elem): e.set_debug_vis(value), ) omni.isaac.ui.ui_utils.add_line_rect_flourish() + # Create a panel for the debug visualization + if isinstance(elem, ManagerLiveVisualizer): + self.ui_window_elements[f"{name}_panel"] = omni.ui.Frame(width=omni.ui.Fraction(1)) + if not elem.set_vis_frame(self.ui_window_elements[f"{name}_panel"]): + print(f"Frame failed to set for ManagerLiveVisualizer: {name}") + + # Add listener for environment selection changes + if isinstance(elem, ManagerLiveVisualizer): + self._ui_listeners.append(elem) + async def _dock_window(self, window_title: str): """Docks the custom UI window to the property window.""" # wait for the window to be created diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/ui/manager_based_rl_env_window.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/ui/manager_based_rl_env_window.py index 3e149e1fef..9fa69c4468 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/ui/manager_based_rl_env_window.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/ui/manager_based_rl_env_window.py @@ -34,5 +34,7 @@ def __init__(self, env: ManagerBasedRLEnv, window_name: str = "IsaacLab"): with self.ui_window_elements["main_vstack"]: with self.ui_window_elements["debug_frame"]: with self.ui_window_elements["debug_vstack"]: - self._create_debug_vis_ui_element("commands", self.env.command_manager) - self._create_debug_vis_ui_element("actions", self.env.action_manager) + self._visualize_manager(title="Commands", class_name="command_manager") + self._visualize_manager(title="Rewards", class_name="reward_manager") + self._visualize_manager(title="Curriculum", class_name="curriculum_manager") + self._visualize_manager(title="Termination", class_name="termination_manager") diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/ui/viewport_camera_controller.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/ui/viewport_camera_controller.py index a1fef9f445..bf7463c9ba 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/ui/viewport_camera_controller.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/ui/viewport_camera_controller.py @@ -156,7 +156,7 @@ def update_view_to_asset_root(self, asset_name: str): # set origin type to asset_root self.cfg.origin_type = "asset_root" # update the camera origins - self.viewer_origin = self._env.scene[self.cfg.asset_name].data.root_pos_w[self.cfg.env_index] + self.viewer_origin = self._env.scene[self.cfg.asset_name].data.root_link_pos_w[self.cfg.env_index] # update the camera view self.update_view_location() diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/utils/__init__.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/utils/__init__.py new file mode 100644 index 0000000000..913e1edb90 --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/utils/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-package for environment utils.""" diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/utils.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/utils/marl.py similarity index 76% rename from source/extensions/omni.isaac.lab/omni/isaac/lab/envs/utils.py rename to source/extensions/omni.isaac.lab/omni/isaac/lab/envs/utils/marl.py index cacbdeaf81..46519048ae 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/utils.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/utils/marl.py @@ -9,9 +9,9 @@ import torch from typing import Any -from .common import ActionType, AgentID, EnvStepReturn, ObsType, StateType, VecEnvObs, VecEnvStepReturn -from .direct_marl_env import DirectMARLEnv -from .direct_rl_env import DirectRLEnv +from ..common import ActionType, AgentID, EnvStepReturn, ObsType, StateType, VecEnvObs, VecEnvStepReturn +from ..direct_marl_env import DirectMARLEnv +from ..direct_rl_env import DirectRLEnv def multi_agent_to_single_agent(env: DirectMARLEnv, state_as_observation: bool = False) -> DirectRLEnv: @@ -39,7 +39,7 @@ def multi_agent_to_single_agent(env: DirectMARLEnv, state_as_observation: bool = Raises: AssertionError: If the environment state cannot be used as observation since it was explicitly defined - as unconstructed (:attr:`DirectMARLEnvCfg.num_states`). + as unconstructed (:attr:`DirectMARLEnvCfg.state_space`). """ class Env(DirectRLEnv): @@ -49,7 +49,7 @@ def __init__(self, env: DirectMARLEnv) -> None: # check if it is possible to use the multi-agent environment state as single-agent observation self._state_as_observation = state_as_observation if self._state_as_observation: - assert self.env.cfg.num_states != 0, ( + assert self.env.cfg.state_space != 0, ( "The environment state cannot be used as observation since it was explicitly defined as" " unconstructed" ) @@ -58,18 +58,17 @@ def __init__(self, env: DirectMARLEnv) -> None: self.cfg = self.env.cfg self.sim = self.env.sim self.scene = self.env.scene - self.num_actions = sum(self.env.cfg.num_actions.values()) - self.num_observations = sum(self.env.cfg.num_observations.values()) - self.num_states = self.env.cfg.num_states self.single_observation_space = gym.spaces.Dict() if self._state_as_observation: self.single_observation_space["policy"] = self.env.state_space else: - self.single_observation_space["policy"] = gym.spaces.Box( - low=-np.inf, high=np.inf, shape=(self.num_observations,) + self.single_observation_space["policy"] = gym.spaces.flatten_space( + gym.spaces.Tuple([self.env.observation_spaces[agent] for agent in self.env.possible_agents]) ) - self.single_action_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.num_actions,)) + self.single_action_space = gym.spaces.flatten_space( + gym.spaces.Tuple([self.env.action_spaces[agent] for agent in self.env.possible_agents]) + ) # batch the spaces for vectorized environments self.observation_space = gym.vector.utils.batch_space( @@ -84,18 +83,25 @@ def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) if self._state_as_observation: obs = {"policy": self.env.state()} # concatenate agents' observations + # FIXME: This implementation assumes the spaces are fundamental ones. Fix it to support composite spaces else: - obs = {"policy": torch.cat([obs[agent] for agent in self.env.possible_agents], dim=-1)} + obs = { + "policy": torch.cat( + [obs[agent].reshape(self.num_envs, -1) for agent in self.env.possible_agents], dim=-1 + ) + } return obs, extras def step(self, action: torch.Tensor) -> VecEnvStepReturn: # split single-agent actions to build the multi-agent ones + # FIXME: This implementation assumes the spaces are fundamental ones. Fix it to support composite spaces index = 0 _actions = {} for agent in self.env.possible_agents: - _actions[agent] = action[:, index : index + self.env.cfg.num_actions[agent]] - index += self.env.cfg.num_actions[agent] + delta = gym.spaces.flatdim(self.env.action_spaces[agent]) + _actions[agent] = action[:, index : index + delta] + index += delta # step the environment obs, rewards, terminated, time_outs, extras = self.env.step(_actions) @@ -104,8 +110,13 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: if self._state_as_observation: obs = {"policy": self.env.state()} # concatenate agents' observations + # FIXME: This implementation assumes the spaces are fundamental ones. Fix it to support composite spaces else: - obs = {"policy": torch.cat([obs[agent] for agent in self.env.possible_agents], dim=-1)} + obs = { + "policy": torch.cat( + [obs[agent].reshape(self.num_envs, -1) for agent in self.env.possible_agents], dim=-1 + ) + } # process environment outputs to return single-agent data rewards = sum(rewards.values()) @@ -147,7 +158,7 @@ def multi_agent_with_one_agent(env: DirectMARLEnv, state_as_observation: bool = Raises: AssertionError: If the environment state cannot be used as observation since it was explicitly defined - as unconstructed (:attr:`DirectMARLEnvCfg.num_states`). + as unconstructed (:attr:`DirectMARLEnvCfg.state_space`). """ class Env(DirectMARLEnv): @@ -157,7 +168,7 @@ def __init__(self, env: DirectMARLEnv) -> None: # check if it is possible to use the multi-agent environment state as agent observation self._state_as_observation = state_as_observation if self._state_as_observation: - assert self.env.cfg.num_states != 0, ( + assert self.env.cfg.state_space != 0, ( "The environment state cannot be used as observation since it was explicitly defined as" " unconstructed" ) @@ -170,13 +181,13 @@ def __init__(self, env: DirectMARLEnv) -> None: self._exported_observation_spaces = {self._agent_id: self.env.state_space} else: self._exported_observation_spaces = { - self._agent_id: gym.spaces.Box( - low=-np.inf, high=np.inf, shape=(sum(self.env.cfg.num_observations.values()),) + self._agent_id: gym.spaces.flatten_space( + gym.spaces.Tuple([self.env.observation_spaces[agent] for agent in self.env.possible_agents]) ) } self._exported_action_spaces = { - self._agent_id: gym.spaces.Box( - low=-np.inf, high=np.inf, shape=(sum(self.env.cfg.num_actions.values()),) + self._agent_id: gym.spaces.flatten_space( + gym.spaces.Tuple([self.env.action_spaces[agent] for agent in self.env.possible_agents]) ) } @@ -208,18 +219,25 @@ def reset( if self._state_as_observation: obs = {self._agent_id: self.env.state()} # concatenate agents' observations + # FIXME: This implementation assumes the spaces are fundamental ones. Fix it to support composite spaces else: - obs = {self._agent_id: torch.cat([obs[agent] for agent in self.env.possible_agents], dim=-1)} + obs = { + self._agent_id: torch.cat( + [obs[agent].reshape(self.num_envs, -1) for agent in self.env.possible_agents], dim=-1 + ) + } return obs, extras def step(self, actions: dict[AgentID, ActionType]) -> EnvStepReturn: # split agent actions to build the multi-agent ones + # FIXME: This implementation assumes the spaces are fundamental ones. Fix it to support composite spaces index = 0 _actions = {} for agent in self.env.possible_agents: - _actions[agent] = actions[self._agent_id][:, index : index + self.env.cfg.num_actions[agent]] - index += self.env.cfg.num_actions[agent] + delta = gym.spaces.flatdim(self.env.action_spaces[agent]) + _actions[agent] = actions[self._agent_id][:, index : index + delta] + index += delta # step the environment obs, rewards, terminated, time_outs, extras = self.env.step(_actions) @@ -228,8 +246,13 @@ def step(self, actions: dict[AgentID, ActionType]) -> EnvStepReturn: if self._state_as_observation: obs = {self._agent_id: self.env.state()} # concatenate agents' observations + # FIXME: This implementation assumes the spaces are fundamental ones. Fix it to support composite spaces else: - obs = {self._agent_id: torch.cat([obs[agent] for agent in self.env.possible_agents], dim=-1)} + obs = { + self._agent_id: torch.cat( + [obs[agent].reshape(self.num_envs, -1) for agent in self.env.possible_agents], dim=-1 + ) + } # process environment outputs to return agent data rewards = {self._agent_id: sum(rewards.values())} diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/utils/spaces.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/utils/spaces.py new file mode 100644 index 0000000000..2a1e30c1ee --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/utils/spaces.py @@ -0,0 +1,221 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import gymnasium as gym +import json +import numpy as np +import torch +from typing import Any + +from ..common import SpaceType + + +def spec_to_gym_space(spec: SpaceType) -> gym.spaces.Space: + """Generate an appropriate Gymnasium space according to the given space specification. + + Args: + spec: Space specification. + + Returns: + Gymnasium space. + + Raises: + ValueError: If the given space specification is not valid/supported. + """ + if isinstance(spec, gym.spaces.Space): + return spec + # fundamental spaces + # Box + elif isinstance(spec, int): + return gym.spaces.Box(low=-np.inf, high=np.inf, shape=(spec,)) + elif isinstance(spec, list) and all(isinstance(x, int) for x in spec): + return gym.spaces.Box(low=-np.inf, high=np.inf, shape=spec) + # Discrete + elif isinstance(spec, set) and len(spec) == 1: + return gym.spaces.Discrete(n=next(iter(spec))) + # MultiDiscrete + elif isinstance(spec, list) and all(isinstance(x, set) and len(x) == 1 for x in spec): + return gym.spaces.MultiDiscrete(nvec=[next(iter(x)) for x in spec]) + # composite spaces + # Tuple + elif isinstance(spec, tuple): + return gym.spaces.Tuple([spec_to_gym_space(x) for x in spec]) + # Dict + elif isinstance(spec, dict): + return gym.spaces.Dict({k: spec_to_gym_space(v) for k, v in spec.items()}) + raise ValueError(f"Unsupported space specification: {spec}") + + +def sample_space(space: gym.spaces.Space, device: str, batch_size: int = -1, fill_value: float | None = None) -> Any: + """Sample a Gymnasium space where the data container are PyTorch tensors. + + Args: + space: Gymnasium space. + device: The device where the tensor should be created. + batch_size: Batch size. If the specified value is greater than zero, a batched space will be created and sampled from it. + fill_value: The value to fill the created tensors with. If None (default value), tensors will keep their random values. + + Returns: + Tensorized sampled space. + """ + + def tensorize(s, x): + if isinstance(s, gym.spaces.Box): + tensor = torch.tensor(x, device=device, dtype=torch.float32).reshape(batch_size, *s.shape) + if fill_value is not None: + tensor.fill_(fill_value) + return tensor + elif isinstance(s, gym.spaces.Discrete): + if isinstance(x, np.ndarray): + tensor = torch.tensor(x, device=device, dtype=torch.int64).reshape(batch_size, 1) + if fill_value is not None: + tensor.fill_(int(fill_value)) + return tensor + elif isinstance(x, np.number) or type(x) in [int, float]: + tensor = torch.tensor([x], device=device, dtype=torch.int64).reshape(batch_size, 1) + if fill_value is not None: + tensor.fill_(int(fill_value)) + return tensor + elif isinstance(s, gym.spaces.MultiDiscrete): + if isinstance(x, np.ndarray): + tensor = torch.tensor(x, device=device, dtype=torch.int64).reshape(batch_size, *s.shape) + if fill_value is not None: + tensor.fill_(int(fill_value)) + return tensor + elif isinstance(s, gym.spaces.Dict): + return {k: tensorize(_s, x[k]) for k, _s in s.items()} + elif isinstance(s, gym.spaces.Tuple): + return tuple([tensorize(_s, v) for _s, v in zip(s, x)]) + + sample = (gym.vector.utils.batch_space(space, batch_size) if batch_size > 0 else space).sample() + return tensorize(space, sample) + + +def serialize_space(space: SpaceType) -> str: + """Serialize a space specification as JSON. + + Args: + space: Space specification. + + Returns: + Serialized JSON representation. + """ + # Gymnasium spaces + if isinstance(space, gym.spaces.Discrete): + return json.dumps({"type": "gymnasium", "space": "Discrete", "n": int(space.n)}) + elif isinstance(space, gym.spaces.Box): + return json.dumps({ + "type": "gymnasium", + "space": "Box", + "low": space.low.tolist(), + "high": space.high.tolist(), + "shape": space.shape, + }) + elif isinstance(space, gym.spaces.MultiDiscrete): + return json.dumps({"type": "gymnasium", "space": "MultiDiscrete", "nvec": space.nvec.tolist()}) + elif isinstance(space, gym.spaces.Tuple): + return json.dumps({"type": "gymnasium", "space": "Tuple", "spaces": tuple(map(serialize_space, space.spaces))}) + elif isinstance(space, gym.spaces.Dict): + return json.dumps( + {"type": "gymnasium", "space": "Dict", "spaces": {k: serialize_space(v) for k, v in space.spaces.items()}} + ) + # Python data types + # Box + elif isinstance(space, int) or (isinstance(space, list) and all(isinstance(x, int) for x in space)): + return json.dumps({"type": "python", "space": "Box", "value": space}) + # Discrete + elif isinstance(space, set) and len(space) == 1: + return json.dumps({"type": "python", "space": "Discrete", "value": next(iter(space))}) + # MultiDiscrete + elif isinstance(space, list) and all(isinstance(x, set) and len(x) == 1 for x in space): + return json.dumps({"type": "python", "space": "MultiDiscrete", "value": [next(iter(x)) for x in space]}) + # composite spaces + # Tuple + elif isinstance(space, tuple): + return json.dumps({"type": "python", "space": "Tuple", "value": [serialize_space(x) for x in space]}) + # Dict + elif isinstance(space, dict): + return json.dumps( + {"type": "python", "space": "Dict", "value": {k: serialize_space(v) for k, v in space.items()}} + ) + raise ValueError(f"Unsupported space ({space})") + + +def deserialize_space(string: str) -> gym.spaces.Space: + """Deserialize a space specification encoded as JSON. + + Args: + string: Serialized JSON representation. + + Returns: + Space specification. + """ + obj = json.loads(string) + # Gymnasium spaces + if obj["type"] == "gymnasium": + if obj["space"] == "Discrete": + return gym.spaces.Discrete(n=obj["n"]) + elif obj["space"] == "Box": + return gym.spaces.Box(low=np.array(obj["low"]), high=np.array(obj["high"]), shape=obj["shape"]) + elif obj["space"] == "MultiDiscrete": + return gym.spaces.MultiDiscrete(nvec=np.array(obj["nvec"])) + elif obj["space"] == "Tuple": + return gym.spaces.Tuple(spaces=tuple(map(deserialize_space, obj["spaces"]))) + elif obj["space"] == "Dict": + return gym.spaces.Dict(spaces={k: deserialize_space(v) for k, v in obj["spaces"].items()}) + else: + raise ValueError(f"Unsupported space ({obj['spaces']})") + # Python data types + elif obj["type"] == "python": + if obj["space"] == "Discrete": + return {obj["value"]} + elif obj["space"] == "Box": + return obj["value"] + elif obj["space"] == "MultiDiscrete": + return [{x} for x in obj["value"]] + elif obj["space"] == "Tuple": + return tuple(map(deserialize_space, obj["value"])) + elif obj["space"] == "Dict": + return {k: deserialize_space(v) for k, v in obj["value"].items()} + else: + raise ValueError(f"Unsupported space ({obj['spaces']})") + else: + raise ValueError(f"Unsupported type ({obj['type']})") + + +def replace_env_cfg_spaces_with_strings(env_cfg: object) -> object: + """Replace spaces objects with their serialized JSON representations in an environment config. + + Args: + env_cfg: Environment config instance. + + Returns: + Environment config instance with spaces replaced if any. + """ + for attr in ["observation_space", "action_space", "state_space"]: + if hasattr(env_cfg, attr): + setattr(env_cfg, attr, serialize_space(getattr(env_cfg, attr))) + for attr in ["observation_spaces", "action_spaces"]: + if hasattr(env_cfg, attr): + setattr(env_cfg, attr, {k: serialize_space(v) for k, v in getattr(env_cfg, attr).items()}) + return env_cfg + + +def replace_strings_with_env_cfg_spaces(env_cfg: object) -> object: + """Replace spaces objects with their serialized JSON representations in an environment config. + + Args: + env_cfg: Environment config instance. + + Returns: + Environment config instance with spaces replaced if any. + """ + for attr in ["observation_space", "action_space", "state_space"]: + if hasattr(env_cfg, attr): + setattr(env_cfg, attr, deserialize_space(getattr(env_cfg, attr))) + for attr in ["observation_spaces", "action_spaces"]: + if hasattr(env_cfg, attr): + setattr(env_cfg, attr, {k: deserialize_space(v) for k, v in getattr(env_cfg, attr).items()}) + return env_cfg diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/__init__.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/__init__.py index d06f832b52..cf8bafe687 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/__init__.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/__init__.py @@ -23,10 +23,12 @@ ManagerTermBaseCfg, ObservationGroupCfg, ObservationTermCfg, + RecorderTermCfg, RewardTermCfg, TerminationTermCfg, ) from .observation_manager import ObservationManager +from .recorder_manager import DatasetExportMode, RecorderManager, RecorderManagerBaseCfg, RecorderTerm from .reward_manager import RewardManager from .scene_entity_cfg import SceneEntityCfg from .termination_manager import TerminationManager diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/action_manager.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/action_manager.py index 56a7ff92dc..30ed9d41ef 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/action_manager.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/action_manager.py @@ -106,6 +106,7 @@ def set_debug_vis(self, debug_vis: bool) -> bool: # check if debug visualization is supported if not self.has_debug_vis_implementation: return False + # toggle debug visualization objects self._set_debug_vis_impl(debug_vis) # toggle debug visualization handles @@ -181,12 +182,21 @@ def __init__(self, cfg: object, env: ManagerBasedEnv): Args: cfg: The configuration object or dictionary (``dict[str, ActionTermCfg]``). env: The environment instance. + + Raises: + ValueError: If the configuration is None. """ + # check if config is None + if cfg is None: + raise ValueError("Action manager configuration is None. Please provide a valid configuration.") + + # call the base class constructor (this prepares the terms) super().__init__(cfg, env) # create buffers to store actions self._action = torch.zeros((self.num_envs, self.total_action_dim), device=self.device) self._prev_action = torch.zeros_like(self._action) + # check if any term has debug visualization implemented self.cfg.debug_vis = False for term in self._terms.values(): self.cfg.debug_vis |= term.cfg.debug_vis @@ -253,7 +263,26 @@ def has_debug_vis_implementation(self) -> bool: Operations. """ - def set_debug_vis(self, debug_vis: bool) -> bool: + def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: + """Returns the active terms as iterable sequence of tuples. + + The first element of the tuple is the name of the term and the second element is the raw value(s) of the term. + + Args: + env_idx: The specific environment to pull the active terms from. + + Returns: + The active terms. + """ + terms = [] + idx = 0 + for name, term in self._terms.items(): + term_actions = self._action[env_idx, idx : idx + term.action_dim].cpu() + terms.append((name, term_actions.tolist())) + idx += term.action_dim + return terms + + def set_debug_vis(self, debug_vis: bool): """Sets whether to visualize the action data. Args: debug_vis: Whether to visualize the action data. @@ -334,8 +363,7 @@ def get_term(self, name: str) -> ActionTerm: """ def _prepare_terms(self): - """Prepares a list of action terms.""" - # parse action terms from the config + # create buffers to parse and store terms self._term_names: list[str] = list() self._terms: dict[str, ActionTerm] = dict() @@ -344,6 +372,7 @@ def _prepare_terms(self): cfg_items = self.cfg.items() else: cfg_items = self.cfg.__dict__.items() + # parse action terms from the config for term_name, term_cfg in cfg_items: # check if term config is None if term_cfg is None: diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/command_manager.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/command_manager.py index 5cf7e929ac..2e1a414ba0 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/command_manager.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/command_manager.py @@ -132,10 +132,7 @@ def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float]: # resolve the environment IDs if env_ids is None: env_ids = slice(None) - # set the command counter to zero - self.command_counter[env_ids] = 0 - # resample the command - self._resample(env_ids) + # add logging metrics extras = {} for metric_name, metric_value in self.metrics.items(): @@ -143,6 +140,12 @@ def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float]: extras[metric_name] = torch.mean(metric_value[env_ids]).item() # reset the metric value metric_value[env_ids] = 0.0 + + # set the command counter to zero + self.command_counter[env_ids] = 0 + # resample the command + self._resample(env_ids) + return extras def compute(self, dt: float): @@ -175,8 +178,8 @@ def _resample(self, env_ids: Sequence[int]): Args: env_ids: The list of environment IDs to resample. """ - # resample the time left before resampling if len(env_ids) != 0: + # resample the time left before resampling self.time_left[env_ids] = self.time_left[env_ids].uniform_(*self.cfg.resampling_time_range) # increment the command counter self.command_counter[env_ids] += 1 @@ -243,12 +246,17 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv): cfg: The configuration object or dictionary (``dict[str, CommandTermCfg]``). env: The environment instance. """ + # create buffers to parse and store terms + self._terms: dict[str, CommandTerm] = dict() + + # call the base class constructor (this prepares the terms) super().__init__(cfg, env) # store the commands self._commands = dict() - self.cfg.debug_vis = False - for term in self._terms.values(): - self.cfg.debug_vis |= term.cfg.debug_vis + if self.cfg: + self.cfg.debug_vis = False + for term in self._terms.values(): + self.cfg.debug_vis |= term.cfg.debug_vis def __str__(self) -> str: """Returns: A string representation for the command manager.""" @@ -291,7 +299,26 @@ def has_debug_vis_implementation(self) -> bool: Operations. """ - def set_debug_vis(self, debug_vis: bool) -> bool: + def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: + """Returns the active terms as iterable sequence of tuples. + + The first element of the tuple is the name of the term and the second element is the raw value(s) of the term. + + Args: + env_idx: The specific environment to pull the active terms from. + + Returns: + The active terms. + """ + + terms = [] + idx = 0 + for name, term in self._terms.items(): + terms.append((name, term.command[env_idx].cpu().tolist())) + idx += term.command.shape[1] + return terms + + def set_debug_vis(self, debug_vis: bool): """Sets whether to visualize the command data. Args: @@ -371,10 +398,6 @@ def get_term(self, name: str) -> CommandTerm: """ def _prepare_terms(self): - """Prepares a list of command terms.""" - # parse command terms from the config - self._terms: dict[str, CommandTerm] = dict() - # check if config is dict already if isinstance(self.cfg, dict): cfg_items = self.cfg.items() @@ -393,5 +416,8 @@ def _prepare_terms(self): ) # create the action term term = term_cfg.class_type(term_cfg, self._env) + # sanity check if term is valid type + if not isinstance(term, CommandTerm): + raise TypeError(f"Returned object for the term '{term_name}' is not of type CommandType.") # add class to dict self._terms[term_name] = term diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/curriculum_manager.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/curriculum_manager.py index b9bef068bf..9e316cccf8 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/curriculum_manager.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/curriculum_manager.py @@ -44,7 +44,14 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv): TypeError: If curriculum term is not of type :class:`CurriculumTermCfg`. ValueError: If curriculum term configuration does not satisfy its function signature. """ + # create buffers to parse and store terms + self._term_names: list[str] = list() + self._term_cfgs: list[CurriculumTermCfg] = list() + self._class_term_cfgs: list[CurriculumTermCfg] = list() + + # call the base class constructor (this will parse the terms config) super().__init__(cfg, env) + # prepare logging self._curriculum_state = dict() for term_name in self._term_names: @@ -131,16 +138,45 @@ def compute(self, env_ids: Sequence[int] | None = None): state = term_cfg.func(self._env, env_ids, **term_cfg.params) self._curriculum_state[name] = state + def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: + """Returns the active terms as iterable sequence of tuples. + + The first element of the tuple is the name of the term and the second element is the raw value(s) of the term. + + Args: + env_idx: The specific environment to pull the active terms from. + + Returns: + The active terms. + """ + + terms = [] + + for term_name, term_state in self._curriculum_state.items(): + if term_state is not None: + # deal with dict + data = [] + + if isinstance(term_state, dict): + # each key is a separate state to log + for key, value in term_state.items(): + if isinstance(value, torch.Tensor): + value = value.item() + terms[term_name].append(value) + else: + # log directly if not a dict + if isinstance(term_state, torch.Tensor): + term_state = term_state.item() + data.append(term_state) + terms.append((term_name, data)) + + return terms + """ Helper functions. """ def _prepare_terms(self): - # parse remaining curriculum terms and decimate their information - self._term_names: list[str] = list() - self._term_cfgs: list[CurriculumTermCfg] = list() - self._class_term_cfgs: list[CurriculumTermCfg] = list() - # check if config is dict already if isinstance(self.cfg, dict): cfg_items = self.cfg.items() diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/event_manager.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/event_manager.py index 9843164ba0..9209fe1d4a 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/event_manager.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/event_manager.py @@ -62,6 +62,12 @@ def __init__(self, cfg: object, env: ManagerBasedEnv): cfg: A configuration object or dictionary (``dict[str, EventTermCfg]``). env: An environment object. """ + # create buffers to parse and store terms + self._mode_term_names: dict[str, list[str]] = dict() + self._mode_term_cfgs: dict[str, list[EventTermCfg]] = dict() + self._mode_class_term_cfgs: dict[str, list[EventTermCfg]] = dict() + + # call the base class (this will parse the terms config) super().__init__(cfg, env) def __str__(self) -> str: @@ -294,11 +300,6 @@ def get_term_cfg(self, term_name: str) -> EventTermCfg: """ def _prepare_terms(self): - """Prepares a list of event functions.""" - # parse remaining event terms and decimate their information - self._mode_term_names: dict[str, list[str]] = dict() - self._mode_term_cfgs: dict[str, list[EventTermCfg]] = dict() - self._mode_class_term_cfgs: dict[str, list[EventTermCfg]] = dict() # buffer to store the time left for "interval" mode # if interval is global, then it is a single value, otherwise it is per environment self._interval_term_time_left: list[torch.Tensor] = list() diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/manager_base.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/manager_base.py index 2bc9236cab..d30545a391 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/manager_base.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/manager_base.py @@ -120,14 +120,15 @@ def __init__(self, cfg: object, env: ManagerBasedEnv): """Initialize the manager. Args: - cfg: The configuration object. + cfg: The configuration object. If None, the manager is initialized without any terms. env: The environment instance. """ # store the inputs self.cfg = copy.deepcopy(cfg) self._env = env # parse config to create terms information - self._prepare_terms() + if self.cfg: + self._prepare_terms() """ Properties. @@ -192,6 +193,16 @@ def find_terms(self, name_keys: str | Sequence[str]) -> list[str]: # return the matching names return string_utils.resolve_matching_names(name_keys, list_of_strings)[1] + def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: + """Returns the active terms as iterable sequence of tuples. + + The first element of the tuple is the name of the term and the second element is the raw value(s) of the term. + + Returns: + The active terms. + """ + raise NotImplementedError + """ Implementation specific. """ diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/manager_term_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/manager_term_cfg.py index 9a2250e48b..19c54f4e70 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/manager_term_cfg.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/manager_term_cfg.py @@ -22,6 +22,7 @@ from .action_manager import ActionTerm from .command_manager import CommandTerm from .manager_base import ManagerTermBase + from .recorder_manager import RecorderTerm @configclass @@ -51,6 +52,22 @@ class ManagerTermBaseCfg: """ +## +# Recorder manager. +## + + +@configclass +class RecorderTermCfg: + """Configuration for an recorder term.""" + + class_type: type[RecorderTerm] = MISSING + """The associated recorder term class. + + The class should inherit from :class:`omni.isaac.lab.managers.action_manager.RecorderTerm`. + """ + + ## # Action manager. ## @@ -76,6 +93,9 @@ class for more details. debug_vis: bool = False """Whether to visualize debug information. Defaults to False.""" + clip: dict[str, tuple] | None = None + """Clip range for the action (dict of regex expressions). Defaults to None.""" + ## # Command manager. @@ -152,9 +172,26 @@ class ObservationTermCfg(ManagerTermBaseCfg): """The clipping range for the observation after adding noise. Defaults to None, in which case no clipping is applied.""" - scale: float | None = None + scale: tuple[float, ...] | float | None = None """The scale to apply to the observation after clipping. Defaults to None, - in which case no scaling is applied (same as setting scale to :obj:`1`).""" + in which case no scaling is applied (same as setting scale to :obj:`1`). + + We leverage PyTorch broadcasting to scale the observation tensor with the provided value. If a tuple is provided, + please make sure the length of the tuple matches the dimensions of the tensor outputted from the term. + """ + + history_length: int = 0 + """Number of past observations to store in the observation buffers. Defaults to 0, meaning no history. + + Observation history initializes to empty, but is filled with the first append after reset or initialization. Subsequent history + only adds a single entry to the history buffer. If flatten_history_dim is set to True, the source data of shape + (N, H, D, ...) where N is the batch dimension and H is the history length will be reshaped to a 2D tensor of shape + (N, H*D*...). Otherwise, the data will be returned as is. + """ + + flatten_history_dim: bool = True + """Whether or not the observation manager should flatten history-based observation terms to a 2D (N, D) tensor. + Defaults to True.""" @configclass @@ -177,6 +214,22 @@ class ObservationGroupCfg: Otherwise, no corruption is applied. """ + history_length: int | None = None + """Number of past observation to store in the observation buffers for all observation terms in group. + + This parameter will override :attr:`ObservationTermCfg.history_length` if set. Defaults to None. If None, each + terms history will be controlled on a per term basis. See :class:`ObservationTermCfg` for details on history_length + implementation. + """ + + flatten_history_dim: bool = True + """Flag to flatten history-based observation terms to a 2D (num_env, D) tensor for all observation terms in group. + Defaults to True. + + This parameter will override all :attr:`ObservationTermCfg.flatten_history_dim` in the group if + ObservationGroupCfg.history_length is set. + """ + ## # Event manager diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/observation_manager.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/observation_manager.py index 58ae0f55f3..1fc7660d79 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/observation_manager.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/observation_manager.py @@ -8,12 +8,14 @@ from __future__ import annotations import inspect +import numpy as np import torch from collections.abc import Sequence from prettytable import PrettyTable from typing import TYPE_CHECKING from omni.isaac.lab.utils import modifiers +from omni.isaac.lab.utils.buffers import CircularBuffer from .manager_base import ManagerBase, ManagerTermBase from .manager_term_cfg import ObservationGroupCfg, ObservationTermCfg @@ -45,6 +47,11 @@ class ObservationManager(ManagerBase): concatenated. In this case, please set the :attr:`ObservationGroupCfg.concatenate_terms` attribute in the group configuration to False. + Observations can also have history. This means a running history is updated per sim step. History can be controlled + per :class:`ObservationTermCfg` (See the :attr:`ObservationTermCfg.history_length` and + :attr:`ObservationTermCfg.flatten_history_dim`). History can also be controlled via :class:`ObservationGroupCfg` + where group configuration overwrites per term configuration if set. History follows an oldest to newest ordering. + The observation manager can be used to compute observations for all the groups or for a specific group. The observations are computed by calling the registered functions for each term in the group. The functions are called in the order of the terms in the group. The functions are expected to return a tensor with shape @@ -63,9 +70,15 @@ def __init__(self, cfg: object, env: ManagerBasedEnv): env: The environment instance. Raises: + ValueError: If the configuration is None. RuntimeError: If the shapes of the observation terms in a group are not compatible for concatenation and the :attr:`~ObservationGroupCfg.concatenate_terms` attribute is set to True. """ + # check that cfg is not None + if cfg is None: + raise ValueError("Observation manager configuration is None. Please provide a valid configuration.") + + # call the base class constructor (this will parse the terms config) super().__init__(cfg, env) # compute combined vector for obs group @@ -87,6 +100,9 @@ def __init__(self, cfg: object, env: ManagerBasedEnv): else: self._group_obs_dim[group_name] = group_term_dims + # Stores the latest observations. + self._obs_buffer: dict[str, torch.Tensor | dict[str, torch.Tensor]] | None = None + def __str__(self) -> str: """Returns: A string representation for the observation manager.""" msg = f" contains {len(self._group_obs_term_names)} groups.\n" @@ -117,6 +133,43 @@ def __str__(self) -> str: return msg + def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: + """Returns the active terms as iterable sequence of tuples. + + The first element of the tuple is the name of the term and the second element is the raw value(s) of the term. + + Args: + env_idx: The specific environment to pull the active terms from. + + Returns: + The active terms. + """ + terms = [] + + if self._obs_buffer is None: + self.compute() + obs_buffer: dict[str, torch.Tensor | dict[str, torch.Tensor]] = self._obs_buffer + + for group_name, _ in self._group_obs_dim.items(): + if not self.group_obs_concatenate[group_name]: + for name, term in obs_buffer[group_name].items(): + terms.append((group_name + "-" + name, term[env_idx].cpu().tolist())) + continue + + idx = 0 + # add info for each term + data = obs_buffer[group_name] + for name, shape in zip( + self._group_obs_term_names[group_name], + self._group_obs_term_dim[group_name], + ): + data_length = np.prod(shape) + term = data[env_idx, idx : idx + data_length] + terms.append((group_name + "-" + name, term.cpu().tolist())) + idx += data_length + + return terms + """ Properties. """ @@ -168,12 +221,17 @@ def group_obs_concatenate(self) -> dict[str, bool]: def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float]: # call all terms that are classes - for group_cfg in self._group_obs_class_term_cfgs.values(): + for group_name, group_cfg in self._group_obs_class_term_cfgs.items(): for term_cfg in group_cfg: term_cfg.func.reset(env_ids=env_ids) + # reset terms with history + for term_name in self._group_obs_term_names[group_name]: + if term_name in self._group_obs_term_history_buffer[group_name]: + self._group_obs_term_history_buffer[group_name][term_name].reset(batch_ids=env_ids) # call all modifiers that are classes for mod in self._group_obs_class_modifiers: mod.reset(env_ids=env_ids) + # nothing to log here return {} @@ -194,6 +252,9 @@ def compute(self) -> dict[str, torch.Tensor | dict[str, torch.Tensor]]: for group_name in self._group_obs_term_names: obs_buffer[group_name] = self.compute_group(group_name) # otherwise return a dict with observations of all groups + + # Cache the observations. + self._obs_buffer = obs_buffer return obs_buffer def compute_group(self, group_name: str) -> torch.Tensor | dict[str, torch.Tensor]: @@ -242,7 +303,7 @@ def compute_group(self, group_name: str) -> torch.Tensor | dict[str, torch.Tenso obs_terms = zip(group_term_names, self._group_obs_term_cfgs[group_name]) # evaluate terms: compute, add noise, clip, scale, custom modifiers - for name, term_cfg in obs_terms: + for term_name, term_cfg in obs_terms: # compute term's value obs: torch.Tensor = term_cfg.func(self._env, **term_cfg.params).clone() # apply post-processing @@ -253,10 +314,19 @@ def compute_group(self, group_name: str) -> torch.Tensor | dict[str, torch.Tenso obs = term_cfg.noise.func(obs, term_cfg.noise) if term_cfg.clip: obs = obs.clip_(min=term_cfg.clip[0], max=term_cfg.clip[1]) - if term_cfg.scale: + if term_cfg.scale is not None: obs = obs.mul_(term_cfg.scale) - # add value to list - group_obs[name] = obs + # Update the history buffer if observation term has history enabled + if term_cfg.history_length > 0: + self._group_obs_term_history_buffer[group_name][term_name].append(obs) + if term_cfg.flatten_history_dim: + group_obs[term_name] = self._group_obs_term_history_buffer[group_name][term_name].buffer.reshape( + self._env.num_envs, -1 + ) + else: + group_obs[term_name] = self._group_obs_term_history_buffer[group_name][term_name].buffer + else: + group_obs[term_name] = obs # concatenate all observations in the group together if self._group_obs_concatenate[group_name]: @@ -277,7 +347,7 @@ def _prepare_terms(self): self._group_obs_term_cfgs: dict[str, list[ObservationTermCfg]] = dict() self._group_obs_class_term_cfgs: dict[str, list[ObservationTermCfg]] = dict() self._group_obs_concatenate: dict[str, bool] = dict() - + self._group_obs_term_history_buffer: dict[str, dict] = dict() # create a list to store modifiers that are classes # we store it as a separate list to only call reset on them and prevent unnecessary calls self._group_obs_class_modifiers: list[modifiers.ModifierBase] = list() @@ -303,6 +373,7 @@ def _prepare_terms(self): self._group_obs_term_dim[group_name] = list() self._group_obs_term_cfgs[group_name] = list() self._group_obs_class_term_cfgs[group_name] = list() + group_entry_history_buffer: dict[str, CircularBuffer] = dict() # read common config for the group self._group_obs_concatenate[group_name] = group_cfg.concatenate_terms # check if config is dict already @@ -311,9 +382,9 @@ def _prepare_terms(self): else: group_cfg_items = group_cfg.__dict__.items() # iterate over all the terms in each group - for term_name, term_cfg in group_cfg.__dict__.items(): + for term_name, term_cfg in group_cfg_items: # skip non-obs settings - if term_name in ["enable_corruption", "concatenate_terms"]: + if term_name in ["enable_corruption", "concatenate_terms", "history_length", "flatten_history_dim"]: continue # check for non config if term_cfg is None: @@ -329,14 +400,45 @@ def _prepare_terms(self): # check noise settings if not group_cfg.enable_corruption: term_cfg.noise = None + # check group history params and override terms + if group_cfg.history_length is not None: + term_cfg.history_length = group_cfg.history_length + term_cfg.flatten_history_dim = group_cfg.flatten_history_dim # add term config to list to list self._group_obs_term_names[group_name].append(term_name) self._group_obs_term_cfgs[group_name].append(term_cfg) - # call function the first time to fill up dimensions obs_dims = tuple(term_cfg.func(self._env, **term_cfg.params).shape) + # create history buffers and calculate history term dimensions + if term_cfg.history_length > 0: + group_entry_history_buffer[term_name] = CircularBuffer( + max_len=term_cfg.history_length, batch_size=self._env.num_envs, device=self._env.device + ) + old_dims = list(obs_dims) + old_dims.insert(1, term_cfg.history_length) + obs_dims = tuple(old_dims) + if term_cfg.flatten_history_dim: + obs_dims = (obs_dims[0], np.prod(obs_dims[1:])) + self._group_obs_term_dim[group_name].append(obs_dims[1:]) + # if scale is set, check if single float or tuple + if term_cfg.scale is not None: + if not isinstance(term_cfg.scale, (float, int, tuple)): + raise TypeError( + f"Scale for observation term '{term_name}' in group '{group_name}'" + f" is not of type float, int or tuple. Received: '{type(term_cfg.scale)}'." + ) + if isinstance(term_cfg.scale, tuple) and len(term_cfg.scale) != obs_dims[1]: + raise ValueError( + f"Scale for observation term '{term_name}' in group '{group_name}'" + f" does not match the dimensions of the observation. Expected: {obs_dims[1]}" + f" but received: {len(term_cfg.scale)}." + ) + + # cast the scale into torch tensor + term_cfg.scale = torch.tensor(term_cfg.scale, dtype=torch.float, device=self._env.device) + # prepare modifiers for each observation if term_cfg.modifiers is not None: # initialize list of modifiers for term @@ -388,3 +490,5 @@ def _prepare_terms(self): self._group_obs_class_term_cfgs[group_name].append(term_cfg) # call reset (in-case above call to get obs dims changed the state) term_cfg.func.reset() + # add history buffers for each group + self._group_obs_term_history_buffer[group_name] = group_entry_history_buffer diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/recorder_manager.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/recorder_manager.py new file mode 100644 index 0000000000..5958e7d73c --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/recorder_manager.py @@ -0,0 +1,442 @@ +# Copyright (c) 2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Recorder manager for recording data produced from the given world.""" + +from __future__ import annotations + +import enum +import os +import torch +from collections.abc import Sequence +from prettytable import PrettyTable +from typing import TYPE_CHECKING + +from omni.isaac.lab.utils import configclass +from omni.isaac.lab.utils.datasets import EpisodeData, HDF5DatasetFileHandler + +from .manager_base import ManagerBase, ManagerTermBase +from .manager_term_cfg import RecorderTermCfg + +if TYPE_CHECKING: + from omni.isaac.lab.envs import ManagerBasedEnv + + +class DatasetExportMode(enum.IntEnum): + """The mode to handle episode exports.""" + + EXPORT_ALL = 0 # Export all episodes to a single dataset file + EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES = 1 # Export succeeded and failed episodes in separate files + EXPORT_SUCCEEDED_ONLY = 2 # Export only succeeded episodes to a single dataset file + + +@configclass +class RecorderManagerBaseCfg: + """Base class for configuring recorder manager terms.""" + + dataset_file_handler_class_type: type = HDF5DatasetFileHandler + + dataset_export_dir_path: str = "/tmp/isaaclab/logs" + """The directory path where the recorded datasets are exported.""" + + dataset_filename: str = "dataset" + """Dataset file name without file extension.""" + + dataset_export_mode: DatasetExportMode = DatasetExportMode.EXPORT_ALL + """The mode to handle episode exports.""" + + +class RecorderTerm(ManagerTermBase): + """Base class for recorder terms. + + The recorder term is responsible for recording data at various stages of the environment's lifecycle. + A recorder term is comprised of four user-defined callbacks to record data in the corresponding stages: + + * Pre-reset recording: This callback is invoked at the beginning of `env.reset()` before the reset is effective. + * Post-reset recording: This callback is invoked at the end of `env.reset()`. + * Pre-step recording: This callback is invoked at the beginning of `env.step()`, after the step action is processed + and before the action is applied by the action manager. + * Post-step recording: This callback is invoked at the end of `env.step()` when all the managers are processed. + """ + + def __init__(self, cfg: RecorderTermCfg, env: ManagerBasedEnv): + """Initialize the recorder term. + + Args: + cfg: The configuration object. + env: The environment instance. + """ + # call the base class constructor + super().__init__(cfg, env) + + """ + User-defined callbacks. + """ + + def record_pre_reset(self, env_ids: Sequence[int] | None) -> tuple[str | None, torch.Tensor | dict | None]: + """Record data at the beginning of env.reset() before reset is effective. + + Args: + env_ids: The environment ids. All environments should be considered when set to None. + + Returns: + A tuple of key and value to be recorded. + The key can contain nested keys separated by '/'. For example, "obs/joint_pos" would add the given + value under ['obs']['policy'] in the underlying dictionary in the recorded episode data. + The value can be a tensor or a nested dictionary of tensors. The shape of a tensor in the value + is (env_ids, ...). + """ + return None, None + + def record_post_reset(self, env_ids: Sequence[int] | None) -> tuple[str | None, torch.Tensor | dict | None]: + """Record data at the end of env.reset(). + + Args: + env_ids: The environment ids. All environments should be considered when set to None. + + Returns: + A tuple of key and value to be recorded. + Please refer to the `record_pre_reset` function for more details. + """ + return None, None + + def record_pre_step(self) -> tuple[str | None, torch.Tensor | dict | None]: + """Record data in the beginning of env.step() after action is cached/processed in the ActionManager. + + Returns: + A tuple of key and value to be recorded. + Please refer to the `record_pre_reset` function for more details. + """ + return None, None + + def record_post_step(self) -> tuple[str | None, torch.Tensor | dict | None]: + """Record data at the end of env.step() when all the managers are processed. + + Returns: + A tuple of key and value to be recorded. + Please refer to the `record_pre_reset` function for more details. + """ + return None, None + + +class RecorderManager(ManagerBase): + """Manager for recording data from recorder terms.""" + + def __init__(self, cfg: object, env: ManagerBasedEnv): + """Initialize the recorder manager. + + Args: + cfg: The configuration object or dictionary (``dict[str, RecorderTermCfg]``). + env: The environment instance. + """ + self._term_names: list[str] = list() + self._terms: dict[str, RecorderTerm] = dict() + + # Do nothing if cfg is None or an empty dict + if not cfg: + return + + super().__init__(cfg, env) + + # Do nothing if no active recorder terms are provided + if len(self.active_terms) == 0: + return + + if not isinstance(cfg, RecorderManagerBaseCfg): + raise TypeError("Configuration for the recorder manager is not of type RecorderManagerBaseCfg.") + + # create episode data buffer indexed by environment id + self._episodes: dict[int, EpisodeData] = dict() + for env_id in range(env.num_envs): + self._episodes[env_id] = EpisodeData() + + env_name = getattr(env.cfg, "env_name", None) + + self._dataset_file_handler = cfg.dataset_file_handler_class_type() + self._dataset_file_handler.create( + os.path.join(cfg.dataset_export_dir_path, cfg.dataset_filename), env_name=env_name + ) + + self._failed_episode_dataset_file_handler = None + if cfg.dataset_export_mode == DatasetExportMode.EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES: + self._failed_episode_dataset_file_handler = cfg.dataset_file_handler_class_type() + self._failed_episode_dataset_file_handler.create( + os.path.join(cfg.dataset_export_dir_path, f"{cfg.dataset_filename}_failed"), env_name=env_name + ) + + def __str__(self) -> str: + """Returns: A string representation for recorder manager.""" + msg = f" contains {len(self._term_names)} active terms.\n" + # create table for term information + table = PrettyTable() + table.title = "Active Recorder Terms" + table.field_names = ["Index", "Name"] + # set alignment of table columns + table.align["Name"] = "l" + # add info on each term + for index, name in enumerate(self._term_names): + table.add_row([index, name]) + # convert table to string + msg += table.get_string() + msg += "\n" + return msg + + def __del__(self): + """Destructor for recorder.""" + # Do nothing if no active recorder terms are provided + if len(self.active_terms) == 0: + return + + if self._dataset_file_handler is not None: + self._dataset_file_handler.close() + + """ + Properties. + """ + + @property + def active_terms(self) -> list[str]: + """Name of active recorder terms.""" + return self._term_names + + """ + Operations. + """ + + def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]: + """Resets the recorder data. + + Args: + env_ids: The environment ids. Defaults to None, in which case + all environments are considered. + + Returns: + An empty dictionary. + """ + # Do nothing if no active recorder terms are provided + if len(self.active_terms) == 0: + return {} + + # resolve environment ids + if env_ids is None: + env_ids = list(range(self._env.num_envs)) + if isinstance(env_ids, torch.Tensor): + env_ids = env_ids.tolist() + + for term in self._terms.values(): + term.reset(env_ids=env_ids) + + for env_id in env_ids: + self._episodes[env_id] = EpisodeData() + + # nothing to log here + return {} + + def get_episode(self, env_id: int) -> EpisodeData: + """Returns the episode data for the given environment id. + + Args: + env_id: The environment id. + + Returns: + The episode data for the given environment id. + """ + return self._episodes.get(env_id, EpisodeData()) + + def add_to_episodes(self, key: str, value: torch.Tensor | dict, env_ids: Sequence[int] | None = None): + """Adds the given key-value pair to the episodes for the given environment ids. + + Args: + key: The key of the given value to be added to the episodes. The key can contain nested keys + separated by '/'. For example, "obs/joint_pos" would add the given value under ['obs']['policy'] + in the underlying dictionary in the episode data. + value: The value to be added to the episodes. The value can be a tensor or a nested dictionary of tensors. + The shape of a tensor in the value is (env_ids, ...). + env_ids: The environment ids. Defaults to None, in which case all environments are considered. + """ + # Do nothing if no active recorder terms are provided + if len(self.active_terms) == 0: + return + + # resolve environment ids + if key is None: + return + if env_ids is None: + env_ids = list(range(self._env.num_envs)) + if isinstance(env_ids, torch.Tensor): + env_ids = env_ids.tolist() + + if isinstance(value, dict): + for sub_key, sub_value in value.items(): + self.add_to_episodes(f"{key}/{sub_key}", sub_value, env_ids) + return + + for value_index, env_id in enumerate(env_ids): + if env_id not in self._episodes: + self._episodes[env_id] = EpisodeData() + self._episodes[env_id].env_id = env_id + self._episodes[env_id].add(key, value[value_index]) + + def set_success_to_episodes(self, env_ids: Sequence[int] | None, success_values: torch.Tensor): + """Sets the task success values to the episodes for the given environment ids. + + Args: + env_ids: The environment ids. Defaults to None, in which case all environments are considered. + success_values: The task success values to be set to the episodes. The shape of the tensor is (env_ids, 1). + """ + # Do nothing if no active recorder terms are provided + if len(self.active_terms) == 0: + return + + # resolve environment ids + if env_ids is None: + env_ids = list(range(self._env.num_envs)) + if isinstance(env_ids, torch.Tensor): + env_ids = env_ids.tolist() + + for value_index, env_id in enumerate(env_ids): + self._episodes[env_id].success = success_values[value_index].item() + + def record_pre_step(self) -> None: + """Trigger recorder terms for pre-step functions.""" + # Do nothing if no active recorder terms are provided + if len(self.active_terms) == 0: + return + + for term in self._terms.values(): + key, value = term.record_pre_step() + self.add_to_episodes(key, value) + + def record_post_step(self) -> None: + """Trigger recorder terms for post-step functions.""" + # Do nothing if no active recorder terms are provided + if len(self.active_terms) == 0: + return + + for term in self._terms.values(): + key, value = term.record_post_step() + self.add_to_episodes(key, value) + + def record_pre_reset(self, env_ids: Sequence[int] | None) -> None: + """Trigger recorder terms for pre-reset functions. + + Args: + env_ids: The environment ids in which a reset is triggered. + """ + # Do nothing if no active recorder terms are provided + if len(self.active_terms) == 0: + return + + if env_ids is None: + env_ids = list(range(self._env.num_envs)) + if isinstance(env_ids, torch.Tensor): + env_ids = env_ids.tolist() + + for term in self._terms.values(): + key, value = term.record_pre_reset(env_ids) + self.add_to_episodes(key, value, env_ids) + + # Set task success values for the relevant episodes + success_results = torch.zeros(len(env_ids), dtype=bool, device=self._env.device) + # Check success indicator from termination terms + if "success" in self._env.termination_manager.active_terms: + success_results |= self._env.termination_manager.get_term("success")[env_ids] + self.set_success_to_episodes(env_ids, success_results) + + self.export_episodes(env_ids) + + def record_post_reset(self, env_ids: Sequence[int] | None) -> None: + """Trigger recorder terms for post-reset functions. + + Args: + env_ids: The environment ids in which a reset is triggered. + """ + # Do nothing if no active recorder terms are provided + if len(self.active_terms) == 0: + return + + for term in self._terms.values(): + key, value = term.record_post_reset(env_ids) + self.add_to_episodes(key, value, env_ids) + + def export_episodes(self, env_ids: Sequence[int] | None = None) -> None: + """Concludes and exports the episodes for the given environment ids. + + Args: + env_ids: The environment ids. Defaults to None, in which case + all environments are considered. + """ + # Do nothing if no active recorder terms are provided + if len(self.active_terms) == 0: + return + + if env_ids is None: + env_ids = list(range(self._env.num_envs)) + if isinstance(env_ids, torch.Tensor): + env_ids = env_ids.tolist() + + # Export episode data through dataset exporter + need_to_flush = False + for env_id in env_ids: + if env_id in self._episodes and not self._episodes[env_id].is_empty(): + episode_succeeded = self._episodes[env_id].success + target_dataset_file_handler = None + if (self.cfg.dataset_export_mode == DatasetExportMode.EXPORT_ALL) or ( + self.cfg.dataset_export_mode == DatasetExportMode.EXPORT_SUCCEEDED_ONLY and episode_succeeded + ): + target_dataset_file_handler = self._dataset_file_handler + elif self.cfg.dataset_export_mode == DatasetExportMode.EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES: + if episode_succeeded: + target_dataset_file_handler = self._dataset_file_handler + else: + target_dataset_file_handler = self._failed_episode_dataset_file_handler + if target_dataset_file_handler is not None: + target_dataset_file_handler.write_episode(self._episodes[env_id]) + need_to_flush = True + # Reset the episode buffer for the given environment after export + self._episodes[env_id] = EpisodeData() + + if need_to_flush: + self._dataset_file_handler.flush() + if self._failed_episode_dataset_file_handler is not None: + self._failed_episode_dataset_file_handler.flush() + + """ + Helper functions. + """ + + def _prepare_terms(self): + """Prepares a list of recorder terms.""" + # check if config is dict already + if isinstance(self.cfg, dict): + cfg_items = self.cfg.items() + else: + cfg_items = self.cfg.__dict__.items() + for term_name, term_cfg in cfg_items: + # skip non-term settings + if term_name in [ + "dataset_file_handler_class_type", + "dataset_filename", + "dataset_export_dir_path", + "dataset_export_mode", + ]: + continue + # check if term config is None + if term_cfg is None: + continue + # check valid type + if not isinstance(term_cfg, RecorderTermCfg): + raise TypeError( + f"Configuration for the term '{term_name}' is not of type RecorderTermCfg." + f" Received: '{type(term_cfg)}'." + ) + # create the recorder term + term = term_cfg.class_type(term_cfg, self._env) + # sanity check if term is valid type + if not isinstance(term, RecorderTerm): + raise TypeError(f"Returned object for the term '{term_name}' is not of type RecorderTerm.") + # add term name and parameters + self._term_names.append(term_name) + self._terms[term_name] = term diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/reward_manager.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/reward_manager.py index c10bc12ec5..49369ca58a 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/reward_manager.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/reward_manager.py @@ -47,6 +47,12 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv): cfg: The configuration object or dictionary (``dict[str, RewardTermCfg]``). env: The environment instance. """ + # create buffers to parse and store terms + self._term_names: list[str] = list() + self._term_cfgs: list[RewardTermCfg] = list() + self._class_term_cfgs: list[RewardTermCfg] = list() + + # call the base class constructor (this will parse the terms config) super().__init__(cfg, env) # prepare extra info to store individual reward term information self._episode_sums = dict() @@ -55,6 +61,9 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv): # create buffer for managing reward per environment self._reward_buf = torch.zeros(self.num_envs, dtype=torch.float, device=self.device) + # Buffer which stores the current step reward for each term for each environment + self._step_reward = torch.zeros((self.num_envs, len(self._term_names)), dtype=torch.float, device=self.device) + def __str__(self) -> str: """Returns: A string representation for reward manager.""" msg = f" contains {len(self._term_names)} active terms.\n" @@ -142,6 +151,9 @@ def compute(self, dt: float) -> torch.Tensor: # update episodic sum self._episode_sums[name] += value + # Update current reward for this step. + self._step_reward[:, self._term_names.index(name)] = value / dt + return self._reward_buf """ @@ -180,17 +192,27 @@ def get_term_cfg(self, term_name: str) -> RewardTermCfg: # return the configuration return self._term_cfgs[self._term_names.index(term_name)] + def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: + """Returns the active terms as iterable sequence of tuples. + + The first element of the tuple is the name of the term and the second element is the raw value(s) of the term. + + Args: + env_idx: The specific environment to pull the active terms from. + + Returns: + The active terms. + """ + terms = [] + for idx, name in enumerate(self._term_names): + terms.append((name, [self._step_reward[env_idx, idx].cpu().item()])) + return terms + """ Helper functions. """ def _prepare_terms(self): - """Prepares a list of reward functions.""" - # parse remaining reward terms and decimate their information - self._term_names: list[str] = list() - self._term_cfgs: list[RewardTermCfg] = list() - self._class_term_cfgs: list[RewardTermCfg] = list() - # check if config is dict already if isinstance(self.cfg, dict): cfg_items = self.cfg.items() diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/scene_entity_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/scene_entity_cfg.py index 17095c6615..3b801dec19 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/scene_entity_cfg.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/scene_entity_cfg.py @@ -7,7 +7,7 @@ from dataclasses import MISSING -from omni.isaac.lab.assets import Articulation, RigidObject +from omni.isaac.lab.assets import Articulation, RigidObject, RigidObjectCollection from omni.isaac.lab.scene import InteractiveScene from omni.isaac.lab.utils import configclass @@ -78,16 +78,34 @@ class for more details. manager. """ + object_collection_names: str | list[str] | None = None + """The names of the objects in the rigid object collection required by the term. Defaults to None. + + The names can be either names or a regular expression matching the object names in the collection. + + These are converted to object indices on initialization of the manager and passed to the term + function as a list of object indices under :attr:`object_collection_ids`. + """ + + object_collection_ids: list[int] | slice = slice(None) + """The indices of the objects from the rigid object collection required by the term. Defaults to slice(None), + which means all the objects in the collection. + + If :attr:`object_collection_names` is specified, this is filled in automatically on initialization of the manager. + """ + preserve_order: bool = False - """Whether to preserve indices ordering to match with that in the specified joint or body names. Defaults to False. + """Whether to preserve indices ordering to match with that in the specified joint, body, or object collection names. + Defaults to False. - If False, the ordering of the indices are sorted in ascending order (i.e. the ordering in the entity's joints - or bodies). Otherwise, the indices are preserved in the order of the specified joint and body names. + If False, the ordering of the indices are sorted in ascending order (i.e. the ordering in the entity's joints, + bodies, or object in the object collection). Otherwise, the indices are preserved in the order of the specified + joint, body, or object collection names. For more details, see the :meth:`omni.isaac.lab.utils.string.resolve_matching_names` function. .. note:: - This attribute is only used when :attr:`joint_names` or :attr:`body_names` are specified. + This attribute is only used when :attr:`joint_names`, :attr:`body_names`, or :attr:`object_collection_names` are specified. """ @@ -106,6 +124,7 @@ def resolve(self, scene: InteractiveScene): ValueError: If both ``joint_names`` and ``joint_ids`` are specified and are not consistent. ValueError: If both ``fixed_tendon_names`` and ``fixed_tendon_ids`` are specified and are not consistent. ValueError: If both ``body_names`` and ``body_ids`` are specified and are not consistent. + ValueError: If both ``object_collection_names`` and ``object_collection_ids`` are specified and are not consistent. """ # check if the entity is valid if self.name not in scene.keys(): @@ -120,6 +139,9 @@ def resolve(self, scene: InteractiveScene): # convert body names to indices based on regex self._resolve_body_names(scene) + # convert object collection names to indices based on regex + self._resolve_object_collection_names(scene) + def _resolve_joint_names(self, scene: InteractiveScene): # convert joint names to indices based on regex if self.joint_names is not None or self.joint_ids != slice(None): @@ -228,3 +250,36 @@ def _resolve_body_names(self, scene: InteractiveScene): if isinstance(self.body_ids, int): self.body_ids = [self.body_ids] self.body_names = [entity.body_names[i] for i in self.body_ids] + + def _resolve_object_collection_names(self, scene: InteractiveScene): + # convert object names to indices based on regex + if self.object_collection_names is not None or self.object_collection_ids != slice(None): + entity: RigidObjectCollection = scene[self.name] + # -- if both are not their default values, check if they are valid + if self.object_collection_names is not None and self.object_collection_ids != slice(None): + if isinstance(self.object_collection_names, str): + self.object_collection_names = [self.object_collection_names] + if isinstance(self.object_collection_ids, int): + self.object_collection_ids = [self.object_collection_ids] + object_ids, _ = entity.find_objects(self.object_collection_names, preserve_order=self.preserve_order) + object_names = [entity.object_names[i] for i in self.object_collection_ids] + if object_ids != self.object_collection_ids or object_names != self.object_collection_names: + raise ValueError( + "Both 'object_collection_names' and 'object_collection_ids' are specified, and are not" + " consistent.\n\tfrom object collection names:" + f" {self.object_collection_names} [{object_ids}]\n\tfrom object collection ids:" + f" {object_names} [{self.object_collection_ids}]\nHint: Use either 'object_collection_names' or" + " 'object_collection_ids' to avoid confusion." + ) + # -- from object names to object indices + elif self.object_collection_names is not None: + if isinstance(self.object_collection_names, str): + self.object_collection_names = [self.object_collection_names] + self.object_collection_ids, _ = entity.find_objects( + self.object_collection_names, preserve_order=self.preserve_order + ) + # -- from object indices to object names + elif self.object_collection_ids != slice(None): + if isinstance(self.object_collection_ids, int): + self.object_collection_ids = [self.object_collection_ids] + self.object_collection_names = [entity.object_names[i] for i in self.object_collection_ids] diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/termination_manager.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/termination_manager.py index bbd1924048..e19ed64fd2 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/termination_manager.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/managers/termination_manager.py @@ -53,6 +53,12 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv): cfg: The configuration object or dictionary (``dict[str, TerminationTermCfg]``). env: An environment object. """ + # create buffers to parse and store terms + self._term_names: list[str] = list() + self._term_cfgs: list[TerminationTermCfg] = list() + self._class_term_cfgs: list[TerminationTermCfg] = list() + + # call the base class constructor (this will parse the terms config) super().__init__(cfg, env) # prepare extra info to store individual termination term information self._term_dones = dict() @@ -178,6 +184,22 @@ def get_term(self, name: str) -> torch.Tensor: """ return self._term_dones[name] + def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: + """Returns the active terms as iterable sequence of tuples. + + The first element of the tuple is the name of the term and the second element is the raw value(s) of the term. + + Args: + env_idx: The specific environment to pull the active terms from. + + Returns: + The active terms. + """ + terms = [] + for key in self._term_dones.keys(): + terms.append((key, [self._term_dones[key][env_idx].float().cpu().item()])) + return terms + """ Operations - Term settings. """ @@ -219,12 +241,6 @@ def get_term_cfg(self, term_name: str) -> TerminationTermCfg: """ def _prepare_terms(self): - """Prepares a list of termination functions.""" - # parse remaining termination terms and decimate their information - self._term_names: list[str] = list() - self._term_cfgs: list[TerminationTermCfg] = list() - self._class_term_cfgs: list[TerminationTermCfg] = list() - # check if config is dict already if isinstance(self.cfg, dict): cfg_items = self.cfg.items() diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/scene/interactive_scene.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/scene/interactive_scene.py index 85775ba425..25dd4bcdb0 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/scene/interactive_scene.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/scene/interactive_scene.py @@ -22,6 +22,8 @@ DeformableObjectCfg, RigidObject, RigidObjectCfg, + RigidObjectCollection, + RigidObjectCollectionCfg, ) from omni.isaac.lab.sensors import ContactSensorCfg, FrameTransformerCfg, SensorBase, SensorBaseCfg from omni.isaac.lab.terrains import TerrainImporter, TerrainImporterCfg @@ -104,6 +106,8 @@ def __init__(self, cfg: InteractiveSceneCfg): Args: cfg: The configuration class for the scene. """ + # check that the config is valid + cfg.validate() # store inputs self.cfg = cfg # initialize scene elements @@ -111,6 +115,7 @@ def __init__(self, cfg: InteractiveSceneCfg): self._articulations = dict() self._deformable_objects = dict() self._rigid_objects = dict() + self._rigid_object_collections = dict() self._sensors = dict() self._extras = dict() # obtain the current stage @@ -307,6 +312,11 @@ def rigid_objects(self) -> dict[str, RigidObject]: """A dictionary of rigid objects in the scene.""" return self._rigid_objects + @property + def rigid_object_collections(self) -> dict[str, RigidObjectCollection]: + """A dictionary of rigid object collections in the scene.""" + return self._rigid_object_collections + @property def sensors(self) -> dict[str, SensorBase]: """A dictionary of the sensors in the scene, such as cameras and contact reporters.""" @@ -331,6 +341,56 @@ def extras(self) -> dict[str, XFormPrimView]: """ return self._extras + @property + def state(self) -> dict[str, dict[str, dict[str, torch.Tensor]]]: + """Returns the state of the scene entities. + + Returns: + A dictionary of the state of the scene entities. + """ + return self.get_state(is_relative=False) + + def get_state(self, is_relative: bool = False) -> dict[str, dict[str, dict[str, torch.Tensor]]]: + """Returns the state of the scene entities. + + Args: + is_relative: If set to True, the state is considered relative to the environment origins. + + Returns: + A dictionary of the state of the scene entities. + """ + state = dict() + # articulations + state["articulation"] = dict() + for asset_name, articulation in self._articulations.items(): + asset_state = dict() + asset_state["root_pose"] = articulation.data.root_link_state_w[:, :7].clone() + if is_relative: + asset_state["root_pose"][:, :3] -= self.env_origins + asset_state["root_velocity"] = articulation.data.root_com_vel_w.clone() + asset_state["joint_position"] = articulation.data.joint_pos.clone() + asset_state["joint_velocity"] = articulation.data.joint_vel.clone() + state["articulation"][asset_name] = asset_state + # deformable objects + state["deformable_object"] = dict() + for asset_name, deformable_object in self._deformable_objects.items(): + asset_state = dict() + asset_state["nodal_position"] = deformable_object.data.nodal_pos_w.clone() + if is_relative: + asset_state["nodal_position"][:, :3] -= self.env_origins + asset_state["nodal_velocity"] = deformable_object.data.nodal_vel_w.clone() + state["deformable_object"][asset_name] = asset_state + # rigid objects + state["rigid_object"] = dict() + for asset_name, rigid_object in self._rigid_objects.items(): + asset_state = dict() + asset_state["root_pose"] = rigid_object.data.root_link_state_w[:, :7].clone() + if is_relative: + asset_state["root_pose"][:, :3] -= self.env_origins + asset_state["root_velocity"] = rigid_object.data.root_com_vel_w.clone() + state["rigid_object"][asset_name] = asset_state + return state + """ Operations. """ @@ -349,10 +409,64 @@ def reset(self, env_ids: Sequence[int] | None = None): deformable_object.reset(env_ids) for rigid_object in self._rigid_objects.values(): rigid_object.reset(env_ids) + for rigid_object_collection in self._rigid_object_collections.values(): + rigid_object_collection.reset(env_ids) # -- sensors for sensor in self._sensors.values(): sensor.reset(env_ids) + def reset_to( + self, + state: dict[str, dict[str, dict[str, torch.Tensor]]], + env_ids: Sequence[int] | None = None, + is_relative: bool = False, + ): + """Resets the scene entities to the given state. + + Args: + state: The state to reset the scene entities to. + env_ids: The indices of the environments to reset. + Defaults to None (all instances). + is_relative: If set to True, the state is considered relative to the environment origins. + """ + if env_ids is None: + env_ids = slice(None) + # articulations + for asset_name, articulation in self._articulations.items(): + asset_state = state["articulation"][asset_name] + # root state + root_pose = asset_state["root_pose"].clone() + if is_relative: + root_pose[:, :3] += self.env_origins[env_ids] + root_velocity = asset_state["root_velocity"].clone() + articulation.write_root_link_pose_to_sim(root_pose, env_ids=env_ids) + articulation.write_root_com_velocity_to_sim(root_velocity, env_ids=env_ids) + # joint state + joint_position = asset_state["joint_position"].clone() + joint_velocity = asset_state["joint_velocity"].clone() + articulation.write_joint_state_to_sim(joint_position, joint_velocity, env_ids=env_ids) + articulation.set_joint_position_target(joint_position, env_ids=env_ids) + articulation.set_joint_velocity_target(joint_velocity, env_ids=env_ids) + # deformable objects + for asset_name, deformable_object in self._deformable_objects.items(): + asset_state = state["deformable_object"][asset_name] + nodal_position = asset_state["nodal_position"].clone() + if is_relative: + nodal_position[:, :3] += self.env_origins[env_ids] + nodal_velocity = asset_state["nodal_velocity"].clone() + deformable_object.write_nodal_pos_to_sim(nodal_position, env_ids=env_ids) + deformable_object.write_nodal_velocity_to_sim(nodal_velocity, env_ids=env_ids) + # rigid objects + for asset_name, rigid_object in self._rigid_objects.items(): + asset_state = state["rigid_object"][asset_name] + root_pose = asset_state["root_pose"].clone() + if is_relative: + root_pose[:, :3] += self.env_origins[env_ids] + root_velocity = asset_state["root_velocity"].clone() + rigid_object.write_root_link_pose_to_sim(root_pose, env_ids=env_ids) + rigid_object.write_root_com_velocity_to_sim(root_velocity, env_ids=env_ids) + self.write_data_to_sim() + def write_data_to_sim(self): """Writes the data of the scene entities to the simulation.""" # -- assets @@ -362,6 +476,8 @@ def write_data_to_sim(self): deformable_object.write_data_to_sim() for rigid_object in self._rigid_objects.values(): rigid_object.write_data_to_sim() + for rigid_object_collection in self._rigid_object_collections.values(): + rigid_object_collection.write_data_to_sim() def update(self, dt: float) -> None: """Update the scene entities. @@ -376,6 +492,8 @@ def update(self, dt: float) -> None: deformable_object.update(dt) for rigid_object in self._rigid_objects.values(): rigid_object.update(dt) + for rigid_object_collection in self._rigid_object_collections.values(): + rigid_object_collection.update(dt) # -- sensors for sensor in self._sensors.values(): sensor.update(dt, force_recompute=not self.cfg.lazy_sensor_update) @@ -395,6 +513,7 @@ def keys(self) -> list[str]: self._articulations, self._deformable_objects, self._rigid_objects, + self._rigid_object_collections, self._sensors, self._extras, ]: @@ -420,6 +539,7 @@ def __getitem__(self, key: str) -> Any: self._articulations, self._deformable_objects, self._rigid_objects, + self._rigid_object_collections, self._sensors, self._extras, ]: @@ -452,7 +572,8 @@ def _add_entities_from_cfg(self): if asset_name in InteractiveSceneCfg.__dataclass_fields__ or asset_cfg is None: continue # resolve regex - asset_cfg.prim_path = asset_cfg.prim_path.format(ENV_REGEX_NS=self.env_regex_ns) + if hasattr(asset_cfg, "prim_path"): + asset_cfg.prim_path = asset_cfg.prim_path.format(ENV_REGEX_NS=self.env_regex_ns) # create asset if isinstance(asset_cfg, TerrainImporterCfg): # terrains are special entities since they define environment origins @@ -465,6 +586,14 @@ def _add_entities_from_cfg(self): self._deformable_objects[asset_name] = asset_cfg.class_type(asset_cfg) elif isinstance(asset_cfg, RigidObjectCfg): self._rigid_objects[asset_name] = asset_cfg.class_type(asset_cfg) + elif isinstance(asset_cfg, RigidObjectCollectionCfg): + for rigid_object_cfg in asset_cfg.rigid_objects.values(): + rigid_object_cfg.prim_path = rigid_object_cfg.prim_path.format(ENV_REGEX_NS=self.env_regex_ns) + self._rigid_object_collections[asset_name] = asset_cfg.class_type(asset_cfg) + for rigid_object_cfg in asset_cfg.rigid_objects.values(): + if hasattr(rigid_object_cfg, "collision_group") and rigid_object_cfg.collision_group == -1: + asset_paths = sim_utils.find_matching_prim_paths(rigid_object_cfg.prim_path) + self._global_prim_paths += asset_paths elif isinstance(asset_cfg, SensorBaseCfg): # Update target frame path(s)' regex name space for FrameTransformer if isinstance(asset_cfg, FrameTransformerCfg): diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/__init__.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/__init__.py index 7ec74d5047..72f1a292a6 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/__init__.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/__init__.py @@ -30,12 +30,15 @@ +---------------------+---------------------------+---------------------------------------------------------------+ | Frame Transformer | /World/robot/base | Leaf exists and is a physics body (Articulation / Rigid Body) | +---------------------+---------------------------+---------------------------------------------------------------+ +| Imu | /World/robot/base | Leaf exists and is a physics body (Rigid Body) | ++---------------------+---------------------------+---------------------------------------------------------------+ """ from .camera import * # noqa: F401, F403 from .contact_sensor import * # noqa: F401, F403 from .frame_transformer import * # noqa: F401 +from .imu import * # noqa: F401, F403 from .ray_caster import * # noqa: F401, F403 from .sensor_base import SensorBase # noqa: F401 from .sensor_base_cfg import SensorBaseCfg # noqa: F401 diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/camera.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/camera.py index 49bbbc7ca0..091a8c1e6e 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/camera.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/camera.py @@ -9,10 +9,10 @@ import re import torch from collections.abc import Sequence -from tensordict import TensorDict from typing import TYPE_CHECKING, Any, Literal import carb +import omni.isaac.core.utils.stage as stage_utils import omni.kit.commands import omni.usd from omni.isaac.core.prims import XFormPrimView @@ -21,11 +21,14 @@ import omni.isaac.lab.sim as sim_utils from omni.isaac.lab.utils import to_camel_case from omni.isaac.lab.utils.array import convert_to_torch -from omni.isaac.lab.utils.math import quat_from_matrix +from omni.isaac.lab.utils.math import ( + convert_camera_frame_orientation_convention, + create_rotation_matrix_from_view, + quat_from_matrix, +) from ..sensor_base import SensorBase from .camera_data import CameraData -from .utils import convert_orientation_convention, create_rotation_matrix_from_view if TYPE_CHECKING: from .camera_cfg import CameraCfg @@ -116,7 +119,9 @@ def __init__(self, cfg: CameraCfg): if self.cfg.spawn is not None: # compute the rotation offset rot = torch.tensor(self.cfg.offset.rot, dtype=torch.float32).unsqueeze(0) - rot_offset = convert_orientation_convention(rot, origin=self.cfg.offset.convention, target="opengl") + rot_offset = convert_camera_frame_orientation_convention( + rot, origin=self.cfg.offset.convention, target="opengl" + ) rot_offset = rot_offset.squeeze(0).numpy() # ensure vertical aperture is set, otherwise replace with default for squared pixels if self.cfg.spawn.vertical_aperture is None: @@ -150,7 +155,7 @@ def __str__(self) -> str: # message for class return ( f"Camera @ '{self.cfg.prim_path}': \n" - f"\tdata types : {self.data.output.sorted_keys} \n" + f"\tdata types : {list(self.data.output.keys())} \n" f"\tsemantic filter : {self.cfg.semantic_filter}\n" f"\tcolorize semantic segm. : {self.cfg.colorize_semantic_segmentation}\n" f"\tcolorize instance segm. : {self.cfg.colorize_instance_segmentation}\n" @@ -289,7 +294,7 @@ def set_world_poses( - :obj:`"ros"` - forward axis: +Z - up axis -Y - Offset is applied in the ROS convention - :obj:`"world"` - forward axis: +X - up axis +Z - Offset is applied in the World Frame convention - See :meth:`omni.isaac.lab.sensors.camera.utils.convert_orientation_convention` for more details + See :meth:`omni.isaac.lab.sensors.camera.utils.convert_camera_frame_orientation_convention` for more details on the conventions. Args: @@ -318,7 +323,7 @@ def set_world_poses( orientations = torch.from_numpy(orientations).to(device=self._device) elif not isinstance(orientations, torch.Tensor): orientations = torch.tensor(orientations, device=self._device) - orientations = convert_orientation_convention(orientations, origin=convention, target="opengl") + orientations = convert_camera_frame_orientation_convention(orientations, origin=convention, target="opengl") # set the pose self._view.set_world_poses(positions, orientations, env_ids) @@ -339,8 +344,10 @@ def set_world_poses_from_view( # resolve env_ids if env_ids is None: env_ids = self._ALL_INDICES + # get up axis of current stage + up_axis = stage_utils.get_stage_up_axis() # set camera poses using the view - orientations = quat_from_matrix(create_rotation_matrix_from_view(eyes, targets, device=self._device)) + orientations = quat_from_matrix(create_rotation_matrix_from_view(eyes, targets, up_axis, device=self._device)) self._view.set_world_poses(eyes, orientations, env_ids) """ @@ -489,7 +496,7 @@ def _update_buffers_impl(self, env_ids: Sequence[int]): self._update_poses(env_ids) # -- read the data from annotator registry # check if buffer is called for the first time. If so then, allocate the memory - if len(self._data.output.sorted_keys) == 0: + if len(self._data.output) == 0: # this is the first time buffer is called # it allocates memory for all the sensors self._create_annotator_data() @@ -506,6 +513,19 @@ def _update_buffers_impl(self, env_ids: Sequence[int]): self._data.output[name][index] = data # add info to output self._data.info[index][name] = info + # NOTE: The `distance_to_camera` annotator returns the distance to the camera optical center. However, + # the replicator depth clipping is applied w.r.t. to the image plane which may result in values + # larger than the clipping range in the output. We apply an additional clipping to ensure values + # are within the clipping range for all the annotators. + if name == "distance_to_camera": + self._data.output[name][self._data.output[name] > self.cfg.spawn.clipping_range[1]] = torch.inf + # apply defined clipping behavior + if ( + name == "distance_to_camera" or name == "distance_to_image_plane" + ) and self.cfg.depth_clipping_behavior != "none": + self._data.output[name][torch.isinf(self._data.output[name])] = ( + 0.0 if self.cfg.depth_clipping_behavior == "zero" else self.cfg.spawn.clipping_range[1] + ) """ Private Helpers @@ -544,7 +564,7 @@ def _create_buffers(self): # lazy allocation of data dictionary # since the size of the output data is not known in advance, we leave it as None # the memory will be allocated when the buffer() function is called for the first time. - self._data.output = TensorDict({}, batch_size=self._view.count, device=self.device) + self._data.output = {} self._data.info = [{name: None for name in self.cfg.data_types} for _ in range(self._view.count)] def _update_intrinsic_matrices(self, env_ids: Sequence[int]): @@ -596,7 +616,9 @@ def _update_poses(self, env_ids: Sequence[int]): # get the poses from the view poses, quat = self._view.get_world_poses(env_ids) self._data.pos_w[env_ids] = poses - self._data.quat_w_world[env_ids] = convert_orientation_convention(quat, origin="opengl", target="world") + self._data.quat_w_world[env_ids] = convert_camera_frame_orientation_convention( + quat, origin="opengl", target="world" + ) def _create_annotator_data(self): """Create the buffers to store the annotator data. @@ -622,6 +644,19 @@ def _create_annotator_data(self): self._data.info[index][name] = info # concatenate the data along the batch dimension self._data.output[name] = torch.stack(data_all_cameras, dim=0) + # NOTE: `distance_to_camera` and `distance_to_image_plane` are not both clipped to the maximum defined + # in the clipping range. The clipping is applied only to `distance_to_image_plane` and then both + # outputs are only clipped where the values in `distance_to_image_plane` exceed the threshold. To + # have a unified behavior between all cameras, we clip both outputs to the maximum value defined. + if name == "distance_to_camera": + self._data.output[name][self._data.output[name] > self.cfg.spawn.clipping_range[1]] = torch.inf + # clip the data if needed + if ( + name == "distance_to_camera" or name == "distance_to_image_plane" + ) and self.cfg.depth_clipping_behavior != "none": + self._data.output[name][torch.isinf(self._data.output[name])] = ( + 0.0 if self.cfg.depth_clipping_behavior == "zero" else self.cfg.spawn.clipping_range[1] + ) def _process_annotator_output(self, name: str, output: Any) -> tuple[torch.tensor, dict | None]: """Process the annotator output. diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/camera_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/camera_cfg.py index ca8367224f..18052d4d16 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/camera_cfg.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/camera_cfg.py @@ -53,6 +53,14 @@ class OffsetCfg: asset is already present in the scene. """ + depth_clipping_behavior: Literal["max", "zero", "none"] = "zero" + """Clipping behavior for the camera for values exceed the maximum value. Defaults to "zero". + + - ``"max"``: Values are clipped to the maximum value. + - ``"zero"``: Values are clipped to zero. + - ``"none``: No clipping is applied. Values will be returned as ``inf``. + """ + data_types: list[str] = ["rgb"] """List of sensor names/types to enable for the camera. Defaults to ["rgb"]. diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/camera_data.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/camera_data.py index d5f25f9e60..540fdb7dcf 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/camera_data.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/camera_data.py @@ -5,10 +5,9 @@ import torch from dataclasses import dataclass -from tensordict import TensorDict from typing import Any -from .utils import convert_orientation_convention +from omni.isaac.lab.utils.math import convert_camera_frame_orientation_convention @dataclass @@ -47,7 +46,7 @@ class CameraData: Shape is (N, 3, 3) where N is the number of sensors. """ - output: TensorDict = None + output: dict[str, torch.Tensor] = None """The retrieved sensor data with sensor types as key. The format of the data is available in the `Replicator Documentation`_. For semantic-based data, @@ -77,7 +76,7 @@ def quat_w_ros(self) -> torch.Tensor: Shape is (N, 4) where N is the number of sensors. """ - return convert_orientation_convention(self.quat_w_world, origin="world", target="ros") + return convert_camera_frame_orientation_convention(self.quat_w_world, origin="world", target="ros") @property def quat_w_opengl(self) -> torch.Tensor: @@ -89,4 +88,4 @@ def quat_w_opengl(self) -> torch.Tensor: Shape is (N, 4) where N is the number of sensors. """ - return convert_orientation_convention(self.quat_w_world, origin="world", target="opengl") + return convert_camera_frame_orientation_convention(self.quat_w_world, origin="world", target="opengl") diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/tiled_camera.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/tiled_camera.py index 1e465a4867..b9cd0efbae 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/tiled_camera.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/tiled_camera.py @@ -9,7 +9,6 @@ import numpy as np import torch from collections.abc import Sequence -from tensordict import TensorDict from typing import TYPE_CHECKING, Any import carb @@ -106,7 +105,7 @@ def __str__(self) -> str: # message for class return ( f"Tiled Camera @ '{self.cfg.prim_path}': \n" - f"\tdata types : {self.data.output.sorted_keys} \n" + f"\tdata types : {list(self.data.output.keys())} \n" f"\tsemantic filter : {self.cfg.semantic_filter}\n" f"\tcolorize semantic segm. : {self.cfg.colorize_semantic_segmentation}\n" f"\tcolorize instance segm. : {self.cfg.colorize_instance_segmentation}\n" @@ -280,6 +279,22 @@ def _update_buffers_impl(self, env_ids: Sequence[int]): if data_type == "rgba" and "rgb" in self.cfg.data_types: self._data.output["rgb"] = self._data.output["rgba"][..., :3] + # NOTE: The `distance_to_camera` annotator returns the distance to the camera optical center. However, + # the replicator depth clipping is applied w.r.t. to the image plane which may result in values + # larger than the clipping range in the output. We apply an additional clipping to ensure values + # are within the clipping range for all the annotators. + if data_type == "distance_to_camera": + self._data.output[data_type][ + self._data.output[data_type] > self.cfg.spawn.clipping_range[1] + ] = torch.inf + # apply defined clipping behavior + if ( + data_type == "distance_to_camera" or data_type == "distance_to_image_plane" or data_type == "depth" + ) and self.cfg.depth_clipping_behavior != "none": + self._data.output[data_type][torch.isinf(self._data.output[data_type])] = ( + 0.0 if self.cfg.depth_clipping_behavior == "zero" else self.cfg.spawn.clipping_range[1] + ) + """ Private Helpers """ @@ -372,7 +387,7 @@ def _create_buffers(self): (self._view.count, self.cfg.height, self.cfg.width, 1), device=self.device, dtype=torch.int32 ).contiguous() - self._data.output = TensorDict(data_dict, batch_size=self._view.count, device=self.device) + self._data.output = data_dict self._data.info = dict() def _tiled_image_shape(self) -> tuple[int, int]: diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/utils.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/utils.py index 4e8c4e63c1..f15f537e7f 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/utils.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/utils.py @@ -8,16 +8,11 @@ # needed to import for allowing type-hinting: torch.device | str | None from __future__ import annotations -import math import numpy as np import torch -import torch.nn.functional as F from collections.abc import Sequence -from typing import Literal -import omni.isaac.core.utils.stage as stage_utils import warp as wp -from pxr import UsdGeom import omni.isaac.lab.utils.math as math_utils from omni.isaac.lab.utils.array import TensorData, convert_to_torch @@ -262,143 +257,6 @@ def create_pointcloud_from_rgbd( return points_xyz, points_rgb -def convert_orientation_convention( - orientation: torch.Tensor, - origin: Literal["opengl", "ros", "world"] = "opengl", - target: Literal["opengl", "ros", "world"] = "ros", -) -> torch.Tensor: - r"""Converts a quaternion representing a rotation from one convention to another. - - In USD, the camera follows the ``"opengl"`` convention. Thus, it is always in **Y up** convention. - This means that the camera is looking down the -Z axis with the +Y axis pointing up , and +X axis pointing right. - However, in ROS, the camera is looking down the +Z axis with the +Y axis pointing down, and +X axis pointing right. - Thus, the camera needs to be rotated by :math:`180^{\circ}` around the X axis to follow the ROS convention. - - .. math:: - - T_{ROS} = \begin{bmatrix} 1 & 0 & 0 & 0 \\ 0 & -1 & 0 & 0 \\ 0 & 0 & -1 & 0 \\ 0 & 0 & 0 & 1 \end{bmatrix} T_{USD} - - On the other hand, the typical world coordinate system is with +X pointing forward, +Y pointing left, - and +Z pointing up. The camera can also be set in this convention by rotating the camera by :math:`90^{\circ}` - around the X axis and :math:`-90^{\circ}` around the Y axis. - - .. math:: - - T_{WORLD} = \begin{bmatrix} 0 & 0 & -1 & 0 \\ -1 & 0 & 0 & 0 \\ 0 & 1 & 0 & 0 \\ 0 & 0 & 0 & 1 \end{bmatrix} T_{USD} - - Thus, based on their application, cameras follow different conventions for their orientation. This function - converts a quaternion from one convention to another. - - Possible conventions are: - - - :obj:`"opengl"` - forward axis: -Z - up axis +Y - Offset is applied in the OpenGL (Usd.Camera) convention - - :obj:`"ros"` - forward axis: +Z - up axis -Y - Offset is applied in the ROS convention - - :obj:`"world"` - forward axis: +X - up axis +Z - Offset is applied in the World Frame convention - - Args: - orientation: Quaternion of form `(w, x, y, z)` with shape (..., 4) in source convention - origin: Convention to convert to. Defaults to "ros". - target: Convention to convert from. Defaults to "opengl". - - Returns: - Quaternion of form `(w, x, y, z)` with shape (..., 4) in target convention - """ - if target == origin: - return orientation.clone() - - # -- unify input type - if origin == "ros": - # convert from ros to opengl convention - rotm = math_utils.matrix_from_quat(orientation) - rotm[:, :, 2] = -rotm[:, :, 2] - rotm[:, :, 1] = -rotm[:, :, 1] - # convert to opengl convention - quat_gl = math_utils.quat_from_matrix(rotm) - elif origin == "world": - # convert from world (x forward and z up) to opengl convention - rotm = math_utils.matrix_from_quat(orientation) - rotm = torch.matmul( - rotm, - math_utils.matrix_from_euler( - torch.tensor([math.pi / 2, -math.pi / 2, 0], device=orientation.device), "XYZ" - ), - ) - # convert to isaac-sim convention - quat_gl = math_utils.quat_from_matrix(rotm) - else: - quat_gl = orientation - - # -- convert to target convention - if target == "ros": - # convert from opengl to ros convention - rotm = math_utils.matrix_from_quat(quat_gl) - rotm[:, :, 2] = -rotm[:, :, 2] - rotm[:, :, 1] = -rotm[:, :, 1] - return math_utils.quat_from_matrix(rotm) - elif target == "world": - # convert from opengl to world (x forward and z up) convention - rotm = math_utils.matrix_from_quat(quat_gl) - rotm = torch.matmul( - rotm, - math_utils.matrix_from_euler( - torch.tensor([math.pi / 2, -math.pi / 2, 0], device=orientation.device), "XYZ" - ).T, - ) - return math_utils.quat_from_matrix(rotm) - else: - return quat_gl.clone() - - -# @torch.jit.script -def create_rotation_matrix_from_view( - eyes: torch.Tensor, - targets: torch.Tensor, - device: str = "cpu", -) -> torch.Tensor: - """ - This function takes a vector ''eyes'' which specifies the location - of the camera in world coordinates and the vector ''targets'' which - indicate the position of the object. - The output is a rotation matrix representing the transformation - from world coordinates -> view coordinates. - - The inputs camera_position and targets can each be a - - 3 element tuple/list - - torch tensor of shape (1, 3) - - torch tensor of shape (N, 3) - - Args: - eyes: position of the camera in world coordinates - targets: position of the object in world coordinates - - The vectors are broadcast against each other so they all have shape (N, 3). - - Returns: - R: (N, 3, 3) batched rotation matrices - - Reference: - Based on PyTorch3D (https://github.com/facebookresearch/pytorch3d/blob/eaf0709d6af0025fe94d1ee7cec454bc3054826a/pytorch3d/renderer/cameras.py#L1635-L1685) - """ - up_axis_token = stage_utils.get_stage_up_axis() - if up_axis_token == UsdGeom.Tokens.y: - up_axis = torch.tensor((0, 1, 0), device=device, dtype=torch.float32).repeat(eyes.shape[0], 1) - elif up_axis_token == UsdGeom.Tokens.z: - up_axis = torch.tensor((0, 0, 1), device=device, dtype=torch.float32).repeat(eyes.shape[0], 1) - else: - raise ValueError(f"Invalid up axis: {up_axis_token}") - - # get rotation matrix in opengl format (-Z forward, +Y up) - z_axis = -F.normalize(targets - eyes, eps=1e-5) - x_axis = F.normalize(torch.cross(up_axis, z_axis, dim=1), eps=1e-5) - y_axis = F.normalize(torch.cross(z_axis, x_axis, dim=1), eps=1e-5) - is_close = torch.isclose(x_axis, torch.tensor(0.0), atol=5e-3).all(dim=1, keepdim=True) - if is_close.any(): - replacement = F.normalize(torch.cross(y_axis, z_axis, dim=1), eps=1e-5) - x_axis = torch.where(is_close, replacement, x_axis) - R = torch.cat((x_axis[:, None, :], y_axis[:, None, :], z_axis[:, None, :]), dim=1) - return R.transpose(1, 2) - - def save_images_to_file(images: torch.Tensor, file_path: str): """Save images to file. diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/imu/__init__.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/imu/__init__.py new file mode 100644 index 0000000000..c1d2015e85 --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/imu/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +""" +Imu Sensor +""" + +from .imu import Imu +from .imu_cfg import ImuCfg +from .imu_data import ImuData diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/imu/imu.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/imu/imu.py new file mode 100644 index 0000000000..feb434a81f --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/imu/imu.py @@ -0,0 +1,243 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import omni.isaac.core.utils.stage as stage_utils +import omni.physics.tensors.impl.api as physx +from pxr import UsdPhysics + +import omni.isaac.lab.sim as sim_utils +import omni.isaac.lab.utils.math as math_utils +from omni.isaac.lab.markers import VisualizationMarkers + +from ..sensor_base import SensorBase +from .imu_data import ImuData + +if TYPE_CHECKING: + from .imu_cfg import ImuCfg + + +class Imu(SensorBase): + """The Inertia Measurement Unit (IMU) sensor. + + The sensor can be attached to any :class:`RigidObject` or :class:`Articulation` in the scene. The sensor provides complete state information. + The sensor is primarily used to provide the linear acceleration and angular velocity of the object in the body frame. The sensor also provides + the position and orientation of the object in the world frame and the angular acceleration and linear velocity in the body frame. The extra + data outputs are useful for simulating with or comparing against "perfect" state estimation. + + .. note:: + + We are computing the accelerations using numerical differentiation from the velocities. Consequently, the + IMU sensor accuracy depends on the chosen phsyx timestep. For a sufficient accuracy, we recommend to keep the + timestep at least as 200Hz. + + .. note:: + + It is suggested to use the OffsetCfg to define an IMU frame relative to a rigid body prim defined at the root of + a :class:`RigidObject` or a prim that is defined by a non-fixed joint in an :class:`Articulation` (except for the + root of a fixed based articulation). The use frames with fixed joints and small mass/inertia to emulate a transform + relative to a body frame can result in lower performance and accuracy. + + """ + + cfg: ImuCfg + """The configuration parameters.""" + + def __init__(self, cfg: ImuCfg): + """Initializes the Imu sensor. + + Args: + cfg: The configuration parameters. + """ + # initialize base class + super().__init__(cfg) + # Create empty variables for storing output data + self._data = ImuData() + + def __str__(self) -> str: + """Returns: A string containing information about the instance.""" + return ( + f"Imu sensor @ '{self.cfg.prim_path}': \n" + f"\tview type : {self._view.__class__}\n" + f"\tupdate period (s) : {self.cfg.update_period}\n" + f"\tnumber of sensors : {self._view.count}\n" + ) + + """ + Properties + """ + + @property + def data(self) -> ImuData: + # update sensors if needed + self._update_outdated_buffers() + # return the data + return self._data + + @property + def num_instances(self) -> int: + return self._view.count + + """ + Operations + """ + + def reset(self, env_ids: Sequence[int] | None = None): + # reset the timestamps + super().reset(env_ids) + # resolve None + if env_ids is None: + env_ids = slice(None) + # reset accumulative data buffers + self._data.quat_w[env_ids] = 0.0 + self._data.lin_vel_b[env_ids] = 0.0 + self._data.ang_vel_b[env_ids] = 0.0 + self._data.lin_acc_b[env_ids] = 0.0 + self._data.ang_acc_b[env_ids] = 0.0 + + def update(self, dt: float, force_recompute: bool = False): + # save timestamp + self._dt = dt + # execute updating + super().update(dt, force_recompute) + + """ + Implementation. + """ + + def _initialize_impl(self): + """Initializes the sensor handles and internal buffers. + + This function creates handles and registers the provided data types with the replicator registry to + be able to access the data from the sensor. It also initializes the internal buffers to store the data. + + Raises: + RuntimeError: If the imu prim is not a RigidBodyPrim + """ + # Initialize parent class + super()._initialize_impl() + # create simulation view + self._physics_sim_view = physx.create_simulation_view(self._backend) + self._physics_sim_view.set_subspace_roots("/") + # check if the prim at path is a rigid prim + prim = sim_utils.find_first_matching_prim(self.cfg.prim_path) + if prim is None: + raise RuntimeError(f"Failed to find a prim at path expression: {self.cfg.prim_path}") + # check if it is a RigidBody Prim + if prim.HasAPI(UsdPhysics.RigidBodyAPI): + self._view = self._physics_sim_view.create_rigid_body_view(self.cfg.prim_path.replace(".*", "*")) + else: + raise RuntimeError(f"Failed to find a RigidBodyAPI for the prim paths: {self.cfg.prim_path}") + + # Create internal buffers + self._initialize_buffers_impl() + + def _update_buffers_impl(self, env_ids: Sequence[int]): + """Fills the buffers of the sensor data.""" + # check if self._dt is set (this is set in the update function) + if not hasattr(self, "_dt"): + raise RuntimeError( + "The update function must be called before the data buffers are accessed the first time." + ) + # default to all sensors + if len(env_ids) == self._num_envs: + env_ids = slice(None) + # obtain the poses of the sensors + pos_w, quat_w = self._view.get_transforms()[env_ids].split([3, 4], dim=-1) + quat_w = math_utils.convert_quat(quat_w, to="wxyz") + + # store the poses + self._data.pos_w[env_ids] = pos_w + math_utils.quat_rotate(quat_w, self._offset_pos_b[env_ids]) + self._data.quat_w[env_ids] = math_utils.quat_mul(quat_w, self._offset_quat_b[env_ids]) + + # get the offset from COM to link origin + com_pos_b = self._view.get_coms().to(self.device).split([3, 4], dim=-1)[0] + + # obtain the velocities of the link COM + lin_vel_w, ang_vel_w = self._view.get_velocities()[env_ids].split([3, 3], dim=-1) + # if an offset is present or the COM does not agree with the link origin, the linear velocity has to be + # transformed taking the angular velocity into account + lin_vel_w += torch.linalg.cross( + ang_vel_w, math_utils.quat_rotate(quat_w, self._offset_pos_b[env_ids] - com_pos_b[env_ids]), dim=-1 + ) + + # numerical derivative + lin_acc_w = (lin_vel_w - self._prev_lin_vel_w[env_ids]) / self._dt + self._gravity_bias_w[env_ids] + ang_acc_w = (ang_vel_w - self._prev_ang_vel_w[env_ids]) / self._dt + # store the velocities + self._data.lin_vel_b[env_ids] = math_utils.quat_rotate_inverse(self._data.quat_w[env_ids], lin_vel_w) + self._data.ang_vel_b[env_ids] = math_utils.quat_rotate_inverse(self._data.quat_w[env_ids], ang_vel_w) + # store the accelerations + self._data.lin_acc_b[env_ids] = math_utils.quat_rotate_inverse(self._data.quat_w[env_ids], lin_acc_w) + self._data.ang_acc_b[env_ids] = math_utils.quat_rotate_inverse(self._data.quat_w[env_ids], ang_acc_w) + + self._prev_lin_vel_w[env_ids] = lin_vel_w + self._prev_ang_vel_w[env_ids] = ang_vel_w + + def _initialize_buffers_impl(self): + """Create buffers for storing data.""" + # data buffers + self._data.pos_w = torch.zeros(self._view.count, 3, device=self._device) + self._data.quat_w = torch.zeros(self._view.count, 4, device=self._device) + self._data.quat_w[:, 0] = 1.0 + self._data.lin_vel_b = torch.zeros_like(self._data.pos_w) + self._data.ang_vel_b = torch.zeros_like(self._data.pos_w) + self._data.lin_acc_b = torch.zeros_like(self._data.pos_w) + self._data.ang_acc_b = torch.zeros_like(self._data.pos_w) + self._prev_lin_vel_w = torch.zeros_like(self._data.pos_w) + self._prev_ang_vel_w = torch.zeros_like(self._data.pos_w) + + # store sensor offset transformation + self._offset_pos_b = torch.tensor(list(self.cfg.offset.pos), device=self._device).repeat(self._view.count, 1) + self._offset_quat_b = torch.tensor(list(self.cfg.offset.rot), device=self._device).repeat(self._view.count, 1) + # set gravity bias + self._gravity_bias_w = torch.tensor(list(self.cfg.gravity_bias), device=self._device).repeat( + self._view.count, 1 + ) + + def _set_debug_vis_impl(self, debug_vis: bool): + # set visibility of markers + # note: parent only deals with callbacks. not their visibility + if debug_vis: + # create markers if necessary for the first tome + if not hasattr(self, "acceleration_visualizer"): + self.acceleration_visualizer = VisualizationMarkers(self.cfg.visualizer_cfg) + # set their visibility to true + self.acceleration_visualizer.set_visibility(True) + else: + if hasattr(self, "acceleration_visualizer"): + self.acceleration_visualizer.set_visibility(False) + + def _debug_vis_callback(self, event): + # safely return if view becomes invalid + # note: this invalidity happens because of isaac sim view callbacks + if self._view is None: + return + # get marker location + # -- base state + base_pos_w = self._data.pos_w.clone() + base_pos_w[:, 2] += 0.5 + # -- resolve the scales + default_scale = self.acceleration_visualizer.cfg.markers["arrow"].scale + arrow_scale = torch.tensor(default_scale, device=self.device).repeat(self._data.lin_acc_b.shape[0], 1) + # get up axis of current stage + up_axis = stage_utils.get_stage_up_axis() + # arrow-direction + quat_opengl = math_utils.quat_from_matrix( + math_utils.create_rotation_matrix_from_view( + self._data.pos_w, + self._data.pos_w + math_utils.quat_rotate(self._data.quat_w, self._data.lin_acc_b), + up_axis=up_axis, + device=self._device, + ) + ) + quat_w = math_utils.convert_camera_frame_orientation_convention(quat_opengl, "opengl", "world") + # display markers + self.acceleration_visualizer.visualize(base_pos_w, quat_w, arrow_scale) diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/imu/imu_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/imu/imu_cfg.py new file mode 100644 index 0000000000..74e38f9e6e --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/imu/imu_cfg.py @@ -0,0 +1,46 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from omni.isaac.lab.markers import VisualizationMarkersCfg +from omni.isaac.lab.markers.config import RED_ARROW_X_MARKER_CFG +from omni.isaac.lab.utils import configclass + +from ..sensor_base_cfg import SensorBaseCfg +from .imu import Imu + + +@configclass +class ImuCfg(SensorBaseCfg): + """Configuration for an Inertial Measurement Unit (IMU) sensor.""" + + class_type: type = Imu + + @configclass + class OffsetCfg: + """The offset pose of the sensor's frame from the sensor's parent frame.""" + + pos: tuple[float, float, float] = (0.0, 0.0, 0.0) + """Translation w.r.t. the parent frame. Defaults to (0.0, 0.0, 0.0).""" + + rot: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0) + """Quaternion rotation (w, x, y, z) w.r.t. the parent frame. Defaults to (1.0, 0.0, 0.0, 0.0).""" + + offset: OffsetCfg = OffsetCfg() + """The offset pose of the sensor's frame from the sensor's parent frame. Defaults to identity.""" + + visualizer_cfg: VisualizationMarkersCfg = RED_ARROW_X_MARKER_CFG.replace(prim_path="/Visuals/Command/velocity_goal") + """The configuration object for the visualization markers. Defaults to RED_ARROW_X_MARKER_CFG. + + This attribute is only used when debug visualization is enabled. + """ + gravity_bias: tuple[float, float, float] = (0.0, 0.0, 9.81) + """The linear acceleration bias applied to the linear acceleration in the world frame (x,y,z). + + Imu sensors typically output a positive gravity acceleration in opposition to the direction of gravity. This + config parameter allows users to subtract that bias if set to (0.,0.,0.). By default this is set to (0.0,0.0,9.81) + which results in a positive acceleration reading in the world Z. + """ diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/imu/imu_data.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/imu/imu_data.py new file mode 100644 index 0000000000..37920cd841 --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/imu/imu_data.py @@ -0,0 +1,50 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from dataclasses import dataclass + + +@dataclass +class ImuData: + """Data container for the Imu sensor.""" + + pos_w: torch.Tensor = None + """Position of the sensor origin in world frame. + + Shape is (N, 3), where ``N`` is the number of environments. + """ + + quat_w: torch.Tensor = None + """Orientation of the sensor origin in quaternion ``(w, x, y, z)`` in world frame. + + Shape is (N, 4), where ``N`` is the number of environments. + """ + + lin_vel_b: torch.Tensor = None + """IMU frame angular velocity relative to the world expressed in IMU frame. + + Shape is (N, 3), where ``N`` is the number of environments. + """ + + ang_vel_b: torch.Tensor = None + """IMU frame angular velocity relative to the world expressed in IMU frame. + + Shape is (N, 3), where ``N`` is the number of environments. + """ + + lin_acc_b: torch.Tensor = None + """IMU frame linear acceleration relative to the world expressed in IMU frame. + + Shape is (N, 3), where ``N`` is the number of environments. + """ + + ang_acc_b: torch.Tensor = None + """IMU frame angular acceleration relative to the world expressed in IMU frame. + + Shape is (N, 3), where ``N`` is the number of environments. + """ diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/ray_caster/ray_caster.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/ray_caster/ray_caster.py index 8be96c358d..04e019f6f9 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/ray_caster/ray_caster.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/ray_caster/ray_caster.py @@ -9,7 +9,7 @@ import re import torch from collections.abc import Sequence -from typing import TYPE_CHECKING, ClassVar +from typing import TYPE_CHECKING import omni.log import omni.physics.tensors.impl.api as physx @@ -48,14 +48,6 @@ class RayCaster(SensorBase): cfg: RayCasterCfg """The configuration parameters.""" - meshes: ClassVar[dict[str, wp.Mesh]] = {} - """The warp meshes available for raycasting. - - The keys correspond to the prim path for the meshes, and values are the corresponding warp Mesh objects. - - Note: - We store a global dictionary of all warp meshes to prevent re-loading the mesh for different ray-cast sensor instances. - """ def __init__(self, cfg: RayCasterCfg): """Initializes the ray-caster object. @@ -77,6 +69,8 @@ def __init__(self, cfg: RayCasterCfg): super().__init__(cfg) # Create empty variables for storing output data self._data = RayCasterData() + # the warp meshes used for raycasting. + self.meshes: dict[str, wp.Mesh] = {} def __str__(self) -> str: """Returns: A string containing information about the instance.""" @@ -84,7 +78,7 @@ def __str__(self) -> str: f"Ray-caster @ '{self.cfg.prim_path}': \n" f"\tview type : {self._view.__class__}\n" f"\tupdate period (s) : {self.cfg.update_period}\n" - f"\tnumber of meshes : {len(RayCaster.meshes)}\n" + f"\tnumber of meshes : {len(self.meshes)}\n" f"\tnumber of sensors : {self._view.count}\n" f"\tnumber of rays/sensor: {self.num_rays}\n" f"\ttotal number of rays : {self.num_rays * self._view.count}" @@ -163,10 +157,6 @@ def _initialize_warp_meshes(self): # read prims to ray-cast for mesh_prim_path in self.cfg.mesh_prim_paths: - # check if mesh already casted into warp mesh - if mesh_prim_path in RayCaster.meshes: - continue - # check if the prim is a plane - handle PhysX plane as a special case # if a plane exists then we need to create an infinite mesh that is a plane mesh_prim = sim_utils.get_first_matching_child_prim( @@ -185,6 +175,9 @@ def _initialize_warp_meshes(self): mesh_prim = UsdGeom.Mesh(mesh_prim) # read the vertices and faces points = np.asarray(mesh_prim.GetPointsAttr().Get()) + transform_matrix = np.array(omni.usd.get_world_transform_matrix(mesh_prim)).T + points = np.matmul(points, transform_matrix[:3, :3].T) + points += transform_matrix[:3, 3] indices = np.asarray(mesh_prim.GetFaceVertexIndicesAttr().Get()) wp_mesh = convert_to_warp_mesh(points, indices, device=self.device) # print info @@ -197,10 +190,10 @@ def _initialize_warp_meshes(self): # print info omni.log.info(f"Created infinite plane mesh prim: {mesh_prim.GetPath()}.") # add the warp mesh to the list - RayCaster.meshes[mesh_prim_path] = wp_mesh + self.meshes[mesh_prim_path] = wp_mesh # throw an error if no meshes are found - if all([mesh_prim_path not in RayCaster.meshes for mesh_prim_path in self.cfg.mesh_prim_paths]): + if all([mesh_prim_path not in self.meshes for mesh_prim_path in self.cfg.mesh_prim_paths]): raise RuntimeError( f"No meshes found for ray-casting! Please check the mesh prim paths: {self.cfg.mesh_prim_paths}" ) @@ -263,7 +256,7 @@ def _update_buffers_impl(self, env_ids: Sequence[int]): ray_starts_w, ray_directions_w, max_dist=self.cfg.max_distance, - mesh=RayCaster.meshes[self.cfg.mesh_prim_paths[0]], + mesh=self.meshes[self.cfg.mesh_prim_paths[0]], )[0] def _set_debug_vis_impl(self, debug_vis: bool): diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/ray_caster/ray_caster_camera.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/ray_caster/ray_caster_camera.py index 745033003c..27a4d323b4 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/ray_caster/ray_caster_camera.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/ray_caster/ray_caster_camera.py @@ -7,15 +7,14 @@ import torch from collections.abc import Sequence -from tensordict import TensorDict from typing import TYPE_CHECKING, ClassVar, Literal +import omni.isaac.core.utils.stage as stage_utils import omni.physics.tensors.impl.api as physx from omni.isaac.core.prims import XFormPrimView import omni.isaac.lab.utils.math as math_utils from omni.isaac.lab.sensors.camera import CameraData -from omni.isaac.lab.sensors.camera.utils import convert_orientation_convention, create_rotation_matrix_from_view from omni.isaac.lab.utils.warp import raycast_mesh from .ray_caster import RayCaster @@ -87,7 +86,7 @@ def __str__(self) -> str: f"Ray-Caster-Camera @ '{self.cfg.prim_path}': \n" f"\tview type : {self._view.__class__}\n" f"\tupdate period (s) : {self.cfg.update_period}\n" - f"\tnumber of meshes : {len(RayCaster.meshes)}\n" + f"\tnumber of meshes : {len(self.meshes)}\n" f"\tnumber of sensors : {self._view.count}\n" f"\tnumber of rays/sensor: {self.num_rays}\n" f"\ttotal number of rays : {self.num_rays * self._view.count}\n" @@ -170,7 +169,7 @@ def set_world_poses( - :obj:`"ros"` - forward axis: +Z - up axis -Y - Offset is applied in the ROS convention - :obj:`"world"` - forward axis: +X - up axis +Z - Offset is applied in the World Frame convention - See :meth:`omni.isaac.lab.sensors.camera.utils.convert_orientation_convention` for more details + See :meth:`omni.isaac.lab.utils.maths.convert_camera_frame_orientation_convention` for more details on the conventions. Args: @@ -196,7 +195,9 @@ def set_world_poses( self._offset_pos[env_ids] = math_utils.quat_apply(math_utils.quat_inv(quat_w), pos_offset_world_frame) if orientations is not None: # convert rotation matrix from input convention to world - quat_w_set = convert_orientation_convention(orientations, origin=convention, target="world") + quat_w_set = math_utils.convert_camera_frame_orientation_convention( + orientations, origin=convention, target="world" + ) self._offset_quat[env_ids] = math_utils.quat_mul(math_utils.quat_inv(quat_w), quat_w_set) # update the data @@ -218,8 +219,12 @@ def set_world_poses_from_view( RuntimeError: If the camera prim is not set. Need to call :meth:`initialize` method first. NotImplementedError: If the stage up-axis is not "Y" or "Z". """ + # get up axis of current stage + up_axis = stage_utils.get_stage_up_axis() # camera position and rotation in opengl convention - orientations = math_utils.quat_from_matrix(create_rotation_matrix_from_view(eyes, targets, device=self._device)) + orientations = math_utils.quat_from_matrix( + math_utils.create_rotation_matrix_from_view(eyes, targets, up_axis=up_axis, device=self._device) + ) self.set_world_poses(eyes, orientations, env_ids, convention="opengl") """ @@ -243,7 +248,7 @@ def _initialize_rays_impl(self): # create buffer to store ray hits self.ray_hits_w = torch.zeros(self._view.count, self.num_rays, 3, device=self._device) # set offsets - quat_w = convert_orientation_convention( + quat_w = math_utils.convert_camera_frame_orientation_convention( torch.tensor([self.cfg.offset.rot], device=self._device), origin=self.cfg.offset.convention, target="world" ) self._offset_quat = quat_w.repeat(self._view.count, 1) @@ -275,7 +280,7 @@ def _update_buffers_impl(self, env_ids: Sequence[int]): self.ray_hits_w, ray_depth, ray_normal, _ = raycast_mesh( ray_starts_w, ray_directions_w, - mesh=RayCasterCamera.meshes[self.cfg.mesh_prim_paths[0]], + mesh=self.meshes[self.cfg.mesh_prim_paths[0]], max_dist=1e6, return_distance=any( [name in self.cfg.data_types for name in ["distance_to_image_plane", "distance_to_camera"]] @@ -292,14 +297,23 @@ def _update_buffers_impl(self, env_ids: Sequence[int]): ) )[:, :, 0] # apply the maximum distance after the transformation - distance_to_image_plane = torch.clip(distance_to_image_plane, max=self.cfg.max_distance) + if self.cfg.depth_clipping_behavior == "max": + distance_to_image_plane = torch.clip(distance_to_image_plane, max=self.cfg.max_distance) + distance_to_image_plane[torch.isnan(distance_to_image_plane)] = self.cfg.max_distance + elif self.cfg.depth_clipping_behavior == "zero": + distance_to_image_plane[distance_to_image_plane > self.cfg.max_distance] = 0.0 + distance_to_image_plane[torch.isnan(distance_to_image_plane)] = 0.0 self._data.output["distance_to_image_plane"][env_ids] = distance_to_image_plane.view( -1, *self.image_shape, 1 ) + if "distance_to_camera" in self.cfg.data_types: - self._data.output["distance_to_camera"][env_ids] = torch.clip( - ray_depth.view(-1, *self.image_shape, 1), max=self.cfg.max_distance - ) + if self.cfg.depth_clipping_behavior == "max": + ray_depth = torch.clip(ray_depth, max=self.cfg.max_distance) + elif self.cfg.depth_clipping_behavior == "zero": + ray_depth[ray_depth > self.cfg.max_distance] = 0.0 + self._data.output["distance_to_camera"][env_ids] = ray_depth.view(-1, *self.image_shape, 1) + if "normals" in self.cfg.data_types: self._data.output["normals"][env_ids] = ray_normal.view(-1, *self.image_shape, 3) @@ -341,7 +355,7 @@ def _create_buffers(self): self._data.image_shape = self.image_shape # -- output data # create the buffers to store the annotator data. - self._data.output = TensorDict({}, batch_size=self._view.count, device=self.device) + self._data.output = {} self._data.info = [{name: None for name in self.cfg.data_types}] * self._view.count for name in self.cfg.data_types: if name in ["distance_to_image_plane", "distance_to_camera"]: diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/ray_caster/ray_caster_camera_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/ray_caster/ray_caster_camera_cfg.py index d770031794..9654f4f3b7 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/ray_caster/ray_caster_camera_cfg.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/ray_caster/ray_caster_camera_cfg.py @@ -46,6 +46,15 @@ class OffsetCfg: data_types: list[str] = ["distance_to_image_plane"] """List of sensor names/types to enable for the camera. Defaults to ["distance_to_image_plane"].""" + depth_clipping_behavior: Literal["max", "zero", "none"] = "zero" + """Clipping behavior for the camera for values exceed the maximum value. Defaults to "zero". + + - ``"max"``: Values are clipped to the maximum value. + - ``"zero"``: Values are clipped to zero. + - ``"none``: No clipping is applied. Values will be returned as ``inf`` for ``distance_to_camera`` and ``nan`` + for ``distance_to_image_plane`` data type. + """ + pattern_cfg: PinholeCameraPatternCfg = MISSING """The pattern that defines the local ray starting positions and directions in a pinhole camera pattern.""" diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/sensor_base.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/sensor_base.py index 8f0d2d5092..b87e209d0a 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/sensor_base.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/sensor_base.py @@ -48,6 +48,8 @@ def __init__(self, cfg: SensorBaseCfg): # check that config is valid if cfg.history_length < 0: raise ValueError(f"History length must be greater than 0! Received: {cfg.history_length}") + # check that the config is valid + cfg.validate() # store inputs self.cfg = cfg # flag for whether the sensor is initialized diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/__init__.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/__init__.py index b722e083ea..2228a2322d 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/__init__.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/__init__.py @@ -28,7 +28,7 @@ from .converters import * # noqa: F401, F403 from .schemas import * # noqa: F401, F403 -from .simulation_cfg import PhysxCfg, SimulationCfg # noqa: F401, F403 +from .simulation_cfg import PhysxCfg, RenderCfg, SimulationCfg # noqa: F401, F403 from .simulation_context import SimulationContext, build_simulation_context # noqa: F401, F403 from .spawners import * # noqa: F401, F403 from .utils import * # noqa: F401, F403 diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/converters/asset_converter_base.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/converters/asset_converter_base.py index 34523f5396..8da8ca2f80 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/converters/asset_converter_base.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/converters/asset_converter_base.py @@ -54,6 +54,8 @@ def __init__(self, cfg: AssetConverterBaseCfg): Raises: ValueError: When provided asset file does not exist. """ + # check that the config is valid + cfg.validate() # check if the asset file exists if not check_file_path(cfg.asset_path): raise ValueError(f"The asset path does not exist: {cfg.asset_path}") diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/converters/mesh_converter.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/converters/mesh_converter.py index 8a42138644..ecaaa4098b 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/converters/mesh_converter.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/converters/mesh_converter.py @@ -10,7 +10,7 @@ import omni.kit.commands import omni.usd from omni.isaac.core.utils.extensions import enable_extension -from pxr import Usd, UsdGeom, UsdPhysics, UsdUtils +from pxr import Gf, Tf, Usd, UsdGeom, UsdPhysics, UsdUtils from omni.isaac.lab.sim.converters.asset_converter_base import AssetConverterBase from omni.isaac.lab.sim.converters.mesh_converter_cfg import MeshConverterCfg @@ -64,12 +64,13 @@ def __init__(self, cfg: MeshConverterCfg): def _convert_asset(self, cfg: MeshConverterCfg): """Generate USD from OBJ, STL or FBX. - It stores the asset in the following format: + The USD file has Y-up axis and is scaled to meters. + The asset hierarchy is arranged as follows: - /file_name (default prim) - |- /geometry <- Made instanceable if requested - |- /Looks - |- /mesh + .. code-block:: none + mesh_file_basename (default prim) + |- /geometry/Looks + |- /geometry/mesh Args: cfg: The configuration for conversion of mesh to USD. @@ -81,17 +82,37 @@ def _convert_asset(self, cfg: MeshConverterCfg): mesh_file_basename, mesh_file_format = os.path.basename(cfg.asset_path).split(".") mesh_file_format = mesh_file_format.lower() + # Check if mesh_file_basename is a valid USD identifier + if not Tf.IsValidIdentifier(mesh_file_basename): + # Correct the name to a valid identifier and update the basename + mesh_file_basename_original = mesh_file_basename + mesh_file_basename = Tf.MakeValidIdentifier(mesh_file_basename) + omni.log.warn( + f"Input file name '{mesh_file_basename_original}' is an invalid identifier for the mesh prim path." + f" Renaming it to '{mesh_file_basename}' for the conversion." + ) + # Convert USD asyncio.get_event_loop().run_until_complete( - self._convert_mesh_to_usd( - in_file=cfg.asset_path, out_file=self.usd_path, prim_path=f"/{mesh_file_basename}" - ) + self._convert_mesh_to_usd(in_file=cfg.asset_path, out_file=self.usd_path) ) + # Create a new stage, set Z up and meters per unit + temp_stage = Usd.Stage.CreateInMemory() + UsdGeom.SetStageUpAxis(temp_stage, UsdGeom.Tokens.z) + UsdGeom.SetStageMetersPerUnit(temp_stage, 1.0) + UsdPhysics.SetStageKilogramsPerUnit(temp_stage, 1.0) + # Add mesh to stage + base_prim = temp_stage.DefinePrim(f"/{mesh_file_basename}", "Xform") + prim = temp_stage.DefinePrim(f"/{mesh_file_basename}/geometry", "Xform") + prim.GetReferences().AddReference(self.usd_path) + temp_stage.SetDefaultPrim(base_prim) + temp_stage.Export(self.usd_path) + # Open converted USD stage - # note: This opens a new stage and does not use the stage created earlier by the user - # create a new stage stage = Usd.Stage.Open(self.usd_path) - # add USD to stage cache + # Need to reload the stage to get the new prim structure, otherwise it can be taken from the cache + stage.Reload() + # Add USD to stage cache stage_id = UsdUtils.StageCache.Get().Insert(stage) # Get the default prim (which is the root prim) -- "/{mesh_file_basename}" xform_prim = stage.GetDefaultPrim() @@ -111,6 +132,32 @@ def _convert_asset(self, cfg: MeshConverterCfg): ) # Delete the old Xform and make the new Xform the default prim stage.SetDefaultPrim(xform_prim) + # Apply default Xform rotation to mesh -> enable to set rotation and scale + omni.kit.commands.execute( + "CreateDefaultXformOnPrimCommand", + prim_path=xform_prim.GetPath(), + **{"stage": stage}, + ) + + # Apply translation, rotation, and scale to the Xform + geom_xform = UsdGeom.Xform(geom_prim) + geom_xform.ClearXformOpOrder() + + # Remove any existing rotation attributes + rotate_attr = geom_prim.GetAttribute("xformOp:rotateXYZ") + if rotate_attr: + geom_prim.RemoveProperty(rotate_attr.GetName()) + + # translation + translate_op = geom_xform.AddTranslateOp(UsdGeom.XformOp.PrecisionDouble) + translate_op.Set(Gf.Vec3d(*cfg.translation)) + # rotation + orient_op = geom_xform.AddOrientOp(UsdGeom.XformOp.PrecisionDouble) + orient_op.Set(Gf.Quatd(*cfg.rotation)) + # scale + scale_op = geom_xform.AddScaleOp(UsdGeom.XformOp.PrecisionDouble) + scale_op.Set(Gf.Vec3d(*cfg.scale)) + # Handle instanceable # Create a new Xform prim that will be the prototype prim if cfg.make_instanceable: @@ -148,28 +195,18 @@ def _convert_asset(self, cfg: MeshConverterCfg): """ @staticmethod - async def _convert_mesh_to_usd( - in_file: str, out_file: str, prim_path: str = "/World", load_materials: bool = True - ) -> bool: + async def _convert_mesh_to_usd(in_file: str, out_file: str, load_materials: bool = True) -> bool: """Convert mesh from supported file types to USD. This function uses the Omniverse Asset Converter extension to convert a mesh file to USD. It is an asynchronous function and should be called using `asyncio.get_event_loop().run_until_complete()`. The converted asset is stored in the USD format in the specified output file. - The USD file has Y-up axis and is scaled to meters. - - The asset hierarchy is arranged as follows: - - .. code-block:: none - prim_path (default prim) - |- /geometry/Looks - |- /geometry/mesh + The USD file has Y-up axis and is scaled to cm. Args: in_file: The file to convert. out_file: The path to store the output file. - prim_path: The prim path of the mesh. load_materials: Set to True to enable attaching materials defined in the input file to the generated USD mesh. Defaults to True. @@ -177,11 +214,9 @@ async def _convert_mesh_to_usd( True if the conversion succeeds. """ enable_extension("omni.kit.asset_converter") - enable_extension("omni.usd.metrics.assembler") import omni.kit.asset_converter import omni.usd - from omni.metrics.assembler.core import get_metrics_assembler_interface # Create converter context converter_context = omni.kit.asset_converter.AssetConverterContext() @@ -202,29 +237,9 @@ async def _convert_mesh_to_usd( # Create converter task instance = omni.kit.asset_converter.get_instance() - out_file_non_metric = out_file.replace(".usd", "_non_metric.usd") - task = instance.create_converter_task(in_file, out_file_non_metric, None, converter_context) + task = instance.create_converter_task(in_file, out_file, None, converter_context) # Start conversion task and wait for it to finish - success = True - while True: - success = await task.wait_until_finished() - if not success: - await asyncio.sleep(0.1) - else: - break - - temp_stage = Usd.Stage.CreateInMemory() - UsdGeom.SetStageUpAxis(temp_stage, UsdGeom.Tokens.z) - UsdGeom.SetStageMetersPerUnit(temp_stage, 1.0) - UsdPhysics.SetStageKilogramsPerUnit(temp_stage, 1.0) - - base_prim = temp_stage.DefinePrim(prim_path, "Xform") - prim = temp_stage.DefinePrim(f"{prim_path}/geometry", "Xform") - prim.GetReferences().AddReference(out_file_non_metric) - cache = UsdUtils.StageCache.Get() - cache.Insert(temp_stage) - stage_id = cache.GetId(temp_stage).ToLongInt() - get_metrics_assembler_interface().resolve_stage(stage_id) - temp_stage.SetDefaultPrim(base_prim) - temp_stage.Export(out_file) + success = await task.wait_until_finished() + if not success: + raise RuntimeError(f"Failed to convert {in_file} to USD. Error: {task.get_error_message()}") return success diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/converters/mesh_converter_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/converters/mesh_converter_cfg.py index 2466048447..92d5a611e2 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/converters/mesh_converter_cfg.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/converters/mesh_converter_cfg.py @@ -12,21 +12,21 @@ class MeshConverterCfg(AssetConverterBaseCfg): """The configuration class for MeshConverter.""" - mass_props: schemas_cfg.MassPropertiesCfg = None + mass_props: schemas_cfg.MassPropertiesCfg | None = None """Mass properties to apply to the USD. Defaults to None. Note: If None, then no mass properties will be added. """ - rigid_props: schemas_cfg.RigidBodyPropertiesCfg = None + rigid_props: schemas_cfg.RigidBodyPropertiesCfg | None = None """Rigid body properties to apply to the USD. Defaults to None. Note: If None, then no rigid body properties will be added. """ - collision_props: schemas_cfg.CollisionPropertiesCfg = None + collision_props: schemas_cfg.CollisionPropertiesCfg | None = None """Collision properties to apply to the USD. Defaults to None. Note: @@ -42,3 +42,12 @@ class MeshConverterCfg(AssetConverterBaseCfg): "none" causes no collision mesh to be added. """ + + translation: tuple[float, float, float] = (0.0, 0.0, 0.0) + """The translation of the mesh to the origin. Defaults to (0.0, 0.0, 0.0).""" + + rotation: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0) + """The rotation of the mesh in quaternion format (w, x, y, z). Defaults to (1.0, 0.0, 0.0, 0.0).""" + + scale: tuple[float, float, float] = (1.0, 1.0, 1.0) + """The scale of the mesh. Defaults to (1.0, 1.0, 1.0).""" diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/converters/urdf_converter_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/converters/urdf_converter_cfg.py index 06e13b3d01..1c7ebb855a 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/converters/urdf_converter_cfg.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/converters/urdf_converter_cfg.py @@ -30,7 +30,7 @@ class UrdfConverterCfg(AssetConverterBaseCfg): """Decompose a convex mesh into smaller pieces for a closer fit. Defaults to False.""" fix_base: bool = MISSING - """Create a fix joint to the root/base link. Defaults to True.""" + """Create a fix joint to the root/base link.""" merge_fixed_joints: bool = False """Consolidate links that are connected by fixed joints. Defaults to False.""" diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/simulation_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/simulation_cfg.py index 0705f5264d..a75d1f028b 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/simulation_cfg.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/simulation_cfg.py @@ -152,6 +152,54 @@ class PhysxCfg: """Size of particle contacts stream buffer allocated in pinned host memory. Default is 2 ** 20.""" +@configclass +class RenderCfg: + """Configuration for Omniverse RTX Renderer. + + These parameters are used to configure the Omniverse RTX Renderer. + For more information, see the `Omniverse RTX Renderer documentation`_. + + .. _Omniverse RTX Renderer documentation: https://docs.omniverse.nvidia.com/materials-and-rendering/latest/rtx-renderer.html + """ + + enable_translucency: bool = False + """Enables translucency for specular transmissive surfaces such as glass at the cost of some performance. Default is False.""" + + enable_reflections: bool = False + """Enables reflections at the cost of some performance. Default is False.""" + + enable_global_illumination: bool = False + """Enables Diffused Global Illumination at the cost of some performance. Default is False.""" + + antialiasing_mode: Literal["Off", "FXAA", "DLSS", "TAA", "DLAA"] = "DLSS" + """Selects the anti-aliasing mode to use. Defaults to DLSS.""" + + enable_dlssg: bool = False + """"Enables the use of DLSS-G. + DLSS Frame Generation boosts performance by using AI to generate more frames. + DLSS analyzes sequential frames and motion data to create additional high quality frames. + This feature requires an Ada Lovelace architecture GPU. + Enabling this feature also enables additional thread-related activities, which can hurt performance. + Default is False.""" + + dlss_mode: Literal[0, 1, 2, 3] = 0 + """For DLSS anti-aliasing, selects the performance/quality tradeoff mode. + Valid values are 0 (Performance), 1 (Balanced), 2 (Quality), or 3 (Auto). Default is 0.""" + + enable_direct_lighting: bool = True + """Enable direct light contributions from lights.""" + + samples_per_pixel: int = 1 + """Defines the Direct Lighting samples per pixel. + Higher values increase the direct lighting quality at the cost of performance. Default is 1.""" + + enable_shadows: bool = True + """Enables shadows at the cost of performance. When disabled, lights will not cast shadows. Defaults to True.""" + + enable_ambient_occlusion: bool = False + """Enables ambient occlusion at the cost of some performance. Default is False.""" + + @configclass class SimulationCfg: """Configuration for simulation physics.""" @@ -234,3 +282,6 @@ class SimulationCfg: The material is created at the path: ``{physics_prim_path}/defaultMaterial``. """ + + render: RenderCfg = RenderCfg() + """Render settings. Default is RenderCfg().""" diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/simulation_context.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/simulation_context.py index 315ad18d47..d36345a6d3 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/simulation_context.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/simulation_context.py @@ -114,6 +114,8 @@ def __init__(self, cfg: SimulationCfg | None = None): # store input if cfg is None: cfg = SimulationCfg() + # check that the config is valid + cfg.validate() self.cfg = cfg # check that simulation is running if stage_utils.get_current_stage() is None: @@ -152,6 +154,26 @@ def __init__(self, cfg: SimulationCfg | None = None): # flag for whether any GUI will be rendered (local, livestreamed or viewport) self._has_gui = self._local_gui or self._livestream_gui + # apply render settings from render config + carb_settings_iface.set_bool("/rtx/translucency/enabled", self.cfg.render.enable_translucency) + carb_settings_iface.set_bool("/rtx/reflections/enabled", self.cfg.render.enable_reflections) + carb_settings_iface.set_bool("/rtx/indirectDiffuse/enabled", self.cfg.render.enable_global_illumination) + carb_settings_iface.set_bool("/rtx/transient/dlssg/enabled", self.cfg.render.enable_dlssg) + carb_settings_iface.set_int("/rtx/post/dlss/execMode", self.cfg.render.dlss_mode) + carb_settings_iface.set_bool("/rtx/directLighting/enabled", self.cfg.render.enable_direct_lighting) + carb_settings_iface.set_int( + "/rtx/directLighting/sampledLighting/samplesPerPixel", self.cfg.render.samples_per_pixel + ) + carb_settings_iface.set_bool("/rtx/shadows/enabled", self.cfg.render.enable_shadows) + carb_settings_iface.set_bool("/rtx/ambientOcclusion/enabled", self.cfg.render.enable_ambient_occlusion) + # set denoiser mode + try: + import omni.replicator.core as rep + + rep.settings.set_render_rtx_realtime(antialiasing=self.cfg.render.antialiasing_mode) + except Exception: + pass + # store the default render mode if not self._has_gui and not self._offscreen_render: # set default render mode @@ -394,6 +416,14 @@ def get_setting(self, name: str) -> Any: """ return self._settings.get(name) + def forward(self) -> None: + """Updates articulation kinematics and fabric for rendering.""" + if self._fabric_iface is not None: + if self.physics_sim_view is not None and self.is_playing(): + # Update the articulations' link's poses before rendering + self.physics_sim_view.update_articulations_kinematic() + self._update_fabric(0.0, 0.0) + """ Operations - Override (standalone) """ @@ -464,11 +494,7 @@ def render(self, mode: RenderMode | None = None): self.set_setting("/app/player/playSimulations", True) else: # manually flush the fabric data to update Hydra textures - if self._fabric_iface is not None: - if self.physics_sim_view is not None and self.is_playing(): - # Update the articulations' link's poses before rendering - self.physics_sim_view.update_articulations_kinematic() - self._update_fabric(0.0, 0.0) + self.forward() # render the simulation # note: we don't call super().render() anymore because they do above operation inside # and we don't want to do it twice. We may remove it once we drop support for Isaac Sim 2022.2. diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/from_files/from_files_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/from_files/from_files_cfg.py index 775460bff6..9f6555d6fd 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/from_files/from_files_cfg.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/from_files/from_files_cfg.py @@ -72,8 +72,7 @@ class UsdFileCfg(FileCfg): ones. These include: - **Selecting variants**: This is done by specifying the :attr:`variants` parameter. - - **Creating and applying materials**: This is done by specifying the :attr:`visual_material` and - :attr:`physics_material` parameters. + - **Creating and applying materials**: This is done by specifying the :attr:`visual_material` parameter. - **Modifying existing properties**: This is done by specifying the respective properties in the configuration class. For instance, to modify the scale of the imported prim, set the :attr:`scale` parameter. diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/sensors/sensors_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/sensors/sensors_cfg.py index 52cd602b4c..ac36ca9485 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/sensors/sensors_cfg.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/sim/spawners/sensors/sensors_cfg.py @@ -184,17 +184,21 @@ class FisheyeCameraCfg(PinholeCameraCfg): func: Callable = sensors.spawn_camera projection_type: Literal[ - "fisheye_orthographic", "fisheye_equidistant", "fisheye_equisolid", "fisheye_polynomial", "fisheye_spherical" - ] = "fisheye_polynomial" - r"""Type of projection to use for the camera. Defaults to "fisheye_polynomial". + "fisheyePolynomial", + "fisheyeSpherical", + "fisheyeKannalaBrandtK3", + "fisheyeRadTanThinPrism", + "omniDirectionalStereo", + ] = "fisheyePolynomial" + r"""Type of projection to use for the camera. Defaults to "fisheyePolynomial". Available options: - - ``"fisheye_orthographic"``: Fisheye camera model using orthographic correction. - - ``"fisheye_equidistant"``: Fisheye camera model using equidistant correction. - - ``"fisheye_equisolid"``: Fisheye camera model using equisolid correction. - - ``"fisheye_polynomial"``: Fisheye camera model with :math:`360^{\circ}` spherical projection. - - ``"fisheye_spherical"``: Fisheye camera model with :math:`360^{\circ}` full-frame projection. + - ``"fisheyePolynomial"``: Fisheye camera model with :math:`360^{\circ}` spherical projection. + - ``"fisheyeSpherical"``: Fisheye camera model with :math:`360^{\circ}` full-frame projection. + - ``"fisheyeKannalaBrandtK3"``: Fisheye camera model using the Kannala-Brandt K3 distortion model. + - ``"fisheyeRadTanThinPrism"``: Fisheye camera model that combines radial and tangential distortions. + - ``"omniDirectionalStereo"``: Fisheye camera model supporting :math:`360^{\circ}` stereoscopic imaging. """ fisheye_nominal_width: float = 1936.0 diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/terrains/terrain_generator_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/terrains/terrain_generator_cfg.py index 5b692bfd81..4b1e9a077a 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/terrains/terrain_generator_cfg.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/terrains/terrain_generator_cfg.py @@ -88,8 +88,12 @@ class SubTerrainBaseCfg: is 0.7. """ - size: tuple[float, float] = MISSING - """The width (along x) and length (along y) of the terrain (in m).""" + size: tuple[float, float] = (10.0, 10.0) + """The width (along x) and length (along y) of the terrain (in m). Defaults to (10.0, 10.0). + + In case the :class:`~omni.isaac.lab.terrains.TerrainImporterCfg` is used, this parameter gets overridden by + :attr:`omni.isaac.lab.scene.TerrainImporterCfg.size` attribute. + """ flat_patch_sampling: dict[str, FlatPatchSamplingCfg] | None = None """Dictionary of configurations for sampling flat patches on the sub-terrain. Defaults to None, diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/terrains/terrain_importer.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/terrains/terrain_importer.py index 2dbdcbf9e9..e8834ceb52 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/terrains/terrain_importer.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/terrains/terrain_importer.py @@ -67,6 +67,8 @@ def __init__(self, cfg: TerrainImporterCfg): ValueError: If terrain type is 'usd' and no configuration provided for ``usd_path``. ValueError: If terrain type is 'usd' or 'plane' and no configuration provided for ``env_spacing``. """ + # check that the config is valid + cfg.validate() # store inputs self.cfg = cfg self.device = sim_utils.SimulationContext.instance().device # type: ignore diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/terrains/terrain_importer_cfg.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/terrains/terrain_importer_cfg.py index d6aca9419e..c420ed2844 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/terrains/terrain_importer_cfg.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/terrains/terrain_importer_cfg.py @@ -36,8 +36,12 @@ class TerrainImporterCfg: All sub-terrains are imported relative to this prim path. """ - num_envs: int = MISSING - """The number of environment origins to consider.""" + num_envs: int = 1 + """The number of environment origins to consider. Defaults to 1. + + In case, the :class:`~omni.isaac.lab.scene.InteractiveSceneCfg` is used, this parameter gets overridden by + :attr:`omni.isaac.lab.scene.InteractiveSceneCfg.num_envs` attribute. + """ terrain_type: Literal["generator", "plane", "usd"] = "generator" """The type of terrain to generate. Defaults to "generator". diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/__init__.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/__init__.py new file mode 100644 index 0000000000..8891b428a4 --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .image_plot import ImagePlot +from .line_plot import LiveLinePlot +from .manager_live_visualizer import ManagerLiveVisualizer +from .ui_visualizer_base import UiVisualizerBase diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/image_plot.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/image_plot.py new file mode 100644 index 0000000000..4b3cb1d6d0 --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/image_plot.py @@ -0,0 +1,201 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import numpy as np +from matplotlib import cm +from typing import TYPE_CHECKING, Optional + +import carb +import omni +import omni.log + +from .ui_widget_wrapper import UIWidgetWrapper + +if TYPE_CHECKING: + import omni.isaac.ui + import omni.ui + + +class ImagePlot(UIWidgetWrapper): + """An image plot widget to display live data. + + It has the following Layout where the mode frame is only useful for depth images: + +-------------------------------------------------------+ + | containing_frame | + |+-----------------------------------------------------+| + | main_plot_frame | + ||+---------------------------------------------------+|| + ||| plot_frames ||| + ||| ||| + ||| ||| + ||| (Image Plot Data) ||| + ||| ||| + ||| ||| + |||+-------------------------------------------------+||| + ||| mode_frame ||| + ||| ||| + ||| [x][Absolute] [x][Grayscaled] [ ][Colorized] ||| + |+-----------------------------------------------------+| + +-------------------------------------------------------+ + + """ + + def __init__( + self, + image: Optional[np.ndarray] = None, + label: str = "", + widget_height: int = 200, + show_min_max: bool = True, + unit: tuple[float, str] = (1, ""), + ): + """Create an XY plot UI Widget with axis scaling, legends, and support for multiple plots. + + Overlapping data is most accurately plotted when centered in the frame with reasonable axis scaling. + Pressing down the mouse gives the x and y values of each function at an x coordinate. + + Args: + image: Image to display + label: Short descriptive text to the left of the plot + widget_height: Height of the plot in pixels + show_min_max: Whether to show the min and max values of the image + unit: Tuple of (scale, name) for the unit of the image + """ + self._show_min_max = show_min_max + self._unit_scale = unit[0] + self._unit_name = unit[1] + + self._curr_mode = "None" + + self._has_built = False + + self._enabled = True + + self._byte_provider = omni.ui.ByteImageProvider() + if image is None: + carb.log_warn("image is NONE") + image = np.ones((480, 640, 3), dtype=np.uint8) * 255 + image[:, :, 0] = 0 + image[:, :240, 1] = 0 + + # if image is channel first, convert to channel last + if image.ndim == 3 and image.shape[0] in [1, 3, 4]: + image = np.moveaxis(image, 0, -1) + + self._aspect_ratio = image.shape[1] / image.shape[0] + self._widget_height = widget_height + self._label = label + self.update_image(image) + + plot_frame = self._create_ui_widget() + + super().__init__(plot_frame) + + def setEnabled(self, enabled: bool): + self._enabled = enabled + + def update_image(self, image: np.ndarray): + if not self._enabled: + return + + # if image is channel first, convert to channel last + if image.ndim == 3 and image.shape[0] in [1, 3, 4]: + image = np.moveaxis(image, 0, -1) + + height, width = image.shape[:2] + + if self._curr_mode == "Normalization": + image = (image - image.min()) / (image.max() - image.min()) + image = (image * 255).astype(np.uint8) + elif self._curr_mode == "Colorization": + if image.ndim == 3 and image.shape[2] == 3: + omni.log.warn("Colorization mode is only available for single channel images") + else: + image = (image - image.min()) / (image.max() - image.min()) + colormap = cm.get_cmap("jet") + if image.ndim == 3 and image.shape[2] == 1: + image = (colormap(image).squeeze(2) * 255).astype(np.uint8) + else: + image = (colormap(image) * 255).astype(np.uint8) + + # convert image to 4-channel RGBA + if image.ndim == 2 or (image.ndim == 3 and image.shape[2] == 1): + image = np.dstack((image, image, image, np.full((height, width, 1), 255, dtype=np.uint8))) + + elif image.ndim == 3 and image.shape[2] == 3: + image = np.dstack((image, np.full((height, width, 1), 255, dtype=np.uint8))) + + self._byte_provider.set_bytes_data(image.flatten().data, [width, height]) + + def update_min_max(self, image: np.ndarray): + if self._show_min_max and hasattr(self, "_min_max_label"): + non_inf = image[np.isfinite(image)].flatten() + if len(non_inf) > 0: + self._min_max_label.text = self._get_unit_description( + np.min(non_inf), np.max(non_inf), np.median(non_inf) + ) + else: + self._min_max_label.text = self._get_unit_description(0, 0) + + def _create_ui_widget(self): + containing_frame = omni.ui.Frame(build_fn=self._build_widget) + return containing_frame + + def _get_unit_description(self, min_value: float, max_value: float, median_value: float = None): + return ( + f"Min: {min_value * self._unit_scale:.2f} {self._unit_name} Max:" + f" {max_value * self._unit_scale:.2f} {self._unit_name}" + + (f" Median: {median_value * self._unit_scale:.2f} {self._unit_name}" if median_value is not None else "") + ) + + def _build_widget(self): + + with omni.ui.VStack(spacing=3): + with omni.ui.HStack(): + # Write the leftmost label for what this plot is + omni.ui.Label( + self._label, width=omni.isaac.ui.ui_utils.LABEL_WIDTH, alignment=omni.ui.Alignment.LEFT_TOP + ) + with omni.ui.Frame(width=self._aspect_ratio * self._widget_height, height=self._widget_height): + self._base_plot = omni.ui.ImageWithProvider(self._byte_provider) + + if self._show_min_max: + self._min_max_label = omni.ui.Label(self._get_unit_description(0, 0)) + + omni.ui.Spacer(height=8) + self._mode_frame = omni.ui.Frame(build_fn=self._build_mode_frame) + + omni.ui.Spacer(width=5) + self._has_built = True + + def _build_mode_frame(self): + """Build the frame containing the mode selection for the plots. + + This is an internal function to build the frame containing the mode selection for the plots. This function + should only be called from within the build function of a frame. + + The built widget has the following layout: + +-------------------------------------------------------+ + | legends_frame | + ||+---------------------------------------------------+|| + ||| ||| + ||| [x][Series 1] [x][Series 2] [ ][Series 3] ||| + ||| ||| + |||+-------------------------------------------------+||| + |+-----------------------------------------------------+| + +-------------------------------------------------------+ + """ + with omni.ui.HStack(): + with omni.ui.HStack(): + + def _change_mode(value): + self._curr_mode = value + + omni.isaac.ui.ui_utils.dropdown_builder( + label="Mode", + type="dropdown", + items=["Original", "Normalization", "Colorization"], + tooltip="Select a mode", + on_clicked_fn=_change_mode, + ) diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/line_plot.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/line_plot.py new file mode 100644 index 0000000000..682b1f0dcf --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/line_plot.py @@ -0,0 +1,603 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import colorsys +import numpy as np +from typing import TYPE_CHECKING + +import omni +from omni.isaac.core.simulation_context import SimulationContext + +from .ui_widget_wrapper import UIWidgetWrapper + +if TYPE_CHECKING: + import omni.isaac.ui + import omni.ui + + +class LiveLinePlot(UIWidgetWrapper): + """A 2D line plot widget to display live data. + + + This widget is used to display live data in a 2D line plot. It can be used to display multiple series + in the same plot. + + It has the following Layout: + +-------------------------------------------------------+ + | containing_frame | + |+-----------------------------------------------------+| + | main_plot_frame | + ||+---------------------------------------------------+|| + ||| plot_frames + grid lines (Z_stacked) ||| + ||| ||| + ||| ||| + ||| (Live Plot Data) ||| + ||| ||| + ||| ||| + |||+-------------------------------------------------+||| + ||| legends_frame ||| + ||| ||| + ||| [x][Series 1] [x][Series 2] [ ][Series 3] ||| + |||+-------------------------------------------------+||| + ||| limits_frame ||| + ||| ||| + ||| [Y-Limits] [min] [max] [Autoscale] ||| + |||+-------------------------------------------------+||| + ||| filter_frame ||| + ||| ||| + ||| ||| + |+-----------------------------------------------------+| + +-------------------------------------------------------+ + + """ + + def __init__( + self, + y_data: list[list[float]], + y_min: float = -10, + y_max: float = 10, + plot_height: int = 150, + show_legend: bool = True, + legends: list[str] | None = None, + max_datapoints: int = 200, + ): + """Create a new LiveLinePlot widget. + + Args: + y_data: A list of lists of floats containing the data to plot. Each list of floats represents a series in the plot. + y_min: The minimum y value to display. Defaults to -10. + y_max: The maximum y value to display. Defaults to 10. + plot_height: The height of the plot in pixels. Defaults to 150. + show_legend: Whether to display the legend. Defaults to True. + legends: A list of strings containing the legend labels for each series. If None, the default labels are "Series_0", "Series_1", etc. Defaults to None. + max_datapoints: The maximum number of data points to display. If the number of data points exceeds this value, the oldest data points are removed. Defaults to 200. + """ + super().__init__(self._create_ui_widget()) + self.plot_height = plot_height + self.show_legend = show_legend + self._legends = legends if legends is not None else ["Series_" + str(i) for i in range(len(y_data))] + self._y_data = y_data + self._colors = self._get_distinct_hex_colors(len(y_data)) + self._y_min = y_min if y_min is not None else -10 + self._y_max = y_max if y_max is not None else 10 + self._max_data_points = max_datapoints + self._show_legend = show_legend + self._series_visible = [True for _ in range(len(y_data))] + self._plot_frames = [] + self._plots = [] + self._plot_selected_values = [] + self._is_built = False + self._filter_frame = None + self._filter_mode = None + self._last_values = None + self._is_paused = False + + # Gets populated when widget is built + self._main_plot_frame = None + + self._autoscale_model = omni.ui.SimpleBoolModel(True) + + """Properties""" + + @property + def autoscale_mode(self) -> bool: + return self._autoscale_model.as_bool + + @property + def y_data(self) -> list[list[float]]: + """The current data in the plot.""" + return self._y_data + + @property + def y_min(self) -> float: + """The current minimum y value.""" + return self._y_min + + @property + def y_max(self) -> float: + """The current maximum y value.""" + return self._y_max + + @property + def legends(self) -> list[str]: + """The current legend labels.""" + return self._legends + + """ General Functions """ + + def clear(self): + """Clears the plot.""" + self._y_data = [[] for _ in range(len(self._y_data))] + self._last_values = None + + for plt in self._plots: + plt.set_data() + + # self._container_frame.rebuild() + + def add_datapoint(self, y_coords: list[float]): + """Add a data point to the plot. + + The data point is added to the end of the plot. If the number of data points exceeds the maximum number + of data points, the oldest data point is removed. + + ``y_coords`` is assumed to be a list of floats with the same length as the number of series in the plot. + + Args: + y_coords: A list of floats containing the y coordinates of the new data points. + """ + + for idx, y_coord in enumerate(y_coords): + + if len(self._y_data[idx]) > self._max_data_points: + self._y_data[idx] = self._y_data[idx][1:] + + if self._filter_mode == "Lowpass": + if self._last_values is not None: + alpha = 0.8 + y_coord = self._y_data[idx][-1] * alpha + y_coord * (1 - alpha) + elif self._filter_mode == "Integrate": + if self._last_values is not None: + y_coord = self._y_data[idx][-1] + y_coord + elif self._filter_mode == "Derivative": + if self._last_values is not None: + y_coord = (y_coord - self._last_values[idx]) / SimulationContext.instance().get_rendering_dt() + + self._y_data[idx].append(float(y_coord)) + + if self._main_plot_frame is None: + # Widget not built, not visible + return + + # Check if the widget has been built, i.e. the plot references have been created. + if not self._is_built or self._is_paused: + return + + if len(self._y_data) != len(self._plots): + # Plots gotten out of sync, rebuild the widget + self._main_plot_frame.rebuild() + return + + if self.autoscale_mode: + self._rescale_btn_pressed() + + for idx, plt in enumerate(self._plots): + plt.set_data(*self._y_data[idx]) + + self._last_values = y_coords + # Autoscale the y-axis to the current data + + """ + Internal functions for building the UI. + """ + + def _build_stacked_plots(self, grid: bool = True): + """Builds multiple plots stacked on top of each other to display multiple series. + + This is an internal function to build the plots. It should not be called from outside the class and only + from within the build function of a frame. + + The built widget has the following layout: + +-------------------------------------------------------+ + | main_plot_frame | + ||+---------------------------------------------------+|| + ||| ||| + ||| y_max|*******-------------------*******| ||| + ||| |-------*****-----------**--------| ||| + ||| 0|------------**-----***-----------| ||| + ||| |--------------***----------------| ||| + ||| y_min|---------------------------------| ||| + ||| ||| + |||+-------------------------------------------------+||| + + + Args: + grid: Whether to display grid lines. Defaults to True. + """ + + # Reset lists which are populated in the build function + self._plot_frames = [] + + # Define internal builder function + def _build_single_plot(y_data: list[float], color: int, plot_idx: int): + """Build a single plot. + + This is an internal function to build a single plot with the given data and color. This function + should only be called from within the build function of a frame. + + Args: + y_data: The data to plot. + color: The color of the plot. + """ + plot = omni.ui.Plot( + omni.ui.Type.LINE, + self._y_min, + self._y_max, + *y_data, + height=self.plot_height, + style={"color": color, "background_color": 0x0}, + ) + + if len(self._plots) <= plot_idx: + self._plots.append(plot) + self._plot_selected_values.append(omni.ui.SimpleStringModel("")) + else: + self._plots[plot_idx] = plot + + # Begin building the widget + with omni.ui.HStack(): + # Space to the left to add y-axis labels + omni.ui.Spacer(width=20) + + # Built plots for each time series stacked on top of each other + with omni.ui.ZStack(): + # Background rectangle + omni.ui.Rectangle( + height=self.plot_height, + style={ + "background_color": 0x0, + "border_color": omni.ui.color.white, + "border_width": 0.4, + "margin": 0.0, + }, + ) + + # Draw grid lines and labels + if grid: + # Calculate the number of grid lines to display + # Absolute range of the plot + plot_range = self._y_max - self._y_min + grid_resolution = 10 ** np.floor(np.log10(0.5 * plot_range)) + + plot_range /= grid_resolution + + # Fraction of the plot range occupied by the first and last grid line + first_space = (self._y_max / grid_resolution) - np.floor(self._y_max / grid_resolution) + last_space = np.ceil(self._y_min / grid_resolution) - self._y_min / grid_resolution + + # Number of grid lines to display + n_lines = int(plot_range - first_space - last_space) + + plot_resolution = self.plot_height / plot_range + + with omni.ui.VStack(): + omni.ui.Spacer(height=plot_resolution * first_space) + + # Draw grid lines + with omni.ui.VGrid(row_height=plot_resolution): + for grid_line_idx in range(n_lines): + # Create grid line + with omni.ui.ZStack(): + omni.ui.Line( + style={ + "color": 0xAA8A8777, + "background_color": 0x0, + "border_width": 0.4, + }, + alignment=omni.ui.Alignment.CENTER_TOP, + height=0, + ) + with omni.ui.Placer(offset_x=-20): + omni.ui.Label( + f"{(self._y_max - first_space * grid_resolution - grid_line_idx * grid_resolution):.3f}", + width=8, + height=8, + alignment=omni.ui.Alignment.RIGHT_TOP, + style={ + "color": 0xFFFFFFFF, + "font_size": 8, + }, + ) + + # Create plots for each series + for idx, (data, color) in enumerate(zip(self._y_data, self._colors)): + plot_frame = omni.ui.Frame( + build_fn=lambda y_data=data, plot_idx=idx, color=color: _build_single_plot( + y_data, color, plot_idx + ), + ) + plot_frame.visible = self._series_visible[idx] + self._plot_frames.append(plot_frame) + + # Create an invisible frame on top that will give a helpful tooltip + self._tooltip_frame = omni.ui.Plot( + height=self.plot_height, + style={"color": 0xFFFFFFFF, "background_color": 0x0}, + ) + + self._tooltip_frame.set_mouse_pressed_fn(self._mouse_moved_on_plot) + + # Create top label for the y-axis + with omni.ui.Placer(offset_x=-20, offset_y=-8): + omni.ui.Label( + f"{self._y_max:.3f}", + width=8, + height=2, + alignment=omni.ui.Alignment.LEFT_TOP, + style={"color": 0xFFFFFFFF, "font_size": 8}, + ) + + # Create bottom label for the y-axis + with omni.ui.Placer(offset_x=-20, offset_y=self.plot_height): + omni.ui.Label( + f"{self._y_min:.3f}", + width=8, + height=2, + alignment=omni.ui.Alignment.LEFT_BOTTOM, + style={"color": 0xFFFFFFFF, "font_size": 8}, + ) + + def _mouse_moved_on_plot(self, x, y, *args): + # Show a tooltip with x,y and function values + if len(self._y_data) == 0 or len(self._y_data[0]) == 0: + # There is no data in the plots, so do nothing + return + + for idx, plot in enumerate(self._plots): + x_pos = plot.screen_position_x + width = plot.computed_width + + location_x = (x - x_pos) / width + + data = self._y_data[idx] + n_samples = len(data) + selected_sample = int(location_x * n_samples) + value = data[selected_sample] + # save the value in scientific notation + self._plot_selected_values[idx].set_value(f"{value:.3f}") + + def _build_legends_frame(self): + """Build the frame containing the legend for the plots. + + This is an internal function to build the frame containing the legend for the plots. This function + should only be called from within the build function of a frame. + + The built widget has the following layout: + +-------------------------------------------------------+ + | legends_frame | + ||+---------------------------------------------------+|| + ||| ||| + ||| [x][Series 1] [x][Series 2] [ ][Series 3] ||| + ||| ||| + |||+-------------------------------------------------+||| + |+-----------------------------------------------------+| + +-------------------------------------------------------+ + """ + if not self._show_legend: + return + + with omni.ui.HStack(): + omni.ui.Spacer(width=32) + + # Find the longest legend to determine the width of the frame + max_legend = max([len(legend) for legend in self._legends]) + CHAR_WIDTH = 8 + with omni.ui.VGrid( + row_height=omni.isaac.ui.ui_utils.LABEL_HEIGHT, + column_width=max_legend * CHAR_WIDTH + 6, + ): + for idx in range(len(self._y_data)): + with omni.ui.HStack(): + model = omni.ui.SimpleBoolModel() + model.set_value(self._series_visible[idx]) + omni.ui.CheckBox(model=model, tooltip="", width=4) + model.add_value_changed_fn(lambda val, idx=idx: self._change_plot_visibility(idx, val.as_bool)) + omni.ui.Spacer(width=2) + with omni.ui.VStack(): + omni.ui.Label( + self._legends[idx], + width=max_legend * CHAR_WIDTH, + alignment=omni.ui.Alignment.LEFT, + style={"color": self._colors[idx], "font_size": 12}, + ) + omni.ui.StringField( + model=self._plot_selected_values[idx], + width=max_legend * CHAR_WIDTH, + alignment=omni.ui.Alignment.LEFT, + style={"color": self._colors[idx], "font_size": 10}, + read_only=True, + ) + + def _build_limits_frame(self): + """Build the frame containing the controls for the y-axis limits. + + This is an internal function to build the frame containing the controls for the y-axis limits. This function + should only be called from within the build function of a frame. + + The built widget has the following layout: + +-------------------------------------------------------+ + | limits_frame | + ||+---------------------------------------------------+|| + ||| ||| + ||| Limits [min] [max] [Re-Sacle] ||| + ||| Autoscale[x] ||| + ||| ------------------------------------------- ||| + |||+-------------------------------------------------+||| + """ + with omni.ui.VStack(): + with omni.ui.HStack(): + omni.ui.Label( + "Limits", + width=omni.isaac.ui.ui_utils.LABEL_WIDTH, + alignment=omni.ui.Alignment.LEFT_CENTER, + ) + + self.lower_limit_drag = omni.ui.FloatDrag(name="min", enabled=True, alignment=omni.ui.Alignment.CENTER) + y_min_model = self.lower_limit_drag.model + y_min_model.set_value(self._y_min) + y_min_model.add_value_changed_fn(lambda x: self._set_y_min(x.as_float)) + omni.ui.Spacer(width=2) + + self.upper_limit_drag = omni.ui.FloatDrag(name="max", enabled=True, alignment=omni.ui.Alignment.CENTER) + y_max_model = self.upper_limit_drag.model + y_max_model.set_value(self._y_max) + y_max_model.add_value_changed_fn(lambda x: self._set_y_max(x.as_float)) + omni.ui.Spacer(width=2) + + omni.ui.Button( + "Re-Scale", + width=omni.isaac.ui.ui_utils.BUTTON_WIDTH, + clicked_fn=self._rescale_btn_pressed, + alignment=omni.ui.Alignment.LEFT_CENTER, + style=omni.isaac.ui.ui_utils.get_style(), + ) + + omni.ui.CheckBox(model=self._autoscale_model, tooltip="", width=4) + + omni.ui.Line( + style={"color": 0x338A8777}, + width=omni.ui.Fraction(1), + alignment=omni.ui.Alignment.CENTER, + ) + + def _build_filter_frame(self): + """Build the frame containing the filter controls. + + This is an internal function to build the frame containing the filter controls. This function + should only be called from within the build function of a frame. + + The built widget has the following layout: + +-------------------------------------------------------+ + | filter_frame | + ||+---------------------------------------------------+|| + ||| ||| + ||| ||| + ||| ||| + |||+-------------------------------------------------+||| + |+-----------------------------------------------------+| + +-------------------------------------------------------+ + """ + with omni.ui.VStack(): + with omni.ui.HStack(): + + def _filter_changed(value): + self.clear() + self._filter_mode = value + + omni.isaac.ui.ui_utils.dropdown_builder( + label="Filter", + type="dropdown", + items=["None", "Lowpass", "Integrate", "Derivative"], + tooltip="Select a filter", + on_clicked_fn=_filter_changed, + ) + + def _toggle_paused(): + self._is_paused = not self._is_paused + + # Button + omni.ui.Button( + "Play/Pause", + width=omni.isaac.ui.ui_utils.BUTTON_WIDTH, + clicked_fn=_toggle_paused, + alignment=omni.ui.Alignment.LEFT_CENTER, + style=omni.isaac.ui.ui_utils.get_style(), + ) + + def _create_ui_widget(self): + """Create the full UI widget.""" + + def _build_widget(): + self._is_built = False + with omni.ui.VStack(): + self._main_plot_frame = omni.ui.Frame(build_fn=self._build_stacked_plots) + omni.ui.Spacer(height=8) + self._legends_frame = omni.ui.Frame(build_fn=self._build_legends_frame) + omni.ui.Spacer(height=8) + self._limits_frame = omni.ui.Frame(build_fn=self._build_limits_frame) + omni.ui.Spacer(height=8) + self._filter_frame = omni.ui.Frame(build_fn=self._build_filter_frame) + self._is_built = True + + containing_frame = omni.ui.Frame(build_fn=_build_widget) + + return containing_frame + + """ UI Actions Listener Functions """ + + def _change_plot_visibility(self, idx: int, visible: bool): + """Change the visibility of a plot at position idx.""" + self._series_visible[idx] = visible + self._plot_frames[idx].visible = visible + # self._main_plot_frame.rebuild() + + def _set_y_min(self, val: float): + """Update the y-axis minimum.""" + self._y_min = val + self.lower_limit_drag.model.set_value(val) + self._main_plot_frame.rebuild() + + def _set_y_max(self, val: float): + """Update the y-axis maximum.""" + self._y_max = val + self.upper_limit_drag.model.set_value(val) + self._main_plot_frame.rebuild() + + def _rescale_btn_pressed(self): + """Autoscale the y-axis to the current data.""" + if any(self._series_visible): + y_min = np.round( + min([min(y) for idx, y in enumerate(self._y_data) if self._series_visible[idx]]), + 4, + ) + y_max = np.round( + max([max(y) for idx, y in enumerate(self._y_data) if self._series_visible[idx]]), + 4, + ) + if y_min == y_max: + y_max += 1e-4 # Make sure axes don't collapse + + self._y_max = y_max + self._y_min = y_min + + if hasattr(self, "lower_limit_drag") and hasattr(self, "upper_limit_drag"): + self.lower_limit_drag.model.set_value(self._y_min) + self.upper_limit_drag.model.set_value(self._y_max) + + self._main_plot_frame.rebuild() + + """ Helper Functions """ + + def _get_distinct_hex_colors(self, num_colors) -> list[int]: + """ + This function returns a list of distinct colors for plotting. + + Args: + num_colors (int): the number of colors to generate + + Returns: + List[int]: a list of distinct colors in hexadecimal format 0xFFBBGGRR + """ + # Generate equally spaced colors in HSV space + rgb_colors = [ + colorsys.hsv_to_rgb(hue / num_colors, 0.75, 1) for hue in np.linspace(0, num_colors - 1, num_colors) + ] + # Convert to 0-255 RGB values + rgb_colors = [[int(c * 255) for c in rgb] for rgb in rgb_colors] + # Convert to 0xFFBBGGRR format + hex_colors = [0xFF * 16**6 + c[2] * 16**4 + c[1] * 16**2 + c[0] for c in rgb_colors] + return hex_colors diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/manager_live_visualizer.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/manager_live_visualizer.py new file mode 100644 index 0000000000..e6ae4f7310 --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/manager_live_visualizer.py @@ -0,0 +1,297 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import numpy +import weakref +from dataclasses import MISSING +from typing import TYPE_CHECKING + +import carb +import omni.kit.app +from omni.isaac.core.simulation_context import SimulationContext + +from omni.isaac.lab.managers import ManagerBase +from omni.isaac.lab.utils import configclass + +from .image_plot import ImagePlot +from .line_plot import LiveLinePlot +from .ui_visualizer_base import UiVisualizerBase + +if TYPE_CHECKING: + import omni.ui + + +@configclass +class ManagerLiveVisualizerCfg: + "Configuration for ManagerLiveVisualizer" + + debug_vis: bool = False + """Flag used to set status of the live visualizers on startup. Defaults to closed.""" + manager_name: str = MISSING + """Manager name that corresponds to the manager of interest in the ManagerBasedEnv and ManagerBasedRLEnv""" + term_names: list[str] | dict[str, list[str]] | None = None + """Specific term names specified in a Manager config that are chosen to be plotted. Defaults to None. + + If None all terms will be plotted. For managers that utilize Groups (i.e. ObservationGroup) use a dictionary of + {group_names: [term_names]}. + """ + + +class ManagerLiveVisualizer(UiVisualizerBase): + """A interface object used to transfer data from a manager to a UI widget. This class handles the creation of UI + Widgets for selected terms given a ManagerLiveVisualizerCfg. + """ + + def __init__(self, manager: ManagerBase, cfg: ManagerLiveVisualizerCfg = ManagerLiveVisualizerCfg()): + """Initialize ManagerLiveVisualizer. + + Args: + manager: The manager with terms to be plotted. The manager must have a get_active_iterable_terms method. + cfg: The configuration file used to select desired manager terms to be plotted. + """ + + self._manager = manager + self.debug_vis = cfg.debug_vis + self._env_idx: int = 0 + self.cfg = cfg + self._viewer_env_idx = 0 + self._vis_frame: omni.ui.Frame + self._vis_window: omni.ui.Window + + # evaluate chosen terms if no terms provided use all available. + self.term_names = [] + + if self.cfg.term_names is not None: + # extract chosen terms + if isinstance(self.cfg.term_names, list): + for term_name in self.cfg.term_names: + if term_name in self._manager.active_terms: + self.term_names.append(term_name) + else: + carb.log_err( + f"ManagerVisualizer Failure: ManagerTerm ({term_name}) does not exist in" + f" Manager({self.cfg.manager_name})" + ) + + # extract chosen group-terms + elif isinstance(self.cfg.term_names, dict): + # if manager is using groups and terms are saved as a dictionary + if isinstance(self._manager.active_terms, dict): + for group, terms in self.cfg.term_names: + if group in self._manager.active_terms.keys(): + for term_name in terms: + if term_name in self._manager.active_terms[group]: + self.term_names.append(f"{group}-{term_name}") + else: + carb.log_err( + f"ManagerVisualizer Failure: ManagerTerm ({term_name}) does not exist in" + f" Group({group})" + ) + else: + carb.log_err( + f"ManagerVisualizer Failure: Group ({group}) does not exist in" + f" Manager({self.cfg.manager_name})" + ) + else: + carb.log_err( + f"ManagerVisualizer Failure: Manager({self.cfg.manager_name}) does not utilize grouping of" + " terms." + ) + + # + # Implementation checks + # + + @property + def get_vis_frame(self) -> omni.ui.Frame: + """Getter for the UI Frame object tied to this visualizer.""" + return self._vis_frame + + @property + def get_vis_window(self) -> omni.ui.Window: + """Getter for the UI Window object tied to this visualizer.""" + return self._vis_window + + # + # Setters + # + + def set_debug_vis(self, debug_vis: bool): + """Set the debug visualization external facing function. + + Args: + debug_vis: Whether to enable or disable the debug visualization. + """ + self._set_debug_vis_impl(debug_vis) + + # + # Implementations + # + + def _set_env_selection_impl(self, env_idx: int): + """Update the index of the selected environment to display. + + Args: + env_idx: The index of the selected environment. + """ + if env_idx > 0 and env_idx < self._manager.num_envs: + self._env_idx = env_idx + else: + carb.log_warn(f"Environment index is out of range (0,{self._manager.num_envs})") + + def _set_vis_frame_impl(self, frame: omni.ui.Frame): + """Updates the assigned frame that can be used for visualizations. + + Args: + frame: The debug visualization frame. + """ + self._vis_frame = frame + + def _debug_vis_callback(self, event): + """Callback for the debug visualization event.""" + + if not SimulationContext.instance().is_playing(): + # Visualizers have not been created yet. + return + + # get updated data and update visualization + for (_, term), vis in zip( + self._manager.get_active_iterable_terms(env_idx=self._env_idx), self._term_visualizers + ): + if isinstance(vis, LiveLinePlot): + vis.add_datapoint(term) + elif isinstance(vis, ImagePlot): + vis.update_image(numpy.array(term)) + + def _set_debug_vis_impl(self, debug_vis: bool): + """Set the debug visualization implementation. + + Args: + debug_vis: Whether to enable or disable debug visualization. + """ + + if not hasattr(self, "_vis_frame"): + raise RuntimeError("No frame set for debug visualization.") + + # Clear internal visualizers + self._term_visualizers = [] + self._vis_frame.clear() + + if debug_vis: + # if enabled create a subscriber for the post update event if it doesn't exist + if not hasattr(self, "_debug_vis_handle") or self._debug_vis_handle is None: + app_interface = omni.kit.app.get_app_interface() + self._debug_vis_handle = app_interface.get_post_update_event_stream().create_subscription_to_pop( + lambda event, obj=weakref.proxy(self): obj._debug_vis_callback(event) + ) + else: + # if disabled remove the subscriber if it exists + if self._debug_vis_handle is not None: + self._debug_vis_handle.unsubscribe() + self._debug_vis_handle = None + + self._vis_frame.visible = False + return + + self._vis_frame.visible = True + + with self._vis_frame: + with omni.ui.VStack(): + # Add a plot in a collapsible frame for each term available + for name, term in self._manager.get_active_iterable_terms(env_idx=self._env_idx): + if name in self.term_names or len(self.term_names) == 0: + frame = omni.ui.CollapsableFrame( + name, + collapsed=False, + style={"border_color": 0xFF8A8777, "padding": 4}, + ) + with frame: + # create line plot for single or multivariable signals + len_term_shape = len(numpy.array(term).shape) + if len_term_shape <= 2: + plot = LiveLinePlot( + y_data=[[elem] for elem in term], + plot_height=150, + show_legend=True, + ) + self._term_visualizers.append(plot) + # create an image plot for 2d and greater data (i.e. mono and rgb images) + elif len_term_shape == 3: + image = ImagePlot( + image=numpy.array(term), + label=name, + ) + self._term_visualizers.append(image) + else: + carb.log_warn( + f"ManagerLiveVisualizer: Term ({name}) is not a supported data type for" + " visualization." + ) + frame.collapsed = True + + self._debug_vis = debug_vis + + +@configclass +class DefaultManagerBasedEnvLiveVisCfg: + """Default configuration to use for the ManagerBasedEnv. Each chosen manager assumes all terms will be plotted.""" + + action_live_vis = ManagerLiveVisualizerCfg(manager_name="action_manager") + observation_live_vis = ManagerLiveVisualizerCfg(manager_name="observation_manager") + + +@configclass +class DefaultManagerBasedRLEnvLiveVisCfg(DefaultManagerBasedEnvLiveVisCfg): + """Default configuration to use for the ManagerBasedRLEnv. Each chosen manager assumes all terms will be plotted.""" + + curriculum_live_vis = ManagerLiveVisualizerCfg(manager_name="curriculum_manager") + command_live_vis = ManagerLiveVisualizerCfg(manager_name="command_manager") + reward_live_vis = ManagerLiveVisualizerCfg(manager_name="reward_manager") + termination_live_vis = ManagerLiveVisualizerCfg(manager_name="termination_manager") + + +class EnvLiveVisualizer: + """A class to handle all ManagerLiveVisualizers used in an Environment.""" + + def __init__(self, cfg: object, managers: dict[str, ManagerBase]): + """Initialize the EnvLiveVisualizer. + + Args: + cfg: The configuration file containing terms of ManagerLiveVisualizers. + managers: A dictionary of labeled managers. i.e. {"manager_name",manager}. + """ + self.cfg = cfg + self.managers = managers + self._prepare_terms() + + def _prepare_terms(self): + self._manager_visualizers: dict[str, ManagerLiveVisualizer] = dict() + + # check if config is dict already + if isinstance(self.cfg, dict): + cfg_items = self.cfg.items() + else: + cfg_items = self.cfg.__dict__.items() + + for term_name, term_cfg in cfg_items: + # check if term config is None + if term_cfg is None: + continue + # check if term config is viable + if isinstance(term_cfg, ManagerLiveVisualizerCfg): + # find appropriate manager name + manager = self.managers[term_cfg.manager_name] + self._manager_visualizers[term_cfg.manager_name] = ManagerLiveVisualizer(manager=manager, cfg=term_cfg) + else: + raise TypeError( + f"Provided EnvLiveVisualizer term: '{term_name}' is not of type ManagerLiveVisualizerCfg" + ) + + @property + def manager_visualizers(self) -> dict[str, ManagerLiveVisualizer]: + """A dictionary of labeled ManagerLiveVisualizers associated manager name as key.""" + return self._manager_visualizers diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/ui_visualizer_base.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/ui_visualizer_base.py new file mode 100644 index 0000000000..a28eb3983a --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/ui_visualizer_base.py @@ -0,0 +1,148 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import inspect +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import omni.ui + + +class UiVisualizerBase: + """Base Class for components that support debug visualizations that requires access to some UI elements. + + This class provides a set of functions that can be used to assign ui interfaces. + + The following functions are provided: + + * :func:`set_debug_vis`: Assigns a debug visualization interface. This function is called by the main UI + when the checkbox for debug visualization is toggled. + * :func:`set_vis_frame`: Assigns a small frame within the isaac lab tab that can be used to visualize debug + information. Such as e.g. plots or images. It is called by the main UI on startup to create the frame. + * :func:`set_window`: Assigngs the main window that is used by the main UI. This allows the user + to have full controller over all UI elements. But be warned, with great power comes great responsibility. + """ + + """ + Exposed Properties + """ + + @property + def has_debug_vis_implementation(self) -> bool: + """Whether the component has a debug visualization implemented.""" + # check if function raises NotImplementedError + source_code = inspect.getsource(self._set_debug_vis_impl) + return "NotImplementedError" not in source_code + + @property + def has_vis_frame_implementation(self) -> bool: + """Whether the component has a debug visualization implemented.""" + # check if function raises NotImplementedError + source_code = inspect.getsource(self._set_vis_frame_impl) + return "NotImplementedError" not in source_code + + @property + def has_window_implementation(self) -> bool: + """Whether the component has a debug visualization implemented.""" + # check if function raises NotImplementedError + source_code = inspect.getsource(self._set_window_impl) + return "NotImplementedError" not in source_code + + @property + def has_env_selection_implementation(self) -> bool: + """Whether the component has a debug visualization implemented.""" + # check if function raises NotImplementedError + source_code = inspect.getsource(self._set_env_selection_impl) + return "NotImplementedError" not in source_code + + """ + Exposed Setters + """ + + def set_env_selection(self, env_selection: int) -> bool: + """Sets the selected environment id. + + This function is called by the main UI when the user selects a different environment. + + Args: + env_selection: The currently selected environment id. + + Returns: + Whether the environment selection was successfully set. False if the component + does not support environment selection. + """ + # check if environment selection is supported + if not self.has_env_selection_implementation: + return False + # set environment selection + self._set_env_selection_impl(env_selection) + return True + + def set_window(self, window: omni.ui.Window) -> bool: + """Sets the current main ui window. + + This function is called by the main UI when the window is created. It allows the component + to add custom UI elements to the window or to control the window and its elements. + + Args: + window: The ui window. + + Returns: + Whether the window was successfully set. False if the component + does not support this functionality. + """ + # check if window is supported + if not self.has_window_implementation: + return False + # set window + self._set_window_impl(window) + return True + + def set_vis_frame(self, vis_frame: omni.ui.Frame) -> bool: + """Sets the debug visualization frame. + + This function is called by the main UI when the window is created. It allows the component + to modify a small frame within the orbit tab that can be used to visualize debug information. + + Args: + vis_frame: The debug visualization frame. + + Returns: + Whether the debug visualization frame was successfully set. False if the component + does not support debug visualization. + """ + # check if debug visualization is supported + if not self.has_vis_frame_implementation: + return False + # set debug visualization frame + self._set_vis_frame_impl(vis_frame) + return True + + """ + Internal Implementation + """ + + def _set_env_selection_impl(self, env_idx: int): + """Set the environment selection.""" + raise NotImplementedError(f"Environment selection is not implemented for {self.__class__.__name__}.") + + def _set_window_impl(self, window: omni.ui.Window): + """Set the window.""" + raise NotImplementedError(f"Window is not implemented for {self.__class__.__name__}.") + + def _set_debug_vis_impl(self, debug_vis: bool): + """Set debug visualization state.""" + raise NotImplementedError(f"Debug visualization is not implemented for {self.__class__.__name__}.") + + def _set_vis_frame_impl(self, vis_frame: omni.ui.Frame): + """Set debug visualization into visualization objects. + + This function is responsible for creating the visualization objects if they don't exist + and input ``debug_vis`` is True. If the visualization objects exist, the function should + set their visibility into the stage. + """ + raise NotImplementedError(f"Debug visualization is not implemented for {self.__class__.__name__}.") diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/ui_widget_wrapper.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/ui_widget_wrapper.py new file mode 100644 index 0000000000..998f6c6da9 --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/ui_widget_wrapper.py @@ -0,0 +1,51 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# This file has been adapted from _isaac_sim/exts/omni.isaac.ui/omni/isaac/ui/element_wrappers/base_ui_element_wrappers.py + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import omni + +if TYPE_CHECKING: + import omni.ui + + +class UIWidgetWrapper: + """ + Base class for creating wrappers around any subclass of omni.ui.Widget in order to provide an easy interface + for creating and managing specific types of widgets such as state buttons or file pickers. + """ + + def __init__(self, container_frame: omni.ui.Frame): + self._container_frame = container_frame + + @property + def container_frame(self) -> omni.ui.Frame: + return self._container_frame + + @property + def enabled(self) -> bool: + return self.container_frame.enabled + + @enabled.setter + def enabled(self, value: bool): + self.container_frame.enabled = value + + @property + def visible(self) -> bool: + return self.container_frame.visible + + @visible.setter + def visible(self, value: bool): + self.container_frame.visible = value + + def cleanup(self): + """ + Perform any necessary cleanup + """ + pass diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/__init__.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/__init__.py index 23fb50d4c2..1cfa3507fd 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/__init__.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/__init__.py @@ -13,3 +13,4 @@ from .modifiers import * from .string import * from .timer import Timer +from .types import * diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/buffers/circular_buffer.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/buffers/circular_buffer.py index 1617344358..197878a016 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/buffers/circular_buffer.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/buffers/circular_buffer.py @@ -75,6 +75,16 @@ def current_length(self) -> torch.Tensor: """ return torch.minimum(self._num_pushes, self._max_len) + @property + def buffer(self) -> torch.Tensor: + """Complete circular buffer with most recent entry at the end and oldest entry at the beginning. + Returns: + Complete circular buffer with most recent entry at the end and oldest entry at the beginning of dimension 1. The shape is [batch_size, max_length, data.shape[1:]]. + """ + buf = self._buffer.clone() + buf = torch.roll(buf, shifts=self.max_length - self._pointer - 1, dims=0) + return torch.transpose(buf, dim0=0, dim1=1) + """ Operations. """ @@ -89,8 +99,10 @@ def reset(self, batch_ids: Sequence[int] | None = None): if batch_ids is None: batch_ids = slice(None) # reset the number of pushes for the specified batch indices - # note: we don't need to reset the buffer since it will be overwritten. The pointer handles this. self._num_pushes[batch_ids] = 0 + if self._buffer is not None: + # set buffer at batch_id reset indices to 0.0 so that the buffer() getter returns the cleared circular buffer after reset. + self._buffer[:, batch_ids, :] = 0.0 def append(self, data: torch.Tensor): """Append the data to the circular buffer. @@ -106,7 +118,7 @@ def append(self, data: torch.Tensor): if data.shape[0] != self.batch_size: raise ValueError(f"The input data has {data.shape[0]} environments while expecting {self.batch_size}") - # at the fist call, initialize the buffer + # at the first call, initialize the buffer size if self._buffer is None: self._pointer = -1 self._buffer = torch.empty((self.max_length, *data.shape), dtype=data.dtype, device=self._device) @@ -114,7 +126,12 @@ def append(self, data: torch.Tensor): self._pointer = (self._pointer + 1) % self.max_length # add the new data to the last layer self._buffer[self._pointer] = data.to(self._device) - # increment number of number of pushes + # Check for batches with zero pushes and initialize all values in batch to first append + if 0 in self._num_pushes.tolist(): + fill_ids = [i for i, x in enumerate(self._num_pushes.tolist()) if x == 0] + self._num_pushes.tolist().index(0) if 0 in self._num_pushes.tolist() else None + self._buffer[:, fill_ids, :] = data.to(self._device)[fill_ids] + # increment number of number of pushes for all batches self._num_pushes += 1 def __getitem__(self, key: torch.Tensor) -> torch.Tensor: diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/configclass.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/configclass.py index 7c1e1a9291..5bb9c30183 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/configclass.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/configclass.py @@ -14,7 +14,7 @@ from .dict import class_to_dict, update_class_from_dict -_CONFIGCLASS_METHODS = ["to_dict", "from_dict", "replace", "copy"] +_CONFIGCLASS_METHODS = ["to_dict", "from_dict", "replace", "copy", "validate"] """List of class methods added at runtime to dataclass.""" """ @@ -98,6 +98,7 @@ class EnvCfg: setattr(cls, "from_dict", _update_class_from_dict) setattr(cls, "replace", _replace_class_with_kwargs) setattr(cls, "copy", _copy_class) + setattr(cls, "validate", _validate) # wrap around dataclass cls = dataclass(cls, **kwargs) # return wrapped class @@ -240,6 +241,56 @@ class State: cls.__annotations__ = hints +def _validate(obj: object, prefix: str = "") -> list[str]: + """Check the validity of configclass object. + + This function checks if the object is a valid configclass object. A valid configclass object contains no MISSING + entries. + + Args: + obj: The object to check. + prefix: The prefix to add to the missing fields. Defaults to ''. + + Returns: + A list of missing fields. + + Raises: + TypeError: When the object is not a valid configuration object. + """ + missing_fields = [] + + if type(obj) is type(MISSING): + missing_fields.append(prefix) + return missing_fields + elif isinstance(obj, (list, tuple)): + for index, item in enumerate(obj): + current_path = f"{prefix}[{index}]" + missing_fields.extend(_validate(item, prefix=current_path)) + return missing_fields + elif isinstance(obj, dict): + obj_dict = obj + elif hasattr(obj, "__dict__"): + obj_dict = obj.__dict__ + else: + return missing_fields + + for key, value in obj_dict.items(): + # disregard builtin attributes + if key.startswith("__"): + continue + current_path = f"{prefix}.{key}" if prefix else key + missing_fields.extend(_validate(value, prefix=current_path)) + + # raise an error only once at the top-level call + if prefix == "" and missing_fields: + formatted_message = "\n".join(f" - {field}" for field in missing_fields) + raise TypeError( + f"Missing values detected in object {obj.__class__.__name__} for the following" + f" fields:\n{formatted_message}\n" + ) + return missing_fields + + def _process_mutable_types(cls): """Initialize all mutable elements through :obj:`dataclasses.Field` to avoid unnecessary complaints. diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/datasets/__init__.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/datasets/__init__.py new file mode 100644 index 0000000000..40ed823c7d --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/datasets/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) 2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +""" +Submodule for datasets classes and methods. +""" + +from .dataset_file_handler_base import DatasetFileHandlerBase +from .episode_data import EpisodeData +from .hdf5_dataset_file_handler import HDF5DatasetFileHandler diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/datasets/dataset_file_handler_base.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/datasets/dataset_file_handler_base.py new file mode 100644 index 0000000000..4dabba566c --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/datasets/dataset_file_handler_base.py @@ -0,0 +1,58 @@ +# Copyright (c) 2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from .episode_data import EpisodeData + + +class DatasetFileHandlerBase(ABC): + """Abstract class for handling dataset files.""" + + def __init__(self): + """Initializes the dataset file handler.""" + pass + + @abstractmethod + def open(self, file_path: str, mode: str = "r"): + """Open a file.""" + return NotImplementedError + + @abstractmethod + def create(self, file_path: str, env_name: str = None): + """Create a new file.""" + return NotImplementedError + + @abstractmethod + def get_env_name(self) -> str | None: + """Get the environment name.""" + return NotImplementedError + + @abstractmethod + def write_episode(self, episode: EpisodeData): + """Write episode data to the file.""" + return NotImplementedError + + @abstractmethod + def flush(self): + """Flush the file.""" + return NotImplementedError + + @abstractmethod + def close(self): + """Close the file.""" + return NotImplementedError + + @abstractmethod + def load_episode(self, episode_name: str) -> EpisodeData | None: + """Load episode data from the file.""" + return NotImplementedError + + @abstractmethod + def get_num_episodes(self) -> int: + """Get number of episodes in the file.""" + return NotImplementedError diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/datasets/episode_data.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/datasets/episode_data.py new file mode 100644 index 0000000000..57a1c306cb --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/datasets/episode_data.py @@ -0,0 +1,171 @@ +# Copyright (c) 2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch + + +class EpisodeData: + """Class to store episode data.""" + + def __init__(self) -> None: + """Initializes episode data class.""" + self._data = dict() + self._next_action_index = 0 + self._next_state_index = 0 + self._seed = None + self._env_id = None + self._success = None + + @property + def data(self): + """Returns the episode data.""" + return self._data + + @data.setter + def data(self, data: dict): + """Set the episode data.""" + self._data = data + + @property + def seed(self): + """Returns the random number generator seed.""" + return self._seed + + @seed.setter + def seed(self, seed: int): + """Set the random number generator seed.""" + self._seed = seed + + @property + def env_id(self): + """Returns the environment ID.""" + return self._env_id + + @env_id.setter + def env_id(self, env_id: int): + """Set the environment ID.""" + self._env_id = env_id + + @property + def next_action_index(self): + """Returns the next action index.""" + return self._next_action_index + + @next_action_index.setter + def next_action_index(self, index: int): + """Set the next action index.""" + self._next_action_index = index + + @property + def next_state_index(self): + """Returns the next state index.""" + return self._next_state_index + + @next_state_index.setter + def next_state_index(self, index: int): + """Set the next state index.""" + self._next_state_index = index + + @property + def success(self): + """Returns the success value.""" + return self._success + + @success.setter + def success(self, success: bool): + """Set the success value.""" + self._success = success + + def is_empty(self): + """Check if the episode data is empty.""" + return not bool(self._data) + + def add(self, key: str, value: torch.Tensor | dict): + """Add a key-value pair to the dataset. + + The key can be nested by using the "/" character. + For example: "obs/joint_pos". + + Args: + key: The key name. + value: The corresponding value of tensor type or of dict type. + """ + # check datatype + if isinstance(value, dict): + for sub_key, sub_value in value.items(): + self.add(f"{key}/{sub_key}", sub_value) + return + + sub_keys = key.split("/") + current_dataset_pointer = self._data + for sub_key_index in range(len(sub_keys)): + if sub_key_index == len(sub_keys) - 1: + # Add value to the final dict layer + if sub_keys[sub_key_index] not in current_dataset_pointer: + current_dataset_pointer[sub_keys[sub_key_index]] = value.unsqueeze(0).clone() + else: + current_dataset_pointer[sub_keys[sub_key_index]] = torch.cat( + (current_dataset_pointer[sub_keys[sub_key_index]], value.unsqueeze(0)) + ) + break + # key index + if sub_keys[sub_key_index] not in current_dataset_pointer: + current_dataset_pointer[sub_keys[sub_key_index]] = dict() + current_dataset_pointer = current_dataset_pointer[sub_keys[sub_key_index]] + + def get_initial_state(self) -> torch.Tensor | None: + """Get the initial state from the dataset.""" + if "initial_state" not in self._data: + return None + return self._data["initial_state"] + + def get_action(self, action_index) -> torch.Tensor | None: + """Get the action of the specified index from the dataset.""" + if "actions" not in self._data: + return None + if action_index >= len(self._data["actions"]): + return None + return self._data["actions"][action_index] + + def get_next_action(self) -> torch.Tensor | None: + """Get the next action from the dataset.""" + action = self.get_action(self._next_action_index) + if action is not None: + self._next_action_index += 1 + return action + + def get_state(self, state_index) -> dict | None: + """Get the state of the specified index from the dataset.""" + if "states" not in self._data: + return None + + states = self._data["states"] + + def get_state_helper(states, state_index) -> dict | torch.Tensor | None: + if isinstance(states, dict): + output_state = dict() + for key, value in states.items(): + output_state[key] = get_state_helper(value, state_index) + if output_state[key] is None: + return None + elif isinstance(states, torch.Tensor): + if state_index >= len(states): + return None + output_state = states[state_index] + else: + raise ValueError(f"Invalid state type: {type(states)}") + return output_state + + output_state = get_state_helper(states, state_index) + return output_state + + def get_next_state(self) -> dict | None: + """Get the next state from the dataset.""" + state = self.get_state(self._next_state_index) + if state is not None: + self._next_state_index += 1 + return state diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/datasets/hdf5_dataset_file_handler.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/datasets/hdf5_dataset_file_handler.py new file mode 100644 index 0000000000..f399030b5e --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/datasets/hdf5_dataset_file_handler.py @@ -0,0 +1,192 @@ +# Copyright (c) 2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import h5py +import json +import numpy as np +import os +import torch +from collections.abc import Iterable + +from .dataset_file_handler_base import DatasetFileHandlerBase +from .episode_data import EpisodeData + + +class HDF5DatasetFileHandler(DatasetFileHandlerBase): + """HDF5 dataset file handler for storing and loading episode data.""" + + def __init__(self): + """Initializes the HDF5 dataset file handler.""" + self._hdf5_file_stream = None + self._hdf5_data_group = None + self._demo_count = 0 + self._env_args = {} + + def open(self, file_path: str, mode: str = "r"): + """Open an existing dataset file.""" + if self._hdf5_file_stream is not None: + raise RuntimeError("HDF5 dataset file stream is already in use") + self._hdf5_file_stream = h5py.File(file_path, mode) + self._hdf5_data_group = self._hdf5_file_stream["data"] + self._demo_count = len(self._hdf5_data_group) + + def create(self, file_path: str, env_name: str = None): + """Create a new dataset file.""" + if self._hdf5_file_stream is not None: + raise RuntimeError("HDF5 dataset file stream is already in use") + if not file_path.endswith(".hdf5"): + file_path += ".hdf5" + dir_path = os.path.dirname(file_path) + if not os.path.isdir(dir_path): + os.makedirs(dir_path) + self._hdf5_file_stream = h5py.File(file_path, "w") + + # set up a data group in the file + self._hdf5_data_group = self._hdf5_file_stream.create_group("data") + self._hdf5_data_group.attrs["total"] = 0 + self._demo_count = 0 + + # set environment arguments + # the environment type (we use gym environment type) is set to be compatible with robomimic + # Ref: https://github.com/ARISE-Initiative/robomimic/blob/master/robomimic/envs/env_base.py#L15 + env_name = env_name if env_name is not None else "" + self.add_env_args({"env_name": env_name, "type": 2}) + + def __del__(self): + """Destructor for the file handler.""" + self.close() + + """ + Properties + """ + + def add_env_args(self, env_args: dict): + """Add environment arguments to the dataset.""" + self._raise_if_not_initialized() + self._env_args.update(env_args) + self._hdf5_data_group.attrs["env_args"] = json.dumps(self._env_args) + + def set_env_name(self, env_name: str): + """Set the environment name.""" + self._raise_if_not_initialized() + self.add_env_args({"env_name": env_name}) + + def get_env_name(self) -> str | None: + """Get the environment name.""" + self._raise_if_not_initialized() + env_args = json.loads(self._hdf5_data_group.attrs["env_args"]) + if "env_name" in env_args: + return env_args["env_name"] + return None + + def get_episode_names(self) -> Iterable[str]: + """Get the names of the episodes in the file.""" + self._raise_if_not_initialized() + return self._hdf5_data_group.keys() + + def get_num_episodes(self) -> int: + """Get number of episodes in the file.""" + return self._demo_count + + @property + def demo_count(self) -> int: + """The number of demos collected so far.""" + return self._demo_count + + """ + Operations. + """ + + def load_episode(self, episode_name: str, device: str) -> EpisodeData | None: + """Load episode data from the file.""" + self._raise_if_not_initialized() + if episode_name not in self._hdf5_data_group: + return None + episode = EpisodeData() + h5_episode_group = self._hdf5_data_group[episode_name] + + def load_dataset_helper(group): + """Helper method to load dataset that contains recursive dict objects.""" + data = {} + for key in group: + if isinstance(group[key], h5py.Group): + data[key] = load_dataset_helper(group[key]) + else: + # Converting group[key] to numpy array greatly improves the performance + # when converting to torch tensor + data[key] = torch.tensor(np.array(group[key]), device=device) + return data + + episode.data = load_dataset_helper(h5_episode_group) + + if "seed" in h5_episode_group.attrs: + episode.seed = h5_episode_group.attrs["seed"] + + if "success" in h5_episode_group.attrs: + episode.success = h5_episode_group.attrs["success"] + + episode.env_id = self.get_env_name() + + return episode + + def write_episode(self, episode: EpisodeData): + """Add an episode to the dataset. + + Args: + episode: The episode data to add. + """ + self._raise_if_not_initialized() + if episode.is_empty(): + return + + # create episode group based on demo count + h5_episode_group = self._hdf5_data_group.create_group(f"demo_{self._demo_count}") + + # store number of steps taken + if "actions" in episode.data: + h5_episode_group.attrs["num_samples"] = len(episode.data["actions"]) + else: + h5_episode_group.attrs["num_samples"] = 0 + + if episode.seed is not None: + h5_episode_group.attrs["seed"] = episode.seed + + if episode.success is not None: + h5_episode_group.attrs["success"] = episode.success + + def create_dataset_helper(group, key, value): + """Helper method to create dataset that contains recursive dict objects.""" + if isinstance(value, dict): + key_group = group.create_group(key) + for sub_key, sub_value in value.items(): + create_dataset_helper(key_group, sub_key, sub_value) + else: + group.create_dataset(key, data=value.cpu().numpy()) + + for key, value in episode.data.items(): + create_dataset_helper(h5_episode_group, key, value) + + # increment total step counts + self._hdf5_data_group.attrs["total"] += h5_episode_group.attrs["num_samples"] + + # increment total demo counts + self._demo_count += 1 + + def flush(self): + """Flush the episode data to disk.""" + self._raise_if_not_initialized() + + self._hdf5_file_stream.flush() + + def close(self): + """Close the dataset file handler.""" + if self._hdf5_file_stream is not None: + self._hdf5_file_stream.close() + self._hdf5_file_stream = None + + def _raise_if_not_initialized(self): + """Raise an error if the dataset file handler is not initialized.""" + if self._hdf5_file_stream is None: + raise RuntimeError("HDF5 dataset file stream is not initialized") diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/dict.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/dict.py index ca6b4f2a4b..e695207c88 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/dict.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/dict.py @@ -8,6 +8,7 @@ import collections.abc import hashlib import json +import torch from collections.abc import Iterable, Mapping from typing import Any @@ -40,8 +41,15 @@ def class_to_dict(obj: object) -> dict[str, Any]: # convert object to dictionary if isinstance(obj, dict): obj_dict = obj - else: + elif isinstance(obj, torch.Tensor): + # We have to treat torch tensors specially because `torch.tensor.__dict__` returns an empty + # dict, which would mean that a torch.tensor would be stored as an empty dict. Instead we + # want to store it directly as the tensor. + return obj + elif hasattr(obj, "__dict__"): obj_dict = obj.__dict__ + else: + return obj # convert to dictionary data = dict() @@ -55,6 +63,9 @@ def class_to_dict(obj: object) -> dict[str, Any]: # check if attribute is a dictionary elif hasattr(value, "__dict__") or isinstance(value, dict): data[key] = class_to_dict(value) + # check if attribute is a list or tuple + elif isinstance(value, (list, tuple)): + data[key] = type(value)([class_to_dict(v) for v in value]) else: data[key] = value return data diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/math.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/math.py index 1d3c0db81f..98e14f5a1c 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/math.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/math.py @@ -8,6 +8,7 @@ # needed to import for allowing type-hinting: torch.Tensor | np.ndarray from __future__ import annotations +import math import numpy as np import torch import torch.nn.functional @@ -745,7 +746,7 @@ def is_identity_pose(pos: torch.tensor, rot: torch.tensor) -> bool: return torch.allclose(pos, pos_identity) and torch.allclose(rot, rot_identity) -# @torch.jit.script +@torch.jit.script def combine_frame_transforms( t01: torch.Tensor, q01: torch.Tensor, t12: torch.Tensor | None = None, q12: torch.Tensor | None = None ) -> tuple[torch.Tensor, torch.Tensor]: @@ -1418,3 +1419,143 @@ def sample_cylinder( xyz[..., 2].uniform_(h_min, h_max) # return positions return xyz + + +""" +Orientation Conversions +""" + + +def convert_camera_frame_orientation_convention( + orientation: torch.Tensor, + origin: Literal["opengl", "ros", "world"] = "opengl", + target: Literal["opengl", "ros", "world"] = "ros", +) -> torch.Tensor: + r"""Converts a quaternion representing a rotation from one convention to another. + + In USD, the camera follows the ``"opengl"`` convention. Thus, it is always in **Y up** convention. + This means that the camera is looking down the -Z axis with the +Y axis pointing up , and +X axis pointing right. + However, in ROS, the camera is looking down the +Z axis with the +Y axis pointing down, and +X axis pointing right. + Thus, the camera needs to be rotated by :math:`180^{\circ}` around the X axis to follow the ROS convention. + + .. math:: + + T_{ROS} = \begin{bmatrix} 1 & 0 & 0 & 0 \\ 0 & -1 & 0 & 0 \\ 0 & 0 & -1 & 0 \\ 0 & 0 & 0 & 1 \end{bmatrix} T_{USD} + + On the other hand, the typical world coordinate system is with +X pointing forward, +Y pointing left, + and +Z pointing up. The camera can also be set in this convention by rotating the camera by :math:`90^{\circ}` + around the X axis and :math:`-90^{\circ}` around the Y axis. + + .. math:: + + T_{WORLD} = \begin{bmatrix} 0 & 0 & -1 & 0 \\ -1 & 0 & 0 & 0 \\ 0 & 1 & 0 & 0 \\ 0 & 0 & 0 & 1 \end{bmatrix} T_{USD} + + Thus, based on their application, cameras follow different conventions for their orientation. This function + converts a quaternion from one convention to another. + + Possible conventions are: + + - :obj:`"opengl"` - forward axis: -Z - up axis +Y - Offset is applied in the OpenGL (Usd.Camera) convention + - :obj:`"ros"` - forward axis: +Z - up axis -Y - Offset is applied in the ROS convention + - :obj:`"world"` - forward axis: +X - up axis +Z - Offset is applied in the World Frame convention + + Args: + orientation: Quaternion of form `(w, x, y, z)` with shape (..., 4) in source convention. + origin: Convention to convert from. Defaults to "opengl". + target: Convention to convert to. Defaults to "ros". + + Returns: + Quaternion of form `(w, x, y, z)` with shape (..., 4) in target convention + """ + if target == origin: + return orientation.clone() + + # -- unify input type + if origin == "ros": + # convert from ros to opengl convention + rotm = matrix_from_quat(orientation) + rotm[:, :, 2] = -rotm[:, :, 2] + rotm[:, :, 1] = -rotm[:, :, 1] + # convert to opengl convention + quat_gl = quat_from_matrix(rotm) + elif origin == "world": + # convert from world (x forward and z up) to opengl convention + rotm = matrix_from_quat(orientation) + rotm = torch.matmul( + rotm, + matrix_from_euler(torch.tensor([math.pi / 2, -math.pi / 2, 0], device=orientation.device), "XYZ"), + ) + # convert to isaac-sim convention + quat_gl = quat_from_matrix(rotm) + else: + quat_gl = orientation + + # -- convert to target convention + if target == "ros": + # convert from opengl to ros convention + rotm = matrix_from_quat(quat_gl) + rotm[:, :, 2] = -rotm[:, :, 2] + rotm[:, :, 1] = -rotm[:, :, 1] + return quat_from_matrix(rotm) + elif target == "world": + # convert from opengl to world (x forward and z up) convention + rotm = matrix_from_quat(quat_gl) + rotm = torch.matmul( + rotm, + matrix_from_euler(torch.tensor([math.pi / 2, -math.pi / 2, 0], device=orientation.device), "XYZ").T, + ) + return quat_from_matrix(rotm) + else: + return quat_gl.clone() + + +def create_rotation_matrix_from_view( + eyes: torch.Tensor, + targets: torch.Tensor, + up_axis: Literal["Y", "Z"] = "Z", + device: str = "cpu", +) -> torch.Tensor: + """Compute the rotation matrix from world to view coordinates. + + This function takes a vector ''eyes'' which specifies the location + of the camera in world coordinates and the vector ''targets'' which + indicate the position of the object. + The output is a rotation matrix representing the transformation + from world coordinates -> view coordinates. + + The inputs eyes and targets can each be a + - 3 element tuple/list + - torch tensor of shape (1, 3) + - torch tensor of shape (N, 3) + + Args: + eyes: Position of the camera in world coordinates. + targets: Position of the object in world coordinates. + up_axis: The up axis of the camera. Defaults to "Z". + device: The device to create torch tensors on. Defaults to "cpu". + + The vectors are broadcast against each other so they all have shape (N, 3). + + Returns: + R: (N, 3, 3) batched rotation matrices + + Reference: + Based on PyTorch3D (https://github.com/facebookresearch/pytorch3d/blob/eaf0709d6af0025fe94d1ee7cec454bc3054826a/pytorch3d/renderer/cameras.py#L1635-L1685) + """ + if up_axis == "Y": + up_axis_vec = torch.tensor((0, 1, 0), device=device, dtype=torch.float32).repeat(eyes.shape[0], 1) + elif up_axis == "Z": + up_axis_vec = torch.tensor((0, 0, 1), device=device, dtype=torch.float32).repeat(eyes.shape[0], 1) + else: + raise ValueError(f"Invalid up axis: {up_axis}. Valid options are 'Y' and 'Z'.") + + # get rotation matrix in opengl format (-Z forward, +Y up) + z_axis = -torch.nn.functional.normalize(targets - eyes, eps=1e-5) + x_axis = torch.nn.functional.normalize(torch.cross(up_axis_vec, z_axis, dim=1), eps=1e-5) + y_axis = torch.nn.functional.normalize(torch.cross(z_axis, x_axis, dim=1), eps=1e-5) + is_close = torch.isclose(x_axis, torch.tensor(0.0), atol=5e-3).all(dim=1, keepdim=True) + if is_close.any(): + replacement = torch.nn.functional.normalize(torch.cross(y_axis, z_axis, dim=1), eps=1e-5) + x_axis = torch.where(is_close, replacement, x_axis) + R = torch.cat((x_axis[:, None, :], y_axis[:, None, :], z_axis[:, None, :]), dim=1) + return R.transpose(1, 2) diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/types.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/types.py new file mode 100644 index 0000000000..3f2066a192 --- /dev/null +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/types.py @@ -0,0 +1,39 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-module for different data types.""" + +from __future__ import annotations + +import torch +from collections.abc import Sequence +from dataclasses import dataclass + + +@dataclass +class ArticulationActions: + """Data container to store articulation's joints actions. + + This class is used to store the actions of the joints of an articulation. + It is used to store the joint positions, velocities, efforts, and indices. + + If the actions are not provided, the values are set to None. + """ + + joint_positions: torch.Tensor | None = None + """The joint positions of the articulation. Defaults to None.""" + + joint_velocities: torch.Tensor | None = None + """The joint velocities of the articulation. Defaults to None.""" + + joint_efforts: torch.Tensor | None = None + """The joint efforts of the articulation. Defaults to None.""" + + joint_indices: torch.Tensor | Sequence[int] | slice | None = None + """The joint indices of the articulation. Defaults to None. + + If the joint indices are a slice, this indicates that the indices are continuous and correspond + to all the joints of the articulation. We use a slice to make the indexing more efficient. + """ diff --git a/source/extensions/omni.isaac.lab/setup.py b/source/extensions/omni.isaac.lab/setup.py index 3b68d3c84f..06c92832e5 100644 --- a/source/extensions/omni.isaac.lab/setup.py +++ b/source/extensions/omni.isaac.lab/setup.py @@ -22,15 +22,18 @@ "torch==2.4.0", "onnx==1.16.1", # 1.16.2 throws access violation on Windows "prettytable==3.3.0", - "tensordict", "toml", # devices "hidapi", # reinforcement learning - "gymnasium==0.29.0", + "gymnasium", # procedural-generation "trimesh", "pyglet<2", + # image processing + "transformers", + "einops", # needed for transformers, doesn't always auto-install + "warp-lang", ] PYTORCH_INDEX_URL = ["https://download.pytorch.org/whl/cu118"] diff --git a/source/extensions/omni.isaac.lab/test/assets/check_external_force.py b/source/extensions/omni.isaac.lab/test/assets/check_external_force.py index ef94da4975..2ba423d239 100644 --- a/source/extensions/omni.isaac.lab/test/assets/check_external_force.py +++ b/source/extensions/omni.isaac.lab/test/assets/check_external_force.py @@ -99,7 +99,8 @@ def main(): root_state = robot.data.default_root_state.clone() root_state[0, :2] = torch.tensor([0.0, -0.5], device=sim.device) root_state[1, :2] = torch.tensor([0.0, 0.5], device=sim.device) - robot.write_root_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) # reset dof state joint_pos, joint_vel = robot.data.default_joint_pos, robot.data.default_joint_vel robot.write_joint_state_to_sim(joint_pos, joint_vel) diff --git a/source/extensions/omni.isaac.lab/test/assets/check_fixed_base_assets.py b/source/extensions/omni.isaac.lab/test/assets/check_fixed_base_assets.py index b5e43b9016..50b8aeb113 100644 --- a/source/extensions/omni.isaac.lab/test/assets/check_fixed_base_assets.py +++ b/source/extensions/omni.isaac.lab/test/assets/check_fixed_base_assets.py @@ -113,7 +113,8 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula root_state = robot.data.default_root_state.clone() root_state[:, :3] += origins[index] root_state[:, :2] += torch.randn_like(root_state[:, :2]) * 0.25 - robot.write_root_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) # joint state joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() robot.write_joint_state_to_sim(joint_pos, joint_vel) diff --git a/source/extensions/omni.isaac.lab/test/assets/test_articulation.py b/source/extensions/omni.isaac.lab/test/assets/test_articulation.py index da0b9accf6..d4b077f9ab 100644 --- a/source/extensions/omni.isaac.lab/test/assets/test_articulation.py +++ b/source/extensions/omni.isaac.lab/test/assets/test_articulation.py @@ -26,6 +26,7 @@ import omni.isaac.core.utils.prims as prim_utils import omni.isaac.lab.sim as sim_utils +import omni.isaac.lab.utils.math as math_utils import omni.isaac.lab.utils.string as string_utils from omni.isaac.lab.actuators import ImplicitActuatorCfg from omni.isaac.lab.assets import Articulation, ArticulationCfg @@ -42,6 +43,8 @@ def generate_articulation_cfg( articulation_type: Literal["humanoid", "panda", "anymal", "shadow_hand", "single_joint"], stiffness: float | None = 10.0, damping: float | None = 2.0, + vel_limit: float | None = 100.0, + effort_limit: float | None = 400.0, ) -> ArticulationCfg: """Generate an articulation configuration. @@ -72,8 +75,8 @@ def generate_articulation_cfg( actuators={ "joint": ImplicitActuatorCfg( joint_names_expr=[".*"], - effort_limit=400.0, - velocity_limit=100.0, + effort_limit=effort_limit, + velocity_limit=vel_limit, stiffness=0.0, damping=10.0, ), @@ -105,9 +108,9 @@ def generate_articulation( The articulation and environment translations. """ - # Generate translations of 1.5m in x for each articulation + # Generate translations of 2.5 m in x for each articulation translations = torch.zeros(num_articulations, 3, device=device) - translations[:, 0] = torch.arange(num_articulations) * 1.5 + translations[:, 0] = torch.arange(num_articulations) * 2.5 # Create Top-level Xforms, one for each articulation for i in range(num_articulations): @@ -147,8 +150,8 @@ def test_initialization_floating_base_non_root(self): # Check that is fixed base self.assertFalse(articulation.is_fixed_base) # Check buffers that exists and have correct shapes - self.assertEqual(articulation.data.root_pos_w.shape, (num_articulations, 3)) - self.assertEqual(articulation.data.root_quat_w.shape, (num_articulations, 4)) + self.assertEqual(articulation.data.root_link_pos_w.shape, (num_articulations, 3)) + self.assertEqual(articulation.data.root_link_quat_w.shape, (num_articulations, 4)) self.assertEqual(articulation.data.joint_pos.shape, (num_articulations, 21)) # Check some internal physx data for debugging @@ -196,8 +199,8 @@ def test_initialization_floating_base(self): # Check that floating base self.assertFalse(articulation.is_fixed_base) # Check buffers that exists and have correct shapes - self.assertEqual(articulation.data.root_pos_w.shape, (num_articulations, 3)) - self.assertEqual(articulation.data.root_quat_w.shape, (num_articulations, 4)) + self.assertEqual(articulation.data.root_link_pos_w.shape, (num_articulations, 3)) + self.assertEqual(articulation.data.root_link_quat_w.shape, (num_articulations, 4)) self.assertEqual(articulation.data.joint_pos.shape, (num_articulations, 12)) self.assertEqual( articulation.data.default_mass.shape, (num_articulations, articulation.num_bodies) @@ -249,8 +252,8 @@ def test_initialization_fixed_base(self): # Check that fixed base self.assertTrue(articulation.is_fixed_base) # Check buffers that exists and have correct shapes - self.assertEqual(articulation.data.root_pos_w.shape, (num_articulations, 3)) - self.assertEqual(articulation.data.root_quat_w.shape, (num_articulations, 4)) + self.assertEqual(articulation.data.root_link_pos_w.shape, (num_articulations, 3)) + self.assertEqual(articulation.data.root_link_quat_w.shape, (num_articulations, 4)) self.assertEqual(articulation.data.joint_pos.shape, (num_articulations, 9)) self.assertEqual( articulation.data.default_mass.shape, (num_articulations, articulation.num_bodies) @@ -308,8 +311,8 @@ def test_initialization_fixed_base_single_joint(self): # Check that fixed base self.assertTrue(articulation.is_fixed_base) # Check buffers that exists and have correct shapes - self.assertEqual(articulation.data.root_pos_w.shape, (num_articulations, 3)) - self.assertEqual(articulation.data.root_quat_w.shape, (num_articulations, 4)) + self.assertEqual(articulation.data.root_link_pos_w.shape, (num_articulations, 3)) + self.assertEqual(articulation.data.root_link_quat_w.shape, (num_articulations, 4)) self.assertEqual(articulation.data.joint_pos.shape, (num_articulations, 1)) self.assertEqual( articulation.data.default_mass.shape, (num_articulations, articulation.num_bodies) @@ -367,8 +370,8 @@ def test_initialization_hand_with_tendons(self): # Check that fixed base self.assertTrue(articulation.is_fixed_base) # Check buffers that exists and have correct shapes - self.assertTrue(articulation.data.root_pos_w.shape == (num_articulations, 3)) - self.assertTrue(articulation.data.root_quat_w.shape == (num_articulations, 4)) + self.assertTrue(articulation.data.root_link_pos_w.shape == (num_articulations, 3)) + self.assertTrue(articulation.data.root_link_quat_w.shape == (num_articulations, 4)) self.assertTrue(articulation.data.joint_pos.shape == (num_articulations, 24)) self.assertEqual( articulation.data.default_mass.shape, (num_articulations, articulation.num_bodies) @@ -417,8 +420,8 @@ def test_initialization_floating_base_made_fixed_base(self): # Check that is fixed base self.assertTrue(articulation.is_fixed_base) # Check buffers that exists and have correct shapes - self.assertEqual(articulation.data.root_pos_w.shape, (num_articulations, 3)) - self.assertEqual(articulation.data.root_quat_w.shape, (num_articulations, 4)) + self.assertEqual(articulation.data.root_link_pos_w.shape, (num_articulations, 3)) + self.assertEqual(articulation.data.root_link_quat_w.shape, (num_articulations, 4)) self.assertEqual(articulation.data.joint_pos.shape, (num_articulations, 12)) # Check some internal physx data for debugging @@ -472,8 +475,8 @@ def test_initialization_fixed_base_made_floating_base(self): # Check that is floating base self.assertFalse(articulation.is_fixed_base) # Check buffers that exists and have correct shapes - self.assertEqual(articulation.data.root_pos_w.shape, (num_articulations, 3)) - self.assertEqual(articulation.data.root_quat_w.shape, (num_articulations, 4)) + self.assertEqual(articulation.data.root_link_pos_w.shape, (num_articulations, 3)) + self.assertEqual(articulation.data.root_link_quat_w.shape, (num_articulations, 4)) self.assertEqual(articulation.data.joint_pos.shape, (num_articulations, 9)) # Check some internal physx data for debugging @@ -572,6 +575,18 @@ def test_joint_limits(self): torch.testing.assert_close(articulation._data.joint_limits, limits) torch.testing.assert_close(articulation._data.default_joint_pos, default_joint_pos) + # Set new joint limits with indexing + env_ids = torch.arange(1, device=device) + joint_ids = torch.arange(2, device=device) + limits = torch.zeros(env_ids.shape[0], joint_ids.shape[0], 2, device=device) + limits[..., 0] = (torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) + 5.0) * -1.0 + limits[..., 1] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) + 5.0 + articulation.write_joint_limits_to_sim(limits, env_ids=env_ids, joint_ids=joint_ids) + + # Check new limits are in place + torch.testing.assert_close(articulation._data.joint_limits[env_ids][:, joint_ids], limits) + torch.testing.assert_close(articulation._data.default_joint_pos, default_joint_pos) + # Set new joint limits that invalidate default joint pos limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) limits[..., 0] = torch.rand(num_articulations, articulation.num_joints, device=device) * -0.1 @@ -584,6 +599,18 @@ def test_joint_limits(self): ) self.assertTrue(torch.all(within_bounds)) + # Set new joint limits that invalidate default joint pos with indexing + limits = torch.zeros(env_ids.shape[0], joint_ids.shape[0], 2, device=device) + limits[..., 0] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) * -0.1 + limits[..., 1] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) * 0.1 + articulation.write_joint_limits_to_sim(limits, env_ids=env_ids, joint_ids=joint_ids) + + # Check if all values are within the bounds + within_bounds = ( + articulation._data.default_joint_pos[env_ids][:, joint_ids] >= limits[..., 0] + ) & (articulation._data.default_joint_pos[env_ids][:, joint_ids] <= limits[..., 1]) + self.assertTrue(torch.all(within_bounds)) + def test_external_force_on_single_body(self): """Test application of external force on the base of the articulation.""" for num_articulations in (1, 2): @@ -606,7 +633,8 @@ def test_external_force_on_single_body(self): # reset root state root_state = articulation.data.default_root_state.clone() - articulation.write_root_state_to_sim(root_state) + articulation.write_root_link_pose_to_sim(root_state[:, :7]) + articulation.write_root_com_velocity_to_sim(root_state[:, 7:]) # reset dof state joint_pos, joint_vel = ( articulation.data.default_joint_pos, @@ -630,7 +658,7 @@ def test_external_force_on_single_body(self): articulation.update(sim.cfg.dt) # check condition that the articulations have fallen down for i in range(num_articulations): - self.assertLess(articulation.data.root_pos_w[i, 2].item(), 0.2) + self.assertLess(articulation.data.root_link_pos_w[i, 2].item(), 0.2) def test_external_force_on_multiple_bodies(self): """Test application of external force on the legs of the articulation.""" @@ -653,7 +681,12 @@ def test_external_force_on_multiple_bodies(self): # Now we are ready! for _ in range(5): # reset root state - articulation.write_root_state_to_sim(articulation.data.default_root_state.clone()) + articulation.write_root_link_pose_to_sim( + articulation.data.default_root_state.clone()[:, :7] + ) + articulation.write_root_com_velocity_to_sim( + articulation.data.default_root_state.clone()[:, 7:] + ) # reset dof state joint_pos, joint_vel = ( articulation.data.default_joint_pos, @@ -678,7 +711,7 @@ def test_external_force_on_multiple_bodies(self): # check condition for i in range(num_articulations): # since there is a moment applied on the articulation, the articulation should rotate - self.assertTrue(articulation.data.root_ang_vel_w[i, 2].item() > 0.1) + self.assertTrue(articulation.data.root_com_ang_vel_w[i, 2].item() > 0.1) def test_loading_gains_from_usd(self): """Test that gains are loaded from USD file if actuator model has them as None.""" @@ -789,6 +822,43 @@ def test_setting_gains_from_cfg_dict(self): torch.testing.assert_close(articulation.actuators["body"].stiffness, expected_stiffness) torch.testing.assert_close(articulation.actuators["body"].damping, expected_damping) + def test_setting_velocity_limits(self): + """Test that velocity limits are loaded form the configuration correctly.""" + for num_articulations in (1, 2): + for device in ("cuda:0", "cpu"): + for limit in (5.0, None): + with self.subTest(num_articulations=num_articulations, device=device, limit=limit): + with build_simulation_context( + device=device, add_ground_plane=False, auto_add_lighting=True + ) as sim: + articulation_cfg = generate_articulation_cfg( + articulation_type="single_joint", vel_limit=limit, effort_limit=limit + ) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=device + ) + # Play sim + sim.reset() + + if limit is not None: + # Expected gains + expected_velocity_limit = torch.full( + (articulation.num_instances, articulation.num_joints), + limit, + device=articulation.device, + ) + # Check that gains are loaded from USD file + torch.testing.assert_close( + articulation.actuators["joint"].velocity_limit, expected_velocity_limit + ) + torch.testing.assert_close( + articulation.data.joint_velocity_limits, expected_velocity_limit + ) + torch.testing.assert_close( + articulation.root_physx_view.get_dof_max_velocities().to(device), + expected_velocity_limit, + ) + def test_reset(self): """Test that reset method works properly. @@ -859,6 +929,206 @@ def test_apply_joint_command(self): # are not properly tuned assert not torch.allclose(articulation.data.joint_pos, joint_pos) + def test_body_root_link_state(self): + """Test for the root_link_state_w property""" + for num_articulations in (1, 2): + # for num_articulations in ( 2,): + for device in ("cuda:0", "cpu"): + # for device in ("cuda:0",): + for with_offset in [True, False]: + # for with_offset in [True,]: + with self.subTest(num_articulations=num_articulations, device=device, with_offset=with_offset): + with build_simulation_context( + device=device, add_ground_plane=False, auto_add_lighting=True + ) as sim: + articulation_cfg = generate_articulation_cfg(articulation_type="single_joint") + articulation, env_pos = generate_articulation(articulation_cfg, num_articulations, device) + env_idx = torch.tensor([x for x in range(num_articulations)]) + # Check that boundedness of articulation is correct + self.assertEqual(ctypes.c_long.from_address(id(articulation)).value, 1) + # Play sim + sim.reset() + # Check if articulation is initialized + self.assertTrue(articulation.is_initialized) + # Check that fixed base + self.assertTrue(articulation.is_fixed_base) + + # change center of mass offset from link frame + if with_offset: + offset = [0.5, 0.0, 0.0] + else: + offset = [0.0, 0.0, 0.0] + + # create com offsets + num_bodies = articulation.num_bodies + com = articulation.root_physx_view.get_coms() + link_offset = [1.0, 0.0, 0.0] # the offset from CenterPivot to Arm frames + new_com = torch.tensor(offset, device=device).repeat(num_articulations, 1, 1) + com[:, 1, :3] = new_com.squeeze(-2) + articulation.root_physx_view.set_coms(com, env_idx) + + # check they are set + torch.testing.assert_close(articulation.root_physx_view.get_coms(), com) + + for i in range(50): + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + # get state properties + root_state_w = articulation.data.root_state_w + root_link_state_w = articulation.data.root_link_state_w + root_com_state_w = articulation.data.root_com_state_w + body_state_w = articulation.data.body_state_w + body_link_state_w = articulation.data.body_link_state_w + body_com_state_w = articulation.data.body_com_state_w + + if with_offset: + # get joint state + joint_pos = articulation.data.joint_pos.unsqueeze(-1) + joint_vel = articulation.data.joint_vel.unsqueeze(-1) + + # LINK state + # pose + torch.testing.assert_close(root_state_w[..., :7], root_link_state_w[..., :7]) + torch.testing.assert_close(body_state_w[..., :7], body_link_state_w[..., :7]) + + # lin_vel arm + lin_vel_gt = torch.zeros(num_articulations, num_bodies, 3, device=device) + vx = -(link_offset[0]) * joint_vel * torch.sin(joint_pos) + vy = torch.zeros(num_articulations, 1, 1, device=device) + vz = (link_offset[0]) * joint_vel * torch.cos(joint_pos) + lin_vel_gt[:, 1, :] = torch.cat([vx, vy, vz], dim=-1).squeeze(-2) + + # linear velocity of root link should be zero + torch.testing.assert_close( + lin_vel_gt[:, 0, :], root_link_state_w[..., 7:10], atol=1e-3, rtol=1e-1 + ) + # linear velocity of pendulum link should be + torch.testing.assert_close( + lin_vel_gt, body_link_state_w[..., 7:10], atol=1e-3, rtol=1e-1 + ) + + # ang_vel + torch.testing.assert_close(root_state_w[..., 10:], root_link_state_w[..., 10:]) + torch.testing.assert_close(body_state_w[..., 10:], body_link_state_w[..., 10:]) + + # COM state + # position and orientation shouldn't match for the _state_com_w but everything else will + pos_gt = torch.zeros(num_articulations, num_bodies, 3, device=device) + px = (link_offset[0] + offset[0]) * torch.cos(joint_pos) + py = torch.zeros(num_articulations, 1, 1, device=device) + pz = (link_offset[0] + offset[0]) * torch.sin(joint_pos) + pos_gt[:, 1, :] = torch.cat([px, py, pz], dim=-1).squeeze(-2) + pos_gt += env_pos.unsqueeze(-2).repeat(1, num_bodies, 1) + torch.testing.assert_close( + pos_gt[:, 0, :], root_com_state_w[..., :3], atol=1e-3, rtol=1e-1 + ) + torch.testing.assert_close(pos_gt, body_com_state_w[..., :3], atol=1e-3, rtol=1e-1) + + # orientation + com_quat_b = articulation.data.com_quat_b + com_quat_w = math_utils.quat_mul(body_link_state_w[..., 3:7], com_quat_b) + torch.testing.assert_close(com_quat_w, body_com_state_w[..., 3:7]) + torch.testing.assert_close(com_quat_w[:, 0, :], root_com_state_w[..., 3:7]) + + # linear vel, and angular vel + torch.testing.assert_close(root_state_w[..., 7:], root_com_state_w[..., 7:]) + torch.testing.assert_close(body_state_w[..., 7:], body_com_state_w[..., 7:]) + else: + # single joint center of masses are at link frames so they will be the same + torch.testing.assert_close(root_state_w, root_link_state_w) + torch.testing.assert_close(root_state_w, root_com_state_w) + torch.testing.assert_close(body_state_w, body_link_state_w) + torch.testing.assert_close(body_state_w, body_com_state_w) + + def test_write_root_state(self): + """Test the setters for root_state using both the link frame and center of mass as reference frame.""" + for num_articulations in (1, 2): + for device in ("cuda:0", "cpu"): + for with_offset in [True, False]: + for state_location in ("com", "link"): + with self.subTest( + num_articulations=num_articulations, + device=device, + with_offset=with_offset, + state_location=state_location, + ): + with build_simulation_context( + device=device, add_ground_plane=False, auto_add_lighting=True, gravity_enabled=False + ) as sim: + articulation_cfg = generate_articulation_cfg(articulation_type="anymal") + articulation, env_pos = generate_articulation( + articulation_cfg, num_articulations, device + ) + env_idx = torch.tensor([x for x in range(num_articulations)]) + + # Play sim + sim.reset() + + # change center of mass offset from link frame + if with_offset: + offset = torch.tensor([1.0, 0.0, 0.0]).repeat(num_articulations, 1, 1) + else: + offset = torch.tensor([0.0, 0.0, 0.0]).repeat(num_articulations, 1, 1) + + # create com offsets + com = articulation.root_physx_view.get_coms() + new_com = offset + com[:, 0, :3] = new_com.squeeze(-2) + articulation.root_physx_view.set_coms(com, env_idx) + + # check they are set + torch.testing.assert_close(articulation.root_physx_view.get_coms(), com) + + rand_state = torch.zeros_like(articulation.data.root_link_state_w) + rand_state[..., :7] = articulation.data.default_root_state[..., :7] + rand_state[..., :3] += env_pos + # make quaternion a unit vector + rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) + + for i in range(10): + + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + if state_location == "com": + articulation.write_root_com_state_to_sim(rand_state) + elif state_location == "link": + articulation.write_root_link_state_to_sim(rand_state) + + if state_location == "com": + torch.testing.assert_close(rand_state, articulation.data.root_com_state_w) + elif state_location == "link": + torch.testing.assert_close(rand_state, articulation.data.root_link_state_w) + + def test_transform_inverses(self): + """Test if math utilities are true inverses of each other.""" + + pose01 = torch.rand(1, 7) + pose01[:, 3:7] = torch.nn.functional.normalize(pose01[..., 3:7], dim=-1) + + pose12 = torch.rand(1, 7) + pose12[:, 3:7] = torch.nn.functional.normalize(pose12[..., 3:7], dim=-1) + + pos02, quat02 = math_utils.combine_frame_transforms( + pose01[..., :3], pose01[..., 3:7], pose12[:, :3], pose12[:, 3:7] + ) + + pos01, quat01 = math_utils.combine_frame_transforms( + pos02, + quat02, + math_utils.quat_rotate(math_utils.quat_inv(pose12[:, 3:7]), -pose12[:, :3]), + math_utils.quat_inv(pose12[:, 3:7]), + ) + print("") + print(pose01) + print(pos01, quat01) + torch.testing.assert_close(pose01, torch.cat((pos01, quat01), dim=-1)) + if __name__ == "__main__": run_tests() diff --git a/source/extensions/omni.isaac.lab/test/assets/test_rigid_object.py b/source/extensions/omni.isaac.lab/test/assets/test_rigid_object.py index 429264c2c3..06d3087d7a 100644 --- a/source/extensions/omni.isaac.lab/test/assets/test_rigid_object.py +++ b/source/extensions/omni.isaac.lab/test/assets/test_rigid_object.py @@ -24,6 +24,7 @@ import ctypes import torch import unittest +from typing import Literal import omni.isaac.core.utils.prims as prim_utils @@ -31,19 +32,23 @@ from omni.isaac.lab.assets import RigidObject, RigidObjectCfg from omni.isaac.lab.sim import build_simulation_context from omni.isaac.lab.sim.spawners import materials -from omni.isaac.lab.utils.assets import ISAAC_NUCLEUS_DIR -from omni.isaac.lab.utils.math import default_orientation, random_orientation +from omni.isaac.lab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR +from omni.isaac.lab.utils.math import default_orientation, quat_mul, quat_rotate_inverse, random_orientation def generate_cubes_scene( - num_cubes: int = 1, height=1.0, has_api: bool = True, kinematic_enabled: bool = False, device: str = "cuda:0" + num_cubes: int = 1, + height=1.0, + api: Literal["none", "rigid_body", "articulation_root"] = "rigid_body", + kinematic_enabled: bool = False, + device: str = "cuda:0", ) -> tuple[RigidObject, torch.Tensor]: """Generate a scene with the provided number of cubes. Args: num_cubes: Number of cubes to generate. height: Height of the cubes. - has_api: Whether the cubes have a rigid body API on them. + api: The type of API that the cubes should have. kinematic_enabled: Whether the cubes are kinematic. device: Device to use for the simulation. @@ -57,17 +62,25 @@ def generate_cubes_scene( prim_utils.create_prim(f"/World/Table_{i}", "Xform", translation=origin) # Resolve spawn configuration - if has_api: - spawn_cfg = sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", - rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=kinematic_enabled), - ) - else: + if api == "none": # since no rigid body properties defined, this is just a static collider spawn_cfg = sim_utils.CuboidCfg( size=(0.1, 0.1, 0.1), collision_props=sim_utils.CollisionPropertiesCfg(), ) + elif api == "rigid_body": + spawn_cfg = sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", + rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=kinematic_enabled), + ) + elif api == "articulation_root": + spawn_cfg = sim_utils.UsdFileCfg( + usd_path=f"{ISAACLAB_NUCLEUS_DIR}/Tests/RigidObject/Cube/dex_cube_instanceable_with_articulation_root.usd", + rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=kinematic_enabled), + ) + else: + raise ValueError(f"Unknown api: {api}") + # Create rigid object cube_object_cfg = RigidObjectCfg( prim_path="/World/Table_.*/Object", @@ -106,8 +119,8 @@ def test_initialization(self): self.assertEqual(len(cube_object.body_names), 1) # Check buffers that exists and have correct shapes - self.assertEqual(cube_object.data.root_pos_w.shape, (num_cubes, 3)) - self.assertEqual(cube_object.data.root_quat_w.shape, (num_cubes, 4)) + self.assertEqual(cube_object.data.root_link_pos_w.shape, (num_cubes, 3)) + self.assertEqual(cube_object.data.root_link_quat_w.shape, (num_cubes, 4)) self.assertEqual(cube_object.data.default_mass.shape, (num_cubes, 1)) self.assertEqual(cube_object.data.default_inertia.shape, (num_cubes, 9)) @@ -140,8 +153,8 @@ def test_initialization_with_kinematic_enabled(self): self.assertEqual(len(cube_object.body_names), 1) # Check buffers that exists and have correct shapes - self.assertEqual(cube_object.data.root_pos_w.shape, (num_cubes, 3)) - self.assertEqual(cube_object.data.root_quat_w.shape, (num_cubes, 4)) + self.assertEqual(cube_object.data.root_link_pos_w.shape, (num_cubes, 3)) + self.assertEqual(cube_object.data.root_link_quat_w.shape, (num_cubes, 4)) # Simulate physics for _ in range(2): @@ -161,7 +174,27 @@ def test_initialization_with_no_rigid_body(self): with self.subTest(num_cubes=num_cubes, device=device): with build_simulation_context(device=device, auto_add_lighting=True) as sim: # Generate cubes scene - cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, has_api=False, device=device) + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, api="none", device=device) + + # Check that boundedness of rigid object is correct + self.assertEqual(ctypes.c_long.from_address(id(cube_object)).value, 1) + + # Play sim + sim.reset() + + # Check if object is initialized + self.assertFalse(cube_object.is_initialized) + + def test_initialization_with_articulation_root(self): + """Test that initialization fails when an articulation root is found at the provided prim path.""" + for num_cubes in (1, 2): + for device in ("cuda:0", "cpu"): + with self.subTest(num_cubes=num_cubes, device=device): + with build_simulation_context(device=device, auto_add_lighting=True) as sim: + # Generate cubes scene + cube_object, _ = generate_cubes_scene( + num_cubes=num_cubes, api="articulation_root", device=device + ) # Check that boundedness of rigid object is correct self.assertEqual(ctypes.c_long.from_address(id(cube_object)).value, 1) @@ -204,7 +237,8 @@ def test_external_force_on_single_body(self): # need to shift the position of the cubes otherwise they will be on top of each other root_state[:, :3] = origins - cube_object.write_root_state_to_sim(root_state) + cube_object.write_root_link_pose_to_sim(root_state[:, :7]) + cube_object.write_root_com_velocity_to_sim(root_state[:, 7:]) # reset object cube_object.reset() @@ -226,10 +260,10 @@ def test_external_force_on_single_body(self): # First object should still be at the same Z position (1.0) torch.testing.assert_close( - cube_object.data.root_pos_w[0::2, 2], torch.ones(num_cubes // 2, device=sim.device) + cube_object.data.root_link_pos_w[0::2, 2], torch.ones(num_cubes // 2, device=sim.device) ) # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - self.assertTrue(torch.all(cube_object.data.root_pos_w[1::2, 2] < 1.0)) + self.assertTrue(torch.all(cube_object.data.root_link_pos_w[1::2, 2] < 1.0)) def test_set_rigid_object_state(self): """Test setting the state of the rigid object. @@ -255,10 +289,14 @@ def test_set_rigid_object_state(self): # Set each state type individually as they are dependent on each other for state_type_to_randomize in state_types: state_dict = { - "root_pos_w": torch.zeros_like(cube_object.data.root_pos_w, device=sim.device), + "root_pos_w": torch.zeros_like(cube_object.data.root_link_pos_w, device=sim.device), "root_quat_w": default_orientation(num=num_cubes, device=sim.device), - "root_lin_vel_w": torch.zeros_like(cube_object.data.root_lin_vel_w, device=sim.device), - "root_ang_vel_w": torch.zeros_like(cube_object.data.root_ang_vel_w, device=sim.device), + "root_lin_vel_w": torch.zeros_like( + cube_object.data.root_com_lin_vel_w, device=sim.device + ), + "root_ang_vel_w": torch.zeros_like( + cube_object.data.root_com_ang_vel_w, device=sim.device + ), } # Now we are ready! @@ -286,7 +324,8 @@ def test_set_rigid_object_state(self): dim=-1, ) # reset root state - cube_object.write_root_state_to_sim(root_state=root_state) + cube_object.write_root_link_pose_to_sim(root_state[:, :7]) + cube_object.write_root_com_velocity_to_sim(root_state[:, 7:]) sim.step() @@ -322,7 +361,8 @@ def test_reset_rigid_object(self): # Random orientation root_state[:, 3:7] = random_orientation(num=num_cubes, device=sim.device) - cube_object.write_root_state_to_sim(root_state) + cube_object.write_root_link_pose_to_sim(root_state[:, :7]) + cube_object.write_root_com_velocity_to_sim(root_state[:, 7:]) if i % 2 == 0: # reset object @@ -407,7 +447,7 @@ def test_rigid_body_no_friction(self): initial_velocity = torch.zeros((num_cubes, 6), device=sim.cfg.device) initial_velocity[:, 0] = 0.1 - cube_object.write_root_velocity_to_sim(initial_velocity) + cube_object.write_root_com_velocity_to_sim(initial_velocity) # Simulate physics for _ in range(5): @@ -426,175 +466,180 @@ def test_rigid_body_no_friction(self): cube_object.data.root_lin_vel_w, initial_velocity[:, :3], rtol=1e-5, atol=tolerance ) - # def test_rigid_body_with_static_friction(self): - # """Test that static friction applied to rigid object works as expected. - - # This test works by applying a force to the object and checking if the object moves or not based on the - # mu (coefficient of static friction) value set for the object. We set the static friction to be non-zero and - # apply a force to the object. When the force applied is below mu, the object should not move. When the force - # applied is above mu, the object should move. - # """ - # for num_cubes in (1, 2): - # for device in ("cuda:0", "cpu"): - # with self.subTest(num_cubes=num_cubes, device=device): - # with build_simulation_context(device=device, add_ground_plane=True, auto_add_lighting=True) as sim: - # cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=0.03125, device=device) - - # # Create ground plane with no friction - # cfg = sim_utils.GroundPlaneCfg( - # physics_material=materials.RigidBodyMaterialCfg( - # static_friction=0.0, - # dynamic_friction=0.0, - # ) - # ) - # cfg.func("/World/GroundPlane", cfg) - - # # Play sim - # sim.reset() - - # # Set static friction to be non-zero - # static_friction_coefficient = 0.5 - # static_friction = torch.Tensor([[static_friction_coefficient]] * num_cubes) - # dynamic_friction = torch.zeros(num_cubes, 1) - # restitution = torch.FloatTensor(num_cubes, 1).uniform_(0.0, 0.2) - - # cube_object_materials = torch.cat([static_friction, dynamic_friction, restitution], dim=-1) - - # indices = torch.tensor(range(num_cubes), dtype=torch.int) - - # # Add friction to cube - # cube_object.root_physx_view.set_material_properties(cube_object_materials, indices) - - # # 2 cases: force applied is below and above mu - # # below mu: block should not move as the force applied is <= mu - # # above mu: block should move as the force applied is > mu - # for force in "below_mu", "above_mu": - # with self.subTest(force=force): - # external_wrench_b = torch.zeros((num_cubes, 1, 6), device=sim.device) - - # if force == "below_mu": - # external_wrench_b[:, 0, 0] = static_friction_coefficient * 0.999 - # else: - # external_wrench_b[:, 0, 0] = static_friction_coefficient * 1.001 - - # cube_object.set_external_force_and_torque( - # external_wrench_b[..., :3], - # external_wrench_b[..., 3:], - # ) - - # # Get root state - # initial_root_state = cube_object.data.root_state_w - - # # Simulate physics - # for _ in range(10): - # # perform rendering - # sim.step() - # # update object - # cube_object.update(sim.cfg.dt) - - # if force == "below_mu": - # # Assert that the block has not moved - # torch.testing.assert_close( - # cube_object.data.root_state_w, initial_root_state, rtol=1e-5, atol=1e-5 - # ) - # else: - # torch.testing.assert_close( - # cube_object.data.root_state_w, initial_root_state, rtol=1e-5, atol=1e-5 - # ) - - # def test_rigid_body_with_restitution(self): - # """Test that restitution when applied to rigid object works as expected. - - # This test works by dropping a block from a height and checking if the block bounces or not based on the - # restitution value set for the object. We set the restitution to be non-zero and drop the block from a height. - # When the restitution is 0, the block should not bounce. When the restitution is 1, the block should bounce - # with the same energy. When the restitution is between 0 and 1, the block should bounce with less energy. - - # """ - # for num_cubes in (1, 2): - # for device in ("cuda:0", "cpu"): - # with self.subTest(num_cubes=num_cubes, device=device): - # with build_simulation_context(device=device, add_ground_plane=True, auto_add_lighting=True) as sim: - # cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=1.0, device=device) - - # # Create ground plane such that has a restitution of 1.0 (perfectly elastic collision) - # cfg = sim_utils.GroundPlaneCfg( - # physics_material=materials.RigidBodyMaterialCfg( - # restitution=1.0, - # ) - # ) - # cfg.func("/World/GroundPlane", cfg) - - # indices = torch.tensor(range(num_cubes), dtype=torch.int) - - # # Play sim - # sim.reset() - - # # 3 cases: inelastic, partially elastic, elastic - # # inelastic: resitution = 0, block should not bounce - # # partially elastic: 0 <= restitution <= 1, block should bounce with less energy - # # elastic: restitution = 1, block should bounce with same energy - # for expected_collision_type in "inelastic", "partially_elastic", "elastic": - # root_state = torch.zeros(1, 13, device=sim.device) - # root_state[0, 3] = 1.0 # To make orientation a quaternion - # root_state[0, 2] = 0.1 # Set an initial drop height - # root_state[0, 9] = -1.0 # Set an initial downward velocity - - # cube_object.write_root_state_to_sim(root_state=root_state) - - # prev_z_velocity = 0.0 - # curr_z_velocity = 0.0 - - # with self.subTest(expected_collision_type=expected_collision_type): - # # cube_object.reset() - # # Set static friction to be non-zero - # if expected_collision_type == "inelastic": - # restitution_coefficient = 0.0 - # elif expected_collision_type == "partially_elastic": - # restitution_coefficient = 0.5 - # else: - # restitution_coefficient = 1.0 - - # restitution = 0.5 - # static_friction = torch.zeros(num_cubes, 1) - # dynamic_friction = torch.zeros(num_cubes, 1) - # restitution = torch.Tensor([[restitution_coefficient]] * num_cubes) - - # cube_object_materials = torch.cat( - # [static_friction, dynamic_friction, restitution], dim=-1 - # ) - - # # Add friction to cube - # cube_object.root_physx_view.set_material_properties(cube_object_materials, indices) - - # curr_z_velocity = cube_object.data.root_lin_vel_w[:, 2] - - # while torch.all(curr_z_velocity <= 0.0): - # # Simulate physics - # curr_z_velocity = cube_object.data.root_lin_vel_w[:, 2] - - # # perform rendering - # sim.step() - - # # update object - # cube_object.update(sim.cfg.dt) - # if torch.all(curr_z_velocity <= 0.0): - # # Still in the air - # prev_z_velocity = curr_z_velocity - - # # We have made contact with the ground and can verify expected collision type - # # based on how velocity has changed after the collision - # if expected_collision_type == "inelastic": - # # Assert that the block has lost most energy by checking that the z velocity is < 1/2 previous - # # velocity. This is because the floor's resitution means it will bounce back an object that itself - # # has restitution set to 0.0 - # self.assertTrue(torch.all(torch.le(curr_z_velocity / 2, abs(prev_z_velocity)))) - # elif expected_collision_type == "partially_elastic": - # # Assert that the block has lost some energy by checking that the z velocity is less - # self.assertTrue(torch.all(torch.le(abs(curr_z_velocity), abs(prev_z_velocity)))) - # elif expected_collision_type == "elastic": - # # Assert that the block has not lost any energy by checking that the z velocity is the same - # torch.testing.assert_close(abs(curr_z_velocity), abs(prev_z_velocity)) + def test_rigid_body_with_static_friction(self): + """Test that static friction applied to rigid object works as expected. + + This test works by applying a force to the object and checking if the object moves or not based on the + mu (coefficient of static friction) value set for the object. We set the static friction to be non-zero and + apply a force to the object. When the force applied is below mu, the object should not move. When the force + applied is above mu, the object should move. + """ + for num_cubes in (1, 2): + for device in ("cuda", "cpu"): + with self.subTest(num_cubes=num_cubes, device=device): + with build_simulation_context( + device=device, dt=0.01, add_ground_plane=False, auto_add_lighting=True + ) as sim: + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=0.03125, device=device) + + # Create ground plane + static_friction_coefficient = 0.5 + cfg = sim_utils.GroundPlaneCfg( + physics_material=materials.RigidBodyMaterialCfg( + static_friction=static_friction_coefficient, + dynamic_friction=static_friction_coefficient, # This shouldn't be required but is due to a bug in PhysX + ) + ) + cfg.func("/World/GroundPlane", cfg) + + # Play sim + sim.reset() + + # Set static friction to be non-zero + # Dynamic friction also needs to be zero due to a bug in PhysX + static_friction = torch.Tensor([[static_friction_coefficient]] * num_cubes) + dynamic_friction = torch.Tensor([[static_friction_coefficient]] * num_cubes) + restitution = torch.zeros(num_cubes, 1) + + cube_object_materials = torch.cat([static_friction, dynamic_friction, restitution], dim=-1) + + indices = torch.tensor(range(num_cubes), dtype=torch.int) + + # Add friction to cube + cube_object.root_physx_view.set_material_properties(cube_object_materials, indices) + + # let everything settle + for _ in range(100): + sim.step() + cube_object.update(sim.cfg.dt) + cube_object.write_root_com_velocity_to_sim(torch.zeros((num_cubes, 6), device=sim.device)) + cube_mass = cube_object.root_physx_view.get_masses() + gravity_magnitude = abs(sim.cfg.gravity[2]) + # 2 cases: force applied is below and above mu + # below mu: block should not move as the force applied is <= mu + # above mu: block should move as the force applied is > mu + for force in "below_mu", "above_mu": + with self.subTest(force=force): + # set initial velocity to zero + cube_object.write_root_com_velocity_to_sim( + torch.zeros((num_cubes, 6), device=sim.device) + ) + + external_wrench_b = torch.zeros((num_cubes, 1, 6), device=sim.device) + if force == "below_mu": + external_wrench_b[..., 0] = ( + static_friction_coefficient * cube_mass * gravity_magnitude * 0.99 + ) + else: + external_wrench_b[..., 0] = ( + static_friction_coefficient * cube_mass * gravity_magnitude * 1.01 + ) + + cube_object.set_external_force_and_torque( + external_wrench_b[..., :3], + external_wrench_b[..., 3:], + ) + + # Get root state + initial_root_pos = cube_object.data.root_link_pos_w.clone() + # Simulate physics + for _ in range(200): + # apply the wrench + cube_object.write_data_to_sim() + sim.step() + # update object + cube_object.update(sim.cfg.dt) + if force == "below_mu": + # Assert that the block has not moved + torch.testing.assert_close( + cube_object.data.root_link_pos_w, initial_root_pos, rtol=1e-3, atol=1e-3 + ) + if force == "above_mu": + self.assertTrue( + (cube_object.data.root_state_w[..., 0] - initial_root_pos[..., 0] > 0.02).all() + ) + + def test_rigid_body_with_restitution(self): + """Test that restitution when applied to rigid object works as expected. + + This test works by dropping a block from a height and checking if the block bounces or not based on the + restitution value set for the object. We set the restitution to be non-zero and drop the block from a height. + When the restitution is 0, the block should not bounce. When the restitution is between 0 and 1, the block + should bounce with less energy. + """ + for num_cubes in (1, 2): + for device in ("cuda:0", "cpu"): + for expected_collision_type in "partially_elastic", "inelastic": + with self.subTest( + expected_collision_type=expected_collision_type, num_cubes=num_cubes, device=device + ): + with build_simulation_context( + device=device, add_ground_plane=False, auto_add_lighting=True + ) as sim: + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=1.0, device=device) + + # Set static friction to be non-zero + if expected_collision_type == "inelastic": + restitution_coefficient = 0.0 + elif expected_collision_type == "partially_elastic": + restitution_coefficient = 0.5 + + # Create ground plane such that has a restitution of 1.0 (perfectly elastic collision) + cfg = sim_utils.GroundPlaneCfg( + physics_material=materials.RigidBodyMaterialCfg( + restitution=restitution_coefficient, + ) + ) + cfg.func("/World/GroundPlane", cfg) + + indices = torch.tensor(range(num_cubes), dtype=torch.int) + + # Play sim + sim.reset() + + root_state = torch.zeros(num_cubes, 13, device=sim.device) + root_state[:, 3] = 1.0 # To make orientation a quaternion + for i in range(num_cubes): + root_state[i, 1] = 1.0 * i + root_state[:, 2] = 1.0 # Set an initial drop height + root_state[:, 9] = -1.0 # Set an initial downward velocity + + cube_object.write_root_link_pose_to_sim(root_state[:, :7]) + cube_object.write_root_com_velocity_to_sim(root_state[:, 7:]) + + static_friction = torch.zeros(num_cubes, 1) + dynamic_friction = torch.zeros(num_cubes, 1) + restitution = torch.Tensor([[restitution_coefficient]] * num_cubes) + + cube_object_materials = torch.cat([static_friction, dynamic_friction, restitution], dim=-1) + + # Add restitution to cube + cube_object.root_physx_view.set_material_properties(cube_object_materials, indices) + + curr_z_velocity = cube_object.data.root_com_lin_vel_w[:, 2].clone() + + for _ in range(100): + sim.step() + + # update object + cube_object.update(sim.cfg.dt) + curr_z_velocity = cube_object.data.root_com_lin_vel_w[:, 2].clone() + + if expected_collision_type == "inelastic": + # assert that the block has not bounced by checking that the z velocity is less than or equal to 0 + self.assertTrue((curr_z_velocity <= 0.0).all()) + + if torch.all(curr_z_velocity <= 0.0): + # Still in the air + prev_z_velocity = curr_z_velocity + else: + # collision has happened, exit the for loop + break + + if expected_collision_type == "partially_elastic": + # Assert that the block has lost some energy by checking that the z velocity is less + self.assertTrue(torch.all(torch.le(abs(curr_z_velocity), abs(prev_z_velocity)))) + self.assertTrue((curr_z_velocity > 0.0).all()) def test_rigid_body_set_mass(self): """Test getting and setting mass of rigid object.""" @@ -674,6 +719,182 @@ def test_gravity_vec_w(self): # Check the body accelerations are correct torch.testing.assert_close(cube_object.data.body_acc_w, gravity) + def test_body_root_state_properties(self): + """Test the root_com_state_w, root_link_state_w, body_com_state_w, and body_link_state_w properties.""" + for num_cubes in (1, 2): + for device in ("cuda:0", "cpu"): + for with_offset in [True, False]: + with self.subTest(num_cubes=num_cubes, device=device, with_offset=with_offset): + with build_simulation_context( + device=device, gravity_enabled=False, auto_add_lighting=True + ) as sim: + # Create a scene with random cubes + cube_object, env_pos = generate_cubes_scene(num_cubes=num_cubes, height=0.0, device=device) + env_idx = torch.tensor([x for x in range(num_cubes)]) + + # Play sim + sim.reset() + + # Check if cube_object is initialized + self.assertTrue(cube_object.is_initialized) + + # change center of mass offset from link frame + if with_offset: + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_cubes, 1) + else: + offset = torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_cubes, 1) + + com = cube_object.root_physx_view.get_coms() + com[..., :3] = offset.to("cpu") + cube_object.root_physx_view.set_coms(com, env_idx) + + # check ceter of mass has been set + torch.testing.assert_close(cube_object.root_physx_view.get_coms(), com) + + # random z spin velocity + spin_twist = torch.zeros(6, device=device) + spin_twist[5] = torch.randn(1, device=device) + + # Simulate physics + for _ in range(100): + # spin the object around Z axis (com) + cube_object.write_root_com_velocity_to_sim(spin_twist.repeat(num_cubes, 1)) + # perform rendering + sim.step() + # update object + cube_object.update(sim.cfg.dt) + + # get state properties + root_state_w = cube_object.data.root_state_w + root_link_state_w = cube_object.data.root_link_state_w + root_com_state_w = cube_object.data.root_com_state_w + body_state_w = cube_object.data.body_state_w + body_link_state_w = cube_object.data.body_link_state_w + body_com_state_w = cube_object.data.body_com_state_w + + # if offset is [0,0,0] all root_state_%_w will match and all body_%_w will match + if not with_offset: + torch.testing.assert_close(root_state_w, root_com_state_w) + torch.testing.assert_close(root_state_w, root_link_state_w) + torch.testing.assert_close(body_state_w, body_com_state_w) + torch.testing.assert_close(body_state_w, body_link_state_w) + else: + # cubes are spinning around center of mass + # position will not match + # center of mass position will be constant (i.e. spining around com) + torch.testing.assert_close(env_pos + offset, root_com_state_w[..., :3]) + torch.testing.assert_close(env_pos + offset, body_com_state_w[..., :3].squeeze(-2)) + # link position will be moving but should stay constant away from center of mass + root_link_state_pos_rel_com = quat_rotate_inverse( + root_link_state_w[..., 3:7], + root_link_state_w[..., :3] - root_com_state_w[..., :3], + ) + torch.testing.assert_close(-offset, root_link_state_pos_rel_com) + body_link_state_pos_rel_com = quat_rotate_inverse( + body_link_state_w[..., 3:7], + body_link_state_w[..., :3] - body_com_state_w[..., :3], + ) + torch.testing.assert_close(-offset, body_link_state_pos_rel_com.squeeze(-2)) + + # orientation of com will be a constant rotation from link orientation + com_quat_b = cube_object.data.com_quat_b + com_quat_w = quat_mul(body_link_state_w[..., 3:7], com_quat_b) + torch.testing.assert_close(com_quat_w, body_com_state_w[..., 3:7]) + torch.testing.assert_close(com_quat_w.squeeze(-2), root_com_state_w[..., 3:7]) + + # orientation of link will match root state will always match + torch.testing.assert_close(root_state_w[..., 3:7], root_link_state_w[..., 3:7]) + torch.testing.assert_close(body_state_w[..., 3:7], body_link_state_w[..., 3:7]) + + # lin_vel will not match + # center of mass vel will be constant (i.e. spining around com) + torch.testing.assert_close( + torch.zeros_like(root_com_state_w[..., 7:10]), root_com_state_w[..., 7:10] + ) + torch.testing.assert_close( + torch.zeros_like(body_com_state_w[..., 7:10]), body_com_state_w[..., 7:10] + ) + # link frame will be moving, and should be equal to input angular velocity cross offset + lin_vel_rel_root_gt = quat_rotate_inverse( + root_link_state_w[..., 3:7], root_link_state_w[..., 7:10] + ) + lin_vel_rel_body_gt = quat_rotate_inverse( + body_link_state_w[..., 3:7], body_link_state_w[..., 7:10] + ) + lin_vel_rel_gt = torch.linalg.cross( + spin_twist.repeat(num_cubes, 1)[..., 3:], -offset + ) + torch.testing.assert_close( + lin_vel_rel_gt, lin_vel_rel_root_gt, atol=1e-4, rtol=1e-4 + ) + torch.testing.assert_close( + lin_vel_rel_gt, lin_vel_rel_body_gt.squeeze(-2), atol=1e-4, rtol=1e-4 + ) + + # ang_vel will always match + torch.testing.assert_close(root_state_w[..., 10:], root_com_state_w[..., 10:]) + torch.testing.assert_close(root_state_w[..., 10:], root_link_state_w[..., 10:]) + torch.testing.assert_close(body_state_w[..., 10:], body_com_state_w[..., 10:]) + torch.testing.assert_close(body_state_w[..., 10:], body_link_state_w[..., 10:]) + + def test_write_root_state(self): + """Test the setters for root_state using both the link frame and center of mass as reference frame.""" + for num_cubes in (1, 2): + for device in ("cuda:0", "cpu"): + for with_offset in [True, False]: + for state_location in ("com", "link"): + with self.subTest(num_cubes=num_cubes, device=device, with_offset=with_offset): + with build_simulation_context( + device=device, gravity_enabled=False, auto_add_lighting=True + ) as sim: + # Create a scene with random cubes + cube_object, env_pos = generate_cubes_scene( + num_cubes=num_cubes, height=0.0, device=device + ) + env_idx = torch.tensor([x for x in range(num_cubes)]) + + # Play sim + sim.reset() + + # Check if cube_object is initialized + self.assertTrue(cube_object.is_initialized) + + # change center of mass offset from link frame + if with_offset: + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_cubes, 1) + else: + offset = torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_cubes, 1) + + com = cube_object.root_physx_view.get_coms() + com[..., :3] = offset.to("cpu") + cube_object.root_physx_view.set_coms(com, env_idx) + + # check ceter of mass has been set + torch.testing.assert_close(cube_object.root_physx_view.get_coms(), com) + + rand_state = torch.zeros_like(cube_object.data.root_link_state_w) + rand_state[..., :7] = cube_object.data.default_root_state[..., :7] + rand_state[..., :3] += env_pos + # make quaternion a unit vector + rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) + + for i in range(10): + + # perform step + sim.step() + # update buffers + cube_object.update(sim.cfg.dt) + + if state_location == "com": + cube_object.write_root_com_state_to_sim(rand_state) + elif state_location == "link": + cube_object.write_root_link_state_to_sim(rand_state) + + if state_location == "com": + torch.testing.assert_close(rand_state, cube_object.data.root_com_state_w) + elif state_location == "link": + torch.testing.assert_close(rand_state, cube_object.data.root_link_state_w) + if __name__ == "__main__": run_tests() diff --git a/source/extensions/omni.isaac.lab/test/assets/test_rigid_object_collection.py b/source/extensions/omni.isaac.lab/test/assets/test_rigid_object_collection.py new file mode 100644 index 0000000000..44b5a570fe --- /dev/null +++ b/source/extensions/omni.isaac.lab/test/assets/test_rigid_object_collection.py @@ -0,0 +1,695 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# ignore private usage of variables warning +# pyright: reportPrivateUsage=none + + +"""Launch Isaac Sim Simulator first.""" + +from omni.isaac.lab.app import AppLauncher, run_tests + +# Can set this to False to see the GUI for debugging +# This will also add lights to the scene +HEADLESS = True + +# launch omniverse app +app_launcher = AppLauncher(headless=HEADLESS) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import ctypes +import torch +import unittest + +import omni.isaac.core.utils.prims as prim_utils + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.assets import RigidObjectCfg, RigidObjectCollection, RigidObjectCollectionCfg +from omni.isaac.lab.sim import build_simulation_context +from omni.isaac.lab.utils.assets import ISAAC_NUCLEUS_DIR +from omni.isaac.lab.utils.math import default_orientation, quat_mul, quat_rotate_inverse, random_orientation + + +def generate_cubes_scene( + num_envs: int = 1, + num_cubes: int = 1, + height=1.0, + has_api: bool = True, + kinematic_enabled: bool = False, + device: str = "cuda:0", +) -> tuple[RigidObjectCollection, torch.Tensor]: + """Generate a scene with the provided number of cubes. + + Args: + num_envs: Number of envs to generate. + num_cubes: Number of cubes to generate. + height: Height of the cubes. + has_api: Whether the cubes have a rigid body API on them. + kinematic_enabled: Whether the cubes are kinematic. + device: Device to use for the simulation. + + Returns: + A tuple containing the rigid object representing the cubes and the origins of the cubes. + + """ + origins = torch.tensor([(i * 3.0, 0, height) for i in range(num_envs)]).to(device) + # Create Top-level Xforms, one for each cube + for i, origin in enumerate(origins): + prim_utils.create_prim(f"/World/Table_{i}", "Xform", translation=origin) + + # Resolve spawn configuration + if has_api: + spawn_cfg = sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", + rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=kinematic_enabled), + ) + else: + # since no rigid body properties defined, this is just a static collider + spawn_cfg = sim_utils.CuboidCfg( + size=(0.1, 0.1, 0.1), + collision_props=sim_utils.CollisionPropertiesCfg(), + ) + + # create the rigid object configs + cube_config_dict = {} + for i in range(num_cubes): + cube_object_cfg = RigidObjectCfg( + prim_path=f"/World/Table_.*/Object_{i}", + spawn=spawn_cfg, + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 3 * i, height)), + ) + cube_config_dict[f"cube_{i}"] = cube_object_cfg + # create the rigid object collection + cube_object_collection_cfg = RigidObjectCollectionCfg(rigid_objects=cube_config_dict) + cube_object_colection = RigidObjectCollection(cfg=cube_object_collection_cfg) + + return cube_object_colection, origins + + +class TestRigidObjectCollection(unittest.TestCase): + """Test for rigid object collection class.""" + + """ + Tests + """ + + def test_initialization(self): + """Test initialization for prim with rigid body API at the provided prim path.""" + for num_envs in (1, 2): + for num_cubes in (1, 3): + for device in ("cuda:0", "cpu"): + with self.subTest(num_envs=num_envs, num_cubes=num_cubes, device=device): + with build_simulation_context(device=device, auto_add_lighting=True) as sim: + # Generate cubes scene + object_collection, _ = generate_cubes_scene( + num_envs=num_envs, num_cubes=num_cubes, device=device + ) + + # Check that boundedness of rigid object is correct + self.assertEqual(ctypes.c_long.from_address(id(object_collection)).value, 1) + + # Play sim + sim.reset() + + # Check if object is initialized + self.assertTrue(object_collection.is_initialized) + self.assertEqual(len(object_collection.object_names), num_cubes) + + # Check buffers that exists and have correct shapes + self.assertEqual(object_collection.data.object_link_pos_w.shape, (num_envs, num_cubes, 3)) + self.assertEqual(object_collection.data.object_link_quat_w.shape, (num_envs, num_cubes, 4)) + self.assertEqual(object_collection.data.default_mass.shape, (num_envs, num_cubes, 1)) + self.assertEqual(object_collection.data.default_inertia.shape, (num_envs, num_cubes, 9)) + + # Simulate physics + for _ in range(2): + # perform rendering + sim.step() + # update object + object_collection.update(sim.cfg.dt) + + def test_id_conversion(self): + """Test environment and object index conversion to physics view indices.""" + for device in ("cuda:0", "cpu"): + with self.subTest(num_envs=2, num_cubes=3, device=device): + with build_simulation_context(device=device, auto_add_lighting=True) as sim: + # Generate cubes scene + object_collection, _ = generate_cubes_scene(num_envs=2, num_cubes=3, device=device) + + # Play sim + sim.reset() + + expected = [ + torch.tensor([4, 5], device=device, dtype=torch.long), + torch.tensor([4], device=device, dtype=torch.long), + torch.tensor([0, 2, 4], device=device, dtype=torch.long), + torch.tensor([1, 3, 5], device=device, dtype=torch.long), + ] + + view_ids = object_collection._env_obj_ids_to_view_ids( + object_collection._ALL_ENV_INDICES, object_collection._ALL_OBJ_INDICES[None, 2] + ) + self.assertTrue((view_ids == expected[0]).all()) + view_ids = object_collection._env_obj_ids_to_view_ids( + object_collection._ALL_ENV_INDICES[None, 0], object_collection._ALL_OBJ_INDICES[None, 2] + ) + self.assertTrue((view_ids == expected[1]).all()) + view_ids = object_collection._env_obj_ids_to_view_ids( + object_collection._ALL_ENV_INDICES[None, 0], object_collection._ALL_OBJ_INDICES + ) + self.assertTrue((view_ids == expected[2]).all()) + view_ids = object_collection._env_obj_ids_to_view_ids( + object_collection._ALL_ENV_INDICES[None, 1], object_collection._ALL_OBJ_INDICES + ) + self.assertTrue((view_ids == expected[3]).all()) + + def test_initialization_with_kinematic_enabled(self): + """Test that initialization for prim with kinematic flag enabled.""" + for num_envs in (1, 2): + for num_cubes in (1, 3): + for device in ("cuda:0", "cpu"): + with self.subTest(num_envs=num_envs, num_cubes=num_cubes, device=device): + with build_simulation_context(device=device, auto_add_lighting=True) as sim: + # Generate cubes scene + object_collection, origins = generate_cubes_scene( + num_envs=num_envs, num_cubes=num_cubes, kinematic_enabled=True, device=device + ) + + # Check that boundedness of rigid object is correct + self.assertEqual(ctypes.c_long.from_address(id(object_collection)).value, 1) + + # Play sim + sim.reset() + + # Check if object is initialized + self.assertTrue(object_collection.is_initialized) + self.assertEqual(len(object_collection.object_names), num_cubes) + + # Check buffers that exists and have correct shapes + self.assertEqual(object_collection.data.object_link_pos_w.shape, (num_envs, num_cubes, 3)) + self.assertEqual(object_collection.data.object_link_quat_w.shape, (num_envs, num_cubes, 4)) + + # Simulate physics + for _ in range(2): + # perform rendering + sim.step() + # update object + object_collection.update(sim.cfg.dt) + # check that the object is kinematic + default_object_state = object_collection.data.default_object_state.clone() + default_object_state[..., :3] += origins.unsqueeze(1) + torch.testing.assert_close( + object_collection.data.object_link_state_w, default_object_state + ) + + def test_initialization_with_no_rigid_body(self): + """Test that initialization fails when no rigid body is found at the provided prim path.""" + for num_cubes in (1, 2): + for device in ("cuda:0", "cpu"): + with self.subTest(num_cubes=num_cubes, device=device): + with build_simulation_context(device=device, auto_add_lighting=True) as sim: + # Generate cubes scene + object_collection, _ = generate_cubes_scene(num_cubes=num_cubes, has_api=False, device=device) + + # Check that boundedness of rigid object is correct + self.assertEqual(ctypes.c_long.from_address(id(object_collection)).value, 1) + + # Play sim + sim.reset() + + # Check if object is initialized + self.assertFalse(object_collection.is_initialized) + + def test_external_force_on_single_body(self): + """Test application of external force on the base of the object. + + In this test, we apply a force equal to the weight of an object on the base of + one of the objects. We check that the object does not move. For the other object, + we do not apply any force and check that it falls down. + """ + for num_envs in (1, 2): + for num_cubes in (1, 4): + for device in ("cuda:0", "cpu"): + with self.subTest(num_envs=num_envs, num_cubes=num_cubes, device=device): + # Generate cubes scene + with build_simulation_context( + device=device, add_ground_plane=True, auto_add_lighting=True + ) as sim: + object_collection, origins = generate_cubes_scene( + num_envs=num_envs, num_cubes=num_cubes, device=device + ) + + # Play the simulator + sim.reset() + + # Find objects to apply the force + object_ids, object_names = object_collection.find_objects(".*") + + # Sample a force equal to the weight of the object + external_wrench_b = torch.zeros( + object_collection.num_instances, len(object_ids), 6, device=sim.device + ) + # Every 2nd cube should have a force applied to it + external_wrench_b[:, 0::2, 2] = 9.81 * object_collection.data.default_mass[:, 0::2, 0] + + # Now we are ready! + for _ in range(5): + # reset object state + object_state = object_collection.data.default_object_state.clone() + + # need to shift the position of the cubes otherwise they will be on top of each other + object_state[..., :2] += origins.unsqueeze(1)[..., :2] + object_collection.write_object_state_to_sim(object_state) + + # reset object + object_collection.reset() + + # apply force + object_collection.set_external_force_and_torque( + external_wrench_b[..., :3], external_wrench_b[..., 3:], object_ids=object_ids + ) + # perform simulation + for _ in range(10): + # apply action to the object + object_collection.write_data_to_sim() + + # perform step + sim.step() + + # update buffers + object_collection.update(sim.cfg.dt) + + # First object should still be at the same Z position (1.0) + torch.testing.assert_close( + object_collection.data.object_link_pos_w[:, 0::2, 2], + torch.ones_like(object_collection.data.object_pos_w[:, 0::2, 2]), + ) + # Second object should have fallen, so it's Z height should be less than initial height of 1.0 + self.assertTrue(torch.all(object_collection.data.object_link_pos_w[:, 1::2, 2] < 1.0)) + + def test_set_object_state(self): + """Test setting the state of the object. + + In this test, we set the state of the object to a random state and check + that the object is in that state after simulation. We set gravity to zero as + we don't want any external forces acting on the object to ensure state remains static. + """ + for num_envs in (1, 3): + for num_cubes in (1, 2): + for device in ("cuda:0", "cpu"): + with self.subTest(num_envs=num_envs, num_cubes=num_cubes, device=device): + # Turn off gravity for this test as we don't want any external forces acting on the object + # to ensure state remains static + with build_simulation_context( + device=device, gravity_enabled=False, auto_add_lighting=True + ) as sim: + # Generate cubes scene + object_collection, origins = generate_cubes_scene( + num_envs=num_envs, num_cubes=num_cubes, device=device + ) + + # Play the simulator + sim.reset() + + state_types = ["object_pos_w", "object_quat_w", "object_lin_vel_w", "object_ang_vel_w"] + + # Set each state type individually as they are dependent on each other + for state_type_to_randomize in state_types: + state_dict = { + "object_pos_w": torch.zeros_like( + object_collection.data.object_pos_w, device=sim.device + ), + "object_quat_w": default_orientation( + num=num_cubes * num_envs, device=sim.device + ).view(num_envs, num_cubes, 4), + "object_lin_vel_w": torch.zeros_like( + object_collection.data.object_lin_vel_w, device=sim.device + ), + "object_ang_vel_w": torch.zeros_like( + object_collection.data.object_ang_vel_w, device=sim.device + ), + } + + # Now we are ready! + for _ in range(5): + # reset object + object_collection.reset() + + # Set random state + if state_type_to_randomize == "object_quat_w": + state_dict[state_type_to_randomize] = random_orientation( + num=num_cubes * num_envs, device=sim.device + ).view(num_envs, num_cubes, 4) + else: + state_dict[state_type_to_randomize] = torch.randn( + num_envs, num_cubes, 3, device=sim.device + ) + # make sure objects do not overlap + if state_type_to_randomize == "object_pos_w": + state_dict[state_type_to_randomize][..., :2] += origins.unsqueeze(1)[ + ..., :2 + ] + + # perform simulation + for _ in range(5): + object_state = torch.cat( + [ + state_dict["object_pos_w"], + state_dict["object_quat_w"], + state_dict["object_lin_vel_w"], + state_dict["object_ang_vel_w"], + ], + dim=-1, + ) + # reset object state + object_collection.write_object_state_to_sim(object_state=object_state) + + sim.step() + + # assert that set object quantities are equal to the ones set in the state_dict + for key, expected_value in state_dict.items(): + value = getattr(object_collection.data, key) + torch.testing.assert_close(value, expected_value, rtol=1e-5, atol=1e-5) + + object_collection.update(sim.cfg.dt) + + def test_object_state_properties(self): + """Test the object_com_state_w and object_link_state_w properties.""" + for num_envs in (1, 4): + for num_cubes in (1, 2): + for device in ("cuda:0", "cpu"): + for with_offset in [True, False]: + with self.subTest( + num_envs=num_envs, num_cubes=num_cubes, device=device, with_offset=with_offset + ): + with build_simulation_context( + device=device, gravity_enabled=False, auto_add_lighting=True + ) as sim: + # Create a scene with random cubes + cube_object, env_pos = generate_cubes_scene( + num_envs=num_envs, num_cubes=num_cubes, height=0.0, device=device + ) + view_ids = torch.tensor([x for x in range(num_cubes * num_envs)]) + + # Play sim + sim.reset() + + # Check if cube_object is initialized + self.assertTrue(cube_object.is_initialized) + + # change center of mass offset from link frame + if with_offset: + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) + else: + offset = torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) + + com = cube_object.reshape_view_to_data(cube_object.root_physx_view.get_coms()) + com[..., :3] = offset.to("cpu") + cube_object.root_physx_view.set_coms( + cube_object.reshape_data_to_view(com.clone()), view_ids + ) + + # check center of mass has been set + torch.testing.assert_close( + cube_object.reshape_view_to_data(cube_object.root_physx_view.get_coms()), com + ) + + # random z spin velocity + spin_twist = torch.zeros(6, device=device) + spin_twist[5] = torch.randn(1, device=device) + + # initial spawn point + init_com = cube_object.data.object_com_state_w[..., :3] + + # Simulate physics + for i in range(10): + # spin the object around Z axis (com) + cube_object.write_object_com_velocity_to_sim( + spin_twist.repeat(num_envs, num_cubes, 1) + ) + # perform rendering + sim.step() + # update object + cube_object.update(sim.cfg.dt) + + # get state properties + object_state_w = cube_object.data.object_state_w + object_link_state_w = cube_object.data.object_link_state_w + object_com_state_w = cube_object.data.object_com_state_w + + # if offset is [0,0,0] all object_state_%_w will match and all body_%_w will match + if not with_offset: + torch.testing.assert_close(object_state_w, object_com_state_w) + torch.testing.assert_close(object_state_w, object_link_state_w) + else: + # cubes are spinning around center of mass + # position will not match + # center of mass position will be constant (i.e. spinning around com) + torch.testing.assert_close(init_com, object_com_state_w[..., :3]) + + # link position will be moving but should stay constant away from center of mass + object_link_state_pos_rel_com = quat_rotate_inverse( + object_link_state_w[..., 3:7], + object_link_state_w[..., :3] - object_com_state_w[..., :3], + ) + + torch.testing.assert_close(-offset, object_link_state_pos_rel_com) + + # orientation of com will be a constant rotation from link orientation + com_quat_b = cube_object.data.com_quat_b + com_quat_w = quat_mul(object_link_state_w[..., 3:7], com_quat_b) + torch.testing.assert_close(com_quat_w, object_com_state_w[..., 3:7]) + + # orientation of link will match object state will always match + torch.testing.assert_close( + object_state_w[..., 3:7], object_link_state_w[..., 3:7] + ) + + # lin_vel will not match + # center of mass vel will be constant (i.e. spining around com) + torch.testing.assert_close( + torch.zeros_like(object_com_state_w[..., 7:10]), + object_com_state_w[..., 7:10], + ) + + # link frame will be moving, and should be equal to input angular velocity cross offset + lin_vel_rel_object_gt = quat_rotate_inverse( + object_link_state_w[..., 3:7], object_link_state_w[..., 7:10] + ) + lin_vel_rel_gt = torch.linalg.cross( + spin_twist.repeat(num_envs, num_cubes, 1)[..., 3:], -offset + ) + torch.testing.assert_close( + lin_vel_rel_gt, lin_vel_rel_object_gt, atol=1e-4, rtol=1e-3 + ) + + # ang_vel will always match + torch.testing.assert_close( + object_state_w[..., 10:], object_com_state_w[..., 10:] + ) + torch.testing.assert_close( + object_state_w[..., 10:], object_link_state_w[..., 10:] + ) + + def test_write_object_state(self): + """Test the setters for object_state using both the link frame and center of mass as reference frame.""" + for num_envs in (1, 3): + for num_cubes in (1, 2): + for device in ("cuda:0", "cpu"): + for with_offset in [True, False]: + for state_location in ("com", "link"): + with self.subTest( + num_envs=num_envs, num_cubes=num_cubes, device=device, with_offset=with_offset + ): + with build_simulation_context( + device=device, gravity_enabled=False, auto_add_lighting=True + ) as sim: + # Create a scene with random cubes + cube_object, env_pos = generate_cubes_scene( + num_envs=num_envs, num_cubes=num_cubes, height=0.0, device=device + ) + view_ids = torch.tensor([x for x in range(num_cubes * num_cubes)]) + + # Play sim + sim.reset() + + # Check if cube_object is initialized + self.assertTrue(cube_object.is_initialized) + + # change center of mass offset from link frame + if with_offset: + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat( + num_envs, num_cubes, 1 + ) + else: + offset = torch.tensor([0.0, 0.0, 0.0], device=device).repeat( + num_envs, num_cubes, 1 + ) + + com = cube_object.reshape_view_to_data(cube_object.root_physx_view.get_coms()) + com[..., :3] = offset.to("cpu") + cube_object.root_physx_view.set_coms( + cube_object.reshape_data_to_view(com.clone()), view_ids + ) + + # check center of mass has been set + torch.testing.assert_close( + cube_object.reshape_view_to_data(cube_object.root_physx_view.get_coms()), com + ) + + rand_state = torch.zeros_like(cube_object.data.object_link_state_w) + rand_state[..., :7] = cube_object.data.default_object_state[..., :7] + rand_state[..., :3] += cube_object.data.object_link_pos_w + # make quaternion a unit vector + rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) + + for i in range(10): + + # perform step + sim.step() + # update buffers + cube_object.update(sim.cfg.dt) + + if state_location == "com": + cube_object.write_object_com_state_to_sim(rand_state) + elif state_location == "link": + cube_object.write_object_link_state_to_sim(rand_state) + + if state_location == "com": + torch.testing.assert_close(rand_state, cube_object.data.object_com_state_w) + elif state_location == "link": + torch.testing.assert_close(rand_state, cube_object.data.object_link_state_w) + + def test_reset_object_collection(self): + """Test resetting the state of the rigid object.""" + for num_envs in (1, 3): + for num_cubes in (1, 2): + for device in ("cuda:0", "cpu"): + with self.subTest(num_envs=num_envs, num_cubes=num_cubes, device=device): + with build_simulation_context( + device=device, gravity_enabled=True, auto_add_lighting=True + ) as sim: + # Generate cubes scene + object_collection, _ = generate_cubes_scene( + num_envs=num_envs, num_cubes=num_cubes, device=device + ) + + # Play the simulator + sim.reset() + + for i in range(5): + # perform rendering + sim.step() + + # update object + object_collection.update(sim.cfg.dt) + + # Move the object to a random position + object_state = object_collection.data.default_object_state.clone() + object_state[..., :3] = torch.randn(num_envs, num_cubes, 3, device=sim.device) + + # Random orientation + object_state[..., 3:7] = random_orientation(num=num_cubes, device=sim.device) + object_collection.write_object_state_to_sim(object_state) + + if i % 2 == 0: + # reset object + object_collection.reset() + + # Reset should zero external forces and torques + self.assertFalse(object_collection.has_external_wrench) + self.assertEqual(torch.count_nonzero(object_collection._external_force_b), 0) + self.assertEqual(torch.count_nonzero(object_collection._external_torque_b), 0) + + def test_set_material_properties(self): + """Test getting and setting material properties of rigid object.""" + for num_envs in (1, 3): + for num_cubes in (1, 2): + for device in ("cuda:0", "cpu"): + with self.subTest(num_envs=num_envs, num_cubes=num_cubes, device=device): + with build_simulation_context( + device=device, gravity_enabled=True, add_ground_plane=True, auto_add_lighting=True + ) as sim: + # Generate cubes scene + object_collection, _ = generate_cubes_scene( + num_envs=num_envs, num_cubes=num_cubes, device=device + ) + + # Play sim + sim.reset() + + # Set material properties + static_friction = torch.FloatTensor(num_envs, num_cubes, 1).uniform_(0.4, 0.8) + dynamic_friction = torch.FloatTensor(num_envs, num_cubes, 1).uniform_(0.4, 0.8) + restitution = torch.FloatTensor(num_envs, num_cubes, 1).uniform_(0.0, 0.2) + + materials = torch.cat([static_friction, dynamic_friction, restitution], dim=-1) + + indices = torch.tensor(range(num_cubes * num_envs), dtype=torch.int) + # Add friction to cube + object_collection.root_physx_view.set_material_properties( + object_collection.reshape_data_to_view(materials), indices + ) + + # Simulate physics + sim.step() + # update object + object_collection.update(sim.cfg.dt) + + # Get material properties + materials_to_check = object_collection.root_physx_view.get_material_properties() + + # Check if material properties are set correctly + torch.testing.assert_close( + object_collection.reshape_view_to_data(materials_to_check), materials + ) + + def test_gravity_vec_w(self): + """Test that gravity vector direction is set correctly for the rigid object.""" + for num_envs in (1, 3): + for num_cubes in (1, 2): + for device in ("cuda:0", "cpu"): + for gravity_enabled in [True, False]: + with self.subTest( + num_envs=num_envs, num_cubes=num_cubes, device=device, gravity_enabled=gravity_enabled + ): + with build_simulation_context(device=device, gravity_enabled=gravity_enabled) as sim: + # Create a scene with random cubes + object_collection, _ = generate_cubes_scene( + num_envs=num_envs, num_cubes=num_cubes, device=device + ) + + # Obtain gravity direction + if gravity_enabled: + gravity_dir = (0.0, 0.0, -1.0) + else: + gravity_dir = (0.0, 0.0, 0.0) + + # Play sim + sim.reset() + + # Check that gravity is set correctly + self.assertEqual(object_collection.data.GRAVITY_VEC_W[0, 0, 0], gravity_dir[0]) + self.assertEqual(object_collection.data.GRAVITY_VEC_W[0, 0, 1], gravity_dir[1]) + self.assertEqual(object_collection.data.GRAVITY_VEC_W[0, 0, 2], gravity_dir[2]) + + # Simulate physics + for _ in range(2): + sim.step() + # update object + object_collection.update(sim.cfg.dt) + + # Expected gravity value is the acceleration of the body + gravity = torch.zeros(num_envs, num_cubes, 6, device=device) + if gravity_enabled: + gravity[..., 2] = -9.81 + # Check the body accelerations are correct + torch.testing.assert_close(object_collection.data.object_acc_w, gravity) + + +if __name__ == "__main__": + run_tests() diff --git a/source/extensions/omni.isaac.lab/test/controllers/test_differential_ik.py b/source/extensions/omni.isaac.lab/test/controllers/test_differential_ik.py index 4d5877aaa1..afbbb3214c 100644 --- a/source/extensions/omni.isaac.lab/test/controllers/test_differential_ik.py +++ b/source/extensions/omni.isaac.lab/test/controllers/test_differential_ik.py @@ -22,7 +22,14 @@ import omni.isaac.lab.sim as sim_utils from omni.isaac.lab.assets import Articulation from omni.isaac.lab.controllers import DifferentialIKController, DifferentialIKControllerCfg -from omni.isaac.lab.utils.math import compute_pose_error, subtract_frame_transforms + +from omni.isaac.lab.utils.math import ( # isort:skip + compute_pose_error, + matrix_from_quat, + quat_inv, + random_yaw_orientation, + subtract_frame_transforms, +) ## # Pre-defined configs @@ -140,8 +147,8 @@ def _run_ik_controller( ee_pose_b_des = torch.zeros(self.num_envs, diff_ik_controller.action_dim, device=self.sim.device) ee_pose_b_des[:] = self.ee_pose_b_des_set[current_goal_idx] # Compute current pose of the end-effector - ee_pose_w = robot.data.body_state_w[:, ee_frame_idx, 0:7] - root_pose_w = robot.data.root_state_w[:, 0:7] + ee_pose_w = robot.data.body_link_state_w[:, ee_frame_idx, 0:7] + root_pose_w = robot.data.root_link_state_w[:, 0:7] ee_pos_b, ee_quat_b = subtract_frame_transforms( root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_pose_w[:, 0:3], ee_pose_w[:, 3:7] ) @@ -169,6 +176,11 @@ def _run_ik_controller( robot.write_joint_state_to_sim(joint_pos, joint_vel) robot.set_joint_position_target(joint_pos) robot.write_data_to_sim() + # randomize root state yaw, ik should work regardless base rotation + root_state = robot.data.root_link_state_w.clone() + root_state[:, 3:7] = random_yaw_orientation(self.num_envs, self.sim.device) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) robot.reset() # reset actions ee_pose_b_des[:] = self.ee_pose_b_des_set[current_goal_idx] @@ -183,8 +195,12 @@ def _run_ik_controller( # so we MUST skip the first step # obtain quantities from simulation jacobian = robot.root_physx_view.get_jacobians()[:, ee_jacobi_idx, :, arm_joint_ids] - ee_pose_w = robot.data.body_state_w[:, ee_frame_idx, 0:7] - root_pose_w = robot.data.root_state_w[:, 0:7] + ee_pose_w = robot.data.body_link_state_w[:, ee_frame_idx, 0:7] + root_pose_w = robot.data.root_link_state_w[:, 0:7] + base_rot = root_pose_w[:, 3:7] + base_rot_matrix = matrix_from_quat(quat_inv(base_rot)) + jacobian[:, :3, :] = torch.bmm(base_rot_matrix, jacobian[:, :3, :]) + jacobian[:, 3:, :] = torch.bmm(base_rot_matrix, jacobian[:, 3:, :]) joint_pos = robot.data.joint_pos[:, arm_joint_ids] # compute frame in root frame ee_pos_b, ee_quat_b = subtract_frame_transforms( diff --git a/source/extensions/omni.isaac.lab/test/controllers/test_operational_space.py b/source/extensions/omni.isaac.lab/test/controllers/test_operational_space.py new file mode 100644 index 0000000000..bfde34681e --- /dev/null +++ b/source/extensions/omni.isaac.lab/test/controllers/test_operational_space.py @@ -0,0 +1,937 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Launch Isaac Sim Simulator first.""" + +from omni.isaac.lab.app import AppLauncher, run_tests + +# launch omniverse app +simulation_app = AppLauncher(headless=True).app + +"""Rest everything follows.""" + +import torch +import unittest + +import omni.isaac.core.utils.prims as prim_utils +import omni.isaac.core.utils.stage as stage_utils +from omni.isaac.cloner import GridCloner + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.assets import Articulation +from omni.isaac.lab.controllers import OperationalSpaceController, OperationalSpaceControllerCfg +from omni.isaac.lab.markers import VisualizationMarkers +from omni.isaac.lab.markers.config import FRAME_MARKER_CFG +from omni.isaac.lab.sensors import ContactSensor, ContactSensorCfg +from omni.isaac.lab.utils.math import ( + apply_delta_pose, + combine_frame_transforms, + compute_pose_error, + matrix_from_quat, + quat_inv, + quat_rotate_inverse, + subtract_frame_transforms, +) + +## +# Pre-defined configs +## +from omni.isaac.lab_assets import FRANKA_PANDA_CFG # isort:skip + + +class TestOperationalSpaceController(unittest.TestCase): + """Test fixture for checking that Operational Space controller tracks commands properly.""" + + def setUp(self): + """Create a blank new stage for each test.""" + # Wait for spawning + stage_utils.create_new_stage() + # Constants + self.num_envs = 16 + # Load kit helper + sim_cfg = sim_utils.SimulationCfg(dt=0.01) + self.sim = sim_utils.SimulationContext(sim_cfg) + # TODO: Remove this once we have a better way to handle this. + self.sim._app_control_on_stop_handle = None + + # Create a ground plane + cfg = sim_utils.GroundPlaneCfg() + cfg.func("/World/GroundPlane", cfg) + + # Markers + frame_marker_cfg = FRAME_MARKER_CFG.copy() + frame_marker_cfg.markers["frame"].scale = (0.1, 0.1, 0.1) + self.ee_marker = VisualizationMarkers(frame_marker_cfg.replace(prim_path="/Visuals/ee_current")) + self.goal_marker = VisualizationMarkers(frame_marker_cfg.replace(prim_path="/Visuals/ee_goal")) + + light_cfg = sim_utils.DistantLightCfg(intensity=5.0, exposure=10.0) + light_cfg.func( + "/Light", + light_cfg, + translation=[0, 0, 1], + ) + + # Create interface to clone the scene + cloner = GridCloner(spacing=2.0) + cloner.define_base_env("/World/envs") + self.env_prim_paths = cloner.generate_paths("/World/envs/env", self.num_envs) + # create source prim + prim_utils.define_prim(self.env_prim_paths[0], "Xform") + # clone the env xform + self.env_origins = cloner.clone( + source_prim_path=self.env_prim_paths[0], + prim_paths=self.env_prim_paths, + replicate_physics=True, + ) + + self.robot_cfg = FRANKA_PANDA_CFG.replace(prim_path="/World/envs/env_.*/Robot") + self.robot_cfg.actuators["panda_shoulder"].stiffness = 0.0 + self.robot_cfg.actuators["panda_shoulder"].damping = 0.0 + self.robot_cfg.actuators["panda_forearm"].stiffness = 0.0 + self.robot_cfg.actuators["panda_forearm"].damping = 0.0 + self.robot_cfg.spawn.rigid_props.disable_gravity = True + + # Define the ContactSensor + self.contact_forces = None + + # Define the target sets + ee_goal_abs_pos_set_b = torch.tensor( + [ + [0.5, 0.5, 0.7], + [0.5, -0.4, 0.6], + [0.5, 0, 0.5], + ], + device=self.sim.device, + ) + ee_goal_abs_quad_set_b = torch.tensor( + [ + [0.707, 0.0, 0.707, 0.0], + [0.707, 0.707, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + ], + device=self.sim.device, + ) + ee_goal_rel_pos_set = torch.tensor( + [ + [0.2, 0.0, 0.0], + [0.2, 0.2, 0.0], + [0.2, 0.2, -0.2], + ], + device=self.sim.device, + ) + ee_goal_rel_axisangle_set = torch.tensor( + [ + [0.0, torch.pi / 2, 0.0], # for [0.707, 0, 0.707, 0] + [torch.pi / 2, 0.0, 0.0], # for [0.707, 0.707, 0, 0] + [torch.pi, 0.0, 0.0], # for [0.0, 1.0, 0, 0] + ], + device=self.sim.device, + ) + ee_goal_abs_wrench_set_b = torch.tensor( + [ + [0.0, 0.0, 10.0, 0.0, -1.0, 0.0], + [0.0, 10.0, 0.0, 0.0, 0.0, 0.0], + [10.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + device=self.sim.device, + ) + kp_set = torch.tensor( + [ + [200.0, 200.0, 200.0, 200.0, 200.0, 200.0], + [240.0, 240.0, 240.0, 240.0, 240.0, 240.0], + [160.0, 160.0, 160.0, 160.0, 160.0, 160.0], + ], + device=self.sim.device, + ) + d_ratio_set = torch.tensor( + [ + [1.0, 1.0, 1.0, 1.0, 1.0, 1.0], + [1.1, 1.1, 1.1, 1.1, 1.1, 1.1], + [0.9, 0.9, 0.9, 0.9, 0.9, 0.9], + ], + device=self.sim.device, + ) + ee_goal_hybrid_set_b = torch.tensor( + [ + [0.6, 0.2, 0.5, 0.0, 0.707, 0.0, 0.707, 10.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.6, -0.29, 0.6, 0.0, 0.707, 0.0, 0.707, 10.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.6, 0.1, 0.8, 0.0, 0.5774, 0.0, 0.8165, 4.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + device=self.sim.device, + ) + ee_goal_pose_set_tilted_b = torch.tensor( + [ + [0.6, 0.15, 0.3, 0.0, 0.92387953, 0.0, 0.38268343], + [0.6, -0.3, 0.3, 0.0, 0.92387953, 0.0, 0.38268343], + [0.8, 0.0, 0.5, 0.0, 0.92387953, 0.0, 0.38268343], + ], + device=self.sim.device, + ) + ee_goal_wrench_set_tilted_task = torch.tensor( + [ + [0.0, 0.0, 10.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 10.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 10.0, 0.0, 0.0, 0.0], + ], + device=self.sim.device, + ) + + # Define goals for the arm [xyz] + self.target_abs_pos_set_b = ee_goal_abs_pos_set_b.clone() + # Define goals for the arm [xyz + quat_wxyz] + self.target_abs_pose_set_b = torch.cat([ee_goal_abs_pos_set_b, ee_goal_abs_quad_set_b], dim=-1) + # Define goals for the arm [xyz] + self.target_rel_pos_set = ee_goal_rel_pos_set.clone() + # Define goals for the arm [xyz + axis-angle] + self.target_rel_pose_set_b = torch.cat([ee_goal_rel_pos_set, ee_goal_rel_axisangle_set], dim=-1) + # Define goals for the arm [force_xyz + torque_xyz] + self.target_abs_wrench_set = ee_goal_abs_wrench_set_b.clone() + # Define goals for the arm [xyz + quat_wxyz] and variable kp [kp_xyz + kp_rot_xyz] + self.target_abs_pose_variable_kp_set = torch.cat([self.target_abs_pose_set_b, kp_set], dim=-1) + # Define goals for the arm [xyz + quat_wxyz] and the variable imp. [kp_xyz + kp_rot_xyz + d_xyz + d_rot_xyz] + self.target_abs_pose_variable_set = torch.cat([self.target_abs_pose_set_b, kp_set, d_ratio_set], dim=-1) + # Define goals for the arm pose [xyz + quat_wxyz] and wrench [force_xyz + torque_xyz] + self.target_hybrid_set_b = ee_goal_hybrid_set_b.clone() + # Define goals for the arm pose, and wrench, and kp + self.target_hybrid_variable_kp_set = torch.cat([self.target_hybrid_set_b, kp_set], dim=-1) + # Define goals for the arm pose [xyz + quat_wxyz] in root and and wrench [force_xyz + torque_xyz] in task frame + self.target_hybrid_set_tilted = torch.cat([ee_goal_pose_set_tilted_b, ee_goal_wrench_set_tilted_task], dim=-1) + + # Reference frame for targets + self.frame = "root" + + def tearDown(self): + """Stops simulator after each test.""" + # stop simulation + self.sim.stop() + # self.sim.clear() # FIXME: This hangs the test for some reason when LIVESTREAM is not enabled. + self.sim.clear_all_callbacks() + self.sim.clear_instance() + # Make contact_forces None after relevant tests otherwise other tests give warning + self.contact_forces = None + + """ + Test fixtures. + """ + + def test_franka_pose_abs_without_inertial_decoupling(self): + """Test absolute pose control with fixed impedance and without inertial dynamics decoupling.""" + robot = Articulation(cfg=self.robot_cfg) + osc_cfg = OperationalSpaceControllerCfg( + target_types=["pose_abs"], + impedance_mode="fixed", + inertial_dynamics_decoupling=False, + gravity_compensation=False, + motion_stiffness_task=[400.0, 400.0, 400.0, 100.0, 100.0, 100.0], + motion_damping_ratio_task=[5.0, 5.0, 5.0, 0.001, 0.001, 0.001], + ) + osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device) + + self._run_op_space_controller(robot, osc, "panda_hand", ["panda_joint.*"], self.target_abs_pose_set_b) + + def test_franka_pose_abs_with_partial_inertial_decoupling(self): + """Test absolute pose control with fixed impedance and partial inertial dynamics decoupling.""" + robot = Articulation(cfg=self.robot_cfg) + osc_cfg = OperationalSpaceControllerCfg( + target_types=["pose_abs"], + impedance_mode="fixed", + inertial_dynamics_decoupling=True, + partial_inertial_dynamics_decoupling=True, + gravity_compensation=False, + motion_stiffness_task=1000.0, + motion_damping_ratio_task=1.0, + ) + osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device) + + self._run_op_space_controller(robot, osc, "panda_hand", ["panda_joint.*"], self.target_abs_pose_set_b) + + def test_franka_pose_abs_fixed_impedance_with_gravity_compensation(self): + """Test absolute pose control with fixed impedance, gravity compensation, and inertial dynamics decoupling.""" + self.robot_cfg.spawn.rigid_props.disable_gravity = False + robot = Articulation(cfg=self.robot_cfg) + osc_cfg = OperationalSpaceControllerCfg( + target_types=["pose_abs"], + impedance_mode="fixed", + inertial_dynamics_decoupling=True, + partial_inertial_dynamics_decoupling=False, + gravity_compensation=True, + motion_stiffness_task=500.0, + motion_damping_ratio_task=1.0, + ) + osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device) + + self._run_op_space_controller(robot, osc, "panda_hand", ["panda_joint.*"], self.target_abs_pose_set_b) + + def test_franka_pose_abs(self): + """Test absolute pose control with fixed impedance and inertial dynamics decoupling.""" + robot = Articulation(cfg=self.robot_cfg) + osc_cfg = OperationalSpaceControllerCfg( + target_types=["pose_abs"], + impedance_mode="fixed", + inertial_dynamics_decoupling=True, + partial_inertial_dynamics_decoupling=False, + gravity_compensation=False, + motion_stiffness_task=500.0, + motion_damping_ratio_task=1.0, + ) + osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device) + + self._run_op_space_controller(robot, osc, "panda_hand", ["panda_joint.*"], self.target_abs_pose_set_b) + + def test_franka_pose_rel(self): + """Test relative pose control with fixed impedance and inertial dynamics decoupling.""" + robot = Articulation(cfg=self.robot_cfg) + osc_cfg = OperationalSpaceControllerCfg( + target_types=["pose_rel"], + impedance_mode="fixed", + inertial_dynamics_decoupling=True, + partial_inertial_dynamics_decoupling=False, + gravity_compensation=False, + motion_stiffness_task=500.0, + motion_damping_ratio_task=1.0, + ) + osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device) + + self._run_op_space_controller(robot, osc, "panda_hand", ["panda_joint.*"], self.target_rel_pose_set_b) + + def test_franka_pose_abs_variable_impedance(self): + """Test absolute pose control with variable impedance and inertial dynamics decoupling.""" + robot = Articulation(cfg=self.robot_cfg) + osc_cfg = OperationalSpaceControllerCfg( + target_types=["pose_abs"], + impedance_mode="variable", + inertial_dynamics_decoupling=True, + partial_inertial_dynamics_decoupling=False, + gravity_compensation=False, + ) + osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device) + + self._run_op_space_controller(robot, osc, "panda_hand", ["panda_joint.*"], self.target_abs_pose_variable_set) + + def test_franka_wrench_abs_open_loop(self): + """Test open loop absolute force control.""" + robot = Articulation(cfg=self.robot_cfg) + + obstacle_spawn_cfg = sim_utils.CuboidCfg( + size=(0.7, 0.7, 0.01), + collision_props=sim_utils.CollisionPropertiesCfg(), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), opacity=0.1), + rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True), + activate_contact_sensors=True, + ) + obstacle_spawn_cfg.func( + "/World/envs/env_.*/obstacle1", + obstacle_spawn_cfg, + translation=(0.2, 0.0, 0.93), + orientation=(0.9848, 0.0, -0.1736, 0.0), + ) + obstacle_spawn_cfg.func( + "/World/envs/env_.*/obstacle2", + obstacle_spawn_cfg, + translation=(0.2, 0.35, 0.7), + orientation=(0.707, 0.707, 0.0, 0.0), + ) + obstacle_spawn_cfg.func( + "/World/envs/env_.*/obstacle3", + obstacle_spawn_cfg, + translation=(0.55, 0.0, 0.7), + orientation=(0.707, 0.0, 0.707, 0.0), + ) + contact_forces_cfg = ContactSensorCfg( + prim_path="/World/envs/env_.*/obstacle.*", + update_period=0.0, + history_length=50, + debug_vis=False, + force_threshold=0.1, + ) + self.contact_forces = ContactSensor(contact_forces_cfg) + + osc_cfg = OperationalSpaceControllerCfg( + target_types=["wrench_abs"], + motion_control_axes_task=[0, 0, 0, 0, 0, 0], + contact_wrench_control_axes_task=[1, 1, 1, 1, 1, 1], + ) + osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device) + + self._run_op_space_controller(robot, osc, "panda_hand", ["panda_joint.*"], self.target_abs_wrench_set) + + def test_franka_wrench_abs_closed_loop(self): + """Test closed loop absolute force control.""" + robot = Articulation(cfg=self.robot_cfg) + + obstacle_spawn_cfg = sim_utils.CuboidCfg( + size=(0.7, 0.7, 0.01), + collision_props=sim_utils.CollisionPropertiesCfg(), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), opacity=0.1), + rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True), + activate_contact_sensors=True, + ) + obstacle_spawn_cfg.func( + "/World/envs/env_.*/obstacle1", + obstacle_spawn_cfg, + translation=(0.2, 0.0, 0.93), + orientation=(0.9848, 0.0, -0.1736, 0.0), + ) + obstacle_spawn_cfg.func( + "/World/envs/env_.*/obstacle2", + obstacle_spawn_cfg, + translation=(0.2, 0.35, 0.7), + orientation=(0.707, 0.707, 0.0, 0.0), + ) + obstacle_spawn_cfg.func( + "/World/envs/env_.*/obstacle3", + obstacle_spawn_cfg, + translation=(0.55, 0.0, 0.7), + orientation=(0.707, 0.0, 0.707, 0.0), + ) + contact_forces_cfg = ContactSensorCfg( + prim_path="/World/envs/env_.*/obstacle.*", + update_period=0.0, + history_length=2, + debug_vis=False, + force_threshold=0.1, + ) + self.contact_forces = ContactSensor(contact_forces_cfg) + + osc_cfg = OperationalSpaceControllerCfg( + target_types=["wrench_abs"], + contact_wrench_stiffness_task=[ + 0.2, + 0.2, + 0.2, + 0.0, + 0.0, + 0.0, + ], # Zero torque feedback as we cannot contact torque + motion_control_axes_task=[0, 0, 0, 0, 0, 0], + contact_wrench_control_axes_task=[1, 1, 1, 1, 1, 1], + ) + osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device) + + self._run_op_space_controller(robot, osc, "panda_hand", ["panda_joint.*"], self.target_abs_wrench_set) + + def test_franka_hybrid_decoupled_motion(self): + """Test hybrid control with fixed impedance and partial inertial dynamics decoupling.""" + robot = Articulation(cfg=self.robot_cfg) + + obstacle_spawn_cfg = sim_utils.CuboidCfg( + size=(1.0, 1.0, 0.01), + collision_props=sim_utils.CollisionPropertiesCfg(), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), opacity=0.1), + rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True), + activate_contact_sensors=True, + ) + obstacle_spawn_cfg.func( + "/World/envs/env_.*/obstacle1", + obstacle_spawn_cfg, + translation=(self.target_hybrid_set_b[0, 0] + 0.05, 0.0, 0.7), + orientation=(0.707, 0.0, 0.707, 0.0), + ) + contact_forces_cfg = ContactSensorCfg( + prim_path="/World/envs/env_.*/obstacle.*", + update_period=0.0, + history_length=2, + debug_vis=False, + force_threshold=0.1, + ) + self.contact_forces = ContactSensor(contact_forces_cfg) + + osc_cfg = OperationalSpaceControllerCfg( + target_types=["pose_abs", "wrench_abs"], + impedance_mode="fixed", + inertial_dynamics_decoupling=True, + partial_inertial_dynamics_decoupling=True, + gravity_compensation=False, + motion_stiffness_task=300.0, + motion_damping_ratio_task=1.0, + contact_wrench_stiffness_task=[0.1, 0.0, 0.0, 0.0, 0.0, 0.0], + motion_control_axes_task=[0, 1, 1, 1, 1, 1], + contact_wrench_control_axes_task=[1, 0, 0, 0, 0, 0], + ) + osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device) + + self._run_op_space_controller(robot, osc, "panda_leftfinger", ["panda_joint.*"], self.target_hybrid_set_b) + + def test_franka_hybrid_variable_kp_impedance(self): + """Test hybrid control with variable kp impedance and inertial dynamics decoupling.""" + robot = Articulation(cfg=self.robot_cfg) + + obstacle_spawn_cfg = sim_utils.CuboidCfg( + size=(1.0, 1.0, 0.01), + collision_props=sim_utils.CollisionPropertiesCfg(), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), opacity=0.1), + rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True), + activate_contact_sensors=True, + ) + obstacle_spawn_cfg.func( + "/World/envs/env_.*/obstacle1", + obstacle_spawn_cfg, + translation=(self.target_hybrid_set_b[0, 0] + 0.05, 0.0, 0.7), + orientation=(0.707, 0.0, 0.707, 0.0), + ) + contact_forces_cfg = ContactSensorCfg( + prim_path="/World/envs/env_.*/obstacle.*", + update_period=0.0, + history_length=2, + debug_vis=False, + force_threshold=0.1, + ) + self.contact_forces = ContactSensor(contact_forces_cfg) + + osc_cfg = OperationalSpaceControllerCfg( + target_types=["pose_abs", "wrench_abs"], + impedance_mode="variable_kp", + inertial_dynamics_decoupling=True, + partial_inertial_dynamics_decoupling=False, + gravity_compensation=False, + motion_damping_ratio_task=0.8, + contact_wrench_stiffness_task=[0.1, 0.0, 0.0, 0.0, 0.0, 0.0], + motion_control_axes_task=[0, 1, 1, 1, 1, 1], + contact_wrench_control_axes_task=[1, 0, 0, 0, 0, 0], + ) + osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device) + + self._run_op_space_controller( + robot, osc, "panda_leftfinger", ["panda_joint.*"], self.target_hybrid_variable_kp_set + ) + + def test_franka_taskframe_pose_abs(self): + """Test absolute pose control in task frame with fixed impedance and inertial dynamics decoupling.""" + robot = Articulation(cfg=self.robot_cfg) + self.frame = "task" + osc_cfg = OperationalSpaceControllerCfg( + target_types=["pose_abs"], + impedance_mode="fixed", + inertial_dynamics_decoupling=True, + partial_inertial_dynamics_decoupling=False, + gravity_compensation=False, + motion_stiffness_task=500.0, + motion_damping_ratio_task=1.0, + ) + osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device) + + self._run_op_space_controller(robot, osc, "panda_hand", ["panda_joint.*"], self.target_abs_pose_set_b) + + def test_franka_taskframe_pose_rel(self): + """Test relative pose control in task frame with fixed impedance and inertial dynamics decoupling.""" + robot = Articulation(cfg=self.robot_cfg) + self.frame = "task" + osc_cfg = OperationalSpaceControllerCfg( + target_types=["pose_rel"], + impedance_mode="fixed", + inertial_dynamics_decoupling=True, + partial_inertial_dynamics_decoupling=False, + gravity_compensation=False, + motion_stiffness_task=500.0, + motion_damping_ratio_task=1.0, + ) + osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device) + + self._run_op_space_controller(robot, osc, "panda_hand", ["panda_joint.*"], self.target_rel_pose_set_b) + + def test_franka_taskframe_hybrid(self): + """Test hybrid control in task frame with fixed impedance and inertial dynamics decoupling.""" + robot = Articulation(cfg=self.robot_cfg) + self.frame = "task" + + obstacle_spawn_cfg = sim_utils.CuboidCfg( + size=(2.0, 1.5, 0.01), + collision_props=sim_utils.CollisionPropertiesCfg(), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), opacity=0.1), + rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True), + activate_contact_sensors=True, + ) + obstacle_spawn_cfg.func( + "/World/envs/env_.*/obstacle1", + obstacle_spawn_cfg, + translation=(self.target_hybrid_set_tilted[0, 0] + 0.085, 0.0, 0.3), + orientation=(0.9238795325, 0.0, -0.3826834324, 0.0), + ) + contact_forces_cfg = ContactSensorCfg( + prim_path="/World/envs/env_.*/obstacle.*", + update_period=0.0, + history_length=2, + debug_vis=False, + force_threshold=0.1, + ) + self.contact_forces = ContactSensor(contact_forces_cfg) + + osc_cfg = OperationalSpaceControllerCfg( + target_types=["pose_abs", "wrench_abs"], + impedance_mode="fixed", + inertial_dynamics_decoupling=True, + partial_inertial_dynamics_decoupling=False, + gravity_compensation=False, + motion_stiffness_task=400.0, + motion_damping_ratio_task=1.0, + contact_wrench_stiffness_task=[0.0, 0.0, 0.1, 0.0, 0.0, 0.0], + motion_control_axes_task=[1, 1, 0, 1, 1, 1], + contact_wrench_control_axes_task=[0, 0, 1, 0, 0, 0], + ) + osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device) + + self._run_op_space_controller(robot, osc, "panda_leftfinger", ["panda_joint.*"], self.target_hybrid_set_tilted) + + """ + Helper functions + """ + + def _run_op_space_controller( + self, + robot: Articulation, + osc: OperationalSpaceController, + ee_frame_name: str, + arm_joint_names: list[str], + target_set: torch.tensor, + ): + """Run the operational space controller with the given parameters. + + Args: + robot (Articulation): The robot to control. + osc (OperationalSpaceController): The operational space controller. + ee_frame_name (str): The name of the end-effector frame. + arm_joint_names (list[str]): The names of the arm joints. + target_set (torch.tensor): The target set to track. + """ + # Initialize the masks for evaluating target convergence according to selection matrices + self.pos_mask = torch.tensor(osc.cfg.motion_control_axes_task[:3], device=self.sim.device).view(1, 3) + self.rot_mask = torch.tensor(osc.cfg.motion_control_axes_task[3:], device=self.sim.device).view(1, 3) + self.wrench_mask = torch.tensor(osc.cfg.contact_wrench_control_axes_task, device=self.sim.device).view(1, 6) + self.force_mask = self.wrench_mask[:, 0:3] # Take only the force components as we can measure only these + + # Define simulation stepping + sim_dt = self.sim.get_physics_dt() + # Play the simulator + self.sim.reset() + + # Obtain the frame index of the end-effector + ee_frame_idx = robot.find_bodies(ee_frame_name)[0][0] + # Obtain joint indices + arm_joint_ids = robot.find_joints(arm_joint_names)[0] + + # Update existing buffers + # Note: We need to update buffers before the first step for the controller. + robot.update(dt=sim_dt) + + # get the updated states + jacobian_b, mass_matrix, gravity, ee_pose_b, ee_vel_b, root_pose_w, ee_pose_w, ee_force_b = self._update_states( + robot, ee_frame_idx, arm_joint_ids + ) + + # Track the given target command + current_goal_idx = 0 # Current goal index for the arm + command = torch.zeros( + self.num_envs, osc.action_dim, device=self.sim.device + ) # Generic target command, which can be pose, position, force, etc. + ee_target_pose_b = torch.zeros(self.num_envs, 7, device=self.sim.device) # Target pose in the body frame + ee_target_pose_w = torch.zeros( + self.num_envs, 7, device=self.sim.device + ) # Target pose in the world frame (for marker) + + # Set joint efforts to zero + zero_joint_efforts = torch.zeros(self.num_envs, robot.num_joints, device=self.sim.device) + joint_efforts = torch.zeros(self.num_envs, len(arm_joint_ids), device=self.sim.device) + + # Now we are ready! + for count in range(1501): + # reset every 500 steps + if count % 500 == 0: + # check that we converged to the goal + if count > 0: + self._check_convergence(osc, ee_pose_b, ee_target_pose_b, ee_force_b, command) + # reset joint state to default + default_joint_pos = robot.data.default_joint_pos.clone() + default_joint_vel = robot.data.default_joint_vel.clone() + robot.write_joint_state_to_sim(default_joint_pos, default_joint_vel) + robot.set_joint_effort_target(zero_joint_efforts) # Set zero torques in the initial step + robot.write_data_to_sim() + robot.reset() + # reset contact sensor + if self.contact_forces is not None: + self.contact_forces.reset() + # reset target pose + robot.update(sim_dt) + _, _, _, ee_pose_b, _, _, _, _ = self._update_states( + robot, ee_frame_idx, arm_joint_ids + ) # at reset, the jacobians are not updated to the latest state + command, ee_target_pose_b, ee_target_pose_w, current_goal_idx = self._update_target( + osc, root_pose_w, ee_pose_b, target_set, current_goal_idx + ) + # set the osc command + osc.reset() + command, task_frame_pose_b = self._convert_to_task_frame( + osc, command=command, ee_target_pose_b=ee_target_pose_b + ) + osc.set_command( + command=command, current_ee_pose_b=ee_pose_b, current_task_frame_pose_b=task_frame_pose_b + ) + else: + # get the updated states + jacobian_b, mass_matrix, gravity, ee_pose_b, ee_vel_b, root_pose_w, ee_pose_w, ee_force_b = ( + self._update_states(robot, ee_frame_idx, arm_joint_ids) + ) + # compute the joint commands + joint_efforts = osc.compute( + jacobian_b=jacobian_b, + current_ee_pose_b=ee_pose_b, + current_ee_vel_b=ee_vel_b, + current_ee_force_b=ee_force_b, + mass_matrix=mass_matrix, + gravity=gravity, + ) + robot.set_joint_effort_target(joint_efforts, joint_ids=arm_joint_ids) + robot.write_data_to_sim() + + # update marker positions + self.ee_marker.visualize(ee_pose_w[:, 0:3], ee_pose_w[:, 3:7]) + self.goal_marker.visualize(ee_target_pose_w[:, 0:3], ee_target_pose_w[:, 3:7]) + + # perform step + self.sim.step(render=False) + # update buffers + robot.update(sim_dt) + + def _update_states( + self, + robot: Articulation, + ee_frame_idx: int, + arm_joint_ids: list[int], + ): + """Update the states of the robot and obtain the relevant quantities for the operational space controller. + + Args: + robot (Articulation): The robot to control. + ee_frame_idx (int): The index of the end-effector frame. + arm_joint_ids (list[int]): The indices of the arm joints. + + Returns: + jacobian_b (torch.tensor): The Jacobian in the root frame. + mass_matrix (torch.tensor): The mass matrix. + gravity (torch.tensor): The gravity vector. + ee_pose_b (torch.tensor): The end-effector pose in the root frame. + ee_vel_b (torch.tensor): The end-effector velocity in the root frame. + root_pose_w (torch.tensor): The root pose in the world frame. + ee_pose_w (torch.tensor): The end-effector pose in the world frame. + ee_force_b (torch.tensor): The end-effector force in the root frame. + """ + # obtain dynamics related quantities from simulation + ee_jacobi_idx = ee_frame_idx - 1 + jacobian_w = robot.root_physx_view.get_jacobians()[:, ee_jacobi_idx, :, arm_joint_ids] + mass_matrix = robot.root_physx_view.get_mass_matrices()[:, arm_joint_ids, :][:, :, arm_joint_ids] + gravity = robot.root_physx_view.get_generalized_gravity_forces()[:, arm_joint_ids] + # Convert the Jacobian from world to root frame + jacobian_b = jacobian_w.clone() + root_rot_matrix = matrix_from_quat(quat_inv(robot.data.root_link_quat_w)) + jacobian_b[:, :3, :] = torch.bmm(root_rot_matrix, jacobian_b[:, :3, :]) + jacobian_b[:, 3:, :] = torch.bmm(root_rot_matrix, jacobian_b[:, 3:, :]) + + # Compute current pose of the end-effector + root_pose_w = robot.data.root_link_state_w[:, 0:7] + ee_pose_w = robot.data.body_link_state_w[:, ee_frame_idx, 0:7] + ee_pos_b, ee_quat_b = subtract_frame_transforms( + root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_pose_w[:, 0:3], ee_pose_w[:, 3:7] + ) + ee_pose_b = torch.cat([ee_pos_b, ee_quat_b], dim=-1) + + # Compute the current velocity of the end-effector + ee_vel_w = robot.data.body_com_vel_w[:, ee_frame_idx, :] # Extract end-effector velocity in the world frame + root_vel_w = robot.data.root_com_vel_w # Extract root velocity in the world frame + relative_vel_w = ee_vel_w - root_vel_w # Compute the relative velocity in the world frame + ee_lin_vel_b = quat_rotate_inverse( + robot.data.root_link_quat_w, relative_vel_w[:, 0:3] + ) # From world to root frame + ee_ang_vel_b = quat_rotate_inverse(robot.data.root_link_quat_w, relative_vel_w[:, 3:6]) + ee_vel_b = torch.cat([ee_lin_vel_b, ee_ang_vel_b], dim=-1) + + # Calculate the contact force + ee_force_w = torch.zeros(self.num_envs, 3, device=self.sim.device) + if self.contact_forces is not None: # Only modify if it exist + sim_dt = self.sim.get_physics_dt() + self.contact_forces.update(sim_dt) # update contact sensor + # Calculate the contact force by averaging over last four time steps (i.e., to smoothen) and + # taking the max of three surfaces as only one should be the contact of interest + ee_force_w, _ = torch.max(torch.mean(self.contact_forces.data.net_forces_w_history, dim=1), dim=1) + + # This is a simplification, only for the sake of testing. + ee_force_b = ee_force_w + + return jacobian_b, mass_matrix, gravity, ee_pose_b, ee_vel_b, root_pose_w, ee_pose_w, ee_force_b + + def _update_target( + self, + osc: OperationalSpaceController, + root_pose_w: torch.tensor, + ee_pose_b: torch.tensor, + target_set: torch.tensor, + current_goal_idx: int, + ): + """Update the target for the operational space controller. + + Args: + osc (OperationalSpaceController): The operational space controller. + root_pose_w (torch.tensor): The root pose in the world frame. + ee_pose_b (torch.tensor): The end-effector pose in the body frame. + target_set (torch.tensor): The target set to track. + current_goal_idx (int): The current goal index. + + Returns: + command (torch.tensor): The target command. + ee_target_pose_b (torch.tensor): The end-effector target pose in the body frame. + ee_target_pose_w (torch.tensor): The end-effector target pose in the world frame. + next_goal_idx (int): The next goal index. + + Raises: + ValueError: If the target type is undefined. + """ + # update the ee desired command + command = torch.zeros(self.num_envs, osc.action_dim, device=self.sim.device) + command[:] = target_set[current_goal_idx] + + # update the ee desired pose + ee_target_pose_b = torch.zeros(self.num_envs, 7, device=self.sim.device) + for target_type in osc.cfg.target_types: + if target_type == "pose_abs": + ee_target_pose_b[:] = command[:, :7] + elif target_type == "pose_rel": + ee_target_pose_b[:, 0:3], ee_target_pose_b[:, 3:7] = apply_delta_pose( + ee_pose_b[:, :3], ee_pose_b[:, 3:], command[:, :7] + ) + elif target_type == "wrench_abs": + pass # ee_target_pose_b could stay at the root frame for force control, what matters is ee_target_b + else: + raise ValueError("Undefined target_type within _update_target().") + + # update the target desired pose in world frame (for marker) + ee_target_pos_w, ee_target_quat_w = combine_frame_transforms( + root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_target_pose_b[:, 0:3], ee_target_pose_b[:, 3:7] + ) + ee_target_pose_w = torch.cat([ee_target_pos_w, ee_target_quat_w], dim=-1) + + next_goal_idx = (current_goal_idx + 1) % len(target_set) + + return command, ee_target_pose_b, ee_target_pose_w, next_goal_idx + + def _convert_to_task_frame( + self, osc: OperationalSpaceController, command: torch.tensor, ee_target_pose_b: torch.tensor + ): + """Convert the target command to the task frame if required. + + Args: + osc (OperationalSpaceController): The operational space controller. + command (torch.tensor): The target command to convert. + ee_target_pose_b (torch.tensor): The end-effector target pose in the body frame. + + Returns: + command (torch.tensor): The converted target command. + task_frame_pose_b (torch.tensor): The task frame pose in the body frame. + + Raises: + ValueError: If the frame is invalid. + """ + command = command.clone() + task_frame_pose_b = None + if self.frame == "root": + # No need to transform anything if they are already in root frame + pass + elif self.frame == "task": + # Convert target commands from base to the task frame + command = command.clone() + task_frame_pose_b = ee_target_pose_b.clone() + + cmd_idx = 0 + for target_type in osc.cfg.target_types: + if target_type == "pose_abs": + command[:, :3], command[:, 3:7] = subtract_frame_transforms( + task_frame_pose_b[:, :3], task_frame_pose_b[:, 3:], command[:, :3], command[:, 3:7] + ) + cmd_idx += 7 + elif target_type == "pose_rel": + # Compute rotation matrices + R_task_b = matrix_from_quat(task_frame_pose_b[:, 3:]) # Task frame to base frame + R_b_task = R_task_b.mT # Base frame to task frame + # Transform the delta position and orientation from base to task frame + command[:, :3] = (R_b_task @ command[:, :3].unsqueeze(-1)).squeeze(-1) + command[:, 3:7] = (R_b_task @ command[:, 3:7].unsqueeze(-1)).squeeze(-1) + cmd_idx += 6 + elif target_type == "wrench_abs": + # These are already defined in target frame for ee_goal_wrench_set_tilted_task (since it is + # easier), so not transforming + cmd_idx += 6 + else: + raise ValueError("Undefined target_type within _convert_to_task_frame().") + else: + # Raise error for invalid frame + raise ValueError("Invalid frame selection for target setting inside the test_operational_space.") + + return command, task_frame_pose_b + + def _check_convergence( + self, + osc: OperationalSpaceController, + ee_pose_b: torch.tensor, + ee_target_pose_b: torch.tensor, + ee_force_b: torch.tensor, + ee_target_b: torch.tensor, + ): + """Check the convergence to the target. + + Args: + osc (OperationalSpaceController): The operational space controller. + ee_pose_b (torch.tensor): The end-effector pose in the body frame. + ee_target_pose_b (torch.tensor): The end-effector target pose in the body frame. + ee_force_b (torch.tensor): The end-effector force in the body frame. + ee_target_b (torch.tensor): The end-effector target in the body frame. + + Raises: + AssertionError: If the convergence is not achieved. + ValueError: If the target type is undefined. + """ + cmd_idx = 0 + for target_type in osc.cfg.target_types: + if target_type == "pose_abs": + pos_error, rot_error = compute_pose_error( + ee_pose_b[:, 0:3], ee_pose_b[:, 3:7], ee_target_pose_b[:, 0:3], ee_target_pose_b[:, 3:7] + ) + pos_error_norm = torch.norm(pos_error * self.pos_mask, dim=-1) + rot_error_norm = torch.norm(rot_error * self.rot_mask, dim=-1) + # desired error (zer) + des_error = torch.zeros_like(pos_error_norm) + # check convergence + torch.testing.assert_close(pos_error_norm, des_error, rtol=0.0, atol=0.1) + torch.testing.assert_close(rot_error_norm, des_error, rtol=0.0, atol=0.1) + cmd_idx += 7 + elif target_type == "pose_rel": + pos_error, rot_error = compute_pose_error( + ee_pose_b[:, 0:3], ee_pose_b[:, 3:7], ee_target_pose_b[:, 0:3], ee_target_pose_b[:, 3:7] + ) + pos_error_norm = torch.norm(pos_error * self.pos_mask, dim=-1) + rot_error_norm = torch.norm(rot_error * self.rot_mask, dim=-1) + # desired error (zer) + des_error = torch.zeros_like(pos_error_norm) + # check convergence + torch.testing.assert_close(pos_error_norm, des_error, rtol=0.0, atol=0.1) + torch.testing.assert_close(rot_error_norm, des_error, rtol=0.0, atol=0.1) + cmd_idx += 6 + elif target_type == "wrench_abs": + force_target_b = ee_target_b[:, cmd_idx : cmd_idx + 3].clone() + # Convert to base frame if the target was defined in task frame + if self.frame == "task": + task_frame_pose_b = ee_target_pose_b.clone() + R_task_b = matrix_from_quat(task_frame_pose_b[:, 3:]) + force_target_b[:] = (R_task_b @ force_target_b[:].unsqueeze(-1)).squeeze(-1) + force_error = ee_force_b - force_target_b + force_error_norm = torch.norm( + force_error * self.force_mask, dim=-1 + ) # ignore torque part as we cannot measure it + des_error = torch.zeros_like(force_error_norm) + # check convergence: big threshold here as the force control is not precise when the robot moves + torch.testing.assert_close(force_error_norm, des_error, rtol=0.0, atol=1.0) + cmd_idx += 6 + else: + raise ValueError("Undefined target_type within _check_convergence().") + + +if __name__ == "__main__": + run_tests() diff --git a/source/extensions/omni.isaac.lab/test/envs/check_manager_based_env_floating_cube.py b/source/extensions/omni.isaac.lab/test/envs/check_manager_based_env_floating_cube.py index 4282eb5e8c..812bb98f37 100644 --- a/source/extensions/omni.isaac.lab/test/envs/check_manager_based_env_floating_cube.py +++ b/source/extensions/omni.isaac.lab/test/envs/check_manager_based_env_floating_cube.py @@ -128,11 +128,11 @@ def process_actions(self, actions: torch.Tensor): def apply_actions(self): # implement a PD controller to track the target position - pos_error = self._processed_actions - (self._asset.data.root_pos_w - self._env.scene.env_origins) - vel_error = -self._asset.data.root_lin_vel_w + pos_error = self._processed_actions - (self._asset.data.root_link_pos_w - self._env.scene.env_origins) + vel_error = -self._asset.data.root_com_lin_vel_w # set velocity targets self._vel_command[:, :3] = self.p_gain * pos_error + self.d_gain * vel_error - self._asset.write_root_velocity_to_sim(self._vel_command) + self._asset.write_root_com_velocity_to_sim(self._vel_command) @configclass @@ -151,7 +151,7 @@ def base_position(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg) -> torch.Tens """Root linear velocity in the asset's root frame.""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return asset.data.root_pos_w - env.scene.env_origins + return asset.data.root_link_pos_w - env.scene.env_origins ## diff --git a/source/extensions/omni.isaac.lab/test/envs/test_action_state_recorder_term.py b/source/extensions/omni.isaac.lab/test/envs/test_action_state_recorder_term.py new file mode 100644 index 0000000000..d3640f63ba --- /dev/null +++ b/source/extensions/omni.isaac.lab/test/envs/test_action_state_recorder_term.py @@ -0,0 +1,123 @@ +# Copyright (c) 2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Launch Isaac Sim Simulator first.""" + +from omni.isaac.lab.app import AppLauncher, run_tests + +# launch the simulator +app_launcher = AppLauncher(headless=True) +simulation_app = app_launcher.app + + +"""Rest everything follows.""" + +import gymnasium as gym +import shutil +import tempfile +import torch +import unittest +import uuid + +import carb +import omni.usd + +from omni.isaac.lab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg + +import omni.isaac.lab_tasks # noqa: F401 +from omni.isaac.lab_tasks.utils.parse_cfg import parse_env_cfg + + +class TestActionStateRecorderManagerCfg(unittest.TestCase): + """Test cases for ActionStateRecorderManagerCfg recorder terms.""" + + @classmethod + def setUpClass(cls): + # this flag is necessary to prevent a bug where the simulation gets stuck randomly when running the + # test on many environments. + carb_settings_iface = carb.settings.get_settings() + carb_settings_iface.set_bool("/physics/cooking/ujitsoCollisionCooking", False) + + def setUp(self): + # create a temporary directory to store the test datasets + self.temp_dir = tempfile.mkdtemp() + + def tearDown(self): + # delete the temporary directory after the test + shutil.rmtree(self.temp_dir) + + def test_action_state_reocrder_terms(self): + """Check action state recorder terms.""" + for task_name in ["Isaac-Lift-Cube-Franka-v0"]: + for device in ["cuda:0", "cpu"]: + for num_envs in [1, 2]: + with self.subTest(task_name=task_name, device=device): + omni.usd.get_context().new_stage() + + dummy_dataset_filename = f"{uuid.uuid4()}.hdf5" + + # parse configuration + env_cfg = parse_env_cfg(task_name, device=device, num_envs=num_envs) + # set recorder configurations for this test + env_cfg.recorders: ActionStateRecorderManagerCfg = ActionStateRecorderManagerCfg() + env_cfg.recorders.dataset_export_dir_path = self.temp_dir + env_cfg.recorders.dataset_filename = dummy_dataset_filename + + # create environment + env = gym.make(task_name, cfg=env_cfg) + + # reset all environment instances to trigger post-reset recorder callbacks + env.reset() + self.check_initial_state_recorder_term(env) + + # reset only one environment that is not the first one + env.unwrapped.reset(env_ids=torch.tensor([num_envs - 1], device=env.unwrapped.device)) + self.check_initial_state_recorder_term(env) + + # close the environment + env.close() + + def check_initial_state_recorder_term(self, env): + """Check values recorded by the initial state recorder terms. + + Args: + env: Environment instance. + """ + current_state = env.unwrapped.scene.get_state(is_relative=True) + for env_id in range(env.unwrapped.num_envs): + recorded_initial_state = env.unwrapped.recorder_manager.get_episode(env_id).get_initial_state() + are_states_equal, output_log = self.compare_states(recorded_initial_state, current_state, env_id) + self.assertTrue(are_states_equal, msg=output_log) + + def compare_states(self, compared_state, ground_truth_state, ground_truth_env_id) -> (bool, str): + """Compare a state with the given ground_truth. + + Args: + compared_state: State to be compared. + ground_truth_state: Ground truth state. + ground_truth_env_id: Index of the environment in the ground_truth states to be compared. + + Returns: + bool: True if states match, False otherwise. + str: Error log if states don't match. + """ + for asset_type in ["articulation", "rigid_object"]: + for asset_name in ground_truth_state[asset_type].keys(): + for state_name in ground_truth_state[asset_type][asset_name].keys(): + runtime_asset_state = ground_truth_state[asset_type][asset_name][state_name][ground_truth_env_id] + dataset_asset_state = compared_state[asset_type][asset_name][state_name][0] + if len(dataset_asset_state) != len(runtime_asset_state): + return False, f"State shape of {state_name} for asset {asset_name} don't match" + for i in range(len(dataset_asset_state)): + if abs(dataset_asset_state[i] - runtime_asset_state[i]) > 0.01: + return ( + False, + f'State ["{asset_type}"]["{asset_name}"]["{state_name}"][{i}] don\'t match\r\n', + ) + return True, "" + + +if __name__ == "__main__": + run_tests() diff --git a/source/extensions/omni.isaac.lab/test/envs/test_direct_marl_env.py b/source/extensions/omni.isaac.lab/test/envs/test_direct_marl_env.py index ec6cb971c2..eba2895e30 100644 --- a/source/extensions/omni.isaac.lab/test/envs/test_direct_marl_env.py +++ b/source/extensions/omni.isaac.lab/test/envs/test_direct_marl_env.py @@ -21,7 +21,6 @@ """Rest everything follows.""" -import torch import unittest import omni.usd @@ -50,9 +49,10 @@ class EmptyEnvCfg(DirectMARLEnvCfg): # Basic settings decimation = 1 possible_agents = ["agent_0", "agent_1"] - num_actions = {"agent_0": 1, "agent_1": 2} - num_observations = {"agent_0": 3, "agent_1": 4} - num_states = -1 + action_spaces = {"agent_0": 1, "agent_1": 2} + observation_spaces = {"agent_0": 3, "agent_1": 4} + state_space = -1 + episode_length_s = 100.0 return EmptyEnvCfg() @@ -69,8 +69,17 @@ def test_initialization(self): with self.subTest(device=device): # create a new stage omni.usd.get_context().new_stage() - # create environment - env = DirectMARLEnv(cfg=get_empty_base_env_cfg(device=device)) + try: + # create environment + env = DirectMARLEnv(cfg=get_empty_base_env_cfg(device=device)) + except Exception as e: + if "env" in locals() and hasattr(env, "_is_closed"): + env.close() + else: + if hasattr(e, "obj") and hasattr(e.obj, "_is_closed"): + e.obj.close() + self.fail(f"Failed to set-up the DirectMARLEnv environment. Error: {e}") + # check multi-agent config self.assertEqual(env.num_agents, 2) self.assertEqual(env.max_num_agents, 2) @@ -78,12 +87,6 @@ def test_initialization(self): self.assertEqual(env.state_space.shape, (7,)) self.assertEqual(len(env.observation_spaces), 2) self.assertEqual(len(env.action_spaces), 2) - # step environment to verify setup - env.reset() - for _ in range(2): - actions = {"agent_0": torch.rand((1, 1)), "agent_1": torch.rand((1, 2))} - obs, reward, terminated, truncate, info = env.step(actions) - env.state() # close the environment env.close() diff --git a/source/extensions/omni.isaac.lab/test/envs/test_env_rendering_logic.py b/source/extensions/omni.isaac.lab/test/envs/test_env_rendering_logic.py index 24af5c91d6..378a089750 100644 --- a/source/extensions/omni.isaac.lab/test/envs/test_env_rendering_logic.py +++ b/source/extensions/omni.isaac.lab/test/envs/test_env_rendering_logic.py @@ -47,6 +47,7 @@ class EnvCfg(ManagerBasedEnvCfg): """Configuration for the test environment.""" decimation: int = 4 + episode_length_s: float = 100.0 sim: SimulationCfg = SimulationCfg(dt=0.005, render_interval=render_interval) scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=1, env_spacing=1.0) actions: EmptyManagerCfg = EmptyManagerCfg() @@ -63,10 +64,13 @@ class EnvCfg(ManagerBasedRLEnvCfg): """Configuration for the test environment.""" decimation: int = 4 + episode_length_s: float = 100.0 sim: SimulationCfg = SimulationCfg(dt=0.005, render_interval=render_interval) scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=1, env_spacing=1.0) actions: EmptyManagerCfg = EmptyManagerCfg() observations: EmptyManagerCfg = EmptyManagerCfg() + rewards: EmptyManagerCfg = EmptyManagerCfg() + terminations: EmptyManagerCfg = EmptyManagerCfg() return ManagerBasedRLEnv(cfg=EnvCfg()) @@ -79,8 +83,9 @@ class EnvCfg(DirectRLEnvCfg): """Configuration for the test environment.""" decimation: int = 4 - num_actions: int = 0 - num_observations: int = 0 + action_space: int = 0 + observation_space: int = 0 + episode_length_s: float = 100.0 sim: SimulationCfg = SimulationCfg(dt=0.005, render_interval=render_interval) scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=1, env_spacing=1.0) @@ -131,14 +136,21 @@ def test_env_rendering_logic(self): # create a new stage omni.usd.get_context().new_stage() - - # create environment - if env_type == "manager_based_env": - env = create_manager_based_env(render_interval) - elif env_type == "manager_based_rl_env": - env = create_manager_based_rl_env(render_interval) - else: - env = create_direct_rl_env(render_interval) + try: + # create environment + if env_type == "manager_based_env": + env = create_manager_based_env(render_interval) + elif env_type == "manager_based_rl_env": + env = create_manager_based_rl_env(render_interval) + else: + env = create_direct_rl_env(render_interval) + except Exception as e: + if "env" in locals() and hasattr(env, "_is_closed"): + env.close() + else: + if hasattr(e, "obj") and hasattr(e.obj, "_is_closed"): + e.obj.close() + self.fail(f"Failed to set-up the environment {env_type}. Error: {e}") # enable the flag to render the environment # note: this is only done for the unit testing to "fake" camera rendering. diff --git a/source/extensions/omni.isaac.lab/test/envs/test_manager_based_rl_env_ui.py b/source/extensions/omni.isaac.lab/test/envs/test_manager_based_rl_env_ui.py new file mode 100644 index 0000000000..f61af77b8d --- /dev/null +++ b/source/extensions/omni.isaac.lab/test/envs/test_manager_based_rl_env_ui.py @@ -0,0 +1,104 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# ignore private usage of variables warning +# pyright: reportPrivateUsage=none + +from __future__ import annotations + +"""Launch Isaac Sim Simulator first.""" + +from omni.isaac.lab.app import AppLauncher, run_tests + +# Can set this to False to see the GUI for debugging +HEADLESS = True + +# launch omniverse app +app_launcher = AppLauncher(headless=HEADLESS, enable_cameras=True) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import unittest + +import carb +import omni.usd +from omni.isaac.core.utils.extensions import enable_extension + +from omni.isaac.lab.envs import ManagerBasedRLEnv, ManagerBasedRLEnvCfg +from omni.isaac.lab.envs.ui import ManagerBasedRLEnvWindow +from omni.isaac.lab.scene import InteractiveSceneCfg +from omni.isaac.lab.utils import configclass + +enable_extension("omni.isaac.ui") + + +@configclass +class EmptyManagerCfg: + """Empty manager specifications for the environment.""" + + pass + + +@configclass +class EmptySceneCfg(InteractiveSceneCfg): + """Configuration for an empty scene.""" + + pass + + +def get_empty_base_env_cfg(device: str = "cuda:0", num_envs: int = 1, env_spacing: float = 1.0): + """Generate base environment config based on device""" + + @configclass + class EmptyEnvCfg(ManagerBasedRLEnvCfg): + """Configuration for the empty test environment.""" + + # Scene settings + scene: EmptySceneCfg = EmptySceneCfg(num_envs=num_envs, env_spacing=env_spacing) + # Basic settings + actions: EmptyManagerCfg = EmptyManagerCfg() + observations: EmptyManagerCfg = EmptyManagerCfg() + rewards: EmptyManagerCfg = EmptyManagerCfg() + terminations: EmptyManagerCfg = EmptyManagerCfg() + # Define window + ui_window_class_type: ManagerBasedRLEnvWindow = ManagerBasedRLEnvWindow + + def __post_init__(self): + """Post initialization.""" + # step settings + self.decimation = 4 # env step every 4 sim steps: 200Hz / 4 = 50Hz + # simulation settings + self.sim.dt = 0.005 # sim step every 5ms: 200Hz + self.sim.render_interval = self.decimation # render every 4 sim steps + # pass device down from test + self.sim.device = device + # episode length + self.episode_length_s = 5.0 + + return EmptyEnvCfg() + + +class TestManagerBasedRLEnvUI(unittest.TestCase): + """Test for manager-based RL env class UI""" + + """ + Tests + """ + + def test_ui_window(self): + device = "cuda:0" + # override sim setting to enable UI + carb.settings.get_settings().set_bool("/app/window/enabled", True) + # create a new stage + omni.usd.get_context().new_stage() + # create environment + env = ManagerBasedRLEnv(cfg=get_empty_base_env_cfg(device=device)) + # close the environment + env.close() + + +if __name__ == "__main__": + run_tests() diff --git a/source/extensions/omni.isaac.lab/test/envs/test_spaces_utils.py b/source/extensions/omni.isaac.lab/test/envs/test_spaces_utils.py new file mode 100644 index 0000000000..791c61380c --- /dev/null +++ b/source/extensions/omni.isaac.lab/test/envs/test_spaces_utils.py @@ -0,0 +1,175 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# ignore private usage of variables warning +# pyright: reportPrivateUsage=none + +from __future__ import annotations + +"""Launch Isaac Sim Simulator first.""" + +from omni.isaac.lab.app import AppLauncher, run_tests + +# Can set this to False to see the GUI for debugging +HEADLESS = True + +# launch omniverse app +app_launcher = AppLauncher(headless=HEADLESS) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import numpy as np +import torch +import unittest +from gymnasium.spaces import Box, Dict, Discrete, MultiDiscrete, Tuple + +from omni.isaac.lab.envs.utils.spaces import deserialize_space, sample_space, serialize_space, spec_to_gym_space + + +class TestSpacesUtils(unittest.TestCase): + """Test for spaces utils' functions""" + + """ + Tests + """ + + def test_spec_to_gym_space(self): + # fundamental spaces + # Box + space = spec_to_gym_space(1) + self.assertIsInstance(space, Box) + self.assertEqual(space.shape, (1,)) + space = spec_to_gym_space([1, 2, 3, 4, 5]) + self.assertIsInstance(space, Box) + self.assertEqual(space.shape, (1, 2, 3, 4, 5)) + space = spec_to_gym_space(Box(low=-1.0, high=1.0, shape=(1, 2))) + self.assertIsInstance(space, Box) + # Discrete + space = spec_to_gym_space({2}) + self.assertIsInstance(space, Discrete) + self.assertEqual(space.n, 2) + space = spec_to_gym_space(Discrete(2)) + self.assertIsInstance(space, Discrete) + # MultiDiscrete + space = spec_to_gym_space([{1}, {2}, {3}]) + self.assertIsInstance(space, MultiDiscrete) + self.assertEqual(space.nvec.shape, (3,)) + space = spec_to_gym_space(MultiDiscrete(np.array([1, 2, 3]))) + self.assertIsInstance(space, MultiDiscrete) + # composite spaces + # Tuple + space = spec_to_gym_space(([1, 2, 3, 4, 5], {2}, [{1}, {2}, {3}])) + self.assertIsInstance(space, Tuple) + self.assertEqual(len(space), 3) + self.assertIsInstance(space[0], Box) + self.assertIsInstance(space[1], Discrete) + self.assertIsInstance(space[2], MultiDiscrete) + space = spec_to_gym_space(Tuple((Box(-1, 1, shape=(1,)), Discrete(2)))) + self.assertIsInstance(space, Tuple) + # Dict + space = spec_to_gym_space({"box": [1, 2, 3, 4, 5], "discrete": {2}, "multi_discrete": [{1}, {2}, {3}]}) + self.assertIsInstance(space, Dict) + self.assertEqual(len(space), 3) + self.assertIsInstance(space["box"], Box) + self.assertIsInstance(space["discrete"], Discrete) + self.assertIsInstance(space["multi_discrete"], MultiDiscrete) + space = spec_to_gym_space(Dict({"box": Box(-1, 1, shape=(1,)), "discrete": Discrete(2)})) + self.assertIsInstance(space, Dict) + + def test_sample_space(self): + device = "cpu" + # fundamental spaces + # Box + sample = sample_space(Box(low=-1.0, high=1.0, shape=(1, 2)), device, batch_size=1) + self.assertIsInstance(sample, torch.Tensor) + self._check_tensorized(sample, batch_size=1) + # Discrete + sample = sample_space(Discrete(2), device, batch_size=2) + self.assertIsInstance(sample, torch.Tensor) + self._check_tensorized(sample, batch_size=2) + # MultiDiscrete + sample = sample_space(MultiDiscrete(np.array([1, 2, 3])), device, batch_size=3) + self.assertIsInstance(sample, torch.Tensor) + self._check_tensorized(sample, batch_size=3) + # composite spaces + # Tuple + sample = sample_space(Tuple((Box(-1, 1, shape=(1,)), Discrete(2))), device, batch_size=4) + self.assertIsInstance(sample, (tuple, list)) + self._check_tensorized(sample, batch_size=4) + # Dict + sample = sample_space(Dict({"box": Box(-1, 1, shape=(1,)), "discrete": Discrete(2)}), device, batch_size=5) + self.assertIsInstance(sample, dict) + self._check_tensorized(sample, batch_size=5) + + def test_space_serialization_deserialization(self): + # fundamental spaces + # Box + space = 1 + output = deserialize_space(serialize_space(space)) + self.assertEqual(space, output) + space = [1, 2, 3, 4, 5] + output = deserialize_space(serialize_space(space)) + self.assertEqual(space, output) + space = Box(low=-1.0, high=1.0, shape=(1, 2)) + output = deserialize_space(serialize_space(space)) + self.assertIsInstance(output, Box) + self.assertTrue((space.low == output.low).all()) + self.assertTrue((space.high == output.high).all()) + self.assertEqual(space.shape, output.shape) + # Discrete + space = {2} + output = deserialize_space(serialize_space(space)) + self.assertEqual(space, output) + space = Discrete(2) + output = deserialize_space(serialize_space(space)) + self.assertIsInstance(output, Discrete) + self.assertEqual(space.n, output.n) + # MultiDiscrete + space = [{1}, {2}, {3}] + output = deserialize_space(serialize_space(space)) + self.assertEqual(space, output) + space = MultiDiscrete(np.array([1, 2, 3])) + output = deserialize_space(serialize_space(space)) + self.assertIsInstance(output, MultiDiscrete) + self.assertTrue((space.nvec == output.nvec).all()) + # composite spaces + # Tuple + space = ([1, 2, 3, 4, 5], {2}, [{1}, {2}, {3}]) + output = deserialize_space(serialize_space(space)) + self.assertEqual(space, output) + space = Tuple((Box(-1, 1, shape=(1,)), Discrete(2))) + output = deserialize_space(serialize_space(space)) + self.assertIsInstance(output, Tuple) + self.assertEqual(len(output), 2) + self.assertIsInstance(output[0], Box) + self.assertIsInstance(output[1], Discrete) + # Dict + space = {"box": [1, 2, 3, 4, 5], "discrete": {2}, "multi_discrete": [{1}, {2}, {3}]} + output = deserialize_space(serialize_space(space)) + self.assertEqual(space, output) + space = Dict({"box": Box(-1, 1, shape=(1,)), "discrete": Discrete(2)}) + output = deserialize_space(serialize_space(space)) + self.assertIsInstance(output, Dict) + self.assertEqual(len(output), 2) + self.assertIsInstance(output["box"], Box) + self.assertIsInstance(output["discrete"], Discrete) + + """ + Helper functions. + """ + + def _check_tensorized(self, sample, batch_size): + if isinstance(sample, (tuple, list)): + list(map(self._check_tensorized, sample, [batch_size] * len(sample))) + elif isinstance(sample, dict): + list(map(self._check_tensorized, sample.values(), [batch_size] * len(sample))) + else: + self.assertIsInstance(sample, torch.Tensor) + self.assertEqual(sample.shape[0], batch_size) + + +if __name__ == "__main__": + run_tests() diff --git a/source/extensions/omni.isaac.lab/test/managers/test_event_manager.py b/source/extensions/omni.isaac.lab/test/managers/test_event_manager.py index 89f43e7060..5a92b7c28d 100644 --- a/source/extensions/omni.isaac.lab/test/managers/test_event_manager.py +++ b/source/extensions/omni.isaac.lab/test/managers/test_event_manager.py @@ -134,6 +134,15 @@ def test_active_terms(self): self.assertEqual(len(self.event_man.active_terms["reset"]), 1) self.assertEqual(len(self.event_man.active_terms["custom"]), 2) + def test_config_empty(self): + """Test the creation of reward manager with empty config.""" + self.event_man = EventManager(None, self.env) + self.assertEqual(len(self.event_man.active_terms), 0) + + # print the expected string + print() + print(self.event_man) + def test_invalid_event_func_module(self): """Test the handling of invalid event function's module in string representation.""" cfg = { diff --git a/source/extensions/omni.isaac.lab/test/managers/test_observation_manager.py b/source/extensions/omni.isaac.lab/test/managers/test_observation_manager.py index 9b73e2d44a..253fc39228 100644 --- a/source/extensions/omni.isaac.lab/test/managers/test_observation_manager.py +++ b/source/extensions/omni.isaac.lab/test/managers/test_observation_manager.py @@ -131,8 +131,51 @@ class SampleGroupCfg(ObservationGroupCfg): self.obs_man = ObservationManager(cfg, self.env) self.assertEqual(len(self.obs_man.active_terms["policy"]), 5) # print the expected string + obs_man_str = str(self.obs_man) print() - print(self.obs_man) + print(obs_man_str) + obs_man_str_split = obs_man_str.split("|") + term_1_str_index = obs_man_str_split.index(" term_1 ") + term_1_str_shape = obs_man_str_split[term_1_str_index + 1].strip() + self.assertEqual(term_1_str_shape, "(4,)") + + def test_str_with_history(self): + """Test the string representation of the observation manager with history terms.""" + + TERM_1_HISTORY = 5 + + @configclass + class MyObservationManagerCfg: + """Test config class for observation manager.""" + + @configclass + class SampleGroupCfg(ObservationGroupCfg): + """Test config class for policy observation group.""" + + term_1 = ObservationTermCfg(func="__main__:grilled_chicken", scale=10, history_length=TERM_1_HISTORY) + term_2 = ObservationTermCfg(func=grilled_chicken, scale=2) + term_3 = ObservationTermCfg(func=grilled_chicken_with_bbq, scale=5, params={"bbq": True}) + term_4 = ObservationTermCfg( + func=grilled_chicken_with_yoghurt, scale=1.0, params={"hot": False, "bland": 2.0} + ) + term_5 = ObservationTermCfg( + func=grilled_chicken_with_yoghurt_and_bbq, scale=1.0, params={"hot": False, "bland": 2.0} + ) + + policy: ObservationGroupCfg = SampleGroupCfg() + + # create observation manager + cfg = MyObservationManagerCfg() + self.obs_man = ObservationManager(cfg, self.env) + self.assertEqual(len(self.obs_man.active_terms["policy"]), 5) + # print the expected string + obs_man_str = str(self.obs_man) + print() + print(obs_man_str) + obs_man_str_split = obs_man_str.split("|") + term_1_str_index = obs_man_str_split.index(" term_1 ") + term_1_str_shape = obs_man_str_split[term_1_str_index + 1].strip() + self.assertEqual(term_1_str_shape, "(20,)") def test_config_equivalence(self): """Test the equivalence of observation manager created from different config types.""" @@ -244,6 +287,8 @@ class SampleImageGroupCfg(ObservationGroupCfg): def test_compute(self): """Test the observation computation.""" + pos_scale_tuple = (2.0, 3.0, 1.0) + @configclass class MyObservationManagerCfg: """Test config class for observation manager.""" @@ -254,14 +299,14 @@ class PolicyCfg(ObservationGroupCfg): term_1 = ObservationTermCfg(func=grilled_chicken, scale=10) term_2 = ObservationTermCfg(func=grilled_chicken_with_curry, scale=0.0, params={"hot": False}) - term_3 = ObservationTermCfg(func=pos_w_data, scale=2.0) + term_3 = ObservationTermCfg(func=pos_w_data, scale=pos_scale_tuple) term_4 = ObservationTermCfg(func=lin_vel_w_data, scale=1.5) @configclass class CriticCfg(ObservationGroupCfg): - term_1 = ObservationTermCfg(func=pos_w_data, scale=2.0) + term_1 = ObservationTermCfg(func=pos_w_data, scale=pos_scale_tuple) term_2 = ObservationTermCfg(func=lin_vel_w_data, scale=1.5) - term_3 = ObservationTermCfg(func=pos_w_data, scale=2.0) + term_3 = ObservationTermCfg(func=pos_w_data, scale=pos_scale_tuple) term_4 = ObservationTermCfg(func=lin_vel_w_data, scale=1.5) @configclass @@ -289,6 +334,11 @@ class ImageCfg(ObservationGroupCfg): self.assertEqual((self.env.num_envs, 11), obs_policy.shape) self.assertEqual((self.env.num_envs, 12), obs_critic.shape) self.assertEqual((self.env.num_envs, 128, 256, 4), obs_image.shape) + # check that the scales are applied correctly + torch.testing.assert_close( + self.env.data.pos_w * torch.tensor(pos_scale_tuple, device=self.env.device), obs_critic[:, :3] + ) + torch.testing.assert_close(self.env.data.lin_vel_w * 1.5, obs_critic[:, 3:6]) # make sure that the data are the same for same terms # -- within group torch.testing.assert_close(obs_critic[:, 0:3], obs_critic[:, 6:9]) @@ -297,6 +347,157 @@ class ImageCfg(ObservationGroupCfg): torch.testing.assert_close(obs_policy[:, 5:8], obs_critic[:, 0:3]) torch.testing.assert_close(obs_policy[:, 8:11], obs_critic[:, 3:6]) + def test_compute_with_history(self): + """Test the observation computation with history buffers.""" + HISTORY_LENGTH = 5 + + @configclass + class MyObservationManagerCfg: + """Test config class for observation manager.""" + + @configclass + class PolicyCfg(ObservationGroupCfg): + """Test config class for policy observation group.""" + + term_1 = ObservationTermCfg(func=grilled_chicken, history_length=HISTORY_LENGTH) + # total observation size: term_dim (4) * history_len (5) = 20 + term_2 = ObservationTermCfg(func=lin_vel_w_data) + # total observation size: term_dim (3) = 3 + + policy: ObservationGroupCfg = PolicyCfg() + + # create observation manager + cfg = MyObservationManagerCfg() + self.obs_man = ObservationManager(cfg, self.env) + # compute observation using manager + observations = self.obs_man.compute() + # obtain the group observations + obs_policy: torch.Tensor = observations["policy"] + # check the observation shape + self.assertEqual((self.env.num_envs, 23), obs_policy.shape) + # check the observation data + expected_obs_term_1_data = torch.ones(self.env.num_envs, 4 * HISTORY_LENGTH, device=self.env.device) + expected_obs_term_2_data = lin_vel_w_data(self.env) + expected_obs_data_t0 = torch.concat((expected_obs_term_1_data, expected_obs_term_2_data), dim=-1) + print(expected_obs_data_t0, obs_policy) + self.assertTrue(torch.equal(expected_obs_data_t0, obs_policy)) + # test that the history buffer holds previous data + for _ in range(HISTORY_LENGTH): + observations = self.obs_man.compute() + obs_policy = observations["policy"] + expected_obs_term_1_data = torch.ones(self.env.num_envs, 4 * HISTORY_LENGTH, device=self.env.device) + expected_obs_data_t5 = torch.concat((expected_obs_term_1_data, expected_obs_term_2_data), dim=-1) + self.assertTrue(torch.equal(expected_obs_data_t5, obs_policy)) + # test reset + self.obs_man.reset() + observations = self.obs_man.compute() + obs_policy = observations["policy"] + self.assertTrue(torch.equal(expected_obs_data_t0, obs_policy)) + # test reset of specific env ids + reset_env_ids = [2, 4, 16] + self.obs_man.reset(reset_env_ids) + self.assertTrue(torch.equal(expected_obs_data_t0[reset_env_ids], obs_policy[reset_env_ids])) + + def test_compute_with_2d_history(self): + """Test the observation computation with history buffers for 2D observations.""" + HISTORY_LENGTH = 5 + + @configclass + class MyObservationManagerCfg: + """Test config class for observation manager.""" + + @configclass + class FlattenedPolicyCfg(ObservationGroupCfg): + """Test config class for policy observation group.""" + + term_1 = ObservationTermCfg( + func=grilled_chicken_image, params={"bland": 1.0, "channel": 1}, history_length=HISTORY_LENGTH + ) + # total observation size: term_dim (128, 256) * history_len (5) = 163840 + + @configclass + class PolicyCfg(ObservationGroupCfg): + """Test config class for policy observation group.""" + + term_1 = ObservationTermCfg( + func=grilled_chicken_image, + params={"bland": 1.0, "channel": 1}, + history_length=HISTORY_LENGTH, + flatten_history_dim=False, + ) + # total observation size: (5, 128, 256, 1) + + flat_obs_policy: ObservationGroupCfg = FlattenedPolicyCfg() + policy: ObservationGroupCfg = PolicyCfg() + + # create observation manager + cfg = MyObservationManagerCfg() + self.obs_man = ObservationManager(cfg, self.env) + # compute observation using manager + observations = self.obs_man.compute() + # obtain the group observations + obs_policy_flat: torch.Tensor = observations["flat_obs_policy"] + obs_policy: torch.Tensor = observations["policy"] + # check the observation shapes + self.assertEqual((self.env.num_envs, 163840), obs_policy_flat.shape) + self.assertEqual((self.env.num_envs, HISTORY_LENGTH, 128, 256, 1), obs_policy.shape) + + def test_compute_with_group_history(self): + """Test the observation computation with group level history buffer configuration.""" + TERM_HISTORY_LENGTH = 5 + GROUP_HISTORY_LENGTH = 10 + + @configclass + class MyObservationManagerCfg: + """Test config class for observation manager.""" + + @configclass + class PolicyCfg(ObservationGroupCfg): + """Test config class for policy observation group.""" + + history_length = GROUP_HISTORY_LENGTH + # group level history length will override all terms + term_1 = ObservationTermCfg(func=grilled_chicken, history_length=TERM_HISTORY_LENGTH) + # total observation size: term_dim (4) * history_len (5) = 20 + # with override total obs size: term_dim (4) * history_len (10) = 40 + term_2 = ObservationTermCfg(func=lin_vel_w_data) + # total observation size: term_dim (3) = 3 + # with override total obs size: term_dim (3) * history_len (10) = 30 + + policy: ObservationGroupCfg = PolicyCfg() + + # create observation manager + cfg = MyObservationManagerCfg() + self.obs_man = ObservationManager(cfg, self.env) + # compute observation using manager + observations = self.obs_man.compute() + # obtain the group observations + obs_policy: torch.Tensor = observations["policy"] + # check the total observation shape + self.assertEqual((self.env.num_envs, 70), obs_policy.shape) + # check the observation data is initialized properly + expected_obs_term_1_data = torch.ones(self.env.num_envs, 4 * GROUP_HISTORY_LENGTH, device=self.env.device) + expected_obs_term_2_data = lin_vel_w_data(self.env).repeat(1, GROUP_HISTORY_LENGTH) + expected_obs_data_t0 = torch.concat((expected_obs_term_1_data, expected_obs_term_2_data), dim=-1) + self.assertTrue(torch.equal(expected_obs_data_t0, obs_policy)) + # test that the history buffer holds previous data + for _ in range(GROUP_HISTORY_LENGTH): + observations = self.obs_man.compute() + obs_policy = observations["policy"] + expected_obs_term_1_data = torch.ones(self.env.num_envs, 4 * GROUP_HISTORY_LENGTH, device=self.env.device) + expected_obs_term_2_data = lin_vel_w_data(self.env).repeat(1, GROUP_HISTORY_LENGTH) + expected_obs_data_t10 = torch.concat((expected_obs_term_1_data, expected_obs_term_2_data), dim=-1) + self.assertTrue(torch.equal(expected_obs_data_t10, obs_policy)) + # test reset + self.obs_man.reset() + observations = self.obs_man.compute() + obs_policy = observations["policy"] + self.assertTrue(torch.equal(expected_obs_data_t0, obs_policy)) + # test reset of specific env ids + reset_env_ids = [2, 4, 16] + self.obs_man.reset(reset_env_ids) + self.assertTrue(torch.equal(expected_obs_data_t0[reset_env_ids], obs_policy[reset_env_ids])) + def test_invalid_observation_config(self): """Test the invalid observation config.""" diff --git a/source/extensions/omni.isaac.lab/test/managers/test_recorder_manager.py b/source/extensions/omni.isaac.lab/test/managers/test_recorder_manager.py new file mode 100644 index 0000000000..594482d6da --- /dev/null +++ b/source/extensions/omni.isaac.lab/test/managers/test_recorder_manager.py @@ -0,0 +1,168 @@ +# Copyright (c) 2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# needed to import for allowing type-hinting: torch.Tensor | None +from __future__ import annotations + +"""Launch Isaac Sim Simulator first.""" + +from omni.isaac.lab.app import AppLauncher, run_tests + +# launch omniverse app +simulation_app = AppLauncher(headless=True).app + +"""Rest everything follows.""" + +import os +import shutil +import tempfile +import torch +import unittest +import uuid +from collections import namedtuple +from collections.abc import Sequence + +from omni.isaac.lab.envs import ManagerBasedEnv +from omni.isaac.lab.managers import ( + DatasetExportMode, + RecorderManager, + RecorderManagerBaseCfg, + RecorderTerm, + RecorderTermCfg, +) +from omni.isaac.lab.utils import configclass + + +class DummyResetRecorderTerm(RecorderTerm): + """Dummy recorder term that records dummy data.""" + + def __init__(self, cfg: RecorderTermCfg, env: ManagerBasedEnv) -> None: + super().__init__(cfg, env) + + def record_pre_reset(self, env_ids: Sequence[int] | None) -> tuple[str | None, torch.Tensor | None]: + return "record_pre_reset", torch.ones(self._env.num_envs, 2, device=self._env.device) + + def record_post_reset(self, env_ids: Sequence[int] | None) -> tuple[str | None, torch.Tensor | None]: + return "record_post_reset", torch.ones(self._env.num_envs, 3, device=self._env.device) + + +class DummyStepRecorderTerm(RecorderTerm): + """Dummy recorder term that records dummy data.""" + + def __init__(self, cfg: RecorderTermCfg, env: ManagerBasedEnv) -> None: + super().__init__(cfg, env) + + def record_pre_step(self) -> tuple[str | None, torch.Tensor | None]: + return "record_pre_step", torch.ones(self._env.num_envs, 4, device=self._env.device) + + def record_post_step(self) -> tuple[str | None, torch.Tensor | None]: + return "record_post_step", torch.ones(self._env.num_envs, 5, device=self._env.device) + + +@configclass +class DummyRecorderManagerCfg(RecorderManagerBaseCfg): + """Dummy recorder configurations.""" + + @configclass + class DummyResetRecorderTermCfg(RecorderTermCfg): + """Configuration for the dummy reset recorder term.""" + + class_type: type[RecorderTerm] = DummyResetRecorderTerm + + @configclass + class DummyStepRecorderTermCfg(RecorderTermCfg): + """Configuration for the dummy step recorder term.""" + + class_type: type[RecorderTerm] = DummyStepRecorderTerm + + record_reset_term = DummyResetRecorderTermCfg() + record_step_term = DummyStepRecorderTermCfg() + + dataset_export_mode = DatasetExportMode.EXPORT_ALL + + +def create_dummy_env(device: str = "cpu") -> ManagerBasedEnv: + """Create a dummy environment.""" + + class DummyTerminationManager: + active_terms = [] + + dummy_termination_manager = DummyTerminationManager() + return namedtuple("ManagerBasedEnv", ["num_envs", "device", "cfg", "termination_manager"])( + 20, device, dict(), dummy_termination_manager + ) + + +class TestRecorderManager(unittest.TestCase): + """Test cases for various situations with recorder manager.""" + + def setUp(self) -> None: + self.dataset_dir = tempfile.mkdtemp() + + def tearDown(self): + # delete the temporary directory after the test + shutil.rmtree(self.dataset_dir) + + def create_dummy_recorder_manager_cfg(self) -> DummyRecorderManagerCfg: + """Get the dummy recorder manager configurations.""" + cfg = DummyRecorderManagerCfg() + cfg.dataset_export_dir_path = self.dataset_dir + cfg.dataset_filename = f"{uuid.uuid4()}.hdf5" + return cfg + + def test_str(self): + """Test the string representation of the recorder manager.""" + # create recorder manager + cfg = DummyRecorderManagerCfg() + recorder_manager = RecorderManager(cfg, create_dummy_env()) + self.assertEqual(len(recorder_manager.active_terms), 2) + # print the expected string + print() + print(recorder_manager) + + def test_initialize_dataset_file(self): + """Test the initialization of the dataset file.""" + # create recorder manager + cfg = self.create_dummy_recorder_manager_cfg() + _ = RecorderManager(cfg, create_dummy_env()) + + # check if the dataset is created + self.assertTrue(os.path.exists(os.path.join(cfg.dataset_export_dir_path, cfg.dataset_filename))) + + def test_record(self): + """Test the recording of the data.""" + for device in ("cuda:0", "cpu"): + with self.subTest(device=device): + env = create_dummy_env(device) + # create recorder manager + recorder_manager = RecorderManager(self.create_dummy_recorder_manager_cfg(), env) + + # record the step data + recorder_manager.record_pre_step() + recorder_manager.record_post_step() + + recorder_manager.record_pre_step() + recorder_manager.record_post_step() + + # check the recorded data + for env_id in range(env.num_envs): + episode = recorder_manager.get_episode(env_id) + self.assertEqual(episode.data["record_pre_step"].shape, (2, 4)) + self.assertEqual(episode.data["record_post_step"].shape, (2, 5)) + + # Trigger pre-reset callbacks which then export and clean the episode data + recorder_manager.record_pre_reset(env_ids=None) + for env_id in range(env.num_envs): + episode = recorder_manager.get_episode(env_id) + self.assertTrue(episode.is_empty()) + + recorder_manager.record_post_reset(env_ids=None) + for env_id in range(env.num_envs): + episode = recorder_manager.get_episode(env_id) + self.assertEqual(episode.data["record_post_reset"].shape, (1, 3)) + + +if __name__ == "__main__": + run_tests() diff --git a/source/extensions/omni.isaac.lab/test/managers/test_reward_manager.py b/source/extensions/omni.isaac.lab/test/managers/test_reward_manager.py index 381886776f..af5d35d858 100644 --- a/source/extensions/omni.isaac.lab/test/managers/test_reward_manager.py +++ b/source/extensions/omni.isaac.lab/test/managers/test_reward_manager.py @@ -12,6 +12,7 @@ """Rest everything follows.""" +import torch import unittest from collections import namedtuple @@ -123,6 +124,21 @@ def test_compute(self): self.assertEqual(float(rewards[0]), expected_reward) self.assertEqual(tuple(rewards.shape), (self.env.num_envs,)) + def test_config_empty(self): + """Test the creation of reward manager with empty config.""" + self.rew_man = RewardManager(None, self.env) + self.assertEqual(len(self.rew_man.active_terms), 0) + + # print the expected string + print() + print(self.rew_man) + + # compute reward + rewards = self.rew_man.compute(dt=self.env.dt) + + # check all rewards are zero + torch.testing.assert_close(rewards, torch.zeros_like(rewards)) + def test_active_terms(self): """Test the correct reading of active terms.""" cfg = { diff --git a/source/extensions/omni.isaac.lab/test/markers/check_markers_visibility.py b/source/extensions/omni.isaac.lab/test/markers/check_markers_visibility.py index 5aafc68c80..ef15ca5844 100644 --- a/source/extensions/omni.isaac.lab/test/markers/check_markers_visibility.py +++ b/source/extensions/omni.isaac.lab/test/markers/check_markers_visibility.py @@ -98,7 +98,8 @@ def run_simulator( # root state root_state = scene["robot"].data.default_root_state.clone() root_state[:, :3] += scene.env_origins - scene["robot"].write_root_state_to_sim(root_state) + scene["robot"].write_root_link_pose_to_sim(root_state[:, :7]) + scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:]) # set joint positions with some noise joint_pos, joint_vel = ( scene["robot"].data.default_joint_pos.clone(), diff --git a/source/extensions/omni.isaac.lab/test/scene/check_interactive_scene.py b/source/extensions/omni.isaac.lab/test/scene/check_interactive_scene.py index 25b63fc087..8120102d66 100644 --- a/source/extensions/omni.isaac.lab/test/scene/check_interactive_scene.py +++ b/source/extensions/omni.isaac.lab/test/scene/check_interactive_scene.py @@ -130,11 +130,13 @@ def main(): joint_vel = scene.articulations["robot_1"].data.default_joint_vel # -- set root state # -- robot 1 - scene.articulations["robot_1"].write_root_state_to_sim(root_state) + scene.articulations["robot_1"].write_root_link_pose_to_sim(root_state[:, :7]) + scene.articulations["robot_1"].write_root_com_velocity_to_sim(root_state[:, 7:]) scene.articulations["robot_1"].write_joint_state_to_sim(joint_pos, joint_vel) # -- robot 2 root_state[:, 1] += 1.0 - scene.articulations["robot_2"].write_root_state_to_sim(root_state) + scene.articulations["robot_2"].write_root_link_pose_to_sim(root_state[:, :7]) + scene.articulations["robot_2"].write_root_com_velocity_to_sim(root_state[:, 7:]) scene.articulations["robot_2"].write_joint_state_to_sim(joint_pos, joint_vel) # reset buffers scene.reset() diff --git a/source/extensions/omni.isaac.lab/test/scene/test_interactive_scene.py b/source/extensions/omni.isaac.lab/test/scene/test_interactive_scene.py index 4e9d08e54b..70d149fd1e 100644 --- a/source/extensions/omni.isaac.lab/test/scene/test_interactive_scene.py +++ b/source/extensions/omni.isaac.lab/test/scene/test_interactive_scene.py @@ -40,7 +40,7 @@ class MySceneCfg(InteractiveSceneCfg): prim_path="/World/Robot", spawn=sim_utils.UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/Simple/revolute_articulation.usd"), actuators={ - "joint": ImplicitActuatorCfg(), + "joint": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=100.0, damping=1.0), }, ) # rigid object diff --git a/source/extensions/omni.isaac.lab/test/sensors/check_imu_sensor.py b/source/extensions/omni.isaac.lab/test/sensors/check_imu_sensor.py new file mode 100644 index 0000000000..69a0422455 --- /dev/null +++ b/source/extensions/omni.isaac.lab/test/sensors/check_imu_sensor.py @@ -0,0 +1,196 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +""" +Visual test script for the imu sensor from the Orbit framework. +""" + +from __future__ import annotations + +"""Launch Isaac Sim Simulator first.""" + +import argparse + +from omni.isaac.kit import SimulationApp + +# add argparse arguments +parser = argparse.ArgumentParser(description="Imu Test Script") +parser.add_argument("--headless", action="store_true", default=False, help="Force display off at all times.") +parser.add_argument("--num_envs", type=int, default=128, help="Number of environments to clone.") +parser.add_argument( + "--terrain_type", + type=str, + default="generator", + choices=["generator", "usd", "plane"], + help="Type of terrain to import. Can be 'generator' or 'usd' or 'plane'.", +) +args_cli = parser.parse_args() + +# launch omniverse app +config = {"headless": args_cli.headless} +simulation_app = SimulationApp(config) + + +"""Rest everything follows.""" + +import torch +import traceback + +import carb +import omni +from omni.isaac.cloner import GridCloner +from omni.isaac.core.simulation_context import SimulationContext +from omni.isaac.core.utils.viewports import set_camera_view +from pxr import PhysxSchema + +import omni.isaac.lab.sim as sim_utils +import omni.isaac.lab.terrains as terrain_gen +from omni.isaac.lab.assets import RigidObject, RigidObjectCfg +from omni.isaac.lab.sensors.imu import Imu, ImuCfg +from omni.isaac.lab.terrains.config.rough import ROUGH_TERRAINS_CFG +from omni.isaac.lab.terrains.terrain_importer import TerrainImporter +from omni.isaac.lab.utils.assets import ISAAC_NUCLEUS_DIR +from omni.isaac.lab.utils.timer import Timer + + +def design_scene(sim: SimulationContext, num_envs: int = 2048) -> RigidObject: + """Design the scene.""" + # Handler for terrains importing + terrain_importer_cfg = terrain_gen.TerrainImporterCfg( + prim_path="/World/ground", + terrain_type="generator", + terrain_generator=ROUGH_TERRAINS_CFG, + usd_path=f"{ISAAC_NUCLEUS_DIR}/Environments/Terrains/rough_plane.usd", + max_init_terrain_level=None, + num_envs=1, + ) + _ = TerrainImporter(terrain_importer_cfg) + # obtain the current stage + stage = omni.usd.get_context().get_stage() + # Create interface to clone the scene + cloner = GridCloner(spacing=2.0) + cloner.define_base_env("/World/envs") + envs_prim_paths = cloner.generate_paths("/World/envs/env", num_paths=num_envs) + # create source prim + stage.DefinePrim(envs_prim_paths[0], "Xform") + # clone the env xform + cloner.clone(source_prim_path="/World/envs/env_0", prim_paths=envs_prim_paths, replicate_physics=True) + # Define the scene + # -- Light + cfg = sim_utils.DistantLightCfg(intensity=2000) + cfg.func("/World/light", cfg) + # -- Balls + cfg = RigidObjectCfg( + spawn=sim_utils.SphereCfg( + radius=0.25, + rigid_props=sim_utils.RigidBodyPropertiesCfg(), + mass_props=sim_utils.MassPropertiesCfg(mass=0.5), + collision_props=sim_utils.CollisionPropertiesCfg(), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 0.0, 1.0)), + ), + prim_path="/World/envs/env_.*/ball", + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 5.0)), + ) + balls = RigidObject(cfg) + # Clone the scene + # obtain the current physics scene + physics_scene_prim_path = None + for prim in stage.Traverse(): + if prim.HasAPI(PhysxSchema.PhysxSceneAPI): + physics_scene_prim_path = prim.GetPrimPath() + carb.log_info(f"Physics scene prim path: {physics_scene_prim_path}") + break + # filter collisions within each environment instance + cloner.filter_collisions( + physics_scene_prim_path, + "/World/collisions", + envs_prim_paths, + ) + return balls + + +def main(): + """Main function.""" + + # Load kit helper + sim_params = { + "use_gpu": True, + "use_gpu_pipeline": True, + "use_flatcache": True, # deprecated from Isaac Sim 2023.1 onwards + "use_fabric": True, # used from Isaac Sim 2023.1 onwards + "enable_scene_query_support": True, + } + sim = SimulationContext( + physics_dt=1.0 / 60.0, rendering_dt=1.0 / 60.0, sim_params=sim_params, backend="torch", device="cuda:0" + ) + # Set main camera + set_camera_view([0.0, 30.0, 25.0], [0.0, 0.0, -2.5]) + + # Parameters + num_envs = args_cli.num_envs + # Design the scene + balls = design_scene(sim=sim, num_envs=num_envs) + + # Create a ray-caster sensor + imu_cfg = ImuCfg( + prim_path="/World/envs/env_.*/ball", + debug_vis=not args_cli.headless, + ) + # increase scale of the arrows for better visualization + imu_cfg.visualizer_cfg.markers["arrow"].scale = (1.0, 0.2, 0.2) + imu = Imu(cfg=imu_cfg) + + # Play simulator and init the Imu + sim.reset() + + # Print the sensor information + print(imu) + + # Get the ball initial positions + sim.step(render=not args_cli.headless) + balls.update(sim.get_physics_dt()) + ball_initial_positions = balls.data.root_link_pos_w.clone() + ball_initial_orientations = balls.data.root_link_quat_w.clone() + + # Create a counter for resetting the scene + step_count = 0 + # Simulate physics + while simulation_app.is_running(): + # If simulation is stopped, then exit. + if sim.is_stopped(): + break + # If simulation is paused, then skip. + if not sim.is_playing(): + sim.step(render=not args_cli.headless) + continue + # Reset the scene + if step_count % 500 == 0: + # reset ball positions + balls.write_root_link_pose_to_sim(torch.cat([ball_initial_positions, ball_initial_orientations], dim=-1)) + balls.reset() + # reset the sensor + imu.reset() + # reset the counter + step_count = 0 + # Step simulation + sim.step() + # Update the imu sensor + with Timer(f"Imu sensor update with {num_envs}"): + imu.update(dt=sim.get_physics_dt(), force_recompute=True) + # Update counter + step_count += 1 + + +if __name__ == "__main__": + try: + # Run the main function + main() + except Exception as err: + carb.log_error(err) + carb.log_error(traceback.format_exc()) + raise + finally: + # close sim app + simulation_app.close() diff --git a/source/extensions/omni.isaac.lab/test/sensors/test_camera.py b/source/extensions/omni.isaac.lab/test/sensors/test_camera.py index fc6953e166..fad7d5893d 100644 --- a/source/extensions/omni.isaac.lab/test/sensors/test_camera.py +++ b/source/extensions/omni.isaac.lab/test/sensors/test_camera.py @@ -121,7 +121,7 @@ def test_camera_init(self): # update camera camera.update(self.dt) # check image data - for im_data in camera.data.output.to_dict().values(): + for im_data in camera.data.output.values(): self.assertEqual(im_data.shape, (1, self.camera_cfg.height, self.camera_cfg.width, 1)) def test_camera_init_offset(self): @@ -228,7 +228,7 @@ def test_multi_camera_init(self): cam_2.update(self.dt) # check image data for cam in [cam_1, cam_2]: - for im_data in cam.data.output.to_dict().values(): + for im_data in cam.data.output.values(): self.assertEqual(im_data.shape, (1, self.camera_cfg.height, self.camera_cfg.width, 1)) def test_multi_camera_with_different_resolution(self): @@ -393,6 +393,137 @@ def test_intrinsic_matrix(self): torch.testing.assert_close(rs_intrinsic_matrix[0, 0, 0], camera.data.intrinsic_matrices[0, 0, 0]) # torch.testing.assert_close(rs_intrinsic_matrix[0, 1, 1], camera.data.intrinsic_matrices[0, 1, 1]) + def test_depth_clipping(self): + """Test depth clipping. + + .. note:: + + This test is the same for all camera models to enforce the same clipping behavior. + """ + # get camera cfgs + camera_cfg_zero = CameraCfg( + prim_path="/World/CameraZero", + offset=CameraCfg.OffsetCfg(pos=(2.5, 2.5, 6.0), rot=(-0.125, 0.362, 0.873, -0.302), convention="ros"), + spawn=sim_utils.PinholeCameraCfg().from_intrinsic_matrix( + focal_length=38.0, + intrinsic_matrix=[380.08, 0.0, 467.79, 0.0, 380.08, 262.05, 0.0, 0.0, 1.0], + height=540, + width=960, + clipping_range=(0.1, 10), + ), + height=540, + width=960, + data_types=["distance_to_image_plane", "distance_to_camera"], + depth_clipping_behavior="zero", + ) + camera_zero = Camera(camera_cfg_zero) + + camera_cfg_none = copy.deepcopy(camera_cfg_zero) + camera_cfg_none.prim_path = "/World/CameraNone" + camera_cfg_none.depth_clipping_behavior = "none" + camera_none = Camera(camera_cfg_none) + + camera_cfg_max = copy.deepcopy(camera_cfg_zero) + camera_cfg_max.prim_path = "/World/CameraMax" + camera_cfg_max.depth_clipping_behavior = "max" + camera_max = Camera(camera_cfg_max) + + # Play sim + self.sim.reset() + + # note: This is a workaround to ensure that the textures are loaded. + # Check "Known Issues" section in the documentation for more details. + for _ in range(5): + self.sim.step() + + camera_zero.update(self.dt) + camera_none.update(self.dt) + camera_max.update(self.dt) + + # none clipping should contain inf values + self.assertTrue(torch.isinf(camera_none.data.output["distance_to_camera"]).any()) + self.assertTrue(torch.isinf(camera_none.data.output["distance_to_image_plane"]).any()) + self.assertTrue( + camera_none.data.output["distance_to_camera"][ + ~torch.isinf(camera_none.data.output["distance_to_camera"]) + ].min() + >= camera_cfg_zero.spawn.clipping_range[0] + ) + self.assertTrue( + camera_none.data.output["distance_to_camera"][ + ~torch.isinf(camera_none.data.output["distance_to_camera"]) + ].max() + <= camera_cfg_zero.spawn.clipping_range[1] + ) + self.assertTrue( + camera_none.data.output["distance_to_image_plane"][ + ~torch.isinf(camera_none.data.output["distance_to_image_plane"]) + ].min() + >= camera_cfg_zero.spawn.clipping_range[0] + ) + self.assertTrue( + camera_none.data.output["distance_to_image_plane"][ + ~torch.isinf(camera_none.data.output["distance_to_camera"]) + ].max() + <= camera_cfg_zero.spawn.clipping_range[1] + ) + + # zero clipping should result in zero values + self.assertTrue( + torch.all( + camera_zero.data.output["distance_to_camera"][ + torch.isinf(camera_none.data.output["distance_to_camera"]) + ] + == 0.0 + ) + ) + self.assertTrue( + torch.all( + camera_zero.data.output["distance_to_image_plane"][ + torch.isinf(camera_none.data.output["distance_to_image_plane"]) + ] + == 0.0 + ) + ) + self.assertTrue( + camera_zero.data.output["distance_to_camera"][camera_zero.data.output["distance_to_camera"] != 0.0].min() + >= camera_cfg_zero.spawn.clipping_range[0] + ) + self.assertTrue(camera_zero.data.output["distance_to_camera"].max() <= camera_cfg_zero.spawn.clipping_range[1]) + self.assertTrue( + camera_zero.data.output["distance_to_image_plane"][ + camera_zero.data.output["distance_to_image_plane"] != 0.0 + ].min() + >= camera_cfg_zero.spawn.clipping_range[0] + ) + self.assertTrue( + camera_zero.data.output["distance_to_image_plane"].max() <= camera_cfg_zero.spawn.clipping_range[1] + ) + + # max clipping should result in max values + self.assertTrue( + torch.all( + camera_max.data.output["distance_to_camera"][torch.isinf(camera_none.data.output["distance_to_camera"])] + == camera_cfg_zero.spawn.clipping_range[1] + ) + ) + self.assertTrue( + torch.all( + camera_max.data.output["distance_to_image_plane"][ + torch.isinf(camera_none.data.output["distance_to_image_plane"]) + ] + == camera_cfg_zero.spawn.clipping_range[1] + ) + ) + self.assertTrue(camera_max.data.output["distance_to_camera"].min() >= camera_cfg_zero.spawn.clipping_range[0]) + self.assertTrue(camera_max.data.output["distance_to_camera"].max() <= camera_cfg_zero.spawn.clipping_range[1]) + self.assertTrue( + camera_max.data.output["distance_to_image_plane"].min() >= camera_cfg_zero.spawn.clipping_range[0] + ) + self.assertTrue( + camera_max.data.output["distance_to_image_plane"].max() <= camera_cfg_zero.spawn.clipping_range[1] + ) + def test_camera_resolution_all_colorize(self): """Test camera resolution is correctly set for all types with colorization enabled.""" # Add all types @@ -705,7 +836,7 @@ def test_throughput(self): with Timer(f"Time taken for writing data with shape {camera.image_shape} "): # Pack data back into replicator format to save them using its writer rep_output = {"annotators": {}} - camera_data = convert_dict_to_backend(camera.data.output[0].to_dict(), backend="numpy") + camera_data = convert_dict_to_backend({k: v[0] for k, v in camera.data.output.items()}, backend="numpy") for key, data, info in zip(camera_data.keys(), camera_data.values(), camera.data.info[0].values()): if info is not None: rep_output["annotators"][key] = {"render_product": {"data": data, **info}} @@ -719,6 +850,15 @@ def test_throughput(self): for im_data in camera.data.output.values(): self.assertEqual(im_data.shape, (1, camera_cfg.height, camera_cfg.width, 1)) + def test_sensor_print(self): + """Test sensor print is working correctly.""" + # Create sensor + sensor = Camera(cfg=self.camera_cfg) + # Play sim + self.sim.reset() + # print info + print(sensor) + """ Helper functions. """ diff --git a/source/extensions/omni.isaac.lab/test/sensors/test_contact_sensor.py b/source/extensions/omni.isaac.lab/test/sensors/test_contact_sensor.py index 8390ee2269..949456c3f4 100644 --- a/source/extensions/omni.isaac.lab/test/sensors/test_contact_sensor.py +++ b/source/extensions/omni.isaac.lab/test/sensors/test_contact_sensor.py @@ -294,6 +294,27 @@ def test_cube_stack_contact_filtering(self): contact_sensor_2.data.force_matrix_w[:, :, 0], contact_sensor.data.force_matrix_w[:, :, 0] ) + def test_sensor_print(self): + """Test sensor print is working correctly.""" + with build_simulation_context(device="cuda:0", dt=self.sim_dt, add_lighting=False) as sim: + # Spawn things into stage + scene_cfg = ContactSensorSceneCfg(num_envs=1, env_spacing=1.0, lazy_sensor_update=False) + scene_cfg.terrain = FLAT_TERRAIN_CFG.replace(prim_path="/World/ground") + scene_cfg.shape = CUBE_CFG + scene_cfg.contact_sensor = ContactSensorCfg( + prim_path=scene_cfg.shape.prim_path, + track_pose=True, + debug_vis=False, + update_period=0.0, + track_air_time=True, + history_length=3, + ) + scene = InteractiveScene(scene_cfg) + # Play the simulator + sim.reset() + # print info + print(scene.sensors["contact_sensor"]) + """ Internal helpers. """ @@ -384,7 +405,7 @@ def _test_sensor_contact(self, shape: RigidObject, sensor: ContactSensor, mode: duration = self.durations[idx] while current_test_time < duration: # set object states to contact the ground plane - shape.write_root_pose_to_sim(root_pose=test_pose) + shape.write_root_link_pose_to_sim(root_pose=test_pose) # perform simulation step self._perform_sim_step() # increment contact time @@ -411,7 +432,7 @@ def _test_sensor_contact(self, shape: RigidObject, sensor: ContactSensor, mode: dt=duration + self.sim_dt, ) # switch the contact mode for 1 dt step before the next contact test begins. - shape.write_root_pose_to_sim(root_pose=reset_pose) + shape.write_root_link_pose_to_sim(root_pose=reset_pose) # perform simulation step self._perform_sim_step() # set the last air time to 2 sim_dt steps, because last_air_time and last_contact_time diff --git a/source/extensions/omni.isaac.lab/test/sensors/test_frame_transformer.py b/source/extensions/omni.isaac.lab/test/sensors/test_frame_transformer.py index aeab3d9e72..bc30bed2f3 100644 --- a/source/extensions/omni.isaac.lab/test/sensors/test_frame_transformer.py +++ b/source/extensions/omni.isaac.lab/test/sensors/test_frame_transformer.py @@ -175,7 +175,8 @@ def test_frame_transformer_feet_wrt_base(self): joint_vel = scene.articulations["robot"].data.default_joint_vel # -- set root state # -- robot - scene.articulations["robot"].write_root_state_to_sim(root_state) + scene.articulations["robot"].write_root_link_pose_to_sim(root_state[:, :7]) + scene.articulations["robot"].write_root_com_velocity_to_sim(root_state[:, 7:]) scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel) # reset buffers scene.reset() @@ -192,9 +193,9 @@ def test_frame_transformer_feet_wrt_base(self): # check absolute frame transforms in world frame # -- ground-truth - root_pose_w = scene.articulations["robot"].data.root_state_w[:, :7] - feet_pos_w_gt = scene.articulations["robot"].data.body_pos_w[:, feet_indices] - feet_quat_w_gt = scene.articulations["robot"].data.body_quat_w[:, feet_indices] + root_pose_w = scene.articulations["robot"].data.root_link_state_w[:, :7] + feet_pos_w_gt = scene.articulations["robot"].data.body_link_pos_w[:, feet_indices] + feet_quat_w_gt = scene.articulations["robot"].data.body_link_quat_w[:, feet_indices] # -- frame transformer source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w source_quat_w_tf = scene.sensors["frame_transformer"].data.source_quat_w @@ -275,7 +276,8 @@ def test_frame_transformer_feet_wrt_thigh(self): joint_vel = scene.articulations["robot"].data.default_joint_vel # -- set root state # -- robot - scene.articulations["robot"].write_root_state_to_sim(root_state) + scene.articulations["robot"].write_root_link_pose_to_sim(root_state[:, :7]) + scene.articulations["robot"].write_root_com_velocity_to_sim(root_state[:, 7:]) scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel) # reset buffers scene.reset() @@ -292,9 +294,9 @@ def test_frame_transformer_feet_wrt_thigh(self): # check absolute frame transforms in world frame # -- ground-truth - source_pose_w_gt = scene.articulations["robot"].data.body_state_w[:, source_frame_index, :7] - feet_pos_w_gt = scene.articulations["robot"].data.body_pos_w[:, feet_indices] - feet_quat_w_gt = scene.articulations["robot"].data.body_quat_w[:, feet_indices] + source_pose_w_gt = scene.articulations["robot"].data.body_link_state_w[:, source_frame_index, :7] + feet_pos_w_gt = scene.articulations["robot"].data.body_link_pos_w[:, feet_indices] + feet_quat_w_gt = scene.articulations["robot"].data.body_link_quat_w[:, feet_indices] # -- frame transformer source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w source_quat_w_tf = scene.sensors["frame_transformer"].data.source_quat_w @@ -356,7 +358,8 @@ def test_frame_transformer_robot_body_to_external_cube(self): joint_vel = scene.articulations["robot"].data.default_joint_vel # -- set root state # -- robot - scene.articulations["robot"].write_root_state_to_sim(root_state) + scene.articulations["robot"].write_root_link_pose_to_sim(root_state[:, :7]) + scene.articulations["robot"].write_root_com_velocity_to_sim(root_state[:, 7:]) scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel) # reset buffers scene.reset() @@ -373,9 +376,9 @@ def test_frame_transformer_robot_body_to_external_cube(self): # check absolute frame transforms in world frame # -- ground-truth - root_pose_w = scene.articulations["robot"].data.root_state_w[:, :7] - cube_pos_w_gt = scene.rigid_objects["cube"].data.root_state_w[:, :3] - cube_quat_w_gt = scene.rigid_objects["cube"].data.root_state_w[:, 3:7] + root_pose_w = scene.articulations["robot"].data.root_link_state_w[:, :7] + cube_pos_w_gt = scene.rigid_objects["cube"].data.root_link_state_w[:, :3] + cube_quat_w_gt = scene.rigid_objects["cube"].data.root_link_state_w[:, 3:7] # -- frame transformer source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w source_quat_w_tf = scene.sensors["frame_transformer"].data.source_quat_w @@ -447,7 +450,8 @@ def test_frame_transformer_offset_frames(self): root_state[:, :3] += scene.env_origins # -- set root state # -- cube - scene["cube"].write_root_state_to_sim(root_state) + scene["cube"].write_root_link_pose_to_sim(root_state[:, :7]) + scene["cube"].write_root_com_velocity_to_sim(root_state[:, 7:]) # reset buffers scene.reset() @@ -460,8 +464,8 @@ def test_frame_transformer_offset_frames(self): # check absolute frame transforms in world frame # -- ground-truth - cube_pos_w_gt = scene["cube"].data.root_state_w[:, :3] - cube_quat_w_gt = scene["cube"].data.root_state_w[:, 3:7] + cube_pos_w_gt = scene["cube"].data.root_link_state_w[:, :3] + cube_quat_w_gt = scene["cube"].data.root_link_state_w[:, 3:7] # -- frame transformer source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w source_quat_w_tf = scene.sensors["frame_transformer"].data.source_quat_w @@ -534,7 +538,8 @@ def test_frame_transformer_all_bodies(self): joint_vel = scene.articulations["robot"].data.default_joint_vel # -- set root state # -- robot - scene.articulations["robot"].write_root_state_to_sim(root_state) + scene.articulations["robot"].write_root_link_pose_to_sim(root_state[:, :7]) + scene.articulations["robot"].write_root_com_velocity_to_sim(root_state[:, 7:]) scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel) # reset buffers scene.reset() @@ -551,9 +556,9 @@ def test_frame_transformer_all_bodies(self): # check absolute frame transforms in world frame # -- ground-truth - root_pose_w = scene.articulations["robot"].data.root_state_w[:, :7] - bodies_pos_w_gt = scene.articulations["robot"].data.body_pos_w - bodies_quat_w_gt = scene.articulations["robot"].data.body_quat_w + root_pose_w = scene.articulations["robot"].data.root_link_state_w[:, :7] + bodies_pos_w_gt = scene.articulations["robot"].data.body_link_pos_w + bodies_quat_w_gt = scene.articulations["robot"].data.body_link_quat_w # -- frame transformer source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w @@ -579,6 +584,25 @@ def test_frame_transformer_all_bodies(self): torch.testing.assert_close(bodies_pos_source_tf[:, index], body_pos_b) torch.testing.assert_close(bodies_quat_source_tf[:, index], body_quat_b) + def test_sensor_print(self): + """Test sensor print is working correctly.""" + # Spawn things into stage + scene_cfg = MySceneCfg(num_envs=2, env_spacing=5.0, lazy_sensor_update=False) + scene_cfg.frame_transformer = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/base", + target_frames=[ + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/.*", + ), + ], + ) + scene = InteractiveScene(scene_cfg) + + # Play the simulator + self.sim.reset() + # print info + print(scene.sensors["frame_transformer"]) + if __name__ == "__main__": run_tests() diff --git a/source/extensions/omni.isaac.lab/test/sensors/test_imu.py b/source/extensions/omni.isaac.lab/test/sensors/test_imu.py new file mode 100644 index 0000000000..d7b755dedd --- /dev/null +++ b/source/extensions/omni.isaac.lab/test/sensors/test_imu.py @@ -0,0 +1,538 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Launch Isaac Sim Simulator first.""" + +from omni.isaac.lab.app import AppLauncher, run_tests + +# launch omniverse app +app_launcher = AppLauncher(headless=True, enable_cameras=True) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import pathlib +import torch +import unittest + +import omni.isaac.core.utils.stage as stage_utils + +import omni.isaac.lab.sim as sim_utils +import omni.isaac.lab.utils.math as math_utils +from omni.isaac.lab.actuators import ImplicitActuatorCfg +from omni.isaac.lab.assets import ArticulationCfg, RigidObjectCfg +from omni.isaac.lab.markers.config import GREEN_ARROW_X_MARKER_CFG, RED_ARROW_X_MARKER_CFG +from omni.isaac.lab.scene import InteractiveScene, InteractiveSceneCfg +from omni.isaac.lab.sensors.imu import ImuCfg +from omni.isaac.lab.terrains import TerrainImporterCfg +from omni.isaac.lab.utils import configclass + +## +# Pre-defined configs +## +from omni.isaac.lab_assets.anymal import ANYMAL_C_CFG # isort: skip +from omni.isaac.lab.utils.assets import NUCLEUS_ASSET_ROOT_DIR # isort: skip + +# offset of imu_link from base_link on anymal_c +POS_OFFSET = (0.2488, 0.00835, 0.04628) +ROT_OFFSET = (0.7071068, 0, 0, 0.7071068) + +# offset of imu_link from link_1 on simple_2_link +PEND_POS_OFFSET = (0.4, 0.0, 0.1) +PEND_ROT_OFFSET = (0.5, 0.5, 0.5, 0.5) + + +@configclass +class MySceneCfg(InteractiveSceneCfg): + """Example scene configuration.""" + + # terrain - flat terrain plane + terrain = TerrainImporterCfg( + prim_path="/World/ground", + terrain_type="plane", + max_init_terrain_level=None, + ) + + # rigid objects - balls + balls = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/ball", + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 0.5)), + spawn=sim_utils.SphereCfg( + radius=0.25, + rigid_props=sim_utils.RigidBodyPropertiesCfg(), + mass_props=sim_utils.MassPropertiesCfg(mass=0.5), + collision_props=sim_utils.CollisionPropertiesCfg(), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 0.0, 1.0)), + ), + ) + + cube = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/cube", + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, -2.0, 0.5)), + spawn=sim_utils.CuboidCfg( + size=(0.25, 0.25, 0.25), + rigid_props=sim_utils.RigidBodyPropertiesCfg(), + mass_props=sim_utils.MassPropertiesCfg(mass=0.5), + collision_props=sim_utils.CollisionPropertiesCfg(), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 0.0, 1.0)), + ), + ) + + # articulations - robot + robot = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/robot") + pendulum = ArticulationCfg( + prim_path="{ENV_REGEX_NS}/pendulum", + spawn=sim_utils.UrdfFileCfg( + fix_base=True, + merge_fixed_joints=False, + make_instanceable=False, + asset_path=f"{pathlib.Path(__file__).parent.resolve()}/urdfs/simple_2_link.urdf", + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=True, solver_position_iteration_count=4, solver_velocity_iteration_count=0 + ), + ), + init_state=ArticulationCfg.InitialStateCfg(), + actuators={ + "joint_1_act": ImplicitActuatorCfg(joint_names_expr=["joint_.*"], stiffness=0.0, damping=0.3), + }, + ) + # sensors - imu (filled inside unit test) + imu_ball: ImuCfg = ImuCfg( + prim_path="{ENV_REGEX_NS}/ball", + gravity_bias=(0.0, 0.0, 0.0), + ) + imu_cube: ImuCfg = ImuCfg( + prim_path="{ENV_REGEX_NS}/cube", + gravity_bias=(0.0, 0.0, 0.0), + ) + imu_robot_imu_link: ImuCfg = ImuCfg( + prim_path="{ENV_REGEX_NS}/robot/imu_link", + gravity_bias=(0.0, 0.0, 0.0), + ) + imu_robot_base: ImuCfg = ImuCfg( + prim_path="{ENV_REGEX_NS}/robot/base", + offset=ImuCfg.OffsetCfg( + pos=POS_OFFSET, + rot=ROT_OFFSET, + ), + gravity_bias=(0.0, 0.0, 0.0), + ) + + imu_pendulum_imu_link: ImuCfg = ImuCfg( + prim_path="{ENV_REGEX_NS}/pendulum/imu_link", + debug_vis=not app_launcher._headless, + visualizer_cfg=RED_ARROW_X_MARKER_CFG.replace(prim_path="/Visuals/Acceleration/imu_link"), + gravity_bias=(0.0, 0.0, 9.81), + ) + imu_pendulum_base: ImuCfg = ImuCfg( + prim_path="{ENV_REGEX_NS}/pendulum/link_1", + offset=ImuCfg.OffsetCfg( + pos=PEND_POS_OFFSET, + rot=PEND_ROT_OFFSET, + ), + debug_vis=not app_launcher._headless, + visualizer_cfg=GREEN_ARROW_X_MARKER_CFG.replace(prim_path="/Visuals/Acceleration/base"), + gravity_bias=(0.0, 0.0, 9.81), + ) + + def __post_init__(self): + """Post initialization.""" + # change position of the robot + self.robot.init_state.pos = (0.0, 2.0, 1.0) + self.pendulum.init_state.pos = (-1.0, 1.0, 0.5) + + # change asset + self.robot.spawn.usd_path = f"{NUCLEUS_ASSET_ROOT_DIR}/Isaac/Robots/ANYbotics/anymal_c.usd" + # change iterations + self.robot.spawn.articulation_props.solver_position_iteration_count = 32 + self.robot.spawn.articulation_props.solver_velocity_iteration_count = 32 + + +class TestImu(unittest.TestCase): + """Test for Imu sensor.""" + + def setUp(self): + """Create a blank new stage for each test.""" + # Create a new stage + stage_utils.create_new_stage() + # Load simulation context + sim_cfg = sim_utils.SimulationCfg(dt=0.001) + sim_cfg.physx.solver_type = 0 # 0: PGS, 1: TGS --> use PGS for more accurate results + self.sim = sim_utils.SimulationContext(sim_cfg) + # construct scene + scene_cfg = MySceneCfg(num_envs=2, env_spacing=5.0, lazy_sensor_update=False) + self.scene = InteractiveScene(scene_cfg) + # Play the simulator + self.sim.reset() + + def tearDown(self): + """Stops simulator after each test.""" + # clear the stage + self.sim.clear_all_callbacks() + self.sim.clear_instance() + + """ + Tests + """ + + def test_constant_velocity(self): + """Test the Imu sensor with a constant velocity. + + Expected behavior is that the linear and angular are approx the same at every time step as in each step we set + the same velocity and therefore reset the physx buffers.""" + prev_lin_acc_ball = torch.zeros((self.scene.num_envs, 3), dtype=torch.float32, device=self.scene.device) + prev_ang_acc_ball = torch.zeros((self.scene.num_envs, 3), dtype=torch.float32, device=self.scene.device) + prev_lin_acc_cube = torch.zeros((self.scene.num_envs, 3), dtype=torch.float32, device=self.scene.device) + prev_ang_acc_cube = torch.zeros((self.scene.num_envs, 3), dtype=torch.float32, device=self.scene.device) + + for idx in range(200): + # set velocity + self.scene.rigid_objects["balls"].write_root_com_velocity_to_sim( + torch.tensor([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=torch.float32, device=self.scene.device).repeat( + self.scene.num_envs, 1 + ) + ) + self.scene.rigid_objects["cube"].write_root_com_velocity_to_sim( + torch.tensor([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=torch.float32, device=self.scene.device).repeat( + self.scene.num_envs, 1 + ) + ) + # write data to sim + self.scene.write_data_to_sim() + + # perform step + self.sim.step() + # read data from sim + self.scene.update(self.sim.get_physics_dt()) + + if idx > 1: + # check the imu accelerations + torch.testing.assert_close( + self.scene.sensors["imu_ball"].data.lin_acc_b, + prev_lin_acc_ball, + rtol=1e-3, + atol=1e-3, + ) + torch.testing.assert_close( + self.scene.sensors["imu_ball"].data.ang_acc_b, + prev_ang_acc_ball, + rtol=1e-3, + atol=1e-3, + ) + + torch.testing.assert_close( + self.scene.sensors["imu_cube"].data.lin_acc_b, + prev_lin_acc_cube, + rtol=1e-3, + atol=1e-3, + ) + torch.testing.assert_close( + self.scene.sensors["imu_cube"].data.ang_acc_b, + prev_ang_acc_cube, + rtol=1e-3, + atol=1e-3, + ) + + # check the imu velocities + # NOTE: the expected lin_vel_b is the same as the set velocity, as write_root_com_velocity_to_sim is + # setting v_0 (initial velocity) and then a calculation step of v_i = v_0 + a*dt. Consequently, + # the data.lin_vel_b is returning approx. v_i. + torch.testing.assert_close( + self.scene.sensors["imu_ball"].data.lin_vel_b, + torch.tensor( + [[1.0, 0.0, -self.scene.physics_dt * 9.81]], dtype=torch.float32, device=self.scene.device + ).repeat(self.scene.num_envs, 1), + rtol=1e-4, + atol=1e-4, + ) + torch.testing.assert_close( + self.scene.sensors["imu_cube"].data.lin_vel_b, + torch.tensor( + [[1.0, 0.0, -self.scene.physics_dt * 9.81]], dtype=torch.float32, device=self.scene.device + ).repeat(self.scene.num_envs, 1), + rtol=1e-4, + atol=1e-4, + ) + + # update previous values + prev_lin_acc_ball = self.scene.sensors["imu_ball"].data.lin_acc_b.clone() + prev_ang_acc_ball = self.scene.sensors["imu_ball"].data.ang_acc_b.clone() + prev_lin_acc_cube = self.scene.sensors["imu_cube"].data.lin_acc_b.clone() + prev_ang_acc_cube = self.scene.sensors["imu_cube"].data.ang_acc_b.clone() + + def test_constant_acceleration(self): + """Test the Imu sensor with a constant acceleration.""" + for idx in range(100): + # set acceleration + self.scene.rigid_objects["balls"].write_root_com_velocity_to_sim( + torch.tensor([[0.1, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=torch.float32, device=self.scene.device).repeat( + self.scene.num_envs, 1 + ) + * (idx + 1) + ) + # write data to sim + self.scene.write_data_to_sim() + # perform step + self.sim.step() + # read data from sim + self.scene.update(self.sim.get_physics_dt()) + + # skip first step where initial velocity is zero + if idx < 1: + continue + + # check the imu data + torch.testing.assert_close( + self.scene.sensors["imu_ball"].data.lin_acc_b, + math_utils.quat_rotate_inverse( + self.scene.rigid_objects["balls"].data.root_link_quat_w, + torch.tensor([[0.1, 0.0, 0.0]], dtype=torch.float32, device=self.scene.device).repeat( + self.scene.num_envs, 1 + ) + / self.sim.get_physics_dt(), + ), + rtol=1e-4, + atol=1e-4, + ) + + # check the angular velocity + torch.testing.assert_close( + self.scene.sensors["imu_ball"].data.ang_vel_b, + self.scene.rigid_objects["balls"].data.root_com_ang_vel_b, + rtol=1e-4, + atol=1e-4, + ) + + def test_single_dof_pendulum(self): + """Test imu against analytical pendulum problem.""" + + # pendulum length + pend_length = PEND_POS_OFFSET[0] + + # should achieve same results between the two imu sensors on the robot + for idx in range(500): + + # write data to sim + self.scene.write_data_to_sim() + # perform step + self.sim.step() + # read data from sim + self.scene.update(self.sim.get_physics_dt()) + + # get pendulum joint state + joint_pos = self.scene.articulations["pendulum"].data.joint_pos + joint_vel = self.scene.articulations["pendulum"].data.joint_vel + joint_acc = self.scene.articulations["pendulum"].data.joint_acc + + # IMU and base data + imu_data = self.scene.sensors["imu_pendulum_imu_link"].data + base_data = self.scene.sensors["imu_pendulum_base"].data + + # extract imu_link imu_sensor dynamics + lin_vel_w_imu_link = math_utils.quat_rotate(imu_data.quat_w, imu_data.lin_vel_b) + lin_acc_w_imu_link = math_utils.quat_rotate(imu_data.quat_w, imu_data.lin_acc_b) + + # calculate the joint dynamics from the imu_sensor (y axis of imu_link is parallel to joint axis of pendulum) + joint_vel_imu = math_utils.quat_rotate(imu_data.quat_w, imu_data.ang_vel_b)[..., 1].unsqueeze(-1) + joint_acc_imu = math_utils.quat_rotate(imu_data.quat_w, imu_data.ang_acc_b)[..., 1].unsqueeze(-1) + + # calculate analytical solution + vx = -joint_vel * pend_length * torch.sin(joint_pos) + vy = torch.zeros(2, 1, device=self.scene.device) + vz = -joint_vel * pend_length * torch.cos(joint_pos) + gt_linear_vel_w = torch.cat([vx, vy, vz], dim=-1) + + ax = -joint_acc * pend_length * torch.sin(joint_pos) - joint_vel**2 * pend_length * torch.cos(joint_pos) + ay = torch.zeros(2, 1, device=self.scene.device) + az = ( + -joint_acc * pend_length * torch.cos(joint_pos) + + joint_vel**2 * pend_length * torch.sin(joint_pos) + + 9.81 + ) + gt_linear_acc_w = torch.cat([ax, ay, az], dim=-1) + + # skip first step where initial velocity is zero + if idx < 2: + continue + + # compare imu angular velocity with joint velocity + torch.testing.assert_close( + joint_vel, + joint_vel_imu, + rtol=1e-1, + atol=1e-3, + ) + # compare imu angular acceleration with joint acceleration + torch.testing.assert_close( + joint_acc, + joint_acc_imu, + rtol=1e-1, + atol=1e-3, + ) + # compare imu linear velocituy with simple pendulum calculation + torch.testing.assert_close( + gt_linear_vel_w, + lin_vel_w_imu_link, + rtol=1e-1, + atol=1e-3, + ) + # compare imu linear acceleration with simple pendulum calculation + torch.testing.assert_close( + gt_linear_acc_w, + lin_acc_w_imu_link, + rtol=1e-1, + atol=1e0, + ) + + # check the position between offset and imu definition + torch.testing.assert_close( + base_data.pos_w, + imu_data.pos_w, + rtol=1e-5, + atol=1e-5, + ) + + # check the orientation between offset and imu definition + torch.testing.assert_close( + base_data.quat_w, + imu_data.quat_w, + rtol=1e-4, + atol=1e-4, + ) + + # check the angular velocities of the imus between offset and imu definition + torch.testing.assert_close( + base_data.ang_vel_b, + imu_data.ang_vel_b, + rtol=1e-4, + atol=1e-4, + ) + # check the angular acceleration of the imus between offset and imu definition + torch.testing.assert_close( + base_data.ang_acc_b, + imu_data.ang_acc_b, + rtol=1e-4, + atol=1e-4, + ) + + # check the linear velocity of the imus between offset and imu definition + torch.testing.assert_close( + base_data.lin_vel_b, + imu_data.lin_vel_b, + rtol=1e-2, + atol=5e-3, + ) + + # check the linear acceleration of the imus between offset and imu definition + torch.testing.assert_close( + base_data.lin_acc_b, + imu_data.lin_acc_b, + rtol=1e-1, + atol=1e-1, + ) + + def test_offset_calculation(self): + """Test offset configuration argument.""" + # should achieve same results between the two imu sensors on the robot + for idx in range(500): + # set acceleration + self.scene.articulations["robot"].write_root_com_velocity_to_sim( + torch.tensor([[0.05, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=torch.float32, device=self.scene.device).repeat( + self.scene.num_envs, 1 + ) + * (idx + 1) + ) + # write data to sim + self.scene.write_data_to_sim() + # perform step + self.sim.step() + # read data from sim + self.scene.update(self.sim.get_physics_dt()) + + # skip first step where initial velocity is zero + if idx < 1: + continue + + # check the accelerations + torch.testing.assert_close( + self.scene.sensors["imu_robot_base"].data.lin_acc_b, + self.scene.sensors["imu_robot_imu_link"].data.lin_acc_b, + rtol=1e-4, + atol=1e-4, + ) + torch.testing.assert_close( + self.scene.sensors["imu_robot_base"].data.ang_acc_b, + self.scene.sensors["imu_robot_imu_link"].data.ang_acc_b, + rtol=1e-4, + atol=1e-4, + ) + + # check the velocities + torch.testing.assert_close( + self.scene.sensors["imu_robot_base"].data.ang_vel_b, + self.scene.sensors["imu_robot_imu_link"].data.ang_vel_b, + rtol=1e-4, + atol=1e-4, + ) + torch.testing.assert_close( + self.scene.sensors["imu_robot_base"].data.lin_vel_b, + self.scene.sensors["imu_robot_imu_link"].data.lin_vel_b, + rtol=1e-4, + atol=1e-4, + ) + + # check the orientation + torch.testing.assert_close( + self.scene.sensors["imu_robot_base"].data.quat_w, + self.scene.sensors["imu_robot_imu_link"].data.quat_w, + rtol=1e-4, + atol=1e-4, + ) + # check the position + torch.testing.assert_close( + self.scene.sensors["imu_robot_base"].data.pos_w, + self.scene.sensors["imu_robot_imu_link"].data.pos_w, + rtol=1e-4, + atol=1e-4, + ) + + def test_env_ids_propogation(self): + """Test that env_ids argument propagates through update and reset methods""" + self.scene.reset() + + for idx in range(10): + # set acceleration + self.scene.articulations["robot"].write_root_com_velocity_to_sim( + torch.tensor([[0.5, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=torch.float32, device=self.scene.device).repeat( + self.scene.num_envs, 1 + ) + * (idx + 1) + ) + # write data to sim + self.scene.write_data_to_sim() + # perform step + self.sim.step() + # read data from sim + self.scene.update(self.sim.get_physics_dt()) + + # reset scene for env 1 + self.scene.reset(env_ids=[1]) + # read data from sim + self.scene.update(self.sim.get_physics_dt()) + # perform step + self.sim.step() + # read data from sim + self.scene.update(self.sim.get_physics_dt()) + + def test_sensor_print(self): + """Test sensor print is working correctly.""" + # Create sensor + sensor = self.scene.sensors["imu_ball"] + # print info + print(sensor) + + +if __name__ == "__main__": + run_tests() diff --git a/source/extensions/omni.isaac.lab/test/sensors/test_outdated_sensor.py b/source/extensions/omni.isaac.lab/test/sensors/test_outdated_sensor.py new file mode 100644 index 0000000000..0079ece967 --- /dev/null +++ b/source/extensions/omni.isaac.lab/test/sensors/test_outdated_sensor.py @@ -0,0 +1,89 @@ +# Copyright (c) 2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Launch Isaac Sim Simulator first.""" + +from omni.isaac.lab.app import AppLauncher, run_tests + +# launch the simulator +app_launcher = AppLauncher(headless=True, enable_cameras=True) +simulation_app = app_launcher.app + + +"""Rest everything follows.""" + +import gymnasium as gym +import shutil +import tempfile +import torch +import unittest + +import carb +import omni.usd + +import omni.isaac.lab_tasks # noqa: F401 +from omni.isaac.lab_tasks.utils.parse_cfg import parse_env_cfg + + +class TestFrameTransformerAfterReset(unittest.TestCase): + """Test cases for checking FrameTransformer values after reset.""" + + @classmethod + def setUpClass(cls): + # this flag is necessary to prevent a bug where the simulation gets stuck randomly when running the + # test on many environments. + carb_settings_iface = carb.settings.get_settings() + carb_settings_iface.set_bool("/physics/cooking/ujitsoCollisionCooking", False) + + def setUp(self): + # create a temporary directory to store the test datasets + self.temp_dir = tempfile.mkdtemp() + + def tearDown(self): + # delete the temporary directory after the test + shutil.rmtree(self.temp_dir) + + def test_action_state_reocrder_terms(self): + """Check FrameTransformer values after reset.""" + for task_name in ["Isaac-Stack-Cube-Franka-IK-Rel-v0"]: + for device in ["cuda:0", "cpu"]: + for num_envs in [1, 2]: + with self.subTest(task_name=task_name, device=device): + omni.usd.get_context().new_stage() + + # parse configuration + env_cfg = parse_env_cfg(task_name, device=device, num_envs=num_envs) + + # create environment + env = gym.make(task_name, cfg=env_cfg) + + # disable control on stop + env.unwrapped.sim._app_control_on_stop_handle = None # type: ignore + + # reset environment + obs = env.reset()[0] + + # get the end effector position after the reset + pre_reset_eef_pos = obs["policy"]["eef_pos"].clone() + print(pre_reset_eef_pos) + + # step the environment with idle actions + idle_actions = torch.zeros(env.action_space.shape, device=env.unwrapped.device) + obs = env.step(idle_actions)[0] + + # get the end effector position after the first step + post_reset_eef_pos = obs["policy"]["eef_pos"] + print(post_reset_eef_pos) + + # check if the end effector position is the same after the reset and the first step + print(torch.all(torch.isclose(pre_reset_eef_pos, post_reset_eef_pos))) + self.assertTrue(torch.all(torch.isclose(pre_reset_eef_pos, post_reset_eef_pos))) + + # close the environment + env.close() + + +if __name__ == "__main__": + run_tests() diff --git a/source/extensions/omni.isaac.lab/test/sensors/test_ray_caster_camera.py b/source/extensions/omni.isaac.lab/test/sensors/test_ray_caster_camera.py index b28d023c85..e369eb569d 100644 --- a/source/extensions/omni.isaac.lab/test/sensors/test_ray_caster_camera.py +++ b/source/extensions/omni.isaac.lab/test/sensors/test_ray_caster_camera.py @@ -129,7 +129,7 @@ def test_camera_init(self): # update camera camera.update(self.dt) # check image data - for im_data in camera.data.output.to_dict().values(): + for im_data in camera.data.output.values(): self.assertEqual( im_data.shape, (1, self.camera_cfg.pattern_cfg.height, self.camera_cfg.pattern_cfg.width, 1) ) @@ -147,11 +147,112 @@ def test_camera_resolution(self): self.sim.step() camera.update(self.dt) # access image data and compare shapes - for im_data in camera.data.output.to_dict().values(): + for im_data in camera.data.output.values(): self.assertTrue( im_data.shape == (1, self.camera_cfg.pattern_cfg.height, self.camera_cfg.pattern_cfg.width, 1) ) + def test_depth_clipping(self): + """Test depth clipping. + + .. note:: + + This test is the same for all camera models to enforce the same clipping behavior. + """ + prim_utils.create_prim("/World/CameraZero", "Xform") + prim_utils.create_prim("/World/CameraNone", "Xform") + prim_utils.create_prim("/World/CameraMax", "Xform") + + # get camera cfgs + camera_cfg_zero = RayCasterCameraCfg( + prim_path="/World/CameraZero", + mesh_prim_paths=["/World/defaultGroundPlane"], + offset=RayCasterCameraCfg.OffsetCfg( + pos=(2.5, 2.5, 6.0), rot=(0.9914449, 0.0, 0.1305, 0.0), convention="world" + ), + pattern_cfg=patterns.PinholeCameraPatternCfg().from_intrinsic_matrix( + focal_length=38.0, + intrinsic_matrix=[380.08, 0.0, 467.79, 0.0, 380.08, 262.05, 0.0, 0.0, 1.0], + height=540, + width=960, + ), + max_distance=10.0, + data_types=["distance_to_image_plane", "distance_to_camera"], + depth_clipping_behavior="zero", + ) + camera_zero = RayCasterCamera(camera_cfg_zero) + + camera_cfg_none = copy.deepcopy(camera_cfg_zero) + camera_cfg_none.prim_path = "/World/CameraNone" + camera_cfg_none.depth_clipping_behavior = "none" + camera_none = RayCasterCamera(camera_cfg_none) + + camera_cfg_max = copy.deepcopy(camera_cfg_zero) + camera_cfg_max.prim_path = "/World/CameraMax" + camera_cfg_max.depth_clipping_behavior = "max" + camera_max = RayCasterCamera(camera_cfg_max) + + # Play sim + self.sim.reset() + + camera_zero.update(self.dt) + camera_none.update(self.dt) + camera_max.update(self.dt) + + # none clipping should contain inf values + self.assertTrue(torch.isinf(camera_none.data.output["distance_to_camera"]).any()) + self.assertTrue(torch.isnan(camera_none.data.output["distance_to_image_plane"]).any()) + self.assertTrue( + camera_none.data.output["distance_to_camera"][ + ~torch.isinf(camera_none.data.output["distance_to_camera"]) + ].max() + > camera_cfg_zero.max_distance + ) + self.assertTrue( + camera_none.data.output["distance_to_image_plane"][ + ~torch.isnan(camera_none.data.output["distance_to_image_plane"]) + ].max() + > camera_cfg_zero.max_distance + ) + + # zero clipping should result in zero values + self.assertTrue( + torch.all( + camera_zero.data.output["distance_to_camera"][ + torch.isinf(camera_none.data.output["distance_to_camera"]) + ] + == 0.0 + ) + ) + self.assertTrue( + torch.all( + camera_zero.data.output["distance_to_image_plane"][ + torch.isnan(camera_none.data.output["distance_to_image_plane"]) + ] + == 0.0 + ) + ) + self.assertTrue(camera_zero.data.output["distance_to_camera"].max() <= camera_cfg_zero.max_distance) + self.assertTrue(camera_zero.data.output["distance_to_image_plane"].max() <= camera_cfg_zero.max_distance) + + # max clipping should result in max values + self.assertTrue( + torch.all( + camera_max.data.output["distance_to_camera"][torch.isinf(camera_none.data.output["distance_to_camera"])] + == camera_cfg_zero.max_distance + ) + ) + self.assertTrue( + torch.all( + camera_max.data.output["distance_to_image_plane"][ + torch.isnan(camera_none.data.output["distance_to_image_plane"]) + ] + == camera_cfg_zero.max_distance + ) + ) + self.assertTrue(camera_max.data.output["distance_to_camera"].max() <= camera_cfg_zero.max_distance) + self.assertTrue(camera_max.data.output["distance_to_image_plane"].max() <= camera_cfg_zero.max_distance) + def test_camera_init_offset(self): """Test camera initialization with offset using different conventions.""" # define the same offset in all conventions @@ -289,7 +390,7 @@ def test_multi_camera_init(self): cam_2.update(self.dt) # check image data for cam in [cam_1, cam_2]: - for im_data in cam.data.output.to_dict().values(): + for im_data in cam.data.output.values(): self.assertEqual( im_data.shape, (1, self.camera_cfg.pattern_cfg.height, self.camera_cfg.pattern_cfg.width, 1) ) @@ -392,7 +493,7 @@ def test_throughput(self): with Timer(f"Time taken for writing data with shape {camera.image_shape} "): # Pack data back into replicator format to save them using its writer rep_output = {"annotators": {}} - camera_data = convert_dict_to_backend(camera.data.output[0].to_dict(), backend="numpy") + camera_data = convert_dict_to_backend({k: v[0] for k, v in camera.data.output.items()}, backend="numpy") for key, data, info in zip(camera_data.keys(), camera_data.values(), camera.data.info[0].values()): if info is not None: rep_output["annotators"][key] = {"render_product": {"data": data, **info}} @@ -817,6 +918,15 @@ def test_output_equal_to_usd_camera_when_intrinsics_set(self): atol=1e-4, ) + def test_sensor_print(self): + """Test sensor print is working correctly.""" + # Create sensor + sensor = RayCasterCamera(cfg=self.camera_cfg) + # Play sim + self.sim.reset() + # print info + print(sensor) + if __name__ == "__main__": run_tests() diff --git a/source/extensions/omni.isaac.lab/test/sensors/test_tiled_camera.py b/source/extensions/omni.isaac.lab/test/sensors/test_tiled_camera.py index 269cb4bd81..38b79718b7 100644 --- a/source/extensions/omni.isaac.lab/test/sensors/test_tiled_camera.py +++ b/source/extensions/omni.isaac.lab/test/sensors/test_tiled_camera.py @@ -111,7 +111,7 @@ def test_single_camera_init(self): # update camera camera.update(self.dt) # check image data - for im_type, im_data in camera.data.output.to_dict().items(): + for im_type, im_data in camera.data.output.items(): if im_type == "rgb": self.assertEqual(im_data.shape, (1, self.camera_cfg.height, self.camera_cfg.width, 3)) self.assertGreater((im_data / 255.0).mean().item(), 0.0) @@ -120,6 +120,117 @@ def test_single_camera_init(self): self.assertGreater(im_data.mean().item(), 0.0) del camera + def test_depth_clipping_max(self): + """Test depth max clipping.""" + # get camera cfgs + camera_cfg = TiledCameraCfg( + prim_path="/World/Camera", + offset=TiledCameraCfg.OffsetCfg(pos=(2.5, 2.5, 6.0), rot=(-0.125, 0.362, 0.873, -0.302), convention="ros"), + spawn=sim_utils.PinholeCameraCfg().from_intrinsic_matrix( + focal_length=38.0, + intrinsic_matrix=[380.08, 0.0, 467.79, 0.0, 380.08, 262.05, 0.0, 0.0, 1.0], + height=540, + width=960, + clipping_range=(4.9, 5.0), + ), + height=540, + width=960, + data_types=["depth"], + depth_clipping_behavior="max", + ) + camera = TiledCamera(camera_cfg) + + # Play sim + self.sim.reset() + + # note: This is a workaround to ensure that the textures are loaded. + # Check "Known Issues" section in the documentation for more details. + for _ in range(5): + self.sim.step() + + camera.update(self.dt) + + self.assertTrue(len(camera.data.output["depth"][torch.isinf(camera.data.output["depth"])]) == 0) + self.assertTrue(camera.data.output["depth"].min() >= camera_cfg.spawn.clipping_range[0]) + self.assertTrue(camera.data.output["depth"].max() <= camera_cfg.spawn.clipping_range[1]) + + del camera + + def test_depth_clipping_none(self): + """Test depth none clipping.""" + # get camera cfgs + camera_cfg = TiledCameraCfg( + prim_path="/World/Camera", + offset=TiledCameraCfg.OffsetCfg(pos=(2.5, 2.5, 6.0), rot=(-0.125, 0.362, 0.873, -0.302), convention="ros"), + spawn=sim_utils.PinholeCameraCfg().from_intrinsic_matrix( + focal_length=38.0, + intrinsic_matrix=[380.08, 0.0, 467.79, 0.0, 380.08, 262.05, 0.0, 0.0, 1.0], + height=540, + width=960, + clipping_range=(4.9, 5.0), + ), + height=540, + width=960, + data_types=["depth"], + depth_clipping_behavior="none", + ) + camera = TiledCamera(camera_cfg) + + # Play sim + self.sim.reset() + + # note: This is a workaround to ensure that the textures are loaded. + # Check "Known Issues" section in the documentation for more details. + for _ in range(5): + self.sim.step() + + camera.update(self.dt) + + self.assertTrue(len(camera.data.output["depth"][torch.isinf(camera.data.output["depth"])]) > 0) + self.assertTrue(camera.data.output["depth"].min() >= camera_cfg.spawn.clipping_range[0]) + self.assertTrue( + camera.data.output["depth"][~torch.isinf(camera.data.output["depth"])].max() + <= camera_cfg.spawn.clipping_range[1] + ) + + del camera + + def test_depth_clipping_zero(self): + """Test depth zero clipping.""" + # get camera cfgs + camera_cfg = TiledCameraCfg( + prim_path="/World/Camera", + offset=TiledCameraCfg.OffsetCfg(pos=(2.5, 2.5, 6.0), rot=(-0.125, 0.362, 0.873, -0.302), convention="ros"), + spawn=sim_utils.PinholeCameraCfg().from_intrinsic_matrix( + focal_length=38.0, + intrinsic_matrix=[380.08, 0.0, 467.79, 0.0, 380.08, 262.05, 0.0, 0.0, 1.0], + height=540, + width=960, + clipping_range=(4.9, 5.0), + ), + height=540, + width=960, + data_types=["depth"], + depth_clipping_behavior="zero", + ) + camera = TiledCamera(camera_cfg) + + # Play sim + self.sim.reset() + + # note: This is a workaround to ensure that the textures are loaded. + # Check "Known Issues" section in the documentation for more details. + for _ in range(5): + self.sim.step() + + camera.update(self.dt) + + self.assertTrue(len(camera.data.output["depth"][torch.isinf(camera.data.output["depth"])]) == 0) + self.assertTrue(camera.data.output["depth"].min() == 0.0) + self.assertTrue(camera.data.output["depth"].max() <= camera_cfg.spawn.clipping_range[1]) + + del camera + def test_multi_camera_init(self): """Test multi-camera initialization.""" @@ -162,7 +273,7 @@ def test_multi_camera_init(self): # update camera camera.update(self.dt) # check image data - for im_type, im_data in camera.data.output.to_dict().items(): + for im_type, im_data in camera.data.output.items(): if im_type == "rgb": self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 3)) for i in range(4): @@ -347,7 +458,7 @@ def test_rgba_only_camera(self): # update camera camera.update(self.dt) # check image data - for _, im_data in camera.data.output.to_dict().items(): + for _, im_data in camera.data.output.items(): self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 4)) for i in range(4): self.assertGreater((im_data[i] / 255.0).mean().item(), 0.0) @@ -399,7 +510,7 @@ def test_distance_to_camera_only_camera(self): # update camera camera.update(self.dt) # check image data - for _, im_data in camera.data.output.to_dict().items(): + for _, im_data in camera.data.output.items(): self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 1)) for i in range(4): self.assertGreater((im_data[i]).mean().item(), 0.0) @@ -451,7 +562,7 @@ def test_distance_to_image_plane_only_camera(self): # update camera camera.update(self.dt) # check image data - for _, im_data in camera.data.output.to_dict().items(): + for _, im_data in camera.data.output.items(): self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 1)) for i in range(4): self.assertGreater((im_data[i]).mean().item(), 0.0) @@ -503,7 +614,7 @@ def test_normals_only_camera(self): # update camera camera.update(self.dt) # check image data - for _, im_data in camera.data.output.to_dict().items(): + for _, im_data in camera.data.output.items(): self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 3)) for i in range(4): self.assertGreater((im_data[i]).mean().item(), 0.0) @@ -555,7 +666,7 @@ def test_motion_vectors_only_camera(self): # update camera camera.update(self.dt) # check image data - for _, im_data in camera.data.output.to_dict().items(): + for _, im_data in camera.data.output.items(): self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 2)) for i in range(4): self.assertGreater((im_data[i]).mean().item(), 0.0) @@ -607,7 +718,7 @@ def test_semantic_segmentation_colorize_only_camera(self): # update camera camera.update(self.dt) # check image data - for _, im_data in camera.data.output.to_dict().items(): + for _, im_data in camera.data.output.items(): self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 4)) for i in range(4): self.assertGreater((im_data[i] / 255.0).mean().item(), 0.0) @@ -660,7 +771,7 @@ def test_instance_segmentation_fast_colorize_only_camera(self): # update camera camera.update(self.dt) # check image data - for _, im_data in camera.data.output.to_dict().items(): + for _, im_data in camera.data.output.items(): self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 4)) for i in range(num_cameras): self.assertGreater((im_data[i] / 255.0).mean().item(), 0.0) @@ -713,7 +824,7 @@ def test_instance_id_segmentation_fast_colorize_only_camera(self): # update camera camera.update(self.dt) # check image data - for _, im_data in camera.data.output.to_dict().items(): + for _, im_data in camera.data.output.items(): self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 4)) for i in range(num_cameras): self.assertGreater((im_data[i] / 255.0).mean().item(), 0.0) @@ -767,7 +878,7 @@ def test_semantic_segmentation_non_colorize_only_camera(self): # update camera camera.update(self.dt) # check image data - for _, im_data in camera.data.output.to_dict().items(): + for _, im_data in camera.data.output.items(): self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 1)) for i in range(num_cameras): self.assertGreater(im_data[i].to(dtype=float).mean().item(), 0.0) @@ -822,7 +933,7 @@ def test_instance_segmentation_fast_non_colorize_only_camera(self): # update camera camera.update(self.dt) # check image data - for _, im_data in camera.data.output.to_dict().items(): + for _, im_data in camera.data.output.items(): self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 1)) for i in range(num_cameras): self.assertGreater(im_data[i].to(dtype=float).mean().item(), 0.0) @@ -876,7 +987,7 @@ def test_instance_id_segmentation_fast_non_colorize_only_camera(self): # update camera camera.update(self.dt) # check image data - for _, im_data in camera.data.output.to_dict().items(): + for _, im_data in camera.data.output.items(): self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 1)) for i in range(num_cameras): self.assertGreater(im_data[i].to(dtype=float).mean().item(), 0.0) @@ -941,7 +1052,7 @@ def test_all_annotators_camera(self): # update camera camera.update(self.dt) # check image data - for data_type, im_data in camera.data.output.to_dict().items(): + for data_type, im_data in camera.data.output.items(): if data_type in ["rgb", "normals"]: self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 3)) elif data_type in [ @@ -1039,7 +1150,7 @@ def test_all_annotators_low_resolution_camera(self): # update camera camera.update(self.dt) # check image data - for data_type, im_data in camera.data.output.to_dict().items(): + for data_type, im_data in camera.data.output.items(): if data_type in ["rgb", "normals"]: self.assertEqual(im_data.shape, (num_cameras, camera_cfg.height, camera_cfg.width, 3)) elif data_type in [ @@ -1135,7 +1246,7 @@ def test_all_annotators_non_perfect_square_number_camera(self): # update camera camera.update(self.dt) # check image data - for data_type, im_data in camera.data.output.to_dict().items(): + for data_type, im_data in camera.data.output.items(): if data_type in ["rgb", "normals"]: self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 3)) elif data_type in [ @@ -1201,7 +1312,7 @@ def test_throughput(self): with Timer(f"Time taken for updating camera with shape {camera.image_shape}"): camera.update(self.dt) # Check image data - for im_type, im_data in camera.data.output.to_dict().items(): + for im_type, im_data in camera.data.output.items(): if im_type == "rgb": self.assertEqual(im_data.shape, (1, camera_cfg.height, camera_cfg.width, 3)) self.assertGreater((im_data / 255.0).mean().item(), 0.0) @@ -1304,6 +1415,15 @@ def test_output_equal_to_usd_camera_intrinsics(self): del camera_tiled del camera_usd + def test_sensor_print(self): + """Test sensor print is working correctly.""" + # Create sensor + sensor = TiledCamera(cfg=self.camera_cfg) + # Play sim + self.sim.reset() + # print info + print(sensor) + """ Helper functions. """ diff --git a/source/extensions/omni.isaac.lab/test/sensors/urdfs/simple_2_link.urdf b/source/extensions/omni.isaac.lab/test/sensors/urdfs/simple_2_link.urdf new file mode 100644 index 0000000000..7c09e1b82c --- /dev/null +++ b/source/extensions/omni.isaac.lab/test/sensors/urdfs/simple_2_link.urdf @@ -0,0 +1,82 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/source/extensions/omni.isaac.lab/test/sim/test_mesh_converter.py b/source/extensions/omni.isaac.lab/test/sim/test_mesh_converter.py index 4d67db879c..a77e0ad8f0 100644 --- a/source/extensions/omni.isaac.lab/test/sim/test_mesh_converter.py +++ b/source/extensions/omni.isaac.lab/test/sim/test_mesh_converter.py @@ -12,7 +12,9 @@ """Rest everything follows.""" +import math import os +import random import tempfile import unittest @@ -27,6 +29,16 @@ from omni.isaac.lab.utils.assets import ISAACLAB_NUCLEUS_DIR, retrieve_file_path +def random_quaternion(): + # Generate four random numbers for the quaternion + u1, u2, u3 = random.random(), random.random(), random.random() + w = math.sqrt(1 - u1) * math.sin(2 * math.pi * u2) + x = math.sqrt(1 - u1) * math.cos(2 * math.pi * u2) + y = math.sqrt(u1) * math.sin(2 * math.pi * u3) + z = math.sqrt(u1) * math.cos(2 * math.pi * u3) + return (w, x, y, z) + + class TestMeshConverter(unittest.TestCase): """Test fixture for the MeshConverter class.""" @@ -105,7 +117,12 @@ def test_config_change(self): def test_convert_obj(self): """Convert an OBJ file""" - mesh_config = MeshConverterCfg(asset_path=self.assets["obj"]) + mesh_config = MeshConverterCfg( + asset_path=self.assets["obj"], + scale=(random.uniform(0.1, 2.0), random.uniform(0.1, 2.0), random.uniform(0.1, 2.0)), + translation=(random.uniform(-10.0, 10.0), random.uniform(-10.0, 10.0), random.uniform(-10.0, 10.0)), + rotation=random_quaternion(), + ) mesh_converter = MeshConverter(mesh_config) # check that mesh conversion is successful @@ -113,7 +130,12 @@ def test_convert_obj(self): def test_convert_stl(self): """Convert an STL file""" - mesh_config = MeshConverterCfg(asset_path=self.assets["stl"]) + mesh_config = MeshConverterCfg( + asset_path=self.assets["stl"], + scale=(random.uniform(0.1, 2.0), random.uniform(0.1, 2.0), random.uniform(0.1, 2.0)), + translation=(random.uniform(-10.0, 10.0), random.uniform(-10.0, 10.0), random.uniform(-10.0, 10.0)), + rotation=random_quaternion(), + ) mesh_converter = MeshConverter(mesh_config) # check that mesh conversion is successful @@ -121,12 +143,24 @@ def test_convert_stl(self): def test_convert_fbx(self): """Convert an FBX file""" - mesh_config = MeshConverterCfg(asset_path=self.assets["fbx"]) + mesh_config = MeshConverterCfg( + asset_path=self.assets["fbx"], + scale=(random.uniform(0.1, 2.0), random.uniform(0.1, 2.0), random.uniform(0.1, 2.0)), + translation=(random.uniform(-10.0, 10.0), random.uniform(-10.0, 10.0), random.uniform(-10.0, 10.0)), + rotation=random_quaternion(), + ) mesh_converter = MeshConverter(mesh_config) # check that mesh conversion is successful self._check_mesh_conversion(mesh_converter) + def test_convert_default_xform_transforms(self): + """Convert an OBJ file and check that default xform transforms are applied correctly""" + mesh_config = MeshConverterCfg(asset_path=self.assets["obj"]) + mesh_converter = MeshConverter(mesh_config) + # check that mesh conversion is successful + self._check_mesh_conversion(mesh_converter) + def test_collider_no_approximation(self): """Convert an OBJ file using no approximation""" collision_props = schemas_cfg.CollisionPropertiesCfg(collision_enabled=True) @@ -229,6 +263,15 @@ def _check_mesh_conversion(self, mesh_converter: MeshConverter): units = UsdGeom.GetStageMetersPerUnit(stage) self.assertEqual(units, 1.0) + # Check mesh settings + pos = tuple(prim_utils.get_prim_at_path("/World/Object/geometry").GetAttribute("xformOp:translate").Get()) + self.assertEqual(pos, mesh_converter.cfg.translation) + quat = prim_utils.get_prim_at_path("/World/Object/geometry").GetAttribute("xformOp:orient").Get() + quat = (quat.GetReal(), quat.GetImaginary()[0], quat.GetImaginary()[1], quat.GetImaginary()[2]) + self.assertEqual(quat, mesh_converter.cfg.rotation) + scale = tuple(prim_utils.get_prim_at_path("/World/Object/geometry").GetAttribute("xformOp:scale").Get()) + self.assertEqual(scale, mesh_converter.cfg.scale) + def _check_mesh_collider_settings(self, mesh_converter: MeshConverter): # Check prim can be properly spawned prim_path = "/World/Object" diff --git a/source/extensions/omni.isaac.lab/test/sim/test_simulation_context.py b/source/extensions/omni.isaac.lab/test/sim/test_simulation_context.py index 4a4f308a8f..7829da9c1c 100644 --- a/source/extensions/omni.isaac.lab/test/sim/test_simulation_context.py +++ b/source/extensions/omni.isaac.lab/test/sim/test_simulation_context.py @@ -12,7 +12,6 @@ """Rest everything follows.""" -import ctypes import numpy as np import unittest @@ -94,33 +93,33 @@ def test_headless_mode(self): # check default render mode self.assertEqual(sim.render_mode, sim.RenderMode.NO_GUI_OR_RENDERING) - def test_boundedness(self): - """Test that the boundedness of the simulation context remains constant. - - Note: This test fails right now because Isaac Sim does not handle boundedness correctly. On creation, - it is registering itself to various callbacks and hence the boundedness is more than 1. This may not be - critical for the simulation context since we usually call various clear functions before deleting the - simulation context. - """ - sim = SimulationContext() - # manually set the boundedness to 1? -- this is not possible because of Isaac Sim. - sim.clear_all_callbacks() - sim._stage_open_callback = None - sim._physics_timer_callback = None - sim._event_timer_callback = None - - # check that boundedness of simulation context is correct - sim_ref_count = ctypes.c_long.from_address(id(sim)).value - # reset the simulation - sim.reset() - self.assertEqual(ctypes.c_long.from_address(id(sim)).value, sim_ref_count) - # step the simulation - for _ in range(10): - sim.step() - self.assertEqual(ctypes.c_long.from_address(id(sim)).value, sim_ref_count) - # clear the simulation - sim.clear_instance() - self.assertEqual(ctypes.c_long.from_address(id(sim)).value, sim_ref_count - 1) + # def test_boundedness(self): + # """Test that the boundedness of the simulation context remains constant. + + # Note: This test fails right now because Isaac Sim does not handle boundedness correctly. On creation, + # it is registering itself to various callbacks and hence the boundedness is more than 1. This may not be + # critical for the simulation context since we usually call various clear functions before deleting the + # simulation context. + # """ + # sim = SimulationContext() + # # manually set the boundedness to 1? -- this is not possible because of Isaac Sim. + # sim.clear_all_callbacks() + # sim._stage_open_callback = None + # sim._physics_timer_callback = None + # sim._event_timer_callback = None + + # # check that boundedness of simulation context is correct + # sim_ref_count = ctypes.c_long.from_address(id(sim)).value + # # reset the simulation + # sim.reset() + # self.assertEqual(ctypes.c_long.from_address(id(sim)).value, sim_ref_count) + # # step the simulation + # for _ in range(10): + # sim.step() + # self.assertEqual(ctypes.c_long.from_address(id(sim)).value, sim_ref_count) + # # clear the simulation + # sim.clear_instance() + # self.assertEqual(ctypes.c_long.from_address(id(sim)).value, sim_ref_count - 1) def test_zero_gravity(self): """Test that gravity can be properly disabled.""" diff --git a/source/extensions/omni.isaac.lab/test/sim/test_simulation_render_config.py b/source/extensions/omni.isaac.lab/test/sim/test_simulation_render_config.py new file mode 100644 index 0000000000..cd885a75a0 --- /dev/null +++ b/source/extensions/omni.isaac.lab/test/sim/test_simulation_render_config.py @@ -0,0 +1,154 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + + +"""Launch Isaac Sim Simulator first.""" + +from omni.isaac.lab.app import AppLauncher, run_tests + +# launch omniverse app +app_launcher = AppLauncher(headless=True, enable_cameras=True) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import unittest + +import carb + +from omni.isaac.lab.sim.simulation_cfg import RenderCfg, SimulationCfg +from omni.isaac.lab.sim.simulation_context import SimulationContext + + +class TestSimulationRenderConfig(unittest.TestCase): + """Tests for simulation context render config.""" + + """ + Tests + """ + + def test_render_cfg(self): + """Test that the simulation context is created with the correct render cfg.""" + enable_translucency = True + enable_reflections = True + enable_global_illumination = True + antialiasing_mode = "DLAA" + enable_dlssg = True + dlss_mode = 3 + enable_direct_lighting = True + samples_per_pixel = 4 + enable_shadows = True + enable_ambient_occlusion = True + + render_cfg = RenderCfg( + enable_translucency=enable_translucency, + enable_reflections=enable_reflections, + enable_global_illumination=enable_global_illumination, + antialiasing_mode=antialiasing_mode, + enable_dlssg=enable_dlssg, + dlss_mode=dlss_mode, + enable_direct_lighting=enable_direct_lighting, + samples_per_pixel=samples_per_pixel, + enable_shadows=enable_shadows, + enable_ambient_occlusion=enable_ambient_occlusion, + ) + + cfg = SimulationCfg(render=render_cfg) + + sim = SimulationContext(cfg) + + self.assertEqual(sim.cfg.render.enable_translucency, enable_translucency) + self.assertEqual(sim.cfg.render.enable_reflections, enable_reflections) + self.assertEqual(sim.cfg.render.enable_global_illumination, enable_global_illumination) + self.assertEqual(sim.cfg.render.antialiasing_mode, antialiasing_mode) + self.assertEqual(sim.cfg.render.enable_dlssg, enable_dlssg) + self.assertEqual(sim.cfg.render.dlss_mode, dlss_mode) + self.assertEqual(sim.cfg.render.enable_direct_lighting, enable_direct_lighting) + self.assertEqual(sim.cfg.render.samples_per_pixel, samples_per_pixel) + self.assertEqual(sim.cfg.render.enable_shadows, enable_shadows) + self.assertEqual(sim.cfg.render.enable_ambient_occlusion, enable_ambient_occlusion) + + carb_settings_iface = carb.settings.get_settings() + self.assertEqual(carb_settings_iface.get("/rtx/translucency/enabled"), sim.cfg.render.enable_translucency) + self.assertEqual(carb_settings_iface.get("/rtx/reflections/enabled"), sim.cfg.render.enable_reflections) + self.assertEqual( + carb_settings_iface.get("/rtx/indirectDiffuse/enabled"), sim.cfg.render.enable_global_illumination + ) + self.assertEqual(carb_settings_iface.get("/rtx/transient/dlssg/enabled"), sim.cfg.render.enable_dlssg) + self.assertEqual(carb_settings_iface.get("/rtx/post/dlss/execMode"), sim.cfg.render.dlss_mode) + self.assertEqual(carb_settings_iface.get("/rtx/directLighting/enabled"), sim.cfg.render.enable_direct_lighting) + self.assertEqual( + carb_settings_iface.get("/rtx/directLighting/sampledLighting/samplesPerPixel"), + sim.cfg.render.samples_per_pixel, + ) + self.assertEqual(carb_settings_iface.get("/rtx/shadows/enabled"), sim.cfg.render.enable_shadows) + self.assertEqual( + carb_settings_iface.get("/rtx/ambientOcclusion/enabled"), sim.cfg.render.enable_ambient_occlusion + ) + self.assertEqual(carb_settings_iface.get("/rtx/post/aa/op"), 4) # dlss = 3, dlaa=4 + + def test_render_cfg_defaults(self): + """Test that the simulation context is created with the correct render cfg.""" + enable_translucency = False + enable_reflections = False + enable_global_illumination = False + antialiasing_mode = "DLSS" + enable_dlssg = False + dlss_mode = 0 + enable_direct_lighting = False + samples_per_pixel = 1 + enable_shadows = False + enable_ambient_occlusion = False + + render_cfg = RenderCfg( + enable_translucency=enable_translucency, + enable_reflections=enable_reflections, + enable_global_illumination=enable_global_illumination, + antialiasing_mode=antialiasing_mode, + enable_dlssg=enable_dlssg, + dlss_mode=dlss_mode, + enable_direct_lighting=enable_direct_lighting, + samples_per_pixel=samples_per_pixel, + enable_shadows=enable_shadows, + enable_ambient_occlusion=enable_ambient_occlusion, + ) + + cfg = SimulationCfg(render=render_cfg) + + sim = SimulationContext(cfg) + + self.assertEqual(sim.cfg.render.enable_translucency, enable_translucency) + self.assertEqual(sim.cfg.render.enable_reflections, enable_reflections) + self.assertEqual(sim.cfg.render.enable_global_illumination, enable_global_illumination) + self.assertEqual(sim.cfg.render.antialiasing_mode, antialiasing_mode) + self.assertEqual(sim.cfg.render.enable_dlssg, enable_dlssg) + self.assertEqual(sim.cfg.render.dlss_mode, dlss_mode) + self.assertEqual(sim.cfg.render.enable_direct_lighting, enable_direct_lighting) + self.assertEqual(sim.cfg.render.samples_per_pixel, samples_per_pixel) + self.assertEqual(sim.cfg.render.enable_shadows, enable_shadows) + self.assertEqual(sim.cfg.render.enable_ambient_occlusion, enable_ambient_occlusion) + + carb_settings_iface = carb.settings.get_settings() + self.assertEqual(carb_settings_iface.get("/rtx/translucency/enabled"), sim.cfg.render.enable_translucency) + self.assertEqual(carb_settings_iface.get("/rtx/reflections/enabled"), sim.cfg.render.enable_reflections) + self.assertEqual( + carb_settings_iface.get("/rtx/indirectDiffuse/enabled"), sim.cfg.render.enable_global_illumination + ) + self.assertEqual(carb_settings_iface.get("/rtx/transient/dlssg/enabled"), sim.cfg.render.enable_dlssg) + self.assertEqual(carb_settings_iface.get("/rtx/post/dlss/execMode"), sim.cfg.render.dlss_mode) + self.assertEqual(carb_settings_iface.get("/rtx/directLighting/enabled"), sim.cfg.render.enable_direct_lighting) + self.assertEqual( + carb_settings_iface.get("/rtx/directLighting/sampledLighting/samplesPerPixel"), + sim.cfg.render.samples_per_pixel, + ) + self.assertEqual(carb_settings_iface.get("/rtx/shadows/enabled"), sim.cfg.render.enable_shadows) + self.assertEqual( + carb_settings_iface.get("/rtx/ambientOcclusion/enabled"), sim.cfg.render.enable_ambient_occlusion + ) + self.assertEqual(carb_settings_iface.get("/rtx/post/aa/op"), 3) # dlss = 3, dlaa=4 + + +if __name__ == "__main__": + run_tests() diff --git a/source/extensions/omni.isaac.lab/test/sim/test_spawn_sensors.py b/source/extensions/omni.isaac.lab/test/sim/test_spawn_sensors.py index 328d69367c..8e9f9a556c 100644 --- a/source/extensions/omni.isaac.lab/test/sim/test_spawn_sensors.py +++ b/source/extensions/omni.isaac.lab/test/sim/test_spawn_sensors.py @@ -69,7 +69,7 @@ def test_spawn_pinhole_camera(self): def test_spawn_fisheye_camera(self): """Test spawning a fisheye camera.""" cfg = sim_utils.FisheyeCameraCfg( - projection_type="fisheye_equidistant", + projection_type="fisheyePolynomial", focal_length=5.0, f_stop=10.0, clipping_range=(0.1, 1000.0), diff --git a/source/extensions/omni.isaac.lab/test/utils/test_circular_buffer.py b/source/extensions/omni.isaac.lab/test/utils/test_circular_buffer.py index 8286283a4c..57d75fc341 100644 --- a/source/extensions/omni.isaac.lab/test/utils/test_circular_buffer.py +++ b/source/extensions/omni.isaac.lab/test/utils/test_circular_buffer.py @@ -46,9 +46,31 @@ def test_reset(self): # reset the buffer self.buffer.reset() - # check if the buffer is empty + # check if the buffer has zeros entries self.assertEqual(self.buffer.current_length.tolist(), [0, 0, 0]) + def test_reset_subset(self): + """Test resetting a subset of batches in the circular buffer.""" + data1 = torch.ones((self.batch_size, 2), device=self.device) + data2 = 2.0 * data1.clone() + data3 = 3.0 * data1.clone() + self.buffer.append(data1) + self.buffer.append(data2) + # reset the buffer + reset_batch_id = 1 + self.buffer.reset(batch_ids=[reset_batch_id]) + # check that correct batch is reset + self.assertEqual(self.buffer.current_length.tolist()[reset_batch_id], 0) + # Append new set of data + self.buffer.append(data3) + # check if the correct number of entries are in each batch + expected_length = [3, 3, 3] + expected_length[reset_batch_id] = 1 + self.assertEqual(self.buffer.current_length.tolist(), expected_length) + # check that all entries of the recently reset and appended batch are equal + for i in range(self.max_len): + torch.testing.assert_close(self.buffer.buffer[reset_batch_id, 0], self.buffer.buffer[reset_batch_id, i]) + def test_append_and_retrieve(self): """Test appending and retrieving data from the circular buffer.""" # append some data @@ -121,6 +143,33 @@ def test_key_greater_than_pushes(self): retrieved_data = self.buffer[torch.tensor([5, 5, 5], device=self.device)] self.assertTrue(torch.equal(retrieved_data, data1)) + def test_return_buffer_prop(self): + """Test retrieving the whole buffer for correct size and contents. + Returning the whole buffer should have the shape [batch_size,max_len,data.shape[1:]] + """ + num_overflow = 2 + for i in range(self.buffer.max_length + num_overflow): + data = torch.tensor([[i]], device=self.device).repeat(3, 2) + self.buffer.append(data) + + retrieved_buffer = self.buffer.buffer + # check shape + self.assertTrue(retrieved_buffer.shape == torch.Size([self.buffer.batch_size, self.buffer.max_length, 2])) + # check that batch is first dimension + torch.testing.assert_close(retrieved_buffer[0], retrieved_buffer[1]) + # check oldest + torch.testing.assert_close( + retrieved_buffer[:, 0], torch.tensor([[num_overflow]], device=self.device).repeat(3, 2) + ) + # check most recent + torch.testing.assert_close( + retrieved_buffer[:, -1], + torch.tensor([[self.buffer.max_length + num_overflow - 1]], device=self.device).repeat(3, 2), + ) + # check that it is returned oldest first + for idx in range(self.buffer.max_length - 1): + self.assertTrue(torch.all(torch.le(retrieved_buffer[:, idx], retrieved_buffer[:, idx + 1]))) + if __name__ == "__main__": run_tests() diff --git a/source/extensions/omni.isaac.lab/test/utils/test_configclass.py b/source/extensions/omni.isaac.lab/test/utils/test_configclass.py index 1ee984ce52..bb4b3e5999 100644 --- a/source/extensions/omni.isaac.lab/test/utils/test_configclass.py +++ b/source/extensions/omni.isaac.lab/test/utils/test_configclass.py @@ -19,11 +19,12 @@ import copy import os +import torch import unittest from collections.abc import Callable from dataclasses import MISSING, asdict, field from functools import wraps -from typing import ClassVar +from typing import Any, ClassVar from omni.isaac.lab.utils.configclass import configclass from omni.isaac.lab.utils.dict import class_to_dict, dict_to_md5_hash, update_class_from_dict @@ -85,6 +86,11 @@ def double(x): return 2 * x +@configclass +class ModifierCfg: + params: dict[str, Any] = {"A": 1, "B": 2} + + @configclass class ViewerCfg: eye: list = [7.5, 7.5, 7.5] # field missing on purpose @@ -113,6 +119,7 @@ class BasicDemoCfg: device_id: int = 0 env: EnvCfg = EnvCfg() robot_default_state: RobotDefaultStateCfg = RobotDefaultStateCfg() + list_config = [ModifierCfg(), ModifierCfg(params={"A": 3, "B": 4})] @configclass @@ -128,6 +135,14 @@ def __post_init__(self): self.add_variable = 3 +@configclass +class BasicDemoTorchCfg: + """Dummy configuration class with a torch tensor .""" + + some_number: int = 0 + some_tensor: torch.Tensor = torch.Tensor([1, 2, 3]) + + """ Dummy configuration to check type annotations ordering. """ @@ -329,6 +344,45 @@ class NestedDictAndListCfg: list_1: list[EnvCfg] = [EnvCfg(), EnvCfg()] +""" +Dummy configuration: Missing attributes +""" + + +@configclass +class MissingParentDemoCfg: + """Dummy parent configuration with missing fields.""" + + a: int = MISSING + + @configclass + class InsideClassCfg: + """Inner dummy configuration.""" + + @configclass + class InsideInsideClassCfg: + """Inner inner dummy configuration.""" + + a: str = MISSING + + inside: str = MISSING + inside_dict = {"a": MISSING} + inside_nested_dict = {"a": {"b": "hello", "c": MISSING, "d": InsideInsideClassCfg()}} + inside_tuple = (10, MISSING, 20) + inside_list = [MISSING, MISSING, 2] + + b: InsideClassCfg = InsideClassCfg() + + +@configclass +class MissingChildDemoCfg(MissingParentDemoCfg): + """Dummy child configuration with missing fields.""" + + c: Callable = MISSING + d: int | None = None + e: dict = {} + + """ Test solutions: Basic """ @@ -342,6 +396,7 @@ class NestedDictAndListCfg: "dof_vel": [0.0, 0.0, 0.0, 0.0, 0.0, 1.0], }, "device_id": 0, + "list_config": [{"params": {"A": 1, "B": 2}}, {"params": {"A": 3, "B": 4}}], } basic_demo_cfg_change_correct = { @@ -353,6 +408,7 @@ class NestedDictAndListCfg: "dof_vel": [0.0, 0.0, 0.0, 0.0, 0.0, 1.0], }, "device_id": 0, + "list_config": [{"params": {"A": 1, "B": 2}}, {"params": {"A": 3, "B": 4}}], } basic_demo_cfg_change_with_none_correct = { @@ -364,6 +420,19 @@ class NestedDictAndListCfg: "dof_vel": [0.0, 0.0, 0.0, 0.0, 0.0, 1.0], }, "device_id": 0, + "list_config": [{"params": {"A": 1, "B": 2}}, {"params": {"A": 3, "B": 4}}], +} + +basic_demo_cfg_change_with_tuple_correct = { + "env": {"num_envs": 56, "episode_length": 2000, "viewer": {"eye": [7.5, 7.5, 7.5], "lookat": [0.0, 0.0, 0.0]}}, + "robot_default_state": { + "pos": (0.0, 0.0, 0.0), + "rot": (1.0, 0.0, 0.0, 0.0), + "dof_pos": (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + "dof_vel": [0.0, 0.0, 0.0, 0.0, 0.0, 1.0], + }, + "device_id": 0, + "list_config": [{"params": {"A": -1, "B": -2}}, {"params": {"A": -3, "B": -4}}], } basic_demo_cfg_nested_dict_and_list = { @@ -404,6 +473,22 @@ class NestedDictAndListCfg: "func_in_dict": {"func": "__main__:dummy_function2"}, } +""" +Test solutions: Missing attributes +""" + +validity_expected_fields = [ + "a", + "b.inside", + "b.inside_dict.a", + "b.inside_nested_dict.a.c", + "b.inside_nested_dict.a.d.a", + "b.inside_tuple[1]", + "b.inside_list[0]", + "b.inside_list[1]", + "c", +] + """ Test fixtures. """ @@ -439,9 +524,15 @@ def test_dict_conversion(self): self.assertDictEqual(cfg.to_dict(), basic_demo_cfg_correct) self.assertDictEqual(cfg.env.to_dict(), basic_demo_cfg_correct["env"]) + torch_cfg = BasicDemoTorchCfg() + torch_cfg_dict = torch_cfg.to_dict() + # We have to do a manual check because torch.Tensor does not work with assertDictEqual. + self.assertEqual(torch_cfg_dict["some_number"], 0) + self.assertTrue(torch.all(torch_cfg_dict["some_tensor"] == torch.tensor([1, 2, 3]))) + def test_dict_conversion_order(self): """Tests that order is conserved when converting to dictionary.""" - true_outer_order = ["device_id", "env", "robot_default_state"] + true_outer_order = ["device_id", "env", "robot_default_state", "list_config"] true_env_order = ["num_envs", "episode_length", "viewer"] # create config cfg = BasicDemoCfg() @@ -459,7 +550,7 @@ def test_dict_conversion_order(self): self.assertEqual(label, parsed_value) # check ordering when copied cfg_dict_copied = copy.deepcopy(cfg_dict) - cfg_dict_copied.pop("robot_default_state") + cfg_dict_copied.pop("list_config") # check ordering for label, parsed_value in zip(true_outer_order, cfg_dict_copied.keys()): self.assertEqual(label, parsed_value) @@ -496,6 +587,13 @@ def test_config_update_dict_with_none(self): update_class_from_dict(cfg, cfg_dict) self.assertDictEqual(asdict(cfg), basic_demo_cfg_change_with_none_correct) + def test_config_update_dict_tuple(self): + """Test updating configclass using a dictionary that modifies a tuple.""" + cfg = BasicDemoCfg() + cfg_dict = {"list_config": [{"params": {"A": -1, "B": -2}}, {"params": {"A": -3, "B": -4}}]} + update_class_from_dict(cfg, cfg_dict) + self.assertDictEqual(asdict(cfg), basic_demo_cfg_change_with_tuple_correct) + def test_config_update_nested_dict(self): """Test updating configclass with sub-dictionaries.""" cfg = NestedDictAndListCfg() @@ -888,6 +986,22 @@ def test_config_md5_hash(self): self.assertEqual(md5_hash_1, md5_hash_2) + def test_validity(self): + """Check that invalid configurations raise errors.""" + + cfg = MissingChildDemoCfg() + + with self.assertRaises(TypeError) as context: + cfg.validate() + + # check that the expected missing fields are in the error message + error_message = str(context.exception) + for elem in validity_expected_fields: + self.assertIn(elem, error_message) + + # check that no more than the expected missing fields are in the error message + self.assertEqual(len(error_message.split("\n")) - 2, len(validity_expected_fields)) + if __name__ == "__main__": run_tests() diff --git a/source/extensions/omni.isaac.lab/test/utils/test_episode_data.py b/source/extensions/omni.isaac.lab/test/utils/test_episode_data.py new file mode 100644 index 0000000000..b3cd1849e0 --- /dev/null +++ b/source/extensions/omni.isaac.lab/test/utils/test_episode_data.py @@ -0,0 +1,140 @@ +# Copyright (c) 2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Launch Isaac Sim Simulator first.""" + +from omni.isaac.lab.app import AppLauncher, run_tests + +# launch omniverse app in headless mode +simulation_app = AppLauncher(headless=True).app + +"""Rest everything follows from here.""" + +import torch +import unittest + +from omni.isaac.lab.utils.datasets import EpisodeData + + +class TestEpisodeData(unittest.TestCase): + """Test EpisodeData implementation.""" + + """ + Test cases for EpisodeData class. + """ + + def test_is_empty(self): + """Test checking whether the episode is empty.""" + for device in ("cuda:0", "cpu"): + with self.subTest(device=device): + episode = EpisodeData() + self.assertTrue(episode.is_empty()) + + episode.add("key", torch.tensor([1, 2, 3], device=device)) + self.assertFalse(episode.is_empty()) + + def test_add_tensors(self): + """Test appending tensor data to the episode.""" + for device in ("cuda:0", "cpu"): + with self.subTest(device=device): + dummy_data_0 = torch.tensor([0], device=device) + dummy_data_1 = torch.tensor([1], device=device) + expected_added_data = torch.cat((dummy_data_0.unsqueeze(0), dummy_data_1.unsqueeze(0))) + episode = EpisodeData() + + # test adding data to a key that does not exist + episode.add("key", dummy_data_0) + self.assertTrue(torch.equal(episode.data.get("key"), dummy_data_0.unsqueeze(0))) + + # test adding data to a key that exists + episode.add("key", dummy_data_1) + self.assertTrue(torch.equal(episode.data.get("key"), expected_added_data)) + + # test adding data to a key with "/" in the name + episode.add("first/second", dummy_data_0) + self.assertTrue(torch.equal(episode.data.get("first").get("second"), dummy_data_0.unsqueeze(0))) + + # test adding data to a key with "/" in the name that already exists + episode.add("first/second", dummy_data_1) + self.assertTrue(torch.equal(episode.data.get("first").get("second"), expected_added_data)) + + def test_add_dict_tensors(self): + """Test appending dict data to the episode.""" + for device in ("cuda:0", "cpu"): + with self.subTest(device=device): + dummy_dict_data_0 = { + "key_0": torch.tensor([0], device=device), + "key_1": {"key_1_0": torch.tensor([1], device=device), "key_1_1": torch.tensor([2], device=device)}, + } + dummy_dict_data_1 = { + "key_0": torch.tensor([3], device=device), + "key_1": {"key_1_0": torch.tensor([4], device=device), "key_1_1": torch.tensor([5], device=device)}, + } + + episode = EpisodeData() + + # test adding dict data to a key that does not exist + episode.add("key", dummy_dict_data_0) + self.assertTrue(torch.equal(episode.data.get("key").get("key_0"), torch.tensor([[0]], device=device))) + self.assertTrue( + torch.equal(episode.data.get("key").get("key_1").get("key_1_0"), torch.tensor([[1]], device=device)) + ) + self.assertTrue( + torch.equal(episode.data.get("key").get("key_1").get("key_1_1"), torch.tensor([[2]], device=device)) + ) + + # test adding dict data to a key that exists + episode.add("key", dummy_dict_data_1) + self.assertTrue( + torch.equal(episode.data.get("key").get("key_0"), torch.tensor([[0], [3]], device=device)) + ) + self.assertTrue( + torch.equal( + episode.data.get("key").get("key_1").get("key_1_0"), torch.tensor([[1], [4]], device=device) + ) + ) + self.assertTrue( + torch.equal( + episode.data.get("key").get("key_1").get("key_1_1"), torch.tensor([[2], [5]], device=device) + ) + ) + + def test_get_initial_state(self): + """Test getting the initial state of the episode.""" + for device in ("cuda:0", "cpu"): + with self.subTest(device=device): + dummy_initial_state = torch.tensor([1, 2, 3], device=device) + episode = EpisodeData() + + episode.add("initial_state", dummy_initial_state) + self.assertTrue(torch.equal(episode.get_initial_state(), dummy_initial_state.unsqueeze(0))) + + def test_get_next_action(self): + """Test getting next actions.""" + for device in ("cuda:0", "cpu"): + with self.subTest(device=device): + # dummy actions + action1 = torch.tensor([1, 2, 3], device=device) + action2 = torch.tensor([4, 5, 6], device=device) + action3 = torch.tensor([7, 8, 9], device=device) + + episode = EpisodeData() + self.assertIsNone(episode.get_next_action()) + + episode.add("actions", action1) + episode.add("actions", action2) + episode.add("actions", action3) + + # check if actions are returned in the correct order + self.assertTrue(torch.equal(episode.get_next_action(), action1)) + self.assertTrue(torch.equal(episode.get_next_action(), action2)) + self.assertTrue(torch.equal(episode.get_next_action(), action3)) + + # check if None is returned when all actions are exhausted + self.assertIsNone(episode.get_next_action()) + + +if __name__ == "__main__": + run_tests() diff --git a/source/extensions/omni.isaac.lab/test/utils/test_hdf5_dataset_file_handler.py b/source/extensions/omni.isaac.lab/test/utils/test_hdf5_dataset_file_handler.py new file mode 100644 index 0000000000..64af6a9e0c --- /dev/null +++ b/source/extensions/omni.isaac.lab/test/utils/test_hdf5_dataset_file_handler.py @@ -0,0 +1,129 @@ +# Copyright (c) 2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Launch Isaac Sim Simulator first.""" + +from omni.isaac.lab.app import AppLauncher, run_tests + +# launch omniverse app in headless mode +simulation_app = AppLauncher(headless=True).app + +"""Rest everything follows from here.""" + +import os +import shutil +import tempfile +import torch +import unittest +import uuid + +from omni.isaac.lab.utils.datasets import EpisodeData, HDF5DatasetFileHandler + + +def create_test_episode(device): + """create a test episode with dummy data.""" + test_episode = EpisodeData() + + test_episode.seed = 0 + test_episode.success = True + + test_episode.add("initial_state", torch.tensor([1, 2, 3], device=device)) + + test_episode.add("actions", torch.tensor([1, 2, 3], device=device)) + test_episode.add("actions", torch.tensor([4, 5, 6], device=device)) + test_episode.add("actions", torch.tensor([7, 8, 9], device=device)) + + test_episode.add("obs/policy/term1", torch.tensor([1, 2, 3, 4, 5], device=device)) + test_episode.add("obs/policy/term1", torch.tensor([6, 7, 8, 9, 10], device=device)) + test_episode.add("obs/policy/term1", torch.tensor([11, 12, 13, 14, 15], device=device)) + + return test_episode + + +class TestHDF5DatasetFileHandler(unittest.TestCase): + """Test HDF5 dataset filer handler implementation.""" + + """ + Test cases for HDF5DatasetFileHandler class. + """ + + def setUp(self): + # create a temporary directory to store the test datasets + self.temp_dir = tempfile.mkdtemp() + + def tearDown(self): + # delete the temporary directory after the test + shutil.rmtree(self.temp_dir) + + def test_create_dataset_file(self): + """Test creating a new dataset file.""" + # create a dataset file given a file name with extension + dataset_file_path = os.path.join(self.temp_dir, f"{uuid.uuid4()}.hdf5") + dataset_file_handler = HDF5DatasetFileHandler() + dataset_file_handler.create(dataset_file_path, "test_env_name") + dataset_file_handler.close() + + # check if the dataset is created + self.assertTrue(os.path.exists(dataset_file_path)) + + # create a dataset file given a file name without extension + dataset_file_path = os.path.join(self.temp_dir, f"{uuid.uuid4()}") + dataset_file_handler = HDF5DatasetFileHandler() + dataset_file_handler.create(dataset_file_path, "test_env_name") + dataset_file_handler.close() + + # check if the dataset is created + self.assertTrue(os.path.exists(dataset_file_path + ".hdf5")) + + def test_write_and_load_episode(self): + """Test writing and loading an episode to and from the dataset file.""" + for device in ("cuda:0", "cpu"): + with self.subTest(device=device): + dataset_file_path = os.path.join(self.temp_dir, f"{uuid.uuid4()}.hdf5") + dataset_file_handler = HDF5DatasetFileHandler() + dataset_file_handler.create(dataset_file_path, "test_env_name") + + test_episode = create_test_episode(device) + + # write the episode to the dataset + dataset_file_handler.write_episode(test_episode) + dataset_file_handler.flush() + + self.assertEqual(dataset_file_handler.get_num_episodes(), 1) + + # write the episode again to test writing 2nd episode + dataset_file_handler.write_episode(test_episode) + dataset_file_handler.flush() + + self.assertEqual(dataset_file_handler.get_num_episodes(), 2) + + # close the dataset file to prepare for testing the load function + dataset_file_handler.close() + + # load the episode from the dataset + dataset_file_handler = HDF5DatasetFileHandler() + dataset_file_handler.open(dataset_file_path) + + self.assertEqual(dataset_file_handler.get_env_name(), "test_env_name") + + loaded_episode_names = dataset_file_handler.get_episode_names() + self.assertEqual(len(list(loaded_episode_names)), 2) + + for episode_name in loaded_episode_names: + loaded_episode = dataset_file_handler.load_episode(episode_name, device=device) + self.assertEqual(loaded_episode.env_id, "test_env_name") + self.assertEqual(loaded_episode.seed, test_episode.seed) + self.assertEqual(loaded_episode.success, test_episode.success) + + self.assertTrue(torch.equal(loaded_episode.get_initial_state(), test_episode.get_initial_state())) + + for action in test_episode.data["actions"]: + self.assertTrue(torch.equal(loaded_episode.get_next_action(), action)) + + dataset_file_handler.close() + + +if __name__ == "__main__": + run_tests() diff --git a/source/extensions/omni.isaac.lab/test/utils/test_math.py b/source/extensions/omni.isaac.lab/test/utils/test_math.py index d6d6086290..33a1cb1497 100644 --- a/source/extensions/omni.isaac.lab/test/utils/test_math.py +++ b/source/extensions/omni.isaac.lab/test/utils/test_math.py @@ -195,6 +195,43 @@ def test_quat_error_mag_with_quat_unique(self): torch.testing.assert_close(error_3, error_4) torch.testing.assert_close(error_4, error_1) + def test_convention_converter(self): + """Test convert_camera_frame_orientation_convention to and from ros, opengl, and world conventions.""" + quat_ros = torch.tensor([[-0.17591989, 0.33985114, 0.82047325, -0.42470819]]) + quat_opengl = torch.tensor([[0.33985113, 0.17591988, 0.42470818, 0.82047324]]) + quat_world = torch.tensor([[-0.3647052, -0.27984815, -0.1159169, 0.88047623]]) + + # from ROS + torch.testing.assert_close( + math_utils.convert_camera_frame_orientation_convention(quat_ros, "ros", "opengl"), quat_opengl + ) + torch.testing.assert_close( + math_utils.convert_camera_frame_orientation_convention(quat_ros, "ros", "world"), quat_world + ) + torch.testing.assert_close( + math_utils.convert_camera_frame_orientation_convention(quat_ros, "ros", "ros"), quat_ros + ) + # from OpenGL + torch.testing.assert_close( + math_utils.convert_camera_frame_orientation_convention(quat_opengl, "opengl", "ros"), quat_ros + ) + torch.testing.assert_close( + math_utils.convert_camera_frame_orientation_convention(quat_opengl, "opengl", "world"), quat_world + ) + torch.testing.assert_close( + math_utils.convert_camera_frame_orientation_convention(quat_opengl, "opengl", "opengl"), quat_opengl + ) + # from World + torch.testing.assert_close( + math_utils.convert_camera_frame_orientation_convention(quat_world, "world", "ros"), quat_ros + ) + torch.testing.assert_close( + math_utils.convert_camera_frame_orientation_convention(quat_world, "world", "opengl"), quat_opengl + ) + torch.testing.assert_close( + math_utils.convert_camera_frame_orientation_convention(quat_world, "world", "world"), quat_world + ) + def test_wrap_to_pi(self): """Test wrap_to_pi method.""" # Define test cases diff --git a/source/extensions/omni.isaac.lab_tasks/config/extension.toml b/source/extensions/omni.isaac.lab_tasks/config/extension.toml index 89ca646936..e82e19521f 100644 --- a/source/extensions/omni.isaac.lab_tasks/config/extension.toml +++ b/source/extensions/omni.isaac.lab_tasks/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "0.10.5" +version = "0.10.16" # Description title = "Isaac Lab Environments" diff --git a/source/extensions/omni.isaac.lab_tasks/docs/CHANGELOG.rst b/source/extensions/omni.isaac.lab_tasks/docs/CHANGELOG.rst index 2614630bd7..69a4c04aea 100644 --- a/source/extensions/omni.isaac.lab_tasks/docs/CHANGELOG.rst +++ b/source/extensions/omni.isaac.lab_tasks/docs/CHANGELOG.rst @@ -1,11 +1,115 @@ Changelog --------- +0.10.16 (2024-12-16) +~~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added ``Factory-Direct-v0`` environment as a direct RL env that + implements contact-rich manipulation tasks including peg insertion, + gear meshing, and nut threading. + + +0.10.15 (2024-12-16) +~~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added ``Isaac-Reach-Franka-OSC-v0`` and ``Isaac-Reach-Franka-OSC-Play-v0`` + variations of the manager based reach environment that uses + :class:`omni.isaac.lab.envs.mdp.actions.OperationalSpaceControllerAction`. + + +0.10.14 (2024-12-03) +~~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added ``Isaac-Stack-Cube-Franka-IK-Rel-v0`` and ``Isaac-Stack-Cube-Instance-Randomize-Franka-IK-Rel-v0`` environments + as manager-based RL envs that implement a three cube stacking task. + + +0.10.13 (2024-10-30) +~~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Defined the Gymnasium task entry points with configuration strings instead of class types. + This avoids unnecessary imports and improves the load types. +* Blacklisted ``mdp`` directories during the recursive module search. + + +0.10.12 (2024-10-28) +~~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Changed manager-based vision cartpole environment names from Isaac-Cartpole-RGB-Camera-v0 + and Isaac-Cartpole-Depth-Camera-v0 to Isaac-Cartpole-RGB-v0 and Isaac-Cartpole-Depth-v0 + + +0.10.11 (2024-10-28) +~~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added feature extracted observation cartpole examples. + + +0.10.10 (2024-10-25) +~~~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Fixed issues with defining Gymnasium spaces in Direct workflows due to Hydra/OmegaConf limitations with non-primitive types. + + +0.10.9 (2024-10-22) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Sets curriculum and commands to None in manager-based environment configurations when not needed. + Earlier, this was done by making an empty configuration object, which is now unnecessary. + + +0.10.8 (2024-10-22) +~~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Fixed the wrong selection of body id's in the :meth:`omni.isaac.lab_tasks.manager_based.locomotion.velocity.mdp.rewards.feet_slide` + reward function. This makes sure the right IDs are selected for the bodies. + + +0.10.7 (2024-10-02) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Replace deprecated :attr:`num_observations`, :attr:`num_actions` and :attr:`num_states` in single-agent direct tasks + by :attr:`observation_space`, :attr:`action_space` and :attr:`state_space` respectively. +* Replace deprecated :attr:`num_observations`, :attr:`num_actions` and :attr:`num_states` in multi-agent direct tasks + by :attr:`observation_spaces`, :attr:`action_spaces` and :attr:`state_space` respectively. + + 0.10.6 (2024-09-25) ~~~~~~~~~~~~~~~~~~~ Added ^^^^^ + * Added ``Isaac-Cartpole-RGB-Camera-v0`` and ``Isaac-Cartpole-Depth-Camera-v0`` manager based camera cartpole environments. diff --git a/source/extensions/omni.isaac.lab_tasks/docs/README.md b/source/extensions/omni.isaac.lab_tasks/docs/README.md index fe225ca25a..ce9dfdb77d 100644 --- a/source/extensions/omni.isaac.lab_tasks/docs/README.md +++ b/source/extensions/omni.isaac.lab_tasks/docs/README.md @@ -32,16 +32,26 @@ The environments are then registered in the `omni/isaac/lab_tasks/locomotion/vel ```python gym.register( id="Isaac-Velocity-Rough-Anymal-C-v0", - entry_point="omni.isaac.lab.envs:RLTaskEnv", + entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, - kwargs={"env_cfg_entry_point": f"{__name__}.rough_env_cfg:AnymalCRoughEnvCfg"}, + kwargs={ + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:AnymalCRoughEnvCfg", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_rough_ppo_cfg.yaml", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalCRoughPPORunnerCfg", + "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", + }, ) gym.register( id="Isaac-Velocity-Flat-Anymal-C-v0", - entry_point="omni.isaac.lab.envs:RLTaskEnv", + entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, - kwargs={"env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalCFlatEnvCfg"}, + kwargs={ + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalCFlatEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalCFlatPPORunnerCfg", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_flat_ppo_cfg.yaml", + "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", + }, ) ``` diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/__init__.py index 5dd39864d3..6f06dbcd02 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/__init__.py @@ -25,6 +25,6 @@ from .utils import import_packages # The blacklist is used to prevent importing configs from sub-packages -_BLACKLIST_PKGS = ["utils"] +_BLACKLIST_PKGS = ["utils", ".mdp"] # Import all configs in this package import_packages(__name__, _BLACKLIST_PKGS) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/allegro_hand/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/allegro_hand/__init__.py index 0f4ed4544f..b0734779f5 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/allegro_hand/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/allegro_hand/__init__.py @@ -10,18 +10,19 @@ import gymnasium as gym from . import agents -from .allegro_hand_env_cfg import AllegroHandEnvCfg ## # Register Gym environments. ## +inhand_task_entry = "omni.isaac.lab_tasks.direct.inhand_manipulation" + gym.register( id="Isaac-Repose-Cube-Allegro-Direct-v0", - entry_point="omni.isaac.lab_tasks.direct.inhand_manipulation:InHandManipulationEnv", + entry_point=f"{inhand_task_entry}.inhand_manipulation_env:InHandManipulationEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": AllegroHandEnvCfg, + "env_cfg_entry_point": f"{__name__}.allegro_hand_env_cfg:AllegroHandEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AllegroHandPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/allegro_hand/allegro_hand_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/allegro_hand/allegro_hand_env_cfg.py index b83b6782a6..b5c53a91d3 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/allegro_hand/allegro_hand_env_cfg.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/allegro_hand/allegro_hand_env_cfg.py @@ -22,9 +22,9 @@ class AllegroHandEnvCfg(DirectRLEnvCfg): # env decimation = 4 episode_length_s = 10.0 - num_actions = 16 - num_observations = 124 # (full) - num_states = 0 + action_space = 16 + observation_space = 124 # (full) + state_space = 0 asymmetric_obs = False obs_type = "full" # simulation diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/ant/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/ant/__init__.py index 97a3814e6c..940702a105 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/ant/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/ant/__init__.py @@ -10,7 +10,6 @@ import gymnasium as gym from . import agents -from .ant_env import AntEnv, AntEnvCfg ## # Register Gym environments. @@ -18,10 +17,10 @@ gym.register( id="Isaac-Ant-Direct-v0", - entry_point="omni.isaac.lab_tasks.direct.ant:AntEnv", + entry_point=f"{__name__}.ant_env:AntEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": AntEnvCfg, + "env_cfg_entry_point": f"{__name__}.ant_env:AntEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AntPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/ant/ant_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/ant/ant_env.py index 8bf6d6bcc9..42f57127ee 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/ant/ant_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/ant/ant_env.py @@ -24,9 +24,9 @@ class AntEnvCfg(DirectRLEnvCfg): episode_length_s = 15.0 decimation = 2 action_scale = 0.5 - num_actions = 8 - num_observations = 36 - num_states = 0 + action_space = 8 + observation_space = 36 + state_space = 0 # simulation sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/anymal_c/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/anymal_c/__init__.py index 1a245b7293..8160e044d9 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/anymal_c/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/anymal_c/__init__.py @@ -10,7 +10,6 @@ import gymnasium as gym from . import agents -from .anymal_c_env import AnymalCEnv, AnymalCFlatEnvCfg, AnymalCRoughEnvCfg ## # Register Gym environments. @@ -18,10 +17,10 @@ gym.register( id="Isaac-Velocity-Flat-Anymal-C-Direct-v0", - entry_point="omni.isaac.lab_tasks.direct.anymal_c:AnymalCEnv", + entry_point=f"{__name__}.anymal_c_env:AnymalCEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": AnymalCFlatEnvCfg, + "env_cfg_entry_point": f"{__name__}.anymal_c_env_cfg:AnymalCFlatEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_flat_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalCFlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", @@ -30,10 +29,10 @@ gym.register( id="Isaac-Velocity-Rough-Anymal-C-Direct-v0", - entry_point="omni.isaac.lab_tasks.direct.anymal_c:AnymalCEnv", + entry_point=f"{__name__}.anymal_c_env:AnymalCEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": AnymalCRoughEnvCfg, + "env_cfg_entry_point": f"{__name__}.anymal_c_env_cfg:AnymalCRoughEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_rough_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalCRoughPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/anymal_c/anymal_c_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/anymal_c/anymal_c_env.py index 5490bb0dd3..0d2c9f16fd 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/anymal_c/anymal_c_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/anymal_c/anymal_c_env.py @@ -5,152 +5,15 @@ from __future__ import annotations +import gymnasium as gym import torch -import omni.isaac.lab.envs.mdp as mdp import omni.isaac.lab.sim as sim_utils -from omni.isaac.lab.assets import Articulation, ArticulationCfg -from omni.isaac.lab.envs import DirectRLEnv, DirectRLEnvCfg -from omni.isaac.lab.managers import EventTermCfg as EventTerm -from omni.isaac.lab.managers import SceneEntityCfg -from omni.isaac.lab.scene import InteractiveSceneCfg -from omni.isaac.lab.sensors import ContactSensor, ContactSensorCfg, RayCaster, RayCasterCfg, patterns -from omni.isaac.lab.sim import SimulationCfg -from omni.isaac.lab.terrains import TerrainImporterCfg -from omni.isaac.lab.utils import configclass +from omni.isaac.lab.assets import Articulation +from omni.isaac.lab.envs import DirectRLEnv +from omni.isaac.lab.sensors import ContactSensor, RayCaster -## -# Pre-defined configs -## -from omni.isaac.lab_assets.anymal import ANYMAL_C_CFG # isort: skip -from omni.isaac.lab.terrains.config.rough import ROUGH_TERRAINS_CFG # isort: skip - - -@configclass -class EventCfg: - """Configuration for randomization.""" - - physics_material = EventTerm( - func=mdp.randomize_rigid_body_material, - mode="startup", - params={ - "asset_cfg": SceneEntityCfg("robot", body_names=".*"), - "static_friction_range": (0.8, 0.8), - "dynamic_friction_range": (0.6, 0.6), - "restitution_range": (0.0, 0.0), - "num_buckets": 64, - }, - ) - - add_base_mass = EventTerm( - func=mdp.randomize_rigid_body_mass, - mode="startup", - params={ - "asset_cfg": SceneEntityCfg("robot", body_names="base"), - "mass_distribution_params": (-5.0, 5.0), - "operation": "add", - }, - ) - - -@configclass -class AnymalCFlatEnvCfg(DirectRLEnvCfg): - # env - episode_length_s = 20.0 - decimation = 4 - action_scale = 0.5 - num_actions = 12 - num_observations = 48 - num_states = 0 - - # simulation - sim: SimulationCfg = SimulationCfg( - dt=1 / 200, - render_interval=decimation, - disable_contact_processing=True, - physics_material=sim_utils.RigidBodyMaterialCfg( - friction_combine_mode="multiply", - restitution_combine_mode="multiply", - static_friction=1.0, - dynamic_friction=1.0, - restitution=0.0, - ), - ) - terrain = TerrainImporterCfg( - prim_path="/World/ground", - terrain_type="plane", - 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=0.0, - ), - debug_vis=False, - ) - - # scene - scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=4096, env_spacing=4.0, replicate_physics=True) - - # events - events: EventCfg = EventCfg() - - # robot - robot: ArticulationCfg = ANYMAL_C_CFG.replace(prim_path="/World/envs/env_.*/Robot") - contact_sensor: ContactSensorCfg = ContactSensorCfg( - prim_path="/World/envs/env_.*/Robot/.*", history_length=3, update_period=0.005, track_air_time=True - ) - - # reward scales - lin_vel_reward_scale = 1.0 - yaw_rate_reward_scale = 0.5 - z_vel_reward_scale = -2.0 - ang_vel_reward_scale = -0.05 - joint_torque_reward_scale = -2.5e-5 - joint_accel_reward_scale = -2.5e-7 - action_rate_reward_scale = -0.01 - feet_air_time_reward_scale = 0.5 - undersired_contact_reward_scale = -1.0 - flat_orientation_reward_scale = -5.0 - - -@configclass -class AnymalCRoughEnvCfg(AnymalCFlatEnvCfg): - # env - num_observations = 235 - - terrain = TerrainImporterCfg( - prim_path="/World/ground", - terrain_type="generator", - terrain_generator=ROUGH_TERRAINS_CFG, - max_init_terrain_level=9, - collision_group=-1, - physics_material=sim_utils.RigidBodyMaterialCfg( - friction_combine_mode="multiply", - restitution_combine_mode="multiply", - static_friction=1.0, - dynamic_friction=1.0, - ), - visual_material=sim_utils.MdlFileCfg( - mdl_path="{NVIDIA_NUCLEUS_DIR}/Materials/Base/Architecture/Shingles_01.mdl", - project_uvw=True, - ), - debug_vis=False, - ) - - # we add a height scanner for perceptive locomotion - height_scanner = RayCasterCfg( - prim_path="/World/envs/env_.*/Robot/base", - offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)), - attach_yaw_only=True, - pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=[1.6, 1.0]), - debug_vis=False, - mesh_prim_paths=["/World/ground"], - ) - - # reward scales (override from flat config) - flat_orientation_reward_scale = 0.0 +from .anymal_c_env_cfg import AnymalCFlatEnvCfg, AnymalCRoughEnvCfg class AnymalCEnv(DirectRLEnv): @@ -160,8 +23,10 @@ def __init__(self, cfg: AnymalCFlatEnvCfg | AnymalCRoughEnvCfg, render_mode: str super().__init__(cfg, render_mode, **kwargs) # Joint position command (deviation from default joint positions) - self._actions = torch.zeros(self.num_envs, self.cfg.num_actions, device=self.device) - self._previous_actions = torch.zeros(self.num_envs, self.cfg.num_actions, device=self.device) + self._actions = torch.zeros(self.num_envs, gym.spaces.flatdim(self.single_action_space), device=self.device) + self._previous_actions = torch.zeros( + self.num_envs, gym.spaces.flatdim(self.single_action_space), device=self.device + ) # X/Y linear velocity and yaw angular velocity commands self._commands = torch.zeros(self.num_envs, 3, device=self.device) @@ -224,8 +89,8 @@ def _get_observations(self) -> dict: [ tensor for tensor in ( - self._robot.data.root_lin_vel_b, - self._robot.data.root_ang_vel_b, + self._robot.data.root_com_lin_vel_b, + self._robot.data.root_com_ang_vel_b, self._robot.data.projected_gravity_b, self._commands, self._robot.data.joint_pos - self._robot.data.default_joint_pos, @@ -242,15 +107,17 @@ def _get_observations(self) -> dict: def _get_rewards(self) -> torch.Tensor: # linear velocity tracking - lin_vel_error = torch.sum(torch.square(self._commands[:, :2] - self._robot.data.root_lin_vel_b[:, :2]), dim=1) + lin_vel_error = torch.sum( + torch.square(self._commands[:, :2] - self._robot.data.root_com_lin_vel_b[:, :2]), dim=1 + ) lin_vel_error_mapped = torch.exp(-lin_vel_error / 0.25) # yaw rate tracking - yaw_rate_error = torch.square(self._commands[:, 2] - self._robot.data.root_ang_vel_b[:, 2]) + yaw_rate_error = torch.square(self._commands[:, 2] - self._robot.data.root_com_ang_vel_b[:, 2]) yaw_rate_error_mapped = torch.exp(-yaw_rate_error / 0.25) # z velocity tracking - z_vel_error = torch.square(self._robot.data.root_lin_vel_b[:, 2]) + z_vel_error = torch.square(self._robot.data.root_com_lin_vel_b[:, 2]) # angular velocity x/y - ang_vel_error = torch.sum(torch.square(self._robot.data.root_ang_vel_b[:, :2]), dim=1) + ang_vel_error = torch.sum(torch.square(self._robot.data.root_com_ang_vel_b[:, :2]), dim=1) # joint torques joint_torques = torch.sum(torch.square(self._robot.data.applied_torque), dim=1) # joint acceleration @@ -313,8 +180,8 @@ def _reset_idx(self, env_ids: torch.Tensor | None): joint_vel = self._robot.data.default_joint_vel[env_ids] default_root_state = self._robot.data.default_root_state[env_ids] default_root_state[:, :3] += self._terrain.env_origins[env_ids] - self._robot.write_root_pose_to_sim(default_root_state[:, :7], env_ids) - self._robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) + self._robot.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids) + self._robot.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids) self._robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) # Logging extras = dict() diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/anymal_c/anymal_c_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/anymal_c/anymal_c_env_cfg.py new file mode 100644 index 0000000000..d5276f1404 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/anymal_c/anymal_c_env_cfg.py @@ -0,0 +1,149 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import omni.isaac.lab.envs.mdp as mdp +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.assets import ArticulationCfg +from omni.isaac.lab.envs import DirectRLEnvCfg +from omni.isaac.lab.managers import EventTermCfg as EventTerm +from omni.isaac.lab.managers import SceneEntityCfg +from omni.isaac.lab.scene import InteractiveSceneCfg +from omni.isaac.lab.sensors import ContactSensorCfg, RayCasterCfg, patterns +from omni.isaac.lab.sim import SimulationCfg +from omni.isaac.lab.terrains import TerrainImporterCfg +from omni.isaac.lab.utils import configclass + +## +# Pre-defined configs +## +from omni.isaac.lab_assets.anymal import ANYMAL_C_CFG # isort: skip +from omni.isaac.lab.terrains.config.rough import ROUGH_TERRAINS_CFG # isort: skip + + +@configclass +class EventCfg: + """Configuration for randomization.""" + + physics_material = EventTerm( + func=mdp.randomize_rigid_body_material, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=".*"), + "static_friction_range": (0.8, 0.8), + "dynamic_friction_range": (0.6, 0.6), + "restitution_range": (0.0, 0.0), + "num_buckets": 64, + }, + ) + + add_base_mass = EventTerm( + func=mdp.randomize_rigid_body_mass, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names="base"), + "mass_distribution_params": (-5.0, 5.0), + "operation": "add", + }, + ) + + +@configclass +class AnymalCFlatEnvCfg(DirectRLEnvCfg): + # env + episode_length_s = 20.0 + decimation = 4 + action_scale = 0.5 + action_space = 12 + observation_space = 48 + state_space = 0 + + # simulation + sim: SimulationCfg = SimulationCfg( + dt=1 / 200, + render_interval=decimation, + disable_contact_processing=True, + physics_material=sim_utils.RigidBodyMaterialCfg( + friction_combine_mode="multiply", + restitution_combine_mode="multiply", + static_friction=1.0, + dynamic_friction=1.0, + restitution=0.0, + ), + ) + terrain = TerrainImporterCfg( + prim_path="/World/ground", + terrain_type="plane", + 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=0.0, + ), + debug_vis=False, + ) + + # scene + scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=4096, env_spacing=4.0, replicate_physics=True) + + # events + events: EventCfg = EventCfg() + + # robot + robot: ArticulationCfg = ANYMAL_C_CFG.replace(prim_path="/World/envs/env_.*/Robot") + contact_sensor: ContactSensorCfg = ContactSensorCfg( + prim_path="/World/envs/env_.*/Robot/.*", history_length=3, update_period=0.005, track_air_time=True + ) + + # reward scales + lin_vel_reward_scale = 1.0 + yaw_rate_reward_scale = 0.5 + z_vel_reward_scale = -2.0 + ang_vel_reward_scale = -0.05 + joint_torque_reward_scale = -2.5e-5 + joint_accel_reward_scale = -2.5e-7 + action_rate_reward_scale = -0.01 + feet_air_time_reward_scale = 0.5 + undersired_contact_reward_scale = -1.0 + flat_orientation_reward_scale = -5.0 + + +@configclass +class AnymalCRoughEnvCfg(AnymalCFlatEnvCfg): + # env + observation_space = 235 + + terrain = TerrainImporterCfg( + prim_path="/World/ground", + terrain_type="generator", + terrain_generator=ROUGH_TERRAINS_CFG, + max_init_terrain_level=9, + collision_group=-1, + physics_material=sim_utils.RigidBodyMaterialCfg( + friction_combine_mode="multiply", + restitution_combine_mode="multiply", + static_friction=1.0, + dynamic_friction=1.0, + ), + visual_material=sim_utils.MdlFileCfg( + mdl_path="{NVIDIA_NUCLEUS_DIR}/Materials/Base/Architecture/Shingles_01.mdl", + project_uvw=True, + ), + debug_vis=False, + ) + + # we add a height scanner for perceptive locomotion + height_scanner = RayCasterCfg( + prim_path="/World/envs/env_.*/Robot/base", + offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)), + attach_yaw_only=True, + pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=[1.6, 1.0]), + debug_vis=False, + mesh_prim_paths=["/World/ground"], + ) + + # reward scales (override from flat config) + flat_orientation_reward_scale = 0.0 diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cart_double_pendulum/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cart_double_pendulum/__init__.py index 2b034f4f2e..90d70311d1 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cart_double_pendulum/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cart_double_pendulum/__init__.py @@ -10,7 +10,6 @@ import gymnasium as gym from . import agents -from .cart_double_pendulum_env import CartDoublePendulumEnv, CartDoublePendulumEnvCfg ## # Register Gym environments. @@ -18,10 +17,10 @@ gym.register( id="Isaac-Cart-Double-Pendulum-Direct-v0", - entry_point="omni.isaac.lab_tasks.direct.cart_double_pendulum:CartDoublePendulumEnv", + entry_point=f"{__name__}.cart_double_pendulum_env:CartDoublePendulumEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": CartDoublePendulumEnvCfg, + "env_cfg_entry_point": f"{__name__}.cart_double_pendulum_env:CartDoublePendulumEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", "skrl_ippo_cfg_entry_point": f"{agents.__name__}:skrl_ippo_cfg.yaml", diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py index 0b606fe899..25f9e59c79 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py @@ -27,9 +27,9 @@ class CartDoublePendulumEnvCfg(DirectMARLEnvCfg): decimation = 2 episode_length_s = 5.0 possible_agents = ["cart", "pendulum"] - num_actions = {"cart": 1, "pendulum": 1} - num_observations = {"cart": 4, "pendulum": 3} - num_states = -1 + action_spaces = {"cart": 1, "pendulum": 1} + observation_spaces = {"cart": 4, "pendulum": 3} + state_space = -1 # simulation sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation) @@ -182,8 +182,8 @@ def _reset_idx(self, env_ids: Sequence[int] | None): self.joint_pos[env_ids] = joint_pos self.joint_vel[env_ids] = joint_vel - self.robot.write_root_pose_to_sim(default_root_state[:, :7], env_ids) - self.robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) + self.robot.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids) + self.robot.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids) self.robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/__init__.py index 9f07c14605..035cf7a7a9 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/__init__.py @@ -10,8 +10,6 @@ import gymnasium as gym from . import agents -from .cartpole_camera_env import CartpoleCameraEnv, CartpoleDepthCameraEnvCfg, CartpoleRGBCameraEnvCfg -from .cartpole_env import CartpoleEnv, CartpoleEnvCfg ## # Register Gym environments. @@ -19,10 +17,10 @@ gym.register( id="Isaac-Cartpole-Direct-v0", - entry_point="omni.isaac.lab_tasks.direct.cartpole:CartpoleEnv", + entry_point=f"{__name__}.cartpole_env:CartpoleEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": CartpoleEnvCfg, + "env_cfg_entry_point": f"{__name__}.cartpole_env:CartpoleEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CartpolePPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", @@ -32,10 +30,10 @@ gym.register( id="Isaac-Cartpole-RGB-Camera-Direct-v0", - entry_point="omni.isaac.lab_tasks.direct.cartpole:CartpoleCameraEnv", + entry_point=f"{__name__}.cartpole_camera_env:CartpoleCameraEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": CartpoleRGBCameraEnvCfg, + "env_cfg_entry_point": f"{__name__}.cartpole_camera_env:CartpoleRGBCameraEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_camera_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_camera_ppo_cfg.yaml", }, @@ -43,10 +41,10 @@ gym.register( id="Isaac-Cartpole-Depth-Camera-Direct-v0", - entry_point="omni.isaac.lab_tasks.direct.cartpole:CartpoleCameraEnv", + entry_point=f"{__name__}.cartpole_camera_env:CartpoleCameraEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": CartpoleDepthCameraEnvCfg, + "env_cfg_entry_point": f"{__name__}.cartpole_camera_env:CartpoleDepthCameraEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_camera_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_camera_ppo_cfg.yaml", }, diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/cartpole_camera_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/cartpole_camera_env.py index b2a1b1e303..fea5f5511a 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/cartpole_camera_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/cartpole_camera_env.py @@ -5,9 +5,7 @@ from __future__ import annotations -import gymnasium as gym import math -import numpy as np import torch from collections.abc import Sequence @@ -29,9 +27,6 @@ class CartpoleRGBCameraEnvCfg(DirectRLEnvCfg): decimation = 2 episode_length_s = 5.0 action_scale = 100.0 # [N] - num_actions = 1 - num_channels = 3 - num_states = 0 # simulation sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation) @@ -52,9 +47,13 @@ class CartpoleRGBCameraEnvCfg(DirectRLEnvCfg): width=80, height=80, ) - num_observations = num_channels * tiled_camera.height * tiled_camera.width write_image_to_file = False + # spaces + action_space = 1 + state_space = 0 + observation_space = [tiled_camera.height, tiled_camera.width, 3] + # change viewer settings viewer = ViewerCfg(eye=(20.0, 20.0, 20.0)) @@ -87,9 +86,8 @@ class CartpoleDepthCameraEnvCfg(CartpoleRGBCameraEnvCfg): height=80, ) - # env - num_channels = 1 - num_observations = num_channels * tiled_camera.height * tiled_camera.width + # spaces + observation_space = [tiled_camera.height, tiled_camera.width, 1] class CartpoleCameraEnv(DirectRLEnv): @@ -118,35 +116,6 @@ def close(self): """Cleanup for the environment.""" super().close() - def _configure_gym_env_spaces(self): - """Configure the action and observation spaces for the Gym environment.""" - # observation space (unbounded since we don't impose any limits) - self.num_actions = self.cfg.num_actions - self.num_observations = self.cfg.num_observations - self.num_states = self.cfg.num_states - - # set up spaces - self.single_observation_space = gym.spaces.Dict() - self.single_observation_space["policy"] = gym.spaces.Box( - low=-np.inf, - high=np.inf, - shape=(self.cfg.tiled_camera.height, self.cfg.tiled_camera.width, self.cfg.num_channels), - ) - if self.num_states > 0: - self.single_observation_space["critic"] = gym.spaces.Box( - low=-np.inf, - high=np.inf, - shape=(self.cfg.tiled_camera.height, self.cfg.tiled_camera.width, self.cfg.num_channels), - ) - self.single_action_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.num_actions,)) - - # batch the spaces for vectorized environments - self.observation_space = gym.vector.utils.batch_space(self.single_observation_space, self.num_envs) - self.action_space = gym.vector.utils.batch_space(self.single_action_space, self.num_envs) - - # RL specifics - self.actions = torch.zeros(self.num_envs, self.num_actions, device=self.sim.device) - def _setup_scene(self): """Setup the scene with the cartpole and camera.""" self._cartpole = Articulation(self.cfg.robot_cfg) @@ -156,7 +125,7 @@ def _setup_scene(self): self.scene.clone_environments(copy_from_source=False) self.scene.filter_collisions(global_prim_paths=[]) - # add articultion and sensors to scene + # add articulation and sensors to scene self.scene.articulations["cartpole"] = self._cartpole self.scene.sensors["tiled_camera"] = self._tiled_camera # add lights @@ -230,8 +199,8 @@ def _reset_idx(self, env_ids: Sequence[int] | None): self.joint_pos[env_ids] = joint_pos self.joint_vel[env_ids] = joint_vel - self._cartpole.write_root_pose_to_sim(default_root_state[:, :7], env_ids) - self._cartpole.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) + self._cartpole.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids) + self._cartpole.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids) self._cartpole.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/cartpole_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/cartpole_env.py index 44926e95f9..88b1877d08 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/cartpole_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/cartpole_env.py @@ -27,9 +27,9 @@ class CartpoleEnvCfg(DirectRLEnvCfg): decimation = 2 episode_length_s = 5.0 action_scale = 100.0 # [N] - num_actions = 1 - num_observations = 4 - num_states = 0 + action_space = 1 + observation_space = 4 + state_space = 0 # simulation sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation) @@ -74,7 +74,7 @@ def _setup_scene(self): # clone, filter, and replicate self.scene.clone_environments(copy_from_source=False) self.scene.filter_collisions(global_prim_paths=[]) - # add articultion to scene + # add articulation to scene self.scene.articulations["cartpole"] = self.cartpole # add lights light_cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75)) @@ -143,8 +143,8 @@ def _reset_idx(self, env_ids: Sequence[int] | None): self.joint_pos[env_ids] = joint_pos self.joint_vel[env_ids] = joint_vel - self.cartpole.write_root_pose_to_sim(default_root_state[:, :7], env_ids) - self.cartpole.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) + self.cartpole.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids) + self.cartpole.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids) self.cartpole.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/__init__.py new file mode 100644 index 0000000000..c19c96f708 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/__init__.py @@ -0,0 +1,44 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import gymnasium as gym + +from . import agents +from .factory_env import FactoryEnv +from .factory_env_cfg import FactoryTaskGearMeshCfg, FactoryTaskNutThreadCfg, FactoryTaskPegInsertCfg + +## +# Register Gym environments. +## + +gym.register( + id="Isaac-Factory-PegInsert-Direct-v0", + entry_point="omni.isaac.lab_tasks.direct.factory:FactoryEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": FactoryTaskPegInsertCfg, + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", + }, +) + +gym.register( + id="Isaac-Factory-GearMesh-Direct-v0", + entry_point="omni.isaac.lab_tasks.direct.factory:FactoryEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": FactoryTaskGearMeshCfg, + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", + }, +) + +gym.register( + id="Isaac-Factory-NutThread-Direct-v0", + entry_point="omni.isaac.lab_tasks.direct.factory:FactoryEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": FactoryTaskNutThreadCfg, + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", + }, +) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/agents/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/agents/__init__.py new file mode 100644 index 0000000000..c3ee657052 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/agents/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/agents/rl_games_ppo_cfg.yaml b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/agents/rl_games_ppo_cfg.yaml new file mode 100644 index 0000000000..5494199846 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/agents/rl_games_ppo_cfg.yaml @@ -0,0 +1,118 @@ +params: + seed: 0 + algo: + name: a2c_continuous + + env: + clip_actions: 1.0 + + model: + name: continuous_a2c_logstd + + network: + name: actor_critic + separate: False + + space: + continuous: + mu_activation: None + sigma_activation: None + mu_init: + name: default + sigma_init: + name: const_initializer + val: 0 + fixed_sigma: False + mlp: + units: [512, 128, 64] + activation: elu + d2rl: False + + initializer: + name: default + regularizer: + name: None + + rnn: + name: lstm + units: 1024 + layers: 2 + before_mlp: True + concat_input: True + layer_norm: True + + load_checkpoint: False + load_path: "" + + config: + name: Factory + device: cuda:0 + full_experiment_name: test + env_name: rlgpu + multi_gpu: False + ppo: True + mixed_precision: True + normalize_input: True + normalize_value: True + value_bootstrap: True + num_actors: 128 + reward_shaper: + scale_value: 1.0 + normalize_advantage: True + gamma: 0.995 + tau: 0.95 + learning_rate: 1.0e-4 + lr_schedule: adaptive + schedule_type: standard + kl_threshold: 0.008 + score_to_win: 20000 + max_epochs: 200 + save_best_after: 10 + save_frequency: 100 + print_stats: True + grad_norm: 1.0 + entropy_coef: 0.0 # 0.0001 # 0.0 + truncate_grads: True + e_clip: 0.2 + horizon_length: 128 + minibatch_size: 512 # batch size = num_envs * horizon_length; minibatch_size = batch_size / num_minibatches + mini_epochs: 4 + critic_coef: 2 + clip_value: True + seq_length: 128 + bounds_loss_coef: 0.0001 + + central_value_config: + minibatch_size: 512 + mini_epochs: 4 + learning_rate: 1e-4 + lr_schedule: adaptive + kl_threshold: 0.008 + clip_value: True + normalize_input: True + truncate_grads: True + + network: + name: actor_critic + central_value: True + + mlp: + units: [512, 128, 64] + activation: elu + d2rl: False + + initializer: + name: default + regularizer: + name: None + + rnn: + name: lstm + units: 1024 + layers: 2 + before_mlp: True + concat_input: True + layer_norm: True + + player: + deterministic: False diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/factory_control.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/factory_control.py new file mode 100644 index 0000000000..33efc41ca0 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/factory_control.py @@ -0,0 +1,196 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Factory: control module. + +Imported by base, environment, and task classes. Not directly executed. +""" + +import math +import torch + +import omni.isaac.core.utils.torch as torch_utils + +from omni.isaac.lab.utils.math import axis_angle_from_quat + + +def compute_dof_torque( + cfg, + dof_pos, + dof_vel, + fingertip_midpoint_pos, + fingertip_midpoint_quat, + fingertip_midpoint_linvel, + fingertip_midpoint_angvel, + jacobian, + arm_mass_matrix, + ctrl_target_fingertip_midpoint_pos, + ctrl_target_fingertip_midpoint_quat, + task_prop_gains, + task_deriv_gains, + device, +): + """Compute Franka DOF torque to move fingertips towards target pose.""" + # References: + # 1) https://ethz.ch/content/dam/ethz/special-interest/mavt/robotics-n-intelligent-systems/rsl-dam/documents/RobotDynamics2018/RD_HS2018script.pdf + # 2) Modern Robotics + + num_envs = cfg.scene.num_envs + dof_torque = torch.zeros((num_envs, dof_pos.shape[1]), device=device) + task_wrench = torch.zeros((num_envs, 6), device=device) + + pos_error, axis_angle_error = get_pose_error( + fingertip_midpoint_pos=fingertip_midpoint_pos, + fingertip_midpoint_quat=fingertip_midpoint_quat, + ctrl_target_fingertip_midpoint_pos=ctrl_target_fingertip_midpoint_pos, + ctrl_target_fingertip_midpoint_quat=ctrl_target_fingertip_midpoint_quat, + jacobian_type="geometric", + rot_error_type="axis_angle", + ) + delta_fingertip_pose = torch.cat((pos_error, axis_angle_error), dim=1) + + # Set tau = k_p * task_pos_error - k_d * task_vel_error (building towards eq. 3.96-3.98) + task_wrench_motion = _apply_task_space_gains( + delta_fingertip_pose=delta_fingertip_pose, + fingertip_midpoint_linvel=fingertip_midpoint_linvel, + fingertip_midpoint_angvel=fingertip_midpoint_angvel, + task_prop_gains=task_prop_gains, + task_deriv_gains=task_deriv_gains, + ) + task_wrench += task_wrench_motion + + # Set tau = J^T * tau, i.e., map tau into joint space as desired + jacobian_T = torch.transpose(jacobian, dim0=1, dim1=2) + dof_torque[:, 0:7] = (jacobian_T @ task_wrench.unsqueeze(-1)).squeeze(-1) + + # adapted from https://gitlab-master.nvidia.com/carbon-gym/carbgym/-/blob/b4bbc66f4e31b1a1bee61dbaafc0766bbfbf0f58/python/examples/franka_cube_ik_osc.py#L70-78 + # roboticsproceedings.org/rss07/p31.pdf + + # useful tensors + arm_mass_matrix_inv = torch.inverse(arm_mass_matrix) + jacobian_T = torch.transpose(jacobian, dim0=1, dim1=2) + arm_mass_matrix_task = torch.inverse( + jacobian @ torch.inverse(arm_mass_matrix) @ jacobian_T + ) # ETH eq. 3.86; geometric Jacobian is assumed + j_eef_inv = arm_mass_matrix_task @ jacobian @ arm_mass_matrix_inv + default_dof_pos_tensor = torch.tensor(cfg.ctrl.default_dof_pos_tensor, device=device).repeat((num_envs, 1)) + # nullspace computation + distance_to_default_dof_pos = default_dof_pos_tensor - dof_pos[:, :7] + distance_to_default_dof_pos = (distance_to_default_dof_pos + math.pi) % ( + 2 * math.pi + ) - math.pi # normalize to [-pi, pi] + u_null = cfg.ctrl.kd_null * -dof_vel[:, :7] + cfg.ctrl.kp_null * distance_to_default_dof_pos + u_null = arm_mass_matrix @ u_null.unsqueeze(-1) + torque_null = (torch.eye(7, device=device).unsqueeze(0) - torch.transpose(jacobian, 1, 2) @ j_eef_inv) @ u_null + dof_torque[:, 0:7] += torque_null.squeeze(-1) + + # TODO: Verify it's okay to no longer do gripper control here. + dof_torque = torch.clamp(dof_torque, min=-100.0, max=100.0) + return dof_torque, task_wrench + + +def get_pose_error( + fingertip_midpoint_pos, + fingertip_midpoint_quat, + ctrl_target_fingertip_midpoint_pos, + ctrl_target_fingertip_midpoint_quat, + jacobian_type, + rot_error_type, +): + """Compute task-space error between target Franka fingertip pose and current pose.""" + # Reference: https://ethz.ch/content/dam/ethz/special-interest/mavt/robotics-n-intelligent-systems/rsl-dam/documents/RobotDynamics2018/RD_HS2018script.pdf + + # Compute pos error + pos_error = ctrl_target_fingertip_midpoint_pos - fingertip_midpoint_pos + + # Compute rot error + if jacobian_type == "geometric": # See example 2.9.8; note use of J_g and transformation between rotation vectors + # Compute quat error (i.e., difference quat) + # Reference: https://personal.utdallas.edu/~sxb027100/dock/quat.html + + # Check for shortest path using quaternion dot product. + quat_dot = (ctrl_target_fingertip_midpoint_quat * fingertip_midpoint_quat).sum(dim=1, keepdim=True) + ctrl_target_fingertip_midpoint_quat = torch.where( + quat_dot.expand(-1, 4) >= 0, ctrl_target_fingertip_midpoint_quat, -ctrl_target_fingertip_midpoint_quat + ) + + fingertip_midpoint_quat_norm = torch_utils.quat_mul( + fingertip_midpoint_quat, torch_utils.quat_conjugate(fingertip_midpoint_quat) + )[ + :, 0 + ] # scalar component + fingertip_midpoint_quat_inv = torch_utils.quat_conjugate( + fingertip_midpoint_quat + ) / fingertip_midpoint_quat_norm.unsqueeze(-1) + quat_error = torch_utils.quat_mul(ctrl_target_fingertip_midpoint_quat, fingertip_midpoint_quat_inv) + + # Convert to axis-angle error + axis_angle_error = axis_angle_from_quat(quat_error) + + if rot_error_type == "quat": + return pos_error, quat_error + elif rot_error_type == "axis_angle": + return pos_error, axis_angle_error + + +def _get_delta_dof_pos(delta_pose, ik_method, jacobian, device): + """Get delta Franka DOF position from delta pose using specified IK method.""" + # References: + # 1) https://www.cs.cmu.edu/~15464-s13/lectures/lecture6/iksurvey.pdf + # 2) https://ethz.ch/content/dam/ethz/special-interest/mavt/robotics-n-intelligent-systems/rsl-dam/documents/RobotDynamics2018/RD_HS2018script.pdf (p. 47) + + if ik_method == "pinv": # Jacobian pseudoinverse + k_val = 1.0 + jacobian_pinv = torch.linalg.pinv(jacobian) + delta_dof_pos = k_val * jacobian_pinv @ delta_pose.unsqueeze(-1) + delta_dof_pos = delta_dof_pos.squeeze(-1) + + elif ik_method == "trans": # Jacobian transpose + k_val = 1.0 + jacobian_T = torch.transpose(jacobian, dim0=1, dim1=2) + delta_dof_pos = k_val * jacobian_T @ delta_pose.unsqueeze(-1) + delta_dof_pos = delta_dof_pos.squeeze(-1) + + elif ik_method == "dls": # damped least squares (Levenberg-Marquardt) + lambda_val = 0.1 # 0.1 + jacobian_T = torch.transpose(jacobian, dim0=1, dim1=2) + lambda_matrix = (lambda_val**2) * torch.eye(n=jacobian.shape[1], device=device) + delta_dof_pos = jacobian_T @ torch.inverse(jacobian @ jacobian_T + lambda_matrix) @ delta_pose.unsqueeze(-1) + delta_dof_pos = delta_dof_pos.squeeze(-1) + + elif ik_method == "svd": # adaptive SVD + k_val = 1.0 + U, S, Vh = torch.linalg.svd(jacobian) + S_inv = 1.0 / S + min_singular_value = 1.0e-5 + S_inv = torch.where(S > min_singular_value, S_inv, torch.zeros_like(S_inv)) + jacobian_pinv = ( + torch.transpose(Vh, dim0=1, dim1=2)[:, :, :6] @ torch.diag_embed(S_inv) @ torch.transpose(U, dim0=1, dim1=2) + ) + delta_dof_pos = k_val * jacobian_pinv @ delta_pose.unsqueeze(-1) + delta_dof_pos = delta_dof_pos.squeeze(-1) + + return delta_dof_pos + + +def _apply_task_space_gains( + delta_fingertip_pose, fingertip_midpoint_linvel, fingertip_midpoint_angvel, task_prop_gains, task_deriv_gains +): + """Interpret PD gains as task-space gains. Apply to task-space error.""" + + task_wrench = torch.zeros_like(delta_fingertip_pose) + + # Apply gains to lin error components + lin_error = delta_fingertip_pose[:, 0:3] + task_wrench[:, 0:3] = task_prop_gains[:, 0:3] * lin_error + task_deriv_gains[:, 0:3] * ( + 0.0 - fingertip_midpoint_linvel + ) + + # Apply gains to rot error components + rot_error = delta_fingertip_pose[:, 3:6] + task_wrench[:, 3:6] = task_prop_gains[:, 3:6] * rot_error + task_deriv_gains[:, 3:6] * ( + 0.0 - fingertip_midpoint_angvel + ) + return task_wrench diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/factory_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/factory_env.py new file mode 100644 index 0000000000..42af6e531c --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/factory_env.py @@ -0,0 +1,888 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import numpy as np +import torch + +import carb +import omni.isaac.core.utils.torch as torch_utils + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.assets import Articulation +from omni.isaac.lab.envs import DirectRLEnv +from omni.isaac.lab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane +from omni.isaac.lab.utils.assets import ISAAC_NUCLEUS_DIR +from omni.isaac.lab.utils.math import axis_angle_from_quat + +from . import factory_control as fc +from .factory_env_cfg import OBS_DIM_CFG, STATE_DIM_CFG, FactoryEnvCfg + + +class FactoryEnv(DirectRLEnv): + cfg: FactoryEnvCfg + + def __init__(self, cfg: FactoryEnvCfg, render_mode: str | None = None, **kwargs): + # Update number of obs/states + cfg.observation_space = sum([OBS_DIM_CFG[obs] for obs in cfg.obs_order]) + cfg.state_space = sum([STATE_DIM_CFG[state] for state in cfg.state_order]) + cfg.observation_space += cfg.action_space + cfg.state_space += cfg.action_space + self.cfg_task = cfg.task + + super().__init__(cfg, render_mode, **kwargs) + + self._set_body_inertias() + self._init_tensors() + self._set_default_dynamics_parameters() + self._compute_intermediate_values(dt=self.physics_dt) + + def _set_body_inertias(self): + """Note: this is to account for the asset_options.armature parameter in IGE.""" + inertias = self._robot.root_physx_view.get_inertias() + offset = torch.zeros_like(inertias) + offset[:, :, [0, 4, 8]] += 0.01 + new_inertias = inertias + offset + self._robot.root_physx_view.set_inertias(new_inertias, torch.arange(self.num_envs)) + + def _set_default_dynamics_parameters(self): + """Set parameters defining dynamic interactions.""" + self.default_gains = torch.tensor(self.cfg.ctrl.default_task_prop_gains, device=self.device).repeat( + (self.num_envs, 1) + ) + + self.pos_threshold = torch.tensor(self.cfg.ctrl.pos_action_threshold, device=self.device).repeat( + (self.num_envs, 1) + ) + self.rot_threshold = torch.tensor(self.cfg.ctrl.rot_action_threshold, device=self.device).repeat( + (self.num_envs, 1) + ) + + # Set masses and frictions. + self._set_friction(self._held_asset, self.cfg_task.held_asset_cfg.friction) + self._set_friction(self._fixed_asset, self.cfg_task.fixed_asset_cfg.friction) + self._set_friction(self._robot, self.cfg_task.robot_cfg.friction) + + def _set_friction(self, asset, value): + """Update material properties for a given asset.""" + materials = asset.root_physx_view.get_material_properties() + materials[..., 0] = value # Static friction. + materials[..., 1] = value # Dynamic friction. + env_ids = torch.arange(self.scene.num_envs, device="cpu") + asset.root_physx_view.set_material_properties(materials, env_ids) + + def _init_tensors(self): + """Initialize tensors once.""" + self.identity_quat = ( + torch.tensor([1.0, 0.0, 0.0, 0.0], device=self.device).unsqueeze(0).repeat(self.num_envs, 1) + ) + + # Control targets. + self.ctrl_target_joint_pos = torch.zeros((self.num_envs, self._robot.num_joints), device=self.device) + self.ctrl_target_fingertip_midpoint_pos = torch.zeros((self.num_envs, 3), device=self.device) + self.ctrl_target_fingertip_midpoint_quat = torch.zeros((self.num_envs, 4), device=self.device) + + # Fixed asset. + self.fixed_pos_action_frame = torch.zeros((self.num_envs, 3), device=self.device) + self.fixed_pos_obs_frame = torch.zeros((self.num_envs, 3), device=self.device) + self.init_fixed_pos_obs_noise = torch.zeros((self.num_envs, 3), device=self.device) + + # Held asset + held_base_x_offset = 0.0 + if self.cfg_task.name == "peg_insert": + held_base_z_offset = 0.0 + elif self.cfg_task.name == "gear_mesh": + gear_base_offset = self._get_target_gear_base_offset() + held_base_x_offset = gear_base_offset[0] + held_base_z_offset = gear_base_offset[2] + elif self.cfg_task.name == "nut_thread": + held_base_z_offset = self.cfg_task.fixed_asset_cfg.base_height + else: + raise NotImplementedError("Task not implemented") + + self.held_base_pos_local = torch.tensor([0.0, 0.0, 0.0], device=self.device).repeat((self.num_envs, 1)) + self.held_base_pos_local[:, 0] = held_base_x_offset + self.held_base_pos_local[:, 2] = held_base_z_offset + self.held_base_quat_local = self.identity_quat.clone().detach() + + self.held_base_pos = torch.zeros_like(self.held_base_pos_local) + self.held_base_quat = self.identity_quat.clone().detach() + + # Computer body indices. + self.left_finger_body_idx = self._robot.body_names.index("panda_leftfinger") + self.right_finger_body_idx = self._robot.body_names.index("panda_rightfinger") + self.fingertip_body_idx = self._robot.body_names.index("panda_fingertip_centered") + + # Tensors for finite-differencing. + self.last_update_timestamp = 0.0 # Note: This is for finite differencing body velocities. + self.prev_fingertip_pos = torch.zeros((self.num_envs, 3), device=self.device) + self.prev_fingertip_quat = self.identity_quat.clone() + self.prev_joint_pos = torch.zeros((self.num_envs, 7), device=self.device) + + # Keypoint tensors. + self.target_held_base_pos = torch.zeros((self.num_envs, 3), device=self.device) + self.target_held_base_quat = self.identity_quat.clone().detach() + + offsets = self._get_keypoint_offsets(self.cfg_task.num_keypoints) + self.keypoint_offsets = offsets * self.cfg_task.keypoint_scale + self.keypoints_held = torch.zeros((self.num_envs, self.cfg_task.num_keypoints, 3), device=self.device) + self.keypoints_fixed = torch.zeros_like(self.keypoints_held, device=self.device) + + # Used to compute target poses. + self.fixed_success_pos_local = torch.zeros((self.num_envs, 3), device=self.device) + if self.cfg_task.name == "peg_insert": + self.fixed_success_pos_local[:, 2] = 0.0 + elif self.cfg_task.name == "gear_mesh": + gear_base_offset = self._get_target_gear_base_offset() + self.fixed_success_pos_local[:, 0] = gear_base_offset[0] + self.fixed_success_pos_local[:, 2] = gear_base_offset[2] + elif self.cfg_task.name == "nut_thread": + head_height = self.cfg_task.fixed_asset_cfg.base_height + shank_length = self.cfg_task.fixed_asset_cfg.height + thread_pitch = self.cfg_task.fixed_asset_cfg.thread_pitch + self.fixed_success_pos_local[:, 2] = head_height + shank_length - thread_pitch * 1.5 + else: + raise NotImplementedError("Task not implemented") + + self.ep_succeeded = torch.zeros((self.num_envs,), dtype=torch.long, device=self.device) + self.ep_success_times = torch.zeros((self.num_envs,), dtype=torch.long, device=self.device) + + def _get_keypoint_offsets(self, num_keypoints): + """Get uniformly-spaced keypoints along a line of unit length, centered at 0.""" + keypoint_offsets = torch.zeros((num_keypoints, 3), device=self.device) + keypoint_offsets[:, -1] = torch.linspace(0.0, 1.0, num_keypoints, device=self.device) - 0.5 + + return keypoint_offsets + + def _setup_scene(self): + """Initialize simulation scene.""" + spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg(), translation=(0.0, 0.0, -0.4)) + + # spawn a usd file of a table into the scene + cfg = sim_utils.UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/SeattleLabTable/table_instanceable.usd") + cfg.func( + "/World/envs/env_.*/Table", cfg, translation=(0.55, 0.0, 0.0), orientation=(0.70711, 0.0, 0.0, 0.70711) + ) + + self._robot = Articulation(self.cfg.robot) + self._fixed_asset = Articulation(self.cfg_task.fixed_asset) + self._held_asset = Articulation(self.cfg_task.held_asset) + if self.cfg_task.name == "gear_mesh": + self._small_gear_asset = Articulation(self.cfg_task.small_gear_cfg) + self._large_gear_asset = Articulation(self.cfg_task.large_gear_cfg) + + self.scene.clone_environments(copy_from_source=False) + self.scene.filter_collisions() + + self.scene.articulations["robot"] = self._robot + self.scene.articulations["fixed_asset"] = self._fixed_asset + self.scene.articulations["held_asset"] = self._held_asset + if self.cfg_task.name == "gear_mesh": + self.scene.articulations["small_gear"] = self._small_gear_asset + self.scene.articulations["large_gear"] = self._large_gear_asset + + # add lights + light_cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75)) + light_cfg.func("/World/Light", light_cfg) + + def _compute_intermediate_values(self, dt): + """Get values computed from raw tensors. This includes adding noise.""" + # TODO: A lot of these can probably only be set once? + self.fixed_pos = self._fixed_asset.data.root_link_pos_w - self.scene.env_origins + self.fixed_quat = self._fixed_asset.data.root_link_quat_w + + self.held_pos = self._held_asset.data.root_link_pos_w - self.scene.env_origins + self.held_quat = self._held_asset.data.root_link_quat_w + + self.fingertip_midpoint_pos = ( + self._robot.data.body_link_pos_w[:, self.fingertip_body_idx] - self.scene.env_origins + ) + self.fingertip_midpoint_quat = self._robot.data.body_link_quat_w[:, self.fingertip_body_idx] + self.fingertip_midpoint_linvel = self._robot.data.body_com_lin_vel_w[:, self.fingertip_body_idx] + self.fingertip_midpoint_angvel = self._robot.data.body_com_ang_vel_w[:, self.fingertip_body_idx] + + jacobians = self._robot.root_physx_view.get_jacobians() + + self.left_finger_jacobian = jacobians[:, self.left_finger_body_idx - 1, 0:6, 0:7] + self.right_finger_jacobian = jacobians[:, self.right_finger_body_idx - 1, 0:6, 0:7] + self.fingertip_midpoint_jacobian = (self.left_finger_jacobian + self.right_finger_jacobian) * 0.5 + self.arm_mass_matrix = self._robot.root_physx_view.get_mass_matrices()[:, 0:7, 0:7] + self.joint_pos = self._robot.data.joint_pos.clone() + self.joint_vel = self._robot.data.joint_vel.clone() + + # Finite-differencing results in more reliable velocity estimates. + self.ee_linvel_fd = (self.fingertip_midpoint_pos - self.prev_fingertip_pos) / dt + self.prev_fingertip_pos = self.fingertip_midpoint_pos.clone() + + # Add state differences if velocity isn't being added. + rot_diff_quat = torch_utils.quat_mul( + self.fingertip_midpoint_quat, torch_utils.quat_conjugate(self.prev_fingertip_quat) + ) + rot_diff_quat *= torch.sign(rot_diff_quat[:, 0]).unsqueeze(-1) + rot_diff_aa = axis_angle_from_quat(rot_diff_quat) + self.ee_angvel_fd = rot_diff_aa / dt + self.prev_fingertip_quat = self.fingertip_midpoint_quat.clone() + + joint_diff = self.joint_pos[:, 0:7] - self.prev_joint_pos + self.joint_vel_fd = joint_diff / dt + self.prev_joint_pos = self.joint_pos[:, 0:7].clone() + + # Keypoint tensors. + self.held_base_quat[:], self.held_base_pos[:] = torch_utils.tf_combine( + self.held_quat, self.held_pos, self.held_base_quat_local, self.held_base_pos_local + ) + self.target_held_base_quat[:], self.target_held_base_pos[:] = torch_utils.tf_combine( + self.fixed_quat, self.fixed_pos, self.identity_quat, self.fixed_success_pos_local + ) + + # Compute pos of keypoints on held asset, and fixed asset in world frame + for idx, keypoint_offset in enumerate(self.keypoint_offsets): + self.keypoints_held[:, idx] = torch_utils.tf_combine( + self.held_base_quat, self.held_base_pos, self.identity_quat, keypoint_offset.repeat(self.num_envs, 1) + )[1] + self.keypoints_fixed[:, idx] = torch_utils.tf_combine( + self.target_held_base_quat, + self.target_held_base_pos, + self.identity_quat, + keypoint_offset.repeat(self.num_envs, 1), + )[1] + + self.keypoint_dist = torch.norm(self.keypoints_held - self.keypoints_fixed, p=2, dim=-1).mean(-1) + self.last_update_timestamp = self._robot._data._sim_timestamp + + def _get_observations(self): + """Get actor/critic inputs using asymmetric critic.""" + noisy_fixed_pos = self.fixed_pos_obs_frame + self.init_fixed_pos_obs_noise + + prev_actions = self.actions.clone() + + obs_dict = { + "fingertip_pos": self.fingertip_midpoint_pos, + "fingertip_pos_rel_fixed": self.fingertip_midpoint_pos - noisy_fixed_pos, + "fingertip_quat": self.fingertip_midpoint_quat, + "ee_linvel": self.ee_linvel_fd, + "ee_angvel": self.ee_angvel_fd, + "prev_actions": prev_actions, + } + + state_dict = { + "fingertip_pos": self.fingertip_midpoint_pos, + "fingertip_pos_rel_fixed": self.fingertip_midpoint_pos - self.fixed_pos_obs_frame, + "fingertip_quat": self.fingertip_midpoint_quat, + "ee_linvel": self.fingertip_midpoint_linvel, + "ee_angvel": self.fingertip_midpoint_angvel, + "joint_pos": self.joint_pos[:, 0:7], + "held_pos": self.held_pos, + "held_pos_rel_fixed": self.held_pos - self.fixed_pos_obs_frame, + "held_quat": self.held_quat, + "fixed_pos": self.fixed_pos, + "fixed_quat": self.fixed_quat, + "task_prop_gains": self.task_prop_gains, + "pos_threshold": self.pos_threshold, + "rot_threshold": self.rot_threshold, + "prev_actions": prev_actions, + } + obs_tensors = [obs_dict[obs_name] for obs_name in self.cfg.obs_order + ["prev_actions"]] + obs_tensors = torch.cat(obs_tensors, dim=-1) + state_tensors = [state_dict[state_name] for state_name in self.cfg.state_order + ["prev_actions"]] + state_tensors = torch.cat(state_tensors, dim=-1) + return {"policy": obs_tensors, "critic": state_tensors} + + def _reset_buffers(self, env_ids): + """Reset buffers.""" + self.ep_succeeded[env_ids] = 0 + + def _pre_physics_step(self, action): + """Apply policy actions with smoothing.""" + env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1) + if len(env_ids) > 0: + self._reset_buffers(env_ids) + + self.actions = ( + self.cfg.ctrl.ema_factor * action.clone().to(self.device) + (1 - self.cfg.ctrl.ema_factor) * self.actions + ) + + def close_gripper_in_place(self): + """Keep gripper in current position as gripper closes.""" + actions = torch.zeros((self.num_envs, 6), device=self.device) + ctrl_target_gripper_dof_pos = 0.0 + + # Interpret actions as target pos displacements and set pos target + pos_actions = actions[:, 0:3] * self.pos_threshold + self.ctrl_target_fingertip_midpoint_pos = self.fingertip_midpoint_pos + pos_actions + + # Interpret actions as target rot (axis-angle) displacements + rot_actions = actions[:, 3:6] + + # Convert to quat and set rot target + angle = torch.norm(rot_actions, p=2, dim=-1) + axis = rot_actions / angle.unsqueeze(-1) + + rot_actions_quat = torch_utils.quat_from_angle_axis(angle, axis) + + rot_actions_quat = torch.where( + angle.unsqueeze(-1).repeat(1, 4) > 1.0e-6, + rot_actions_quat, + torch.tensor([1.0, 0.0, 0.0, 0.0], device=self.device).repeat(self.num_envs, 1), + ) + self.ctrl_target_fingertip_midpoint_quat = torch_utils.quat_mul(rot_actions_quat, self.fingertip_midpoint_quat) + + target_euler_xyz = torch.stack(torch_utils.get_euler_xyz(self.ctrl_target_fingertip_midpoint_quat), dim=1) + target_euler_xyz[:, 0] = 3.14159 + target_euler_xyz[:, 1] = 0.0 + + self.ctrl_target_fingertip_midpoint_quat = torch_utils.quat_from_euler_xyz( + roll=target_euler_xyz[:, 0], pitch=target_euler_xyz[:, 1], yaw=target_euler_xyz[:, 2] + ) + + self.ctrl_target_gripper_dof_pos = ctrl_target_gripper_dof_pos + self.generate_ctrl_signals() + + def _apply_action(self): + """Apply actions for policy as delta targets from current position.""" + # Get current yaw for success checking. + _, _, curr_yaw = torch_utils.get_euler_xyz(self.fingertip_midpoint_quat) + self.curr_yaw = torch.where(curr_yaw > np.deg2rad(235), curr_yaw - 2 * np.pi, curr_yaw) + + # Note: We use finite-differenced velocities for control and observations. + # Check if we need to re-compute velocities within the decimation loop. + if self.last_update_timestamp < self._robot._data._sim_timestamp: + self._compute_intermediate_values(dt=self.physics_dt) + + # Interpret actions as target pos displacements and set pos target + pos_actions = self.actions[:, 0:3] * self.pos_threshold + + # Interpret actions as target rot (axis-angle) displacements + rot_actions = self.actions[:, 3:6] + if self.cfg_task.unidirectional_rot: + rot_actions[:, 2] = -(rot_actions[:, 2] + 1.0) * 0.5 # [-1, 0] + rot_actions = rot_actions * self.rot_threshold + + self.ctrl_target_fingertip_midpoint_pos = self.fingertip_midpoint_pos + pos_actions + # To speed up learning, never allow the policy to move more than 5cm away from the base. + delta_pos = self.ctrl_target_fingertip_midpoint_pos - self.fixed_pos_action_frame + pos_error_clipped = torch.clip( + delta_pos, -self.cfg.ctrl.pos_action_bounds[0], self.cfg.ctrl.pos_action_bounds[1] + ) + self.ctrl_target_fingertip_midpoint_pos = self.fixed_pos_action_frame + pos_error_clipped + + # Convert to quat and set rot target + angle = torch.norm(rot_actions, p=2, dim=-1) + axis = rot_actions / angle.unsqueeze(-1) + + rot_actions_quat = torch_utils.quat_from_angle_axis(angle, axis) + rot_actions_quat = torch.where( + angle.unsqueeze(-1).repeat(1, 4) > 1e-6, + rot_actions_quat, + torch.tensor([1.0, 0.0, 0.0, 0.0], device=self.device).repeat(self.num_envs, 1), + ) + self.ctrl_target_fingertip_midpoint_quat = torch_utils.quat_mul(rot_actions_quat, self.fingertip_midpoint_quat) + + target_euler_xyz = torch.stack(torch_utils.get_euler_xyz(self.ctrl_target_fingertip_midpoint_quat), dim=1) + target_euler_xyz[:, 0] = 3.14159 # Restrict actions to be upright. + target_euler_xyz[:, 1] = 0.0 + + self.ctrl_target_fingertip_midpoint_quat = torch_utils.quat_from_euler_xyz( + roll=target_euler_xyz[:, 0], pitch=target_euler_xyz[:, 1], yaw=target_euler_xyz[:, 2] + ) + + self.ctrl_target_gripper_dof_pos = 0.0 + self.generate_ctrl_signals() + + def _set_gains(self, prop_gains, rot_deriv_scale=1.0): + """Set robot gains using critical damping.""" + self.task_prop_gains = prop_gains + self.task_deriv_gains = 2 * torch.sqrt(prop_gains) + self.task_deriv_gains[:, 3:6] /= rot_deriv_scale + + def generate_ctrl_signals(self): + """Get Jacobian. Set Franka DOF position targets (fingers) or DOF torques (arm).""" + self.joint_torque, self.applied_wrench = fc.compute_dof_torque( + cfg=self.cfg, + dof_pos=self.joint_pos, + dof_vel=self.joint_vel, # _fd, + fingertip_midpoint_pos=self.fingertip_midpoint_pos, + fingertip_midpoint_quat=self.fingertip_midpoint_quat, + fingertip_midpoint_linvel=self.ee_linvel_fd, + fingertip_midpoint_angvel=self.ee_angvel_fd, + jacobian=self.fingertip_midpoint_jacobian, + arm_mass_matrix=self.arm_mass_matrix, + ctrl_target_fingertip_midpoint_pos=self.ctrl_target_fingertip_midpoint_pos, + ctrl_target_fingertip_midpoint_quat=self.ctrl_target_fingertip_midpoint_quat, + task_prop_gains=self.task_prop_gains, + task_deriv_gains=self.task_deriv_gains, + device=self.device, + ) + + # set target for gripper joints to use physx's PD controller + self.ctrl_target_joint_pos[:, 7:9] = self.ctrl_target_gripper_dof_pos + self.joint_torque[:, 7:9] = 0.0 + + self._robot.set_joint_position_target(self.ctrl_target_joint_pos) + self._robot.set_joint_effort_target(self.joint_torque) + + def _get_dones(self): + """Update intermediate values used for rewards and observations.""" + self._compute_intermediate_values(dt=self.physics_dt) + time_out = self.episode_length_buf >= self.max_episode_length - 1 + return time_out, time_out + + def _get_curr_successes(self, success_threshold, check_rot=False): + """Get success mask at current timestep.""" + curr_successes = torch.zeros((self.num_envs,), dtype=torch.bool, device=self.device) + + xy_dist = torch.linalg.vector_norm(self.target_held_base_pos[:, 0:2] - self.held_base_pos[:, 0:2], dim=1) + z_disp = self.held_base_pos[:, 2] - self.target_held_base_pos[:, 2] + + is_centered = torch.where(xy_dist < 0.0025, torch.ones_like(curr_successes), torch.zeros_like(curr_successes)) + # Height threshold to target + fixed_cfg = self.cfg_task.fixed_asset_cfg + if self.cfg_task.name == "peg_insert" or self.cfg_task.name == "gear_mesh": + height_threshold = fixed_cfg.height * success_threshold + elif self.cfg_task.name == "nut_thread": + height_threshold = fixed_cfg.thread_pitch * success_threshold + else: + raise NotImplementedError("Task not implemented") + is_close_or_below = torch.where( + z_disp < height_threshold, torch.ones_like(curr_successes), torch.zeros_like(curr_successes) + ) + curr_successes = torch.logical_and(is_centered, is_close_or_below) + + if check_rot: + is_rotated = self.curr_yaw < self.cfg_task.ee_success_yaw + curr_successes = torch.logical_and(curr_successes, is_rotated) + + return curr_successes + + def _get_rewards(self): + """Update rewards and compute success statistics.""" + # Get successful and failed envs at current timestep + check_rot = self.cfg_task.name == "nut_thread" + curr_successes = self._get_curr_successes( + success_threshold=self.cfg_task.success_threshold, check_rot=check_rot + ) + + rew_buf = self._update_rew_buf(curr_successes) + + # Only log episode success rates at the end of an episode. + if torch.any(self.reset_buf): + self.extras["successes"] = torch.count_nonzero(curr_successes) / self.num_envs + + # Get the time at which an episode first succeeds. + first_success = torch.logical_and(curr_successes, torch.logical_not(self.ep_succeeded)) + self.ep_succeeded[curr_successes] = 1 + + first_success_ids = first_success.nonzero(as_tuple=False).squeeze(-1) + self.ep_success_times[first_success_ids] = self.episode_length_buf[first_success_ids] + nonzero_success_ids = self.ep_success_times.nonzero(as_tuple=False).squeeze(-1) + + if len(nonzero_success_ids) > 0: # Only log for successful episodes. + success_times = self.ep_success_times[nonzero_success_ids].sum() / len(nonzero_success_ids) + self.extras["success_times"] = success_times + + self.prev_actions = self.actions.clone() + return rew_buf + + def _update_rew_buf(self, curr_successes): + """Compute reward at current timestep.""" + rew_dict = {} + + # Keypoint rewards. + def squashing_fn(x, a, b): + return 1 / (torch.exp(a * x) + b + torch.exp(-a * x)) + + a0, b0 = self.cfg_task.keypoint_coef_baseline + rew_dict["kp_baseline"] = squashing_fn(self.keypoint_dist, a0, b0) + # a1, b1 = 25, 2 + a1, b1 = self.cfg_task.keypoint_coef_coarse + rew_dict["kp_coarse"] = squashing_fn(self.keypoint_dist, a1, b1) + a2, b2 = self.cfg_task.keypoint_coef_fine + # a2, b2 = 300, 0 + rew_dict["kp_fine"] = squashing_fn(self.keypoint_dist, a2, b2) + + # Action penalties. + rew_dict["action_penalty"] = torch.norm(self.actions, p=2) + rew_dict["action_grad_penalty"] = torch.norm(self.actions - self.prev_actions, p=2, dim=-1) + rew_dict["curr_engaged"] = ( + self._get_curr_successes(success_threshold=self.cfg_task.engage_threshold, check_rot=False).clone().float() + ) + rew_dict["curr_successes"] = curr_successes.clone().float() + + rew_buf = ( + rew_dict["kp_coarse"] + + rew_dict["kp_baseline"] + + rew_dict["kp_fine"] + - rew_dict["action_penalty"] * self.cfg_task.action_penalty_scale + - rew_dict["action_grad_penalty"] * self.cfg_task.action_grad_penalty_scale + + rew_dict["curr_engaged"] + + rew_dict["curr_successes"] + ) + + for rew_name, rew in rew_dict.items(): + self.extras[f"logs_rew_{rew_name}"] = rew.mean() + + return rew_buf + + def _reset_idx(self, env_ids): + """ + We assume all envs will always be reset at the same time. + """ + super()._reset_idx(env_ids) + + self._set_assets_to_default_pose(env_ids) + self._set_franka_to_default_pose(joints=self.cfg.ctrl.reset_joints, env_ids=env_ids) + self.step_sim_no_action() + + self.randomize_initial_state(env_ids) + + def _get_target_gear_base_offset(self): + """Get offset of target gear from the gear base asset.""" + target_gear = self.cfg_task.target_gear + if target_gear == "gear_large": + gear_base_offset = self.cfg_task.fixed_asset_cfg.large_gear_base_offset + elif target_gear == "gear_medium": + gear_base_offset = self.cfg_task.fixed_asset_cfg.medium_gear_base_offset + elif target_gear == "gear_small": + gear_base_offset = self.cfg_task.fixed_asset_cfg.small_gear_base_offset + else: + raise ValueError(f"{target_gear} not valid in this context!") + return gear_base_offset + + def _set_assets_to_default_pose(self, env_ids): + """Move assets to default pose before randomization.""" + held_state = self._held_asset.data.default_root_state.clone()[env_ids] + held_state[:, 0:3] += self.scene.env_origins[env_ids] + held_state[:, 7:] = 0.0 + self._held_asset.write_root_link_pose_to_sim(held_state[:, 0:7], env_ids=env_ids) + self._held_asset.write_root_com_velocity_to_sim(held_state[:, 7:], env_ids=env_ids) + self._held_asset.reset() + + fixed_state = self._fixed_asset.data.default_root_state.clone()[env_ids] + fixed_state[:, 0:3] += self.scene.env_origins[env_ids] + fixed_state[:, 7:] = 0.0 + self._fixed_asset.write_root_link_pose_to_sim(fixed_state[:, 0:7], env_ids=env_ids) + self._fixed_asset.write_root_com_velocity_to_sim(fixed_state[:, 7:], env_ids=env_ids) + self._fixed_asset.reset() + + def set_pos_inverse_kinematics(self, env_ids): + """Set robot joint position using DLS IK.""" + ik_time = 0.0 + while ik_time < 0.25: + # Compute error to target. + pos_error, axis_angle_error = fc.get_pose_error( + fingertip_midpoint_pos=self.fingertip_midpoint_pos[env_ids], + fingertip_midpoint_quat=self.fingertip_midpoint_quat[env_ids], + ctrl_target_fingertip_midpoint_pos=self.ctrl_target_fingertip_midpoint_pos[env_ids], + ctrl_target_fingertip_midpoint_quat=self.ctrl_target_fingertip_midpoint_quat[env_ids], + jacobian_type="geometric", + rot_error_type="axis_angle", + ) + + delta_hand_pose = torch.cat((pos_error, axis_angle_error), dim=-1) + + # Solve DLS problem. + delta_dof_pos = fc._get_delta_dof_pos( + delta_pose=delta_hand_pose, + ik_method="dls", + jacobian=self.fingertip_midpoint_jacobian[env_ids], + device=self.device, + ) + self.joint_pos[env_ids, 0:7] += delta_dof_pos[:, 0:7] + self.joint_vel[env_ids, :] = torch.zeros_like(self.joint_pos[env_ids,]) + + self.ctrl_target_joint_pos[env_ids, 0:7] = self.joint_pos[env_ids, 0:7] + # Update dof state. + self._robot.write_joint_state_to_sim(self.joint_pos, self.joint_vel) + self._robot.set_joint_position_target(self.ctrl_target_joint_pos) + + # Simulate and update tensors. + self.step_sim_no_action() + ik_time += self.physics_dt + + return pos_error, axis_angle_error + + def get_handheld_asset_relative_pose(self): + """Get default relative pose between help asset and fingertip.""" + if self.cfg_task.name == "peg_insert": + held_asset_relative_pos = torch.zeros_like(self.held_base_pos_local) + held_asset_relative_pos[:, 2] = self.cfg_task.held_asset_cfg.height + held_asset_relative_pos[:, 2] -= self.cfg_task.robot_cfg.franka_fingerpad_length + elif self.cfg_task.name == "gear_mesh": + held_asset_relative_pos = torch.zeros_like(self.held_base_pos_local) + gear_base_offset = self._get_target_gear_base_offset() + held_asset_relative_pos[:, 0] += gear_base_offset[0] + held_asset_relative_pos[:, 2] += gear_base_offset[2] + held_asset_relative_pos[:, 2] += self.cfg_task.held_asset_cfg.height / 2.0 * 1.1 + elif self.cfg_task.name == "nut_thread": + held_asset_relative_pos = self.held_base_pos_local + else: + raise NotImplementedError("Task not implemented") + + held_asset_relative_quat = self.identity_quat + if self.cfg_task.name == "nut_thread": + # Rotate along z-axis of frame for default position. + initial_rot_deg = self.cfg_task.held_asset_rot_init + rot_yaw_euler = torch.tensor([0.0, 0.0, initial_rot_deg * np.pi / 180.0], device=self.device).repeat( + self.num_envs, 1 + ) + held_asset_relative_quat = torch_utils.quat_from_euler_xyz( + roll=rot_yaw_euler[:, 0], pitch=rot_yaw_euler[:, 1], yaw=rot_yaw_euler[:, 2] + ) + + return held_asset_relative_pos, held_asset_relative_quat + + def _set_franka_to_default_pose(self, joints, env_ids): + """Return Franka to its default joint position.""" + gripper_width = self.cfg_task.held_asset_cfg.diameter / 2 * 1.25 + joint_pos = self._robot.data.default_joint_pos[env_ids] + joint_pos[:, 7:] = gripper_width # MIMIC + joint_pos[:, :7] = torch.tensor(joints, device=self.device)[None, :] + joint_vel = torch.zeros_like(joint_pos) + joint_effort = torch.zeros_like(joint_pos) + self.ctrl_target_joint_pos[env_ids, :] = joint_pos + print(f"Resetting {len(env_ids)} envs...") + self._robot.set_joint_position_target(self.ctrl_target_joint_pos[env_ids], env_ids=env_ids) + self._robot.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids) + self._robot.reset() + self._robot.set_joint_effort_target(joint_effort, env_ids=env_ids) + + self.step_sim_no_action() + + def step_sim_no_action(self): + """Step the simulation without an action. Used for resets.""" + self.scene.write_data_to_sim() + self.sim.step(render=False) + self.scene.update(dt=self.physics_dt) + self._compute_intermediate_values(dt=self.physics_dt) + + def randomize_initial_state(self, env_ids): + """Randomize initial state and perform any episode-level randomization.""" + # Disable gravity. + physics_sim_view = sim_utils.SimulationContext.instance().physics_sim_view + physics_sim_view.set_gravity(carb.Float3(0.0, 0.0, 0.0)) + + # (1.) Randomize fixed asset pose. + fixed_state = self._fixed_asset.data.default_root_state.clone()[env_ids] + # (1.a.) Position + rand_sample = torch.rand((len(env_ids), 3), dtype=torch.float32, device=self.device) + fixed_pos_init_rand = 2 * (rand_sample - 0.5) # [-1, 1] + fixed_asset_init_pos_rand = torch.tensor( + self.cfg_task.fixed_asset_init_pos_noise, dtype=torch.float32, device=self.device + ) + fixed_pos_init_rand = fixed_pos_init_rand @ torch.diag(fixed_asset_init_pos_rand) + fixed_state[:, 0:3] += fixed_pos_init_rand + self.scene.env_origins[env_ids] + # (1.b.) Orientation + fixed_orn_init_yaw = np.deg2rad(self.cfg_task.fixed_asset_init_orn_deg) + fixed_orn_yaw_range = np.deg2rad(self.cfg_task.fixed_asset_init_orn_range_deg) + rand_sample = torch.rand((len(env_ids), 3), dtype=torch.float32, device=self.device) + fixed_orn_euler = fixed_orn_init_yaw + fixed_orn_yaw_range * rand_sample + fixed_orn_euler[:, 0:2] = 0.0 # Only change yaw. + fixed_orn_quat = torch_utils.quat_from_euler_xyz( + fixed_orn_euler[:, 0], fixed_orn_euler[:, 1], fixed_orn_euler[:, 2] + ) + fixed_state[:, 3:7] = fixed_orn_quat + # (1.c.) Velocity + fixed_state[:, 7:] = 0.0 # vel + # (1.d.) Update values. + self._fixed_asset.write_root_link_pose_to_sim(fixed_state[:, 0:7], env_ids=env_ids) + self._fixed_asset.write_root_com_velocity_to_sim(fixed_state[:, 7:], env_ids=env_ids) + self._fixed_asset.reset() + + # (1.e.) Noisy position observation. + fixed_asset_pos_noise = torch.randn((len(env_ids), 3), dtype=torch.float32, device=self.device) + fixed_asset_pos_rand = torch.tensor(self.cfg.obs_rand.fixed_asset_pos, dtype=torch.float32, device=self.device) + fixed_asset_pos_noise = fixed_asset_pos_noise @ torch.diag(fixed_asset_pos_rand) + self.init_fixed_pos_obs_noise[:] = fixed_asset_pos_noise + + self.step_sim_no_action() + + # Compute the frame on the bolt that would be used as observation: fixed_pos_obs_frame + # For example, the tip of the bolt can be used as the observation frame + fixed_tip_pos_local = torch.zeros_like(self.fixed_pos) + fixed_tip_pos_local[:, 2] += self.cfg_task.fixed_asset_cfg.height + fixed_tip_pos_local[:, 2] += self.cfg_task.fixed_asset_cfg.base_height + if self.cfg_task.name == "gear_mesh": + fixed_tip_pos_local[:, 0] = self._get_target_gear_base_offset()[0] + + _, fixed_tip_pos = torch_utils.tf_combine( + self.fixed_quat, self.fixed_pos, self.identity_quat, fixed_tip_pos_local + ) + self.fixed_pos_obs_frame[:] = fixed_tip_pos + + # (2) Move gripper to randomizes location above fixed asset. Keep trying until IK succeeds. + # (a) get position vector to target + bad_envs = env_ids.clone() + ik_attempt = 0 + + hand_down_quat = torch.zeros((self.num_envs, 4), dtype=torch.float32, device=self.device) + self.hand_down_euler = torch.zeros((self.num_envs, 3), dtype=torch.float32, device=self.device) + while True: + n_bad = bad_envs.shape[0] + + above_fixed_pos = fixed_tip_pos.clone() + above_fixed_pos[:, 2] += self.cfg_task.hand_init_pos[2] + + rand_sample = torch.rand((n_bad, 3), dtype=torch.float32, device=self.device) + above_fixed_pos_rand = 2 * (rand_sample - 0.5) # [-1, 1] + hand_init_pos_rand = torch.tensor(self.cfg_task.hand_init_pos_noise, device=self.device) + above_fixed_pos_rand = above_fixed_pos_rand @ torch.diag(hand_init_pos_rand) + above_fixed_pos[bad_envs] += above_fixed_pos_rand + + # (b) get random orientation facing down + hand_down_euler = ( + torch.tensor(self.cfg_task.hand_init_orn, device=self.device).unsqueeze(0).repeat(n_bad, 1) + ) + + rand_sample = torch.rand((n_bad, 3), dtype=torch.float32, device=self.device) + above_fixed_orn_noise = 2 * (rand_sample - 0.5) # [-1, 1] + hand_init_orn_rand = torch.tensor(self.cfg_task.hand_init_orn_noise, device=self.device) + above_fixed_orn_noise = above_fixed_orn_noise @ torch.diag(hand_init_orn_rand) + hand_down_euler += above_fixed_orn_noise + self.hand_down_euler[bad_envs, ...] = hand_down_euler + hand_down_quat[bad_envs, :] = torch_utils.quat_from_euler_xyz( + roll=hand_down_euler[:, 0], pitch=hand_down_euler[:, 1], yaw=hand_down_euler[:, 2] + ) + + # (c) iterative IK Method + self.ctrl_target_fingertip_midpoint_pos[bad_envs, ...] = above_fixed_pos[bad_envs, ...] + self.ctrl_target_fingertip_midpoint_quat[bad_envs, ...] = hand_down_quat[bad_envs, :] + + pos_error, aa_error = self.set_pos_inverse_kinematics(env_ids=bad_envs) + pos_error = torch.linalg.norm(pos_error, dim=1) > 1e-3 + angle_error = torch.norm(aa_error, dim=1) > 1e-3 + any_error = torch.logical_or(pos_error, angle_error) + bad_envs = bad_envs[any_error.nonzero(as_tuple=False).squeeze(-1)] + + # Check IK succeeded for all envs, otherwise try again for those envs + if bad_envs.shape[0] == 0: + break + + self._set_franka_to_default_pose( + joints=[0.00871, -0.10368, -0.00794, -1.49139, -0.00083, 1.38774, 0.0], env_ids=bad_envs + ) + + ik_attempt += 1 + print(f"IK Attempt: {ik_attempt}\tBad Envs: {bad_envs.shape[0]}") + + self.step_sim_no_action() + + # Add flanking gears after servo (so arm doesn't move them). + if self.cfg_task.name == "gear_mesh" and self.cfg_task.add_flanking_gears: + small_gear_state = self._small_gear_asset.data.default_root_state.clone()[env_ids] + small_gear_state[:, 0:7] = fixed_state[:, 0:7] + small_gear_state[:, 7:] = 0.0 # vel + self._small_gear_asset.write_root_link_pose_to_sim(small_gear_state[:, 0:7], env_ids=env_ids) + self._small_gear_asset.write_root_com_velocity_to_sim(small_gear_state[:, 7:], env_ids=env_ids) + self._small_gear_asset.reset() + + large_gear_state = self._large_gear_asset.data.default_root_state.clone()[env_ids] + large_gear_state[:, 0:7] = fixed_state[:, 0:7] + large_gear_state[:, 7:] = 0.0 # vel + self._large_gear_asset.write_root_link_pose_to_sim(large_gear_state[:, 0:7], env_ids=env_ids) + self._large_gear_asset.write_root_com_velocity_to_sim(large_gear_state[:, 7:], env_ids=env_ids) + self._large_gear_asset.reset() + + # (3) Randomize asset-in-gripper location. + # flip gripper z orientation + flip_z_quat = torch.tensor([0.0, 0.0, 1.0, 0.0], device=self.device).unsqueeze(0).repeat(self.num_envs, 1) + fingertip_flipped_quat, fingertip_flipped_pos = torch_utils.tf_combine( + q1=self.fingertip_midpoint_quat, + t1=self.fingertip_midpoint_pos, + q2=flip_z_quat, + t2=torch.zeros_like(self.fingertip_midpoint_pos), + ) + + # get default gripper in asset transform + held_asset_relative_pos, held_asset_relative_quat = self.get_handheld_asset_relative_pose() + asset_in_hand_quat, asset_in_hand_pos = torch_utils.tf_inverse( + held_asset_relative_quat, held_asset_relative_pos + ) + + translated_held_asset_quat, translated_held_asset_pos = torch_utils.tf_combine( + q1=fingertip_flipped_quat, t1=fingertip_flipped_pos, q2=asset_in_hand_quat, t2=asset_in_hand_pos + ) + + # Add asset in hand randomization + rand_sample = torch.rand((self.num_envs, 3), dtype=torch.float32, device=self.device) + self.held_asset_pos_noise = 2 * (rand_sample - 0.5) # [-1, 1] + if self.cfg_task.name == "gear_mesh": + self.held_asset_pos_noise[:, 2] = -rand_sample[:, 2] # [-1, 0] + + held_asset_pos_noise = torch.tensor(self.cfg_task.held_asset_pos_noise, device=self.device) + self.held_asset_pos_noise = self.held_asset_pos_noise @ torch.diag(held_asset_pos_noise) + translated_held_asset_quat, translated_held_asset_pos = torch_utils.tf_combine( + q1=translated_held_asset_quat, + t1=translated_held_asset_pos, + q2=self.identity_quat, + t2=self.held_asset_pos_noise, + ) + + held_state = self._held_asset.data.default_root_state.clone() + held_state[:, 0:3] = translated_held_asset_pos + self.scene.env_origins + held_state[:, 3:7] = translated_held_asset_quat + held_state[:, 7:] = 0.0 + self._held_asset.write_root_link_pose_to_sim(held_state[:, 0:7]) + self._held_asset.write_root_com_velocity_to_sim(held_state[:, 7:]) + self._held_asset.reset() + + # Close hand + # Set gains to use for quick resets. + reset_task_prop_gains = torch.tensor(self.cfg.ctrl.reset_task_prop_gains, device=self.device).repeat( + (self.num_envs, 1) + ) + reset_rot_deriv_scale = self.cfg.ctrl.reset_rot_deriv_scale + self._set_gains(reset_task_prop_gains, reset_rot_deriv_scale) + + self.step_sim_no_action() + + grasp_time = 0.0 + while grasp_time < 0.25: + self.ctrl_target_joint_pos[env_ids, 7:] = 0.0 # Close gripper. + self.ctrl_target_gripper_dof_pos = 0.0 + self.close_gripper_in_place() + self.step_sim_no_action() + grasp_time += self.sim.get_physics_dt() + + self.prev_joint_pos = self.joint_pos[:, 0:7].clone() + self.prev_fingertip_pos = self.fingertip_midpoint_pos.clone() + self.prev_fingertip_quat = self.fingertip_midpoint_quat.clone() + + # Set initial actions to involve no-movement. Needed for EMA/correct penalties. + self.actions = torch.zeros_like(self.actions) + self.prev_actions = torch.zeros_like(self.actions) + # Back out what actions should be for initial state. + # Relative position to bolt tip. + self.fixed_pos_action_frame[:] = self.fixed_pos_obs_frame + self.init_fixed_pos_obs_noise + + pos_actions = self.fingertip_midpoint_pos - self.fixed_pos_action_frame + pos_action_bounds = torch.tensor(self.cfg.ctrl.pos_action_bounds, device=self.device) + pos_actions = pos_actions @ torch.diag(1.0 / pos_action_bounds) + self.actions[:, 0:3] = self.prev_actions[:, 0:3] = pos_actions + + # Relative yaw to bolt. + unrot_180_euler = torch.tensor([-np.pi, 0.0, 0.0], device=self.device).repeat(self.num_envs, 1) + unrot_quat = torch_utils.quat_from_euler_xyz( + roll=unrot_180_euler[:, 0], pitch=unrot_180_euler[:, 1], yaw=unrot_180_euler[:, 2] + ) + + fingertip_quat_rel_bolt = torch_utils.quat_mul(unrot_quat, self.fingertip_midpoint_quat) + fingertip_yaw_bolt = torch_utils.get_euler_xyz(fingertip_quat_rel_bolt)[-1] + fingertip_yaw_bolt = torch.where( + fingertip_yaw_bolt > torch.pi / 2, fingertip_yaw_bolt - 2 * torch.pi, fingertip_yaw_bolt + ) + fingertip_yaw_bolt = torch.where( + fingertip_yaw_bolt < -torch.pi, fingertip_yaw_bolt + 2 * torch.pi, fingertip_yaw_bolt + ) + + yaw_action = (fingertip_yaw_bolt + np.deg2rad(180.0)) / np.deg2rad(270.0) * 2.0 - 1.0 + self.actions[:, 5] = self.prev_actions[:, 5] = yaw_action + + # Zero initial velocity. + self.ee_angvel_fd[:, :] = 0.0 + self.ee_linvel_fd[:, :] = 0.0 + + # Set initial gains for the episode. + self._set_gains(self.default_gains) + + physics_sim_view.set_gravity(carb.Float3(*self.cfg.sim.gravity)) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/factory_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/factory_env_cfg.py new file mode 100644 index 0000000000..0356e9e434 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/factory_env_cfg.py @@ -0,0 +1,208 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.actuators.actuator_cfg import ImplicitActuatorCfg +from omni.isaac.lab.assets import ArticulationCfg +from omni.isaac.lab.envs import DirectRLEnvCfg +from omni.isaac.lab.scene import InteractiveSceneCfg +from omni.isaac.lab.sim import PhysxCfg, SimulationCfg +from omni.isaac.lab.sim.spawners.materials.physics_materials_cfg import RigidBodyMaterialCfg +from omni.isaac.lab.utils import configclass + +from .factory_tasks_cfg import ASSET_DIR, FactoryTask, GearMesh, NutThread, PegInsert + +OBS_DIM_CFG = { + "fingertip_pos": 3, + "fingertip_pos_rel_fixed": 3, + "fingertip_quat": 4, + "ee_linvel": 3, + "ee_angvel": 3, +} + +STATE_DIM_CFG = { + "fingertip_pos": 3, + "fingertip_pos_rel_fixed": 3, + "fingertip_quat": 4, + "ee_linvel": 3, + "ee_angvel": 3, + "joint_pos": 7, + "held_pos": 3, + "held_pos_rel_fixed": 3, + "held_quat": 4, + "fixed_pos": 3, + "fixed_quat": 4, + "task_prop_gains": 6, + "ema_factor": 1, + "pos_threshold": 3, + "rot_threshold": 3, +} + + +@configclass +class ObsRandCfg: + fixed_asset_pos = [0.001, 0.001, 0.001] + + +@configclass +class CtrlCfg: + ema_factor = 0.2 + + pos_action_bounds = [0.05, 0.05, 0.05] + rot_action_bounds = [1.0, 1.0, 1.0] + + pos_action_threshold = [0.02, 0.02, 0.02] + rot_action_threshold = [0.097, 0.097, 0.097] + + reset_joints = [1.5178e-03, -1.9651e-01, -1.4364e-03, -1.9761, -2.7717e-04, 1.7796, 7.8556e-01] + reset_task_prop_gains = [300, 300, 300, 20, 20, 20] + reset_rot_deriv_scale = 10.0 + default_task_prop_gains = [100, 100, 100, 30, 30, 30] + + # Null space parameters. + default_dof_pos_tensor = [-1.3003, -0.4015, 1.1791, -2.1493, 0.4001, 1.9425, 0.4754] + kp_null = 10.0 + kd_null = 6.3246 + + +@configclass +class FactoryEnvCfg(DirectRLEnvCfg): + decimation = 8 + action_space = 6 + # num_*: will be overwritten to correspond to obs_order, state_order. + observation_space = 21 + state_space = 72 + obs_order: list = ["fingertip_pos_rel_fixed", "fingertip_quat", "ee_linvel", "ee_angvel"] + state_order: list = [ + "fingertip_pos", + "fingertip_quat", + "ee_linvel", + "ee_angvel", + "joint_pos", + "held_pos", + "held_pos_rel_fixed", + "held_quat", + "fixed_pos", + "fixed_quat", + ] + + task_name: str = "peg_insert" # peg_insert, gear_mesh, nut_thread + task: FactoryTask = FactoryTask() + obs_rand: ObsRandCfg = ObsRandCfg() + ctrl: CtrlCfg = CtrlCfg() + + episode_length_s = 10.0 # Probably need to override. + sim: SimulationCfg = SimulationCfg( + device="cuda:0", + dt=1 / 120, + gravity=(0.0, 0.0, -9.81), + physx=PhysxCfg( + solver_type=1, + max_position_iteration_count=192, # Important to avoid interpenetration. + max_velocity_iteration_count=1, + bounce_threshold_velocity=0.2, + friction_offset_threshold=0.01, + friction_correlation_distance=0.00625, + gpu_max_rigid_contact_count=2**23, + gpu_max_rigid_patch_count=2**23, + gpu_max_num_partitions=1, # Important for stable simulation. + ), + physics_material=RigidBodyMaterialCfg( + static_friction=1.0, + dynamic_friction=1.0, + ), + ) + + scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=128, env_spacing=2.0) + + robot = ArticulationCfg( + prim_path="/World/envs/env_.*/Robot", + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ASSET_DIR}/franka_mimic.usd", + activate_contact_sensors=True, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=True, + max_depenetration_velocity=5.0, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=3666.0, + enable_gyroscopic_forces=True, + solver_position_iteration_count=192, + solver_velocity_iteration_count=1, + max_contact_impulse=1e32, + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=False, + solver_position_iteration_count=192, + solver_velocity_iteration_count=1, + ), + collision_props=sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + ), + init_state=ArticulationCfg.InitialStateCfg( + joint_pos={ + "panda_joint1": 0.00871, + "panda_joint2": -0.10368, + "panda_joint3": -0.00794, + "panda_joint4": -1.49139, + "panda_joint5": -0.00083, + "panda_joint6": 1.38774, + "panda_joint7": 0.0, + "panda_finger_joint2": 0.04, + }, + pos=(0.0, 0.0, 0.0), + rot=(1.0, 0.0, 0.0, 0.0), + ), + actuators={ + "panda_arm1": ImplicitActuatorCfg( + joint_names_expr=["panda_joint[1-4]"], + stiffness=0.0, + damping=0.0, + friction=0.0, + armature=0.0, + effort_limit=87, + velocity_limit=124.6, + ), + "panda_arm2": ImplicitActuatorCfg( + joint_names_expr=["panda_joint[5-7]"], + stiffness=0.0, + damping=0.0, + friction=0.0, + armature=0.0, + effort_limit=12, + velocity_limit=149.5, + ), + "panda_hand": ImplicitActuatorCfg( + joint_names_expr=["panda_finger_joint[1-2]"], + effort_limit=40.0, + velocity_limit=0.04, + stiffness=7500.0, + damping=173.0, + friction=0.1, + armature=0.0, + ), + }, + ) + + +@configclass +class FactoryTaskPegInsertCfg(FactoryEnvCfg): + task_name = "peg_insert" + task = PegInsert() + episode_length_s = 10.0 + + +@configclass +class FactoryTaskGearMeshCfg(FactoryEnvCfg): + task_name = "gear_mesh" + task = GearMesh() + episode_length_s = 20.0 + + +@configclass +class FactoryTaskNutThreadCfg(FactoryEnvCfg): + task_name = "nut_thread" + task = NutThread() + episode_length_s = 30.0 diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/factory_tasks_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/factory_tasks_cfg.py new file mode 100644 index 0000000000..643cbc2afb --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/factory_tasks_cfg.py @@ -0,0 +1,448 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.assets import ArticulationCfg +from omni.isaac.lab.utils import configclass +from omni.isaac.lab.utils.assets import ISAACLAB_NUCLEUS_DIR + +ASSET_DIR = f"{ISAACLAB_NUCLEUS_DIR}/Factory" + + +@configclass +class FixedAssetCfg: + usd_path: str = "" + diameter: float = 0.0 + height: float = 0.0 + base_height: float = 0.0 # Used to compute held asset CoM. + friction: float = 0.75 + mass: float = 0.05 + + +@configclass +class HeldAssetCfg: + usd_path: str = "" + diameter: float = 0.0 # Used for gripper width. + height: float = 0.0 + friction: float = 0.75 + mass: float = 0.05 + + +@configclass +class RobotCfg: + robot_usd: str = "" + franka_fingerpad_length: float = 0.017608 + friction: float = 0.75 + + +@configclass +class FactoryTask: + robot_cfg: RobotCfg = RobotCfg() + name: str = "" + duration_s = 5.0 + + fixed_asset_cfg: FixedAssetCfg = FixedAssetCfg() + held_asset_cfg: HeldAssetCfg = HeldAssetCfg() + asset_size: float = 0.0 + + # Robot + hand_init_pos: list = [0.0, 0.0, 0.015] # Relative to fixed asset tip. + hand_init_pos_noise: list = [0.02, 0.02, 0.01] + hand_init_orn: list = [3.1416, 0, 2.356] + hand_init_orn_noise: list = [0.0, 0.0, 1.57] + + # Action + unidirectional_rot: bool = False + + # Fixed Asset (applies to all tasks) + fixed_asset_init_pos_noise: list = [0.05, 0.05, 0.05] + fixed_asset_init_orn_deg: float = 0.0 + fixed_asset_init_orn_range_deg: float = 360.0 + + # Held Asset (applies to all tasks) + held_asset_pos_noise: list = [0.0, 0.006, 0.003] # noise level of the held asset in gripper + held_asset_rot_init: float = -90.0 + + # Reward + ee_success_yaw: float = 0.0 # nut_thread task only. + action_penalty_scale: float = 0.0 + action_grad_penalty_scale: float = 0.0 + # Reward function details can be found in Appendix B of https://arxiv.org/pdf/2408.04587. + # Multi-scale keypoints are used to capture different phases of the task. + # Each reward passes the keypoint distance, x, through a squashing function: + # r(x) = 1/(exp(-ax) + b + exp(ax)). + # Each list defines [a, b] which control the slope and maximum of the squashing function. + num_keypoints: int = 4 + keypoint_scale: float = 0.15 + keypoint_coef_baseline: list = [5, 4] # General movement towards fixed object. + keypoint_coef_coarse: list = [50, 2] # Movement to align the assets. + keypoint_coef_fine: list = [100, 0] # Smaller distances for threading or last-inch insertion. + # Fixed-asset height fraction for which different bonuses are rewarded (see individual tasks). + success_threshold: float = 0.04 + engage_threshold: float = 0.9 + + +@configclass +class Peg8mm(HeldAssetCfg): + usd_path = f"{ASSET_DIR}/factory_peg_8mm.usd" + diameter = 0.007986 + height = 0.050 + mass = 0.019 + + +@configclass +class Hole8mm(FixedAssetCfg): + usd_path = f"{ASSET_DIR}/factory_hole_8mm.usd" + diameter = 0.0081 + height = 0.025 + base_height = 0.0 + + +@configclass +class PegInsert(FactoryTask): + name = "peg_insert" + fixed_asset_cfg = Hole8mm() + held_asset_cfg = Peg8mm() + asset_size = 8.0 + duration_s = 10.0 + + # Robot + hand_init_pos: list = [0.0, 0.0, 0.047] # Relative to fixed asset tip. + hand_init_pos_noise: list = [0.02, 0.02, 0.01] + hand_init_orn: list = [3.1416, 0.0, 0.0] + hand_init_orn_noise: list = [0.0, 0.0, 0.785] + + # Fixed Asset (applies to all tasks) + fixed_asset_init_pos_noise: list = [0.05, 0.05, 0.05] + fixed_asset_init_orn_deg: float = 0.0 + fixed_asset_init_orn_range_deg: float = 360.0 + + # Held Asset (applies to all tasks) + held_asset_pos_noise: list = [0.003, 0.0, 0.003] # noise level of the held asset in gripper + held_asset_rot_init: float = 0.0 + + # Rewards + keypoint_coef_baseline: list = [5, 4] + keypoint_coef_coarse: list = [50, 2] + keypoint_coef_fine: list = [100, 0] + # Fraction of socket height. + success_threshold: float = 0.04 + engage_threshold: float = 0.9 + + fixed_asset: ArticulationCfg = ArticulationCfg( + prim_path="/World/envs/env_.*/FixedAsset", + spawn=sim_utils.UsdFileCfg( + usd_path=fixed_asset_cfg.usd_path, + activate_contact_sensors=True, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + max_depenetration_velocity=5.0, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=3666.0, + enable_gyroscopic_forces=True, + solver_position_iteration_count=192, + solver_velocity_iteration_count=1, + max_contact_impulse=1e32, + ), + mass_props=sim_utils.MassPropertiesCfg(mass=fixed_asset_cfg.mass), + collision_props=sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + ), + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.6, 0.0, 0.05), rot=(1.0, 0.0, 0.0, 0.0), joint_pos={}, joint_vel={} + ), + actuators={}, + ) + held_asset: ArticulationCfg = ArticulationCfg( + prim_path="/World/envs/env_.*/HeldAsset", + spawn=sim_utils.UsdFileCfg( + usd_path=held_asset_cfg.usd_path, + activate_contact_sensors=True, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=True, + max_depenetration_velocity=5.0, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=3666.0, + enable_gyroscopic_forces=True, + solver_position_iteration_count=192, + solver_velocity_iteration_count=1, + max_contact_impulse=1e32, + ), + mass_props=sim_utils.MassPropertiesCfg(mass=held_asset_cfg.mass), + collision_props=sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + ), + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.0, 0.4, 0.1), rot=(1.0, 0.0, 0.0, 0.0), joint_pos={}, joint_vel={} + ), + actuators={}, + ) + + +@configclass +class GearBase(FixedAssetCfg): + usd_path = f"{ASSET_DIR}/factory_gear_base.usd" + height = 0.02 + base_height = 0.005 + small_gear_base_offset = [5.075e-2, 0.0, 0.0] + medium_gear_base_offset = [2.025e-2, 0.0, 0.0] + large_gear_base_offset = [-3.025e-2, 0.0, 0.0] + + +@configclass +class MediumGear(HeldAssetCfg): + usd_path = f"{ASSET_DIR}/factory_gear_medium.usd" + diameter = 0.03 # Used for gripper width. + height: float = 0.03 + mass = 0.012 + + +@configclass +class GearMesh(FactoryTask): + name = "gear_mesh" + fixed_asset_cfg = GearBase() + held_asset_cfg = MediumGear() + target_gear = "gear_medium" + duration_s = 20.0 + + small_gear_usd = f"{ASSET_DIR}/factory_gear_small.usd" + large_gear_usd = f"{ASSET_DIR}/factory_gear_large.usd" + + small_gear_cfg: ArticulationCfg = ArticulationCfg( + prim_path="/World/envs/env_.*/SmallGearAsset", + spawn=sim_utils.UsdFileCfg( + usd_path=small_gear_usd, + activate_contact_sensors=True, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + max_depenetration_velocity=5.0, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=3666.0, + enable_gyroscopic_forces=True, + solver_position_iteration_count=192, + solver_velocity_iteration_count=1, + max_contact_impulse=1e32, + ), + mass_props=sim_utils.MassPropertiesCfg(mass=0.019), + collision_props=sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + ), + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.0, 0.4, 0.1), rot=(1.0, 0.0, 0.0, 0.0), joint_pos={}, joint_vel={} + ), + actuators={}, + ) + + large_gear_cfg: ArticulationCfg = ArticulationCfg( + prim_path="/World/envs/env_.*/LargeGearAsset", + spawn=sim_utils.UsdFileCfg( + usd_path=large_gear_usd, + activate_contact_sensors=True, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + max_depenetration_velocity=5.0, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=3666.0, + enable_gyroscopic_forces=True, + solver_position_iteration_count=192, + solver_velocity_iteration_count=1, + max_contact_impulse=1e32, + ), + mass_props=sim_utils.MassPropertiesCfg(mass=0.019), + collision_props=sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + ), + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.0, 0.4, 0.1), rot=(1.0, 0.0, 0.0, 0.0), joint_pos={}, joint_vel={} + ), + actuators={}, + ) + + # Gears Asset + add_flanking_gears = True + add_flanking_gears_prob = 1.0 + + # Robot + hand_init_pos: list = [0.0, 0.0, 0.035] # Relative to fixed asset tip. + hand_init_pos_noise: list = [0.02, 0.02, 0.01] + hand_init_orn: list = [3.1416, 0, 0.0] + hand_init_orn_noise: list = [0.0, 0.0, 0.785] + + # Fixed Asset (applies to all tasks) + fixed_asset_init_pos_noise: list = [0.05, 0.05, 0.05] + fixed_asset_init_orn_deg: float = 0.0 + fixed_asset_init_orn_range_deg: float = 15.0 + + # Held Asset (applies to all tasks) + held_asset_pos_noise: list = [0.003, 0.0, 0.003] # noise level of the held asset in gripper + held_asset_rot_init: float = -90.0 + + keypoint_coef_baseline: list = [5, 4] + keypoint_coef_coarse: list = [50, 2] + keypoint_coef_fine: list = [100, 0] + # Fraction of gear peg height. + success_threshold: float = 0.05 + engage_threshold: float = 0.9 + + fixed_asset: ArticulationCfg = ArticulationCfg( + prim_path="/World/envs/env_.*/FixedAsset", + spawn=sim_utils.UsdFileCfg( + usd_path=fixed_asset_cfg.usd_path, + activate_contact_sensors=True, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + max_depenetration_velocity=5.0, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=3666.0, + enable_gyroscopic_forces=True, + solver_position_iteration_count=192, + solver_velocity_iteration_count=1, + max_contact_impulse=1e32, + ), + mass_props=sim_utils.MassPropertiesCfg(mass=fixed_asset_cfg.mass), + collision_props=sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + ), + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.6, 0.0, 0.05), rot=(1.0, 0.0, 0.0, 0.0), joint_pos={}, joint_vel={} + ), + actuators={}, + ) + held_asset: ArticulationCfg = ArticulationCfg( + prim_path="/World/envs/env_.*/HeldAsset", + spawn=sim_utils.UsdFileCfg( + usd_path=held_asset_cfg.usd_path, + activate_contact_sensors=True, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=True, + max_depenetration_velocity=5.0, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=3666.0, + enable_gyroscopic_forces=True, + solver_position_iteration_count=192, + solver_velocity_iteration_count=1, + max_contact_impulse=1e32, + ), + mass_props=sim_utils.MassPropertiesCfg(mass=held_asset_cfg.mass), + collision_props=sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + ), + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.0, 0.4, 0.1), rot=(1.0, 0.0, 0.0, 0.0), joint_pos={}, joint_vel={} + ), + actuators={}, + ) + + +@configclass +class NutM16(HeldAssetCfg): + usd_path = f"{ASSET_DIR}/factory_nut_m16.usd" + diameter = 0.024 + height = 0.01 + mass = 0.03 + friction = 0.01 # Additive with the nut means friction is (-0.25 + 0.75)/2 = 0.25 + + +@configclass +class BoltM16(FixedAssetCfg): + usd_path = f"{ASSET_DIR}/factory_bolt_m16.usd" + diameter = 0.024 + height = 0.025 + base_height = 0.01 + thread_pitch = 0.002 + + +@configclass +class NutThread(FactoryTask): + name = "nut_thread" + fixed_asset_cfg = BoltM16() + held_asset_cfg = NutM16() + asset_size = 16.0 + duration_s = 30.0 + + # Robot + hand_init_pos: list = [0.0, 0.0, 0.015] # Relative to fixed asset tip. + hand_init_pos_noise: list = [0.02, 0.02, 0.01] + hand_init_orn: list = [3.1416, 0.0, 1.83] + hand_init_orn_noise: list = [0.0, 0.0, 0.26] + + # Action + unidirectional_rot: bool = True + + # Fixed Asset (applies to all tasks) + fixed_asset_init_pos_noise: list = [0.05, 0.05, 0.05] + fixed_asset_init_orn_deg: float = 120.0 + fixed_asset_init_orn_range_deg: float = 30.0 + + # Held Asset (applies to all tasks) + held_asset_pos_noise: list = [0.0, 0.003, 0.003] # noise level of the held asset in gripper + held_asset_rot_init: float = -90.0 + + # Reward. + ee_success_yaw = 0.0 + keypoint_coef_baseline: list = [100, 2] + keypoint_coef_coarse: list = [500, 2] # 100, 2 + keypoint_coef_fine: list = [1500, 0] # 500, 0 + # Fraction of thread-height. + success_threshold: float = 0.375 + engage_threshold: float = 0.5 + keypoint_scale: float = 0.05 + + fixed_asset: ArticulationCfg = ArticulationCfg( + prim_path="/World/envs/env_.*/FixedAsset", + spawn=sim_utils.UsdFileCfg( + usd_path=fixed_asset_cfg.usd_path, + activate_contact_sensors=True, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + max_depenetration_velocity=5.0, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=3666.0, + enable_gyroscopic_forces=True, + solver_position_iteration_count=192, + solver_velocity_iteration_count=1, + max_contact_impulse=1e32, + ), + mass_props=sim_utils.MassPropertiesCfg(mass=fixed_asset_cfg.mass), + collision_props=sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + ), + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.6, 0.0, 0.05), rot=(1.0, 0.0, 0.0, 0.0), joint_pos={}, joint_vel={} + ), + actuators={}, + ) + held_asset: ArticulationCfg = ArticulationCfg( + prim_path="/World/envs/env_.*/HeldAsset", + spawn=sim_utils.UsdFileCfg( + usd_path=held_asset_cfg.usd_path, + activate_contact_sensors=True, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=True, + max_depenetration_velocity=5.0, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=3666.0, + enable_gyroscopic_forces=True, + solver_position_iteration_count=192, + solver_velocity_iteration_count=1, + max_contact_impulse=1e32, + ), + mass_props=sim_utils.MassPropertiesCfg(mass=held_asset_cfg.mass), + collision_props=sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + ), + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.0, 0.4, 0.1), rot=(1.0, 0.0, 0.0, 0.0), joint_pos={}, joint_vel={} + ), + actuators={}, + ) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/franka_cabinet/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/franka_cabinet/__init__.py index 5f333dc9cf..f91915d3ed 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/franka_cabinet/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/franka_cabinet/__init__.py @@ -9,7 +9,6 @@ import gymnasium as gym from . import agents -from .franka_cabinet_env import FrankaCabinetEnv, FrankaCabinetEnvCfg ## # Register Gym environments. @@ -17,10 +16,10 @@ gym.register( id="Isaac-Franka-Cabinet-Direct-v0", - entry_point="omni.isaac.lab_tasks.direct.franka_cabinet:FrankaCabinetEnv", + entry_point=f"{__name__}.franka_cabinet_env:FrankaCabinetEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": FrankaCabinetEnvCfg, + "env_cfg_entry_point": f"{__name__}.franka_cabinet_env:FrankaCabinetEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaCabinetPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/franka_cabinet/franka_cabinet_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/franka_cabinet/franka_cabinet_env.py index 4eb01953fe..d7234ca386 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/franka_cabinet/franka_cabinet_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/franka_cabinet/franka_cabinet_env.py @@ -28,9 +28,9 @@ class FrankaCabinetEnvCfg(DirectRLEnvCfg): # env episode_length_s = 8.3333 # 500 timesteps decimation = 2 - num_actions = 9 - num_observations = 23 - num_states = 0 + action_space = 9 + observation_space = 23 + state_space = 0 # simulation sim: SimulationCfg = SimulationCfg( @@ -299,8 +299,8 @@ def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]: def _get_rewards(self) -> torch.Tensor: # Refresh the intermediate values after the physics steps self._compute_intermediate_values() - robot_left_finger_pos = self._robot.data.body_pos_w[:, self.left_finger_link_idx] - robot_right_finger_pos = self._robot.data.body_pos_w[:, self.right_finger_link_idx] + robot_left_finger_pos = self._robot.data.body_link_pos_w[:, self.left_finger_link_idx] + robot_right_finger_pos = self._robot.data.body_link_pos_w[:, self.right_finger_link_idx] return self._compute_rewards( self.actions, @@ -372,10 +372,10 @@ def _compute_intermediate_values(self, env_ids: torch.Tensor | None = None): if env_ids is None: env_ids = self._robot._ALL_INDICES - hand_pos = self._robot.data.body_pos_w[env_ids, self.hand_link_idx] - hand_rot = self._robot.data.body_quat_w[env_ids, self.hand_link_idx] - drawer_pos = self._cabinet.data.body_pos_w[env_ids, self.drawer_link_idx] - drawer_rot = self._cabinet.data.body_quat_w[env_ids, self.drawer_link_idx] + hand_pos = self._robot.data.body_link_pos_w[env_ids, self.hand_link_idx] + hand_rot = self._robot.data.body_link_quat_w[env_ids, self.hand_link_idx] + drawer_pos = self._cabinet.data.body_link_pos_w[env_ids, self.drawer_link_idx] + drawer_rot = self._cabinet.data.body_link_quat_w[env_ids, self.drawer_link_idx] ( self.robot_grasp_rot[env_ids], self.robot_grasp_pos[env_ids], diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/humanoid/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/humanoid/__init__.py index 92e8384926..9bd3dfec17 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/humanoid/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/humanoid/__init__.py @@ -10,7 +10,6 @@ import gymnasium as gym from . import agents -from .humanoid_env import HumanoidEnv, HumanoidEnvCfg ## # Register Gym environments. @@ -18,10 +17,10 @@ gym.register( id="Isaac-Humanoid-Direct-v0", - entry_point="omni.isaac.lab_tasks.direct.humanoid:HumanoidEnv", + entry_point=f"{__name__}.humanoid_env:HumanoidEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": HumanoidEnvCfg, + "env_cfg_entry_point": f"{__name__}.humanoid_env:HumanoidEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:HumanoidPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/humanoid/humanoid_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/humanoid/humanoid_env.py index bfaf8f8190..2a4d330e6a 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/humanoid/humanoid_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/humanoid/humanoid_env.py @@ -24,9 +24,9 @@ class HumanoidEnvCfg(DirectRLEnvCfg): episode_length_s = 15.0 decimation = 2 action_scale = 1.0 - num_actions = 21 - num_observations = 75 - num_states = 0 + action_space = 21 + observation_space = 75 + state_space = 0 # simulation sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/inhand_manipulation/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/inhand_manipulation/__init__.py index cdce057355..c3ee657052 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/inhand_manipulation/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/inhand_manipulation/__init__.py @@ -2,5 +2,3 @@ # All rights reserved. # # SPDX-License-Identifier: BSD-3-Clause - -from .inhand_manipulation_env import InHandManipulationEnv diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py index 52f4dc6e3c..8017889c6d 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py @@ -19,8 +19,8 @@ from omni.isaac.lab.utils.math import quat_conjugate, quat_from_angle_axis, quat_mul, sample_uniform, saturate if TYPE_CHECKING: - from omni.isaac.lab_tasks.direct.allegro_hand import AllegroHandEnvCfg - from omni.isaac.lab_tasks.direct.shadow_hand import ShadowHandEnvCfg + from omni.isaac.lab_tasks.direct.allegro_hand.allegro_hand_env_cfg import AllegroHandEnvCfg + from omni.isaac.lab_tasks.direct.shadow_hand.shadow_hand_env_cfg import ShadowHandEnvCfg class InHandManipulationEnv(DirectRLEnv): @@ -84,7 +84,7 @@ def _setup_scene(self): spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg()) # clone and replicate (no need to filter for this environment) self.scene.clone_environments(copy_from_source=False) - # add articultion to scene - we must register to scene to randomize with EventManager + # add articulation to scene - we must register to scene to randomize with EventManager self.scene.articulations["robot"] = self.hand self.scene.rigid_objects["object"] = self.object # add lights @@ -221,7 +221,8 @@ def _reset_idx(self, env_ids: Sequence[int] | None): ) object_default_state[:, 7:] = torch.zeros_like(self.object.data.default_root_state[env_ids, 7:]) - self.object.write_root_state_to_sim(object_default_state, env_ids) + self.object.write_root_link_pose_to_sim(object_default_state[:, :7], env_ids) + self.object.write_root_com_velocity_to_sim(object_default_state[:, 7:], env_ids) # reset hand delta_max = self.hand_dof_upper_limits[env_ids] - self.hand.data.default_joint_pos[env_ids] @@ -260,22 +261,22 @@ def _reset_target_pose(self, env_ids): def _compute_intermediate_values(self): # data for hand - self.fingertip_pos = self.hand.data.body_pos_w[:, self.finger_bodies] - self.fingertip_rot = self.hand.data.body_quat_w[:, self.finger_bodies] + self.fingertip_pos = self.hand.data.body_link_pos_w[:, self.finger_bodies] + self.fingertip_rot = self.hand.data.body_link_quat_w[:, self.finger_bodies] self.fingertip_pos -= self.scene.env_origins.repeat((1, self.num_fingertips)).reshape( self.num_envs, self.num_fingertips, 3 ) - self.fingertip_velocities = self.hand.data.body_vel_w[:, self.finger_bodies] + self.fingertip_velocities = self.hand.data.body_com_vel_w[:, self.finger_bodies] self.hand_dof_pos = self.hand.data.joint_pos self.hand_dof_vel = self.hand.data.joint_vel # data for object - self.object_pos = self.object.data.root_pos_w - self.scene.env_origins - self.object_rot = self.object.data.root_quat_w - self.object_velocities = self.object.data.root_vel_w - self.object_linvel = self.object.data.root_lin_vel_w - self.object_angvel = self.object.data.root_ang_vel_w + self.object_pos = self.object.data.root_link_pos_w - self.scene.env_origins + self.object_rot = self.object.data.root_link_quat_w + self.object_velocities = self.object.data.root_com_vel_w + self.object_linvel = self.object.data.root_com_lin_vel_w + self.object_angvel = self.object.data.root_com_ang_vel_w def compute_reduced_observations(self): # Per https://arxiv.org/pdf/1808.00177.pdf Table 2 diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/locomotion/locomotion_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/locomotion/locomotion_env.py index c5c41a3c48..dc3a5439da 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/locomotion/locomotion_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/locomotion/locomotion_env.py @@ -54,7 +54,7 @@ def _setup_scene(self): # clone, filter, and replicate self.scene.clone_environments(copy_from_source=False) self.scene.filter_collisions(global_prim_paths=[self.cfg.terrain.prim_path]) - # add articultion to scene + # add articulation to scene self.scene.articulations["robot"] = self.robot # add lights light_cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75)) @@ -68,8 +68,8 @@ def _apply_action(self): self.robot.set_joint_effort_target(forces, joint_ids=self._joint_dof_idx) def _compute_intermediate_values(self): - self.torso_position, self.torso_rotation = self.robot.data.root_pos_w, self.robot.data.root_quat_w - self.velocity, self.ang_velocity = self.robot.data.root_lin_vel_w, self.robot.data.root_ang_vel_w + self.torso_position, self.torso_rotation = self.robot.data.root_link_pos_w, self.robot.data.root_link_quat_w + self.velocity, self.ang_velocity = self.robot.data.root_com_lin_vel_w, self.robot.data.root_com_ang_vel_w self.dof_pos, self.dof_vel = self.robot.data.joint_pos, self.robot.data.joint_vel ( @@ -161,8 +161,8 @@ def _reset_idx(self, env_ids: torch.Tensor | None): default_root_state = self.robot.data.default_root_state[env_ids] default_root_state[:, :3] += self.scene.env_origins[env_ids] - self.robot.write_root_pose_to_sim(default_root_state[:, :7], env_ids) - self.robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) + self.robot.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids) + self.robot.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids) self.robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) to_target = self.targets[env_ids] - default_root_state[:, :3] diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/quadcopter/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/quadcopter/__init__.py index bc2bc74ce9..2605b599f0 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/quadcopter/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/quadcopter/__init__.py @@ -10,7 +10,6 @@ import gymnasium as gym from . import agents -from .quadcopter_env import QuadcopterEnv, QuadcopterEnvCfg ## # Register Gym environments. @@ -18,10 +17,10 @@ gym.register( id="Isaac-Quadcopter-Direct-v0", - entry_point="omni.isaac.lab_tasks.direct.quadcopter:QuadcopterEnv", + entry_point=f"{__name__}.quadcopter_env:QuadcopterEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": QuadcopterEnvCfg, + "env_cfg_entry_point": f"{__name__}.quadcopter_env:QuadcopterEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:QuadcopterPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/quadcopter/quadcopter_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/quadcopter/quadcopter_env.py index c6df659ec6..7389975ebc 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/quadcopter/quadcopter_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/quadcopter/quadcopter_env.py @@ -5,6 +5,7 @@ from __future__ import annotations +import gymnasium as gym import torch import omni.isaac.lab.sim as sim_utils @@ -50,9 +51,9 @@ class QuadcopterEnvCfg(DirectRLEnvCfg): # env episode_length_s = 10.0 decimation = 2 - num_actions = 4 - num_observations = 12 - num_states = 0 + action_space = 4 + observation_space = 12 + state_space = 0 debug_vis = True ui_window_class_type = QuadcopterEnvWindow @@ -105,7 +106,7 @@ def __init__(self, cfg: QuadcopterEnvCfg, render_mode: str | None = None, **kwar super().__init__(cfg, render_mode, **kwargs) # Total thrust and moment applied to the base of the quadcopter - self._actions = torch.zeros(self.num_envs, self.cfg.num_actions, device=self.device) + self._actions = torch.zeros(self.num_envs, gym.spaces.flatdim(self.single_action_space), device=self.device) self._thrust = torch.zeros(self.num_envs, 1, 3, device=self.device) self._moment = torch.zeros(self.num_envs, 1, 3, device=self.device) # Goal position @@ -153,12 +154,12 @@ def _apply_action(self): def _get_observations(self) -> dict: desired_pos_b, _ = subtract_frame_transforms( - self._robot.data.root_state_w[:, :3], self._robot.data.root_state_w[:, 3:7], self._desired_pos_w + self._robot.data.root_link_state_w[:, :3], self._robot.data.root_link_state_w[:, 3:7], self._desired_pos_w ) obs = torch.cat( [ - self._robot.data.root_lin_vel_b, - self._robot.data.root_ang_vel_b, + self._robot.data.root_com_lin_vel_b, + self._robot.data.root_com_ang_vel_b, self._robot.data.projected_gravity_b, desired_pos_b, ], @@ -168,9 +169,9 @@ def _get_observations(self) -> dict: return observations def _get_rewards(self) -> torch.Tensor: - lin_vel = torch.sum(torch.square(self._robot.data.root_lin_vel_b), dim=1) - ang_vel = torch.sum(torch.square(self._robot.data.root_ang_vel_b), dim=1) - distance_to_goal = torch.linalg.norm(self._desired_pos_w - self._robot.data.root_pos_w, dim=1) + lin_vel = torch.sum(torch.square(self._robot.data.root_com_lin_vel_b), dim=1) + ang_vel = torch.sum(torch.square(self._robot.data.root_com_ang_vel_b), dim=1) + distance_to_goal = torch.linalg.norm(self._desired_pos_w - self._robot.data.root_link_pos_w, dim=1) distance_to_goal_mapped = 1 - torch.tanh(distance_to_goal / 0.8) rewards = { "lin_vel": lin_vel * self.cfg.lin_vel_reward_scale * self.step_dt, @@ -185,7 +186,9 @@ def _get_rewards(self) -> torch.Tensor: def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]: time_out = self.episode_length_buf >= self.max_episode_length - 1 - died = torch.logical_or(self._robot.data.root_pos_w[:, 2] < 0.1, self._robot.data.root_pos_w[:, 2] > 2.0) + died = torch.logical_or( + self._robot.data.root_link_pos_w[:, 2] < 0.1, self._robot.data.root_link_pos_w[:, 2] > 2.0 + ) return died, time_out def _reset_idx(self, env_ids: torch.Tensor | None): @@ -194,7 +197,7 @@ def _reset_idx(self, env_ids: torch.Tensor | None): # Logging final_distance_to_goal = torch.linalg.norm( - self._desired_pos_w[env_ids] - self._robot.data.root_pos_w[env_ids], dim=1 + self._desired_pos_w[env_ids] - self._robot.data.root_link_pos_w[env_ids], dim=1 ).mean() extras = dict() for key in self._episode_sums.keys(): @@ -225,8 +228,8 @@ def _reset_idx(self, env_ids: torch.Tensor | None): joint_vel = self._robot.data.default_joint_vel[env_ids] default_root_state = self._robot.data.default_root_state[env_ids] default_root_state[:, :3] += self._terrain.env_origins[env_ids] - self._robot.write_root_pose_to_sim(default_root_state[:, :7], env_ids) - self._robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) + self._robot.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids) + self._robot.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids) self._robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) def _set_debug_vis_impl(self, debug_vis: bool): diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand/__init__.py index 07ac1b10f2..73cdc43270 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand/__init__.py @@ -10,19 +10,19 @@ import gymnasium as gym from . import agents -from .shadow_hand_env_cfg import ShadowHandEnvCfg, ShadowHandOpenAIEnvCfg -from .shadow_hand_vision_env import ShadowHandVisionEnvCfg, ShadowHandVisionEnvPlayCfg ## # Register Gym environments. ## +inhand_task_entry = "omni.isaac.lab_tasks.direct.inhand_manipulation" + gym.register( id="Isaac-Repose-Cube-Shadow-Direct-v0", - entry_point="omni.isaac.lab_tasks.direct.inhand_manipulation:InHandManipulationEnv", + entry_point=f"{inhand_task_entry}.inhand_manipulation_env:InHandManipulationEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": ShadowHandEnvCfg, + "env_cfg_entry_point": f"{__name__}.shadow_hand_env_cfg:ShadowHandEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:ShadowHandPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", @@ -31,10 +31,10 @@ gym.register( id="Isaac-Repose-Cube-Shadow-OpenAI-FF-Direct-v0", - entry_point="omni.isaac.lab_tasks.direct.inhand_manipulation:InHandManipulationEnv", + entry_point=f"{inhand_task_entry}.inhand_manipulation_env:InHandManipulationEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": ShadowHandOpenAIEnvCfg, + "env_cfg_entry_point": f"{__name__}.shadow_hand_env_cfg:ShadowHandOpenAIEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_ff_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:ShadowHandAsymFFPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ff_ppo_cfg.yaml", @@ -43,10 +43,10 @@ gym.register( id="Isaac-Repose-Cube-Shadow-OpenAI-LSTM-Direct-v0", - entry_point="omni.isaac.lab_tasks.direct.inhand_manipulation:InHandManipulationEnv", + entry_point=f"{inhand_task_entry}.inhand_manipulation_env:InHandManipulationEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": ShadowHandOpenAIEnvCfg, + "env_cfg_entry_point": f"{__name__}.shadow_hand_env_cfg:ShadowHandOpenAIEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_lstm_cfg.yaml", }, ) @@ -55,10 +55,10 @@ gym.register( id="Isaac-Repose-Cube-Shadow-Vision-Direct-v0", - entry_point="omni.isaac.lab_tasks.direct.shadow_hand.shadow_hand_vision_env:ShadowHandVisionEnv", + entry_point=f"{__name__}.shadow_hand_vision_env:ShadowHandVisionEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": ShadowHandVisionEnvCfg, + "env_cfg_entry_point": f"{__name__}.shadow_hand_vision_env:ShadowHandVisionEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:ShadowHandVisionFFPPORunnerCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_vision_cfg.yaml", }, @@ -66,10 +66,10 @@ gym.register( id="Isaac-Repose-Cube-Shadow-Vision-Direct-Play-v0", - entry_point="omni.isaac.lab_tasks.direct.shadow_hand.shadow_hand_vision_env:ShadowHandVisionEnv", + entry_point=f"{__name__}.shadow_hand_vision_env:ShadowHandVisionEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": ShadowHandVisionEnvPlayCfg, + "env_cfg_entry_point": f"{__name__}.shadow_hand_vision_env:ShadowHandVisionEnvPlayCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:ShadowHandVisionFFPPORunnerCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_vision_cfg.yaml", }, diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand/feature_extractor.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand/feature_extractor.py index 1dbfb39b1a..fc92bbfb1b 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand/feature_extractor.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand/feature_extractor.py @@ -7,7 +7,6 @@ import os import torch import torch.nn as nn - import torchvision from omni.isaac.lab.sensors import save_images_to_file diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand/shadow_hand_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand/shadow_hand_env_cfg.py index f4b8407296..af88124792 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand/shadow_hand_env_cfg.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand/shadow_hand_env_cfg.py @@ -119,9 +119,9 @@ class ShadowHandEnvCfg(DirectRLEnvCfg): # env decimation = 2 episode_length_s = 10.0 - num_actions = 20 - num_observations = 157 # (full) - num_states = 0 + action_space = 20 + observation_space = 157 # (full) + state_space = 0 asymmetric_obs = False obs_type = "full" @@ -232,9 +232,9 @@ class ShadowHandOpenAIEnvCfg(ShadowHandEnvCfg): # env decimation = 3 episode_length_s = 8.0 - num_actions = 20 - num_observations = 42 - num_states = 187 + action_space = 20 + observation_space = 42 + state_space = 187 asymmetric_obs = True obs_type = "openai" # simulation diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand/shadow_hand_vision_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand/shadow_hand_vision_env.py index b025bfb052..8677c16eba 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand/shadow_hand_vision_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand/shadow_hand_vision_env.py @@ -48,8 +48,8 @@ class ShadowHandVisionEnvCfg(ShadowHandEnvCfg): feature_extractor = FeatureExtractorCfg() # env - num_observations = 164 + 27 # state observation + vision CNN embedding - num_states = 187 + 27 # asymettric states + vision CNN embedding + observation_space = 164 + 27 # state observation + vision CNN embedding + state_space = 187 + 27 # asymettric states + vision CNN embedding @configclass @@ -88,7 +88,7 @@ def _setup_scene(self): sem.GetSemanticDataAttr().Set("cube") # clone and replicate (no need to filter for this environment) self.scene.clone_environments(copy_from_source=False) - # add articultion to scene - we must register to scene to randomize with EventManager + # add articulation to scene - we must register to scene to randomize with EventManager self.scene.articulations["robot"] = self.hand self.scene.rigid_objects["object"] = self.object self.scene.sensors["tiled_camera"] = self._tiled_camera diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand_over/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand_over/__init__.py index 0beb0a0d61..8eef9dd395 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand_over/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand_over/__init__.py @@ -10,8 +10,6 @@ import gymnasium as gym from . import agents -from .shadow_hand_over_env import ShadowHandOverEnv -from .shadow_hand_over_env_cfg import ShadowHandOverEnvCfg ## # Register Gym environments. @@ -19,10 +17,10 @@ gym.register( id="Isaac-Shadow-Hand-Over-Direct-v0", - entry_point="omni.isaac.lab_tasks.direct.shadow_hand_over:ShadowHandOverEnv", + entry_point=f"{__name__}.shadow_hand_over_env:ShadowHandOverEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": ShadowHandOverEnvCfg, + "env_cfg_entry_point": f"{__name__}.shadow_hand_over_env_cfg:ShadowHandOverEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", "skrl_ippo_cfg_entry_point": f"{agents.__name__}:skrl_ippo_cfg.yaml", diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py index 01ac2c22d7..622a95e842 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py @@ -322,7 +322,8 @@ def _reset_idx(self, env_ids: Sequence[int] | torch.Tensor | None): ) object_default_state[:, 7:] = torch.zeros_like(self.object.data.default_root_state[env_ids, 7:]) - self.object.write_root_state_to_sim(object_default_state, env_ids) + self.object.write_root_link_pose_to_sim(object_default_state[:, :7], env_ids) + self.object.write_root_com_velocity_to_sim(object_default_state[:, 7:], env_ids) # reset right hand delta_max = self.hand_dof_upper_limits[env_ids] - self.right_hand.data.default_joint_pos[env_ids] @@ -376,33 +377,33 @@ def _reset_target_pose(self, env_ids): def _compute_intermediate_values(self): # data for right hand - self.right_fingertip_pos = self.right_hand.data.body_pos_w[:, self.finger_bodies] - self.right_fingertip_rot = self.right_hand.data.body_quat_w[:, self.finger_bodies] + self.right_fingertip_pos = self.right_hand.data.body_link_pos_w[:, self.finger_bodies] + self.right_fingertip_rot = self.right_hand.data.body_link_quat_w[:, self.finger_bodies] self.right_fingertip_pos -= self.scene.env_origins.repeat((1, self.num_fingertips)).reshape( self.num_envs, self.num_fingertips, 3 ) - self.right_fingertip_velocities = self.right_hand.data.body_vel_w[:, self.finger_bodies] + self.right_fingertip_velocities = self.right_hand.data.body_com_vel_w[:, self.finger_bodies] self.right_hand_dof_pos = self.right_hand.data.joint_pos self.right_hand_dof_vel = self.right_hand.data.joint_vel # data for left hand - self.left_fingertip_pos = self.left_hand.data.body_pos_w[:, self.finger_bodies] - self.left_fingertip_rot = self.left_hand.data.body_quat_w[:, self.finger_bodies] + self.left_fingertip_pos = self.left_hand.data.body_link_pos_w[:, self.finger_bodies] + self.left_fingertip_rot = self.left_hand.data.body_link_quat_w[:, self.finger_bodies] self.left_fingertip_pos -= self.scene.env_origins.repeat((1, self.num_fingertips)).reshape( self.num_envs, self.num_fingertips, 3 ) - self.left_fingertip_velocities = self.left_hand.data.body_vel_w[:, self.finger_bodies] + self.left_fingertip_velocities = self.left_hand.data.body_com_vel_w[:, self.finger_bodies] self.left_hand_dof_pos = self.left_hand.data.joint_pos self.left_hand_dof_vel = self.left_hand.data.joint_vel # data for object - self.object_pos = self.object.data.root_pos_w - self.scene.env_origins - self.object_rot = self.object.data.root_quat_w - self.object_velocities = self.object.data.root_vel_w - self.object_linvel = self.object.data.root_lin_vel_w - self.object_angvel = self.object.data.root_ang_vel_w + self.object_pos = self.object.data.root_link_pos_w - self.scene.env_origins + self.object_rot = self.object.data.root_link_quat_w + self.object_velocities = self.object.data.root_com_vel_w + self.object_linvel = self.object.data.root_com_lin_vel_w + self.object_angvel = self.object.data.root_com_ang_vel_w @torch.jit.script diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand_over/shadow_hand_over_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand_over/shadow_hand_over_env_cfg.py index d6dbb3d6a2..d3a7c33b3f 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand_over/shadow_hand_over_env_cfg.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand_over/shadow_hand_over_env_cfg.py @@ -118,9 +118,9 @@ class ShadowHandOverEnvCfg(DirectMARLEnvCfg): decimation = 2 episode_length_s = 7.5 possible_agents = ["right_hand", "left_hand"] - num_actions = {"right_hand": 20, "left_hand": 20} - num_observations = {"right_hand": 157, "left_hand": 157} - num_states = 290 + action_spaces = {"right_hand": 20, "left_hand": 20} + observation_spaces = {"right_hand": 157, "left_hand": 157} + state_space = 290 # simulation sim: SimulationCfg = SimulationCfg( diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/ant/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/ant/__init__.py index ed3c86eab3..a128cf5116 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/ant/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/ant/__init__.py @@ -9,7 +9,7 @@ import gymnasium as gym -from . import agents, ant_env_cfg +from . import agents ## # Register Gym environments. @@ -20,7 +20,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": ant_env_cfg.AntEnvCfg, + "env_cfg_entry_point": f"{__name__}.ant_env_cfg:AntEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AntPPORunnerCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/ant/ant_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/ant/ant_env_cfg.py index f12a046305..68f091c515 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/ant/ant_env_cfg.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/ant/ant_env_cfg.py @@ -58,14 +58,6 @@ class MySceneCfg(InteractiveSceneCfg): ## -@configclass -class CommandsCfg: - """Command terms for the MDP.""" - - # no commands for this MDP - null = mdp.NullCommandCfg() - - @configclass class ActionsCfg: """Action specifications for the MDP.""" @@ -163,13 +155,6 @@ class TerminationsCfg: torso_height = DoneTerm(func=mdp.root_height_below_minimum, params={"minimum_height": 0.31}) -@configclass -class CurriculumCfg: - """Curriculum terms for the MDP.""" - - pass - - @configclass class AntEnvCfg(ManagerBasedRLEnvCfg): """Configuration for the MuJoCo-style Ant walking environment.""" @@ -179,13 +164,10 @@ class AntEnvCfg(ManagerBasedRLEnvCfg): # 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.""" diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/__init__.py index 7a3070d775..3b0107ecb7 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/__init__.py @@ -10,8 +10,6 @@ import gymnasium as gym from . import agents -from .cartpole_camera_env_cfg import CartpoleDepthCameraEnvCfg, CartpoleRGBCameraEnvCfg -from .cartpole_env_cfg import CartpoleEnvCfg ## # Register Gym environments. @@ -22,7 +20,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": CartpoleEnvCfg, + "env_cfg_entry_point": f"{__name__}.cartpole_env_cfg:CartpoleEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CartpolePPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", @@ -31,21 +29,41 @@ ) gym.register( - id="Isaac-Cartpole-RGB-Camera-v0", + id="Isaac-Cartpole-RGB-v0", entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": CartpoleRGBCameraEnvCfg, + "env_cfg_entry_point": f"{__name__}.cartpole_camera_env_cfg:CartpoleRGBCameraEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_camera_ppo_cfg.yaml", }, ) gym.register( - id="Isaac-Cartpole-Depth-Camera-v0", + id="Isaac-Cartpole-Depth-v0", entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": CartpoleDepthCameraEnvCfg, + "env_cfg_entry_point": f"{__name__}.cartpole_camera_env_cfg:CartpoleDepthCameraEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_camera_ppo_cfg.yaml", }, ) + +gym.register( + id="Isaac-Cartpole-RGB-ResNet18-v0", + entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.cartpole_camera_env_cfg:CartpoleResNet18CameraEnvCfg", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_feature_ppo_cfg.yaml", + }, +) + +gym.register( + id="Isaac-Cartpole-RGB-TheiaTiny-v0", + entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.cartpole_camera_env_cfg:CartpoleTheiaTinyCameraEnvCfg", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_feature_ppo_cfg.yaml", + }, +) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/agents/rl_games_feature_ppo_cfg.yaml b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/agents/rl_games_feature_ppo_cfg.yaml new file mode 100644 index 0000000000..41e265e9f2 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/agents/rl_games_feature_ppo_cfg.yaml @@ -0,0 +1,79 @@ +params: + seed: 42 + + # environment wrapper clipping + env: + # added to the wrapper + clip_observations: 5.0 + # can make custom wrapper? + clip_actions: 1.0 + + algo: + name: a2c_continuous + + model: + name: continuous_a2c_logstd + + # doesn't have this fine grained control but made it close + network: + name: actor_critic + separate: False + space: + continuous: + mu_activation: None + sigma_activation: None + + mu_init: + name: default + sigma_init: + name: const_initializer + val: 0 + fixed_sigma: True + mlp: + units: [256] + activation: elu + d2rl: False + + initializer: + name: default + regularizer: + name: None + + load_checkpoint: False # flag which sets whether to load the checkpoint + load_path: '' # path to the checkpoint to load + + config: + name: cartpole_features + env_name: rlgpu + device: 'cuda:0' + device_name: 'cuda:0' + multi_gpu: False + ppo: True + mixed_precision: False + normalize_input: True + normalize_value: True + value_bootstraop: True + num_actors: -1 # configured from the script (based on num_envs) + reward_shaper: + scale_value: 1.0 + normalize_advantage: True + gamma: 0.99 + tau : 0.95 + learning_rate: 3e-4 + lr_schedule: adaptive + kl_threshold: 0.008 + score_to_win: 20000 + max_epochs: 200 + save_best_after: 50 + save_frequency: 25 + grad_norm: 1.0 + entropy_coef: 0.0 + truncate_grads: True + e_clip: 0.2 + horizon_length: 16 + minibatch_size: 2048 + mini_epochs: 8 + critic_coef: 4 + clip_value: True + seq_length: 4 + bounds_loss_coef: 0.0001 diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/cartpole_camera_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/cartpole_camera_env_cfg.py index ce5a6c90b8..35c828d2cb 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/cartpole_camera_env_cfg.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/cartpole_camera_env_cfg.py @@ -78,7 +78,7 @@ class DepthObservationsCfg: """Observation specifications for the MDP.""" @configclass - class DepthCameraPolicyCfg(RGBObservationsCfg.RGBCameraPolicyCfg): + class DepthCameraPolicyCfg(ObsGroup): """Observations for policy group with depth images.""" image = ObsTerm( @@ -88,6 +88,43 @@ class DepthCameraPolicyCfg(RGBObservationsCfg.RGBCameraPolicyCfg): policy: ObsGroup = DepthCameraPolicyCfg() +@configclass +class ResNet18ObservationCfg: + """Observation specifications for the MDP.""" + + @configclass + class ResNet18FeaturesCameraPolicyCfg(ObsGroup): + """Observations for policy group with features extracted from RGB images with a frozen ResNet18.""" + + image = ObsTerm( + func=mdp.image_features, + params={"sensor_cfg": SceneEntityCfg("tiled_camera"), "data_type": "rgb", "model_name": "resnet18"}, + ) + + policy: ObsGroup = ResNet18FeaturesCameraPolicyCfg() + + +@configclass +class TheiaTinyObservationCfg: + """Observation specifications for the MDP.""" + + @configclass + class TheiaTinyFeaturesCameraPolicyCfg(ObsGroup): + """Observations for policy group with features extracted from RGB images with a frozen Theia-Tiny Transformer""" + + image = ObsTerm( + func=mdp.image_features, + params={ + "sensor_cfg": SceneEntityCfg("tiled_camera"), + "data_type": "rgb", + "model_name": "theia-tiny-patch16-224-cddsv", + "model_device": "cuda:0", + }, + ) + + policy: ObsGroup = TheiaTinyFeaturesCameraPolicyCfg() + + ## # Environment configuration ## @@ -97,13 +134,43 @@ class DepthCameraPolicyCfg(RGBObservationsCfg.RGBCameraPolicyCfg): class CartpoleRGBCameraEnvCfg(CartpoleEnvCfg): """Configuration for the cartpole environment with RGB camera.""" - scene: CartpoleSceneCfg = CartpoleRGBCameraSceneCfg(num_envs=1024, env_spacing=20) + scene: CartpoleRGBCameraSceneCfg = CartpoleRGBCameraSceneCfg(num_envs=1024, env_spacing=20) observations: RGBObservationsCfg = RGBObservationsCfg() + def __post_init__(self): + super().__post_init__() + # remove ground as it obstructs the camera + self.scene.ground = None + # viewer settings + self.viewer.eye = (7.0, 0.0, 2.5) + self.viewer.lookat = (0.0, 0.0, 2.5) + @configclass class CartpoleDepthCameraEnvCfg(CartpoleEnvCfg): """Configuration for the cartpole environment with depth camera.""" - scene: CartpoleSceneCfg = CartpoleDepthCameraSceneCfg(num_envs=1024, env_spacing=20) + scene: CartpoleDepthCameraSceneCfg = CartpoleDepthCameraSceneCfg(num_envs=1024, env_spacing=20) observations: DepthObservationsCfg = DepthObservationsCfg() + + def __post_init__(self): + super().__post_init__() + # remove ground as it obstructs the camera + self.scene.ground = None + # viewer settings + self.viewer.eye = (7.0, 0.0, 2.5) + self.viewer.lookat = (0.0, 0.0, 2.5) + + +@configclass +class CartpoleResNet18CameraEnvCfg(CartpoleRGBCameraEnvCfg): + """Configuration for the cartpole environment with ResNet18 features as observations.""" + + observations: ResNet18ObservationCfg = ResNet18ObservationCfg() + + +@configclass +class CartpoleTheiaTinyCameraEnvCfg(CartpoleRGBCameraEnvCfg): + """Configuration for the cartpole environment with Theia-Tiny features as observations.""" + + observations: TheiaTinyObservationCfg = TheiaTinyObservationCfg() diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/cartpole_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/cartpole_env_cfg.py index 8c92d3d5ae..ca8784d8db 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/cartpole_env_cfg.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/cartpole_env_cfg.py @@ -48,11 +48,6 @@ class CartpoleSceneCfg(InteractiveSceneCfg): prim_path="/World/DomeLight", spawn=sim_utils.DomeLightCfg(color=(0.9, 0.9, 0.9), intensity=500.0), ) - distant_light = AssetBaseCfg( - prim_path="/World/DistantLight", - spawn=sim_utils.DistantLightCfg(color=(0.9, 0.9, 0.9), intensity=2500.0), - init_state=AssetBaseCfg.InitialStateCfg(rot=(0.738, 0.477, 0.477, 0.0)), - ) ## @@ -60,14 +55,6 @@ class CartpoleSceneCfg(InteractiveSceneCfg): ## -@configclass -class CommandsCfg: - """Command terms for the MDP.""" - - # no commands for this MDP - null = mdp.NullCommandCfg() - - @configclass class ActionsCfg: """Action specifications for the MDP.""" @@ -162,13 +149,6 @@ class TerminationsCfg: ) -@configclass -class CurriculumCfg: - """Configuration for the curriculum.""" - - pass - - ## # Environment configuration ## @@ -185,11 +165,8 @@ class CartpoleEnvCfg(ManagerBasedRLEnvCfg): actions: ActionsCfg = ActionsCfg() events: EventCfg = EventCfg() # MDP settings - curriculum: CurriculumCfg = CurriculumCfg() rewards: RewardsCfg = RewardsCfg() terminations: TerminationsCfg = TerminationsCfg() - # No command generator - commands: CommandsCfg = CommandsCfg() # Post initialization def __post_init__(self) -> None: diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/__init__.py index c902b83af5..497901f711 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/__init__.py @@ -9,7 +9,7 @@ import gymnasium as gym -from . import agents, humanoid_env_cfg +from . import agents ## # Register Gym environments. @@ -20,7 +20,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": humanoid_env_cfg.HumanoidEnvCfg, + "env_cfg_entry_point": f"{__name__}.humanoid_env_cfg:HumanoidEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:HumanoidPPORunnerCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/humanoid_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/humanoid_env_cfg.py index f376811a43..e02dd94edb 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/humanoid_env_cfg.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/humanoid_env_cfg.py @@ -102,14 +102,6 @@ class MySceneCfg(InteractiveSceneCfg): ## -@configclass -class CommandsCfg: - """Command terms for the MDP.""" - - # no commands for this MDP - null = mdp.NullCommandCfg() - - @configclass class ActionsCfg: """Action specifications for the MDP.""" @@ -248,13 +240,6 @@ class TerminationsCfg: torso_height = DoneTerm(func=mdp.root_height_below_minimum, params={"minimum_height": 0.8}) -@configclass -class CurriculumCfg: - """Curriculum terms for the MDP.""" - - pass - - @configclass class HumanoidEnvCfg(ManagerBasedRLEnvCfg): """Configuration for the MuJoCo-style Humanoid walking environment.""" @@ -264,13 +249,10 @@ class HumanoidEnvCfg(ManagerBasedRLEnvCfg): # 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.""" diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/mdp/observations.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/mdp/observations.py index 1a65c381a7..e74b872582 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/mdp/observations.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/mdp/observations.py @@ -21,7 +21,7 @@ def base_yaw_roll(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityC # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] # extract euler angles (in world frame) - roll, _, yaw = math_utils.euler_xyz_from_quat(asset.data.root_quat_w) + roll, _, yaw = math_utils.euler_xyz_from_quat(asset.data.root_link_quat_w) # normalize angle to [-pi, pi] roll = torch.atan2(torch.sin(roll), torch.cos(roll)) yaw = torch.atan2(torch.sin(yaw), torch.cos(yaw)) @@ -46,11 +46,11 @@ def base_heading_proj( # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] # compute desired heading direction - to_target_pos = torch.tensor(target_pos, device=env.device) - asset.data.root_pos_w[:, :3] + to_target_pos = torch.tensor(target_pos, device=env.device) - asset.data.root_link_pos_w[:, :3] to_target_pos[:, 2] = 0.0 to_target_dir = math_utils.normalize(to_target_pos) # compute base forward vector - heading_vec = math_utils.quat_rotate(asset.data.root_quat_w, asset.data.FORWARD_VEC_B) + heading_vec = math_utils.quat_rotate(asset.data.root_link_quat_w, asset.data.FORWARD_VEC_B) # compute dot product between heading and target direction heading_proj = torch.bmm(heading_vec.view(env.num_envs, 1, 3), to_target_dir.view(env.num_envs, 3, 1)) @@ -64,10 +64,10 @@ def base_angle_to_target( # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] # compute desired heading direction - to_target_pos = torch.tensor(target_pos, device=env.device) - asset.data.root_pos_w[:, :3] + to_target_pos = torch.tensor(target_pos, device=env.device) - asset.data.root_link_pos_w[:, :3] walk_target_angle = torch.atan2(to_target_pos[:, 1], to_target_pos[:, 0]) # compute base forward vector - _, _, yaw = math_utils.euler_xyz_from_quat(asset.data.root_quat_w) + _, _, yaw = math_utils.euler_xyz_from_quat(asset.data.root_link_quat_w) # normalize angle to target to [-pi, pi] angle_to_target = walk_target_angle - yaw angle_to_target = torch.atan2(torch.sin(angle_to_target), torch.cos(angle_to_target)) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/mdp/rewards.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/mdp/rewards.py index b21be8d31f..913be8c1d2 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/mdp/rewards.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/mdp/rewards.py @@ -53,7 +53,7 @@ def reset(self, env_ids: torch.Tensor): asset: Articulation = self._env.scene["robot"] # compute projection of current heading to desired heading vector target_pos = torch.tensor(self.cfg.params["target_pos"], device=self.device) - to_target_pos = target_pos - asset.data.root_pos_w[env_ids, :3] + to_target_pos = target_pos - asset.data.root_link_pos_w[env_ids, :3] # reward terms self.potentials[env_ids] = -torch.norm(to_target_pos, p=2, dim=-1) / self._env.step_dt self.prev_potentials[env_ids] = self.potentials[env_ids] @@ -68,7 +68,7 @@ def __call__( asset: Articulation = env.scene[asset_cfg.name] # compute vector to target target_pos = torch.tensor(target_pos, device=env.device) - to_target_pos = target_pos - asset.data.root_pos_w[:, :3] + to_target_pos = target_pos - asset.data.root_link_pos_w[:, :3] to_target_pos[:, 2] = 0.0 # update history buffer and compute new potential self.prev_potentials[:] = self.potentials[:] diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/a1/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/a1/__init__.py index 80e9c3750f..fd628c91f5 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/a1/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/a1/__init__.py @@ -5,7 +5,7 @@ import gymnasium as gym -from . import agents, flat_env_cfg, rough_env_cfg +from . import agents ## # Register Gym environments. @@ -16,7 +16,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": flat_env_cfg.UnitreeA1FlatEnvCfg, + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeA1FlatEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeA1FlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -27,7 +27,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": flat_env_cfg.UnitreeA1FlatEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeA1FlatEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeA1FlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -38,7 +38,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": rough_env_cfg.UnitreeA1RoughEnvCfg, + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:UnitreeA1RoughEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeA1RoughPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", }, @@ -49,7 +49,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": rough_env_cfg.UnitreeA1RoughEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:UnitreeA1RoughEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeA1RoughPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", }, diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/anymal_b/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/anymal_b/__init__.py index 65e89a1d9e..3796c67d81 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/anymal_b/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/anymal_b/__init__.py @@ -4,7 +4,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gymnasium as gym -from . import agents, flat_env_cfg, rough_env_cfg +from . import agents ## # Register Gym environments. @@ -15,7 +15,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": flat_env_cfg.AnymalBFlatEnvCfg, + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalBFlatEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalBFlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -26,7 +26,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": flat_env_cfg.AnymalBFlatEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalBFlatEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalBFlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -37,7 +37,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": rough_env_cfg.AnymalBRoughEnvCfg, + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:AnymalBRoughEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalBRoughPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", }, @@ -48,7 +48,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": rough_env_cfg.AnymalBRoughEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:AnymalBRoughEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalBRoughPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", }, diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/anymal_c/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/anymal_c/__init__.py index 1d5aa41c87..770b52337e 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/anymal_c/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/anymal_c/__init__.py @@ -5,7 +5,7 @@ import gymnasium as gym -from . import agents, flat_env_cfg, rough_env_cfg +from . import agents ## # Register Gym environments. @@ -16,7 +16,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": flat_env_cfg.AnymalCFlatEnvCfg, + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalCFlatEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalCFlatPPORunnerCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_flat_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", @@ -28,7 +28,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": flat_env_cfg.AnymalCFlatEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalCFlatEnvCfg_PLAY", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_flat_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalCFlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", @@ -40,7 +40,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": rough_env_cfg.AnymalCRoughEnvCfg, + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:AnymalCRoughEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_rough_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalCRoughPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", @@ -52,7 +52,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": rough_env_cfg.AnymalCRoughEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:AnymalCRoughEnvCfg_PLAY", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_rough_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalCRoughPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/anymal_d/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/anymal_d/__init__.py index 0f6b8a4776..eea14799e9 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/anymal_d/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/anymal_d/__init__.py @@ -5,7 +5,7 @@ import gymnasium as gym -from . import agents, flat_env_cfg, rough_env_cfg +from . import agents ## # Register Gym environments. @@ -16,7 +16,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": flat_env_cfg.AnymalDFlatEnvCfg, + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalDFlatEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDFlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -27,7 +27,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": flat_env_cfg.AnymalDFlatEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalDFlatEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDFlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -38,7 +38,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": rough_env_cfg.AnymalDRoughEnvCfg, + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:AnymalDRoughEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDRoughPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", }, @@ -49,7 +49,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": rough_env_cfg.AnymalDRoughEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:AnymalDRoughEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDRoughPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", }, diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/cassie/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/cassie/__init__.py index 495491f617..69658caf46 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/cassie/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/cassie/__init__.py @@ -5,7 +5,7 @@ import gymnasium as gym -from . import agents, flat_env_cfg, rough_env_cfg +from . import agents ## # Register Gym environments. @@ -16,7 +16,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": flat_env_cfg.CassieFlatEnvCfg, + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:CassieFlatEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CassieFlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -27,7 +27,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": flat_env_cfg.CassieFlatEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:CassieFlatEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CassieFlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -38,7 +38,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": rough_env_cfg.CassieRoughEnvCfg, + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:CassieRoughEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CassieRoughPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", }, @@ -49,7 +49,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": rough_env_cfg.CassieRoughEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:CassieRoughEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CassieRoughPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", }, diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/g1/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/g1/__init__.py index c84c627099..316252f77d 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/g1/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/g1/__init__.py @@ -5,19 +5,18 @@ import gymnasium as gym -from . import agents, flat_env_cfg, rough_env_cfg +from . import agents ## # Register Gym environments. ## - gym.register( id="Isaac-Velocity-Rough-G1-v0", entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": rough_env_cfg.G1RoughEnvCfg, + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:G1RoughEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:G1RoughPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", }, @@ -29,7 +28,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": rough_env_cfg.G1RoughEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:G1RoughEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:G1RoughPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", }, @@ -41,7 +40,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": flat_env_cfg.G1FlatEnvCfg, + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:G1FlatEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:G1FlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -53,7 +52,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": flat_env_cfg.G1FlatEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:G1FlatEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:G1FlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/g1/rough_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/g1/rough_env_cfg.py index 98503f4fb8..0be5d6eba8 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/g1/rough_env_cfg.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/g1/rough_env_cfg.py @@ -5,7 +5,6 @@ from omni.isaac.lab.managers import RewardTermCfg as RewTerm from omni.isaac.lab.managers import SceneEntityCfg -from omni.isaac.lab.managers import TerminationTermCfg as DoneTerm from omni.isaac.lab.utils import configclass import omni.isaac.lab_tasks.manager_based.locomotion.velocity.mdp as mdp @@ -104,21 +103,9 @@ class G1Rewards(RewardsCfg): ) -@configclass -class TerminationsCfg: - """Termination terms for the MDP.""" - - time_out = DoneTerm(func=mdp.time_out, time_out=True) - base_contact = DoneTerm( - func=mdp.illegal_contact, - params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names="torso_link"), "threshold": 1.0}, - ) - - @configclass class G1RoughEnvCfg(LocomotionVelocityRoughEnvCfg): rewards: G1Rewards = G1Rewards() - terminations: TerminationsCfg = TerminationsCfg() def __post_init__(self): # post init of parent @@ -163,6 +150,9 @@ def __post_init__(self): self.commands.base_velocity.ranges.lin_vel_y = (-0.0, 0.0) self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0) + # terminations + self.terminations.base_contact.params["sensor_cfg"].body_names = "torso_link" + @configclass class G1RoughEnvCfg_PLAY(G1RoughEnvCfg): diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/go1/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/go1/__init__.py index 260e68cff1..b30722bc7f 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/go1/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/go1/__init__.py @@ -5,7 +5,7 @@ import gymnasium as gym -from . import agents, flat_env_cfg, rough_env_cfg +from . import agents ## # Register Gym environments. @@ -16,7 +16,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": flat_env_cfg.UnitreeGo1FlatEnvCfg, + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeGo1FlatEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeGo1FlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -27,7 +27,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": flat_env_cfg.UnitreeGo1FlatEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeGo1FlatEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeGo1FlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -38,7 +38,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": rough_env_cfg.UnitreeGo1RoughEnvCfg, + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:UnitreeGo1RoughEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeGo1RoughPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", }, @@ -49,7 +49,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": rough_env_cfg.UnitreeGo1RoughEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:UnitreeGo1RoughEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeGo1RoughPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", }, diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/go2/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/go2/__init__.py index 27b3483956..85071ad547 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/go2/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/go2/__init__.py @@ -5,7 +5,7 @@ import gymnasium as gym -from . import agents, flat_env_cfg, rough_env_cfg +from . import agents ## # Register Gym environments. @@ -16,7 +16,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": flat_env_cfg.UnitreeGo2FlatEnvCfg, + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeGo2FlatEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeGo2FlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -27,7 +27,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": flat_env_cfg.UnitreeGo2FlatEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeGo2FlatEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeGo2FlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -38,7 +38,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": rough_env_cfg.UnitreeGo2RoughEnvCfg, + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:UnitreeGo2RoughEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeGo2RoughPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", }, @@ -49,7 +49,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": rough_env_cfg.UnitreeGo2RoughEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:UnitreeGo2RoughEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeGo2RoughPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", }, diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/h1/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/h1/__init__.py index 5c405ea4ef..b03937ae60 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/h1/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/h1/__init__.py @@ -5,19 +5,18 @@ import gymnasium as gym -from . import agents, flat_env_cfg, rough_env_cfg +from . import agents ## # Register Gym environments. ## - gym.register( id="Isaac-Velocity-Rough-H1-v0", entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": rough_env_cfg.H1RoughEnvCfg, + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:H1RoughEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:H1RoughPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", }, @@ -29,7 +28,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": rough_env_cfg.H1RoughEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:H1RoughEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:H1RoughPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", }, @@ -41,7 +40,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": flat_env_cfg.H1FlatEnvCfg, + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:H1FlatEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:H1FlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -53,7 +52,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": flat_env_cfg.H1FlatEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:H1FlatEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:H1FlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/h1/rough_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/h1/rough_env_cfg.py index 516119908f..e1414cb364 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/h1/rough_env_cfg.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/h1/rough_env_cfg.py @@ -5,7 +5,6 @@ from omni.isaac.lab.managers import RewardTermCfg as RewTerm from omni.isaac.lab.managers import SceneEntityCfg -from omni.isaac.lab.managers import TerminationTermCfg as DoneTerm from omni.isaac.lab.utils import configclass import omni.isaac.lab_tasks.manager_based.locomotion.velocity.mdp as mdp @@ -71,21 +70,9 @@ class H1Rewards(RewardsCfg): ) -@configclass -class TerminationsCfg: - """Termination terms for the MDP.""" - - time_out = DoneTerm(func=mdp.time_out, time_out=True) - base_contact = DoneTerm( - func=mdp.illegal_contact, - params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*torso_link"), "threshold": 1.0}, - ) - - @configclass class H1RoughEnvCfg(LocomotionVelocityRoughEnvCfg): rewards: H1Rewards = H1Rewards() - terminations: TerminationsCfg = TerminationsCfg() def __post_init__(self): # post init of parent @@ -127,6 +114,9 @@ def __post_init__(self): self.commands.base_velocity.ranges.lin_vel_y = (0.0, 0.0) self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0) + # terminations + self.terminations.base_contact.params["sensor_cfg"].body_names = ".*torso_link" + @configclass class H1RoughEnvCfg_PLAY(H1RoughEnvCfg): diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/spot/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/spot/__init__.py index fda9123ea6..a2a852dd68 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/spot/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/spot/__init__.py @@ -5,7 +5,7 @@ import gymnasium as gym -from . import agents, flat_env_cfg +from . import agents ## # Register Gym environments. @@ -16,7 +16,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": flat_env_cfg.SpotFlatEnvCfg, + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:SpotFlatEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:SpotFlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -27,7 +27,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": flat_env_cfg.SpotFlatEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:SpotFlatEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:SpotFlatPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/spot/flat_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/spot/flat_env_cfg.py index 986b10da33..fba2e69ee1 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/spot/flat_env_cfg.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/spot/flat_env_cfg.py @@ -293,13 +293,6 @@ class SpotTerminationsCfg: ) -@configclass -class SpotCurriculumCfg: - """Curriculum terms for the MDP.""" - - pass - - @configclass class SpotFlatEnvCfg(LocomotionVelocityRoughEnvCfg): @@ -312,7 +305,6 @@ class SpotFlatEnvCfg(LocomotionVelocityRoughEnvCfg): rewards: SpotRewardsCfg = SpotRewardsCfg() terminations: SpotTerminationsCfg = SpotTerminationsCfg() events: SpotEventCfg = SpotEventCfg() - curriculum: SpotCurriculumCfg = SpotCurriculumCfg() # Viewer viewer = ViewerCfg(eye=(10.5, 10.5, 0.3), origin_type="world", env_index=0, asset_name="robot") diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/spot/mdp/rewards.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/spot/mdp/rewards.py index d4deab6a6a..f587827f73 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/spot/mdp/rewards.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/spot/mdp/rewards.py @@ -49,7 +49,7 @@ def air_time_reward( 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) + body_vel = torch.linalg.norm(asset.data.root_com_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), @@ -64,7 +64,7 @@ def base_angular_velocity_reward(env: ManagerBasedRLEnv, asset_cfg: SceneEntityC asset: RigidObject = env.scene[asset_cfg.name] # compute the error target = env.command_manager.get_command("base_velocity")[:, 2] - ang_vel_error = torch.linalg.norm((target - asset.data.root_ang_vel_b[:, 2]).unsqueeze(1), dim=1) + ang_vel_error = torch.linalg.norm((target - asset.data.root_com_ang_vel_b[:, 2]).unsqueeze(1), dim=1) return torch.exp(-ang_vel_error / std) @@ -76,7 +76,7 @@ def base_linear_velocity_reward( asset: RigidObject = env.scene[asset_cfg.name] # compute the error target = env.command_manager.get_command("base_velocity")[:, :2] - lin_vel_error = torch.linalg.norm((target - asset.data.root_lin_vel_b[:, :2]), dim=1) + lin_vel_error = torch.linalg.norm((target - asset.data.root_com_lin_vel_b[:, :2]), dim=1) # fixed 1.0 multiple for tracking below the ramp_at_vel value, then scale by the rate above vel_cmd_magnitude = torch.linalg.norm(target, dim=1) velocity_scaling_multiple = torch.clamp(1.0 + ramp_rate * (vel_cmd_magnitude - ramp_at_vel), min=1.0) @@ -148,7 +148,7 @@ def __call__( async_reward = async_reward_0 * async_reward_1 * async_reward_2 * async_reward_3 # only enforce gait if cmd > 0 cmd = torch.norm(env.command_manager.get_command("base_velocity"), dim=1) - body_vel = torch.linalg.norm(self.asset.data.root_lin_vel_b[:, :2], dim=1) + body_vel = torch.linalg.norm(self.asset.data.root_com_lin_vel_b[:, :2], dim=1) return torch.where( torch.logical_or(cmd > 0.0, body_vel > self.velocity_threshold), sync_reward * async_reward, 0.0 ) @@ -182,8 +182,10 @@ def foot_clearance_reward( ) -> 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.norm(asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2], dim=2)) + foot_z_target_error = torch.square(asset.data.body_link_pos_w[:, asset_cfg.body_ids, 2] - target_height) + foot_velocity_tanh = torch.tanh( + tanh_mult * torch.norm(asset.data.body_com_lin_vel_w[:, asset_cfg.body_ids, :2], dim=2) + ) reward = foot_z_target_error * foot_velocity_tanh return torch.exp(-torch.sum(reward, dim=1) / std) @@ -217,8 +219,8 @@ def base_motion_penalty(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> to """Penalize base vertical and roll/pitch velocity""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return 0.8 * torch.square(asset.data.root_lin_vel_b[:, 2]) + 0.2 * torch.sum( - torch.abs(asset.data.root_ang_vel_b[:, :2]), dim=1 + return 0.8 * torch.square(asset.data.root_com_lin_vel_b[:, 2]) + 0.2 * torch.sum( + torch.abs(asset.data.root_com_ang_vel_b[:, :2]), dim=1 ) @@ -243,7 +245,7 @@ def foot_slip_penalty( # 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 - foot_planar_velocity = torch.linalg.norm(asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2], dim=2) + foot_planar_velocity = torch.linalg.norm(asset.data.body_com_lin_vel_w[:, asset_cfg.body_ids, :2], dim=2) reward = is_contact * foot_planar_velocity return torch.sum(reward, dim=1) @@ -263,7 +265,7 @@ def joint_position_penalty( # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] cmd = torch.linalg.norm(env.command_manager.get_command("base_velocity"), dim=1) - body_vel = torch.linalg.norm(asset.data.root_lin_vel_b[:, :2], dim=1) + body_vel = torch.linalg.norm(asset.data.root_com_lin_vel_b[:, :2], dim=1) reward = torch.linalg.norm((asset.data.joint_pos - asset.data.default_joint_pos), dim=1) return torch.where(torch.logical_or(cmd > 0.0, body_vel > velocity_threshold), reward, stand_still_scale * reward) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/curriculums.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/curriculums.py index 38e7b93d10..8092acca2d 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/curriculums.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/curriculums.py @@ -43,7 +43,7 @@ def terrain_levels_vel( terrain: TerrainImporter = env.scene.terrain command = env.command_manager.get_command("base_velocity") # compute the distance the robot walked - distance = torch.norm(asset.data.root_pos_w[env_ids, :2] - env.scene.env_origins[env_ids, :2], dim=1) + distance = torch.norm(asset.data.root_link_pos_w[env_ids, :2] - env.scene.env_origins[env_ids, :2], dim=1) # robots that walked far enough progress to harder terrains move_up = distance > terrain.cfg.terrain_generator.size[0] / 2 # robots that walked less than half of their required distance go to simpler terrains diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/rewards.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/rewards.py index db8ad9a0cc..07fa615f99 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/rewards.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/rewards.py @@ -77,7 +77,8 @@ def feet_slide(env, sensor_cfg: SceneEntityCfg, asset_cfg: SceneEntityCfg = Scen 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 = env.scene[asset_cfg.name] - body_vel = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2] + + body_vel = asset.data.body_com_lin_vel_w[:, asset_cfg.body_ids, :2] reward = torch.sum(body_vel.norm(dim=-1) * contacts, dim=1) return reward @@ -88,7 +89,7 @@ def track_lin_vel_xy_yaw_frame_exp( """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_rotate_inverse(yaw_quat(asset.data.root_quat_w), asset.data.root_lin_vel_w[:, :3]) + vel_yaw = quat_rotate_inverse(yaw_quat(asset.data.root_link_quat_w), asset.data.root_com_lin_vel_w[:, :3]) lin_vel_error = torch.sum( torch.square(env.command_manager.get_command(command_name)[:, :2] - vel_yaw[:, :2]), dim=1 ) @@ -101,5 +102,7 @@ def track_ang_vel_z_world_exp( """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]) + ang_vel_error = torch.square( + env.command_manager.get_command(command_name)[:, 2] - asset.data.root_com_ang_vel_w[:, 2] + ) return torch.exp(-ang_vel_error / std**2) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/terminations.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/terminations.py index 464716deec..cd504f0b6f 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/terminations.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/terminations.py @@ -45,8 +45,8 @@ def terrain_out_of_bounds( asset: RigidObject = env.scene[asset_cfg.name] # check if the agent is out of bounds - x_out_of_bounds = torch.abs(asset.data.root_pos_w[:, 0]) > 0.5 * map_width - distance_buffer - y_out_of_bounds = torch.abs(asset.data.root_pos_w[:, 1]) > 0.5 * map_height - distance_buffer + x_out_of_bounds = torch.abs(asset.data.root_link_pos_w[:, 0]) > 0.5 * map_width - distance_buffer + y_out_of_bounds = torch.abs(asset.data.root_link_pos_w[:, 1]) > 0.5 * map_height - distance_buffer return torch.logical_or(x_out_of_bounds, y_out_of_bounds) else: raise ValueError("Received unsupported terrain type, must be either 'plane' or 'generator'.") diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/cabinet/cabinet_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/cabinet/cabinet_env_cfg.py index 7faa94601a..56c7e5d2b8 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/cabinet/cabinet_env_cfg.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/cabinet/cabinet_env_cfg.py @@ -123,13 +123,6 @@ class CabinetSceneCfg(InteractiveSceneCfg): ## -@configclass -class CommandsCfg: - """Command terms for the MDP.""" - - null_command = mdp.NullCommandCfg() - - @configclass class ActionsCfg: """Action specifications for the MDP.""" @@ -267,7 +260,6 @@ class CabinetEnvCfg(ManagerBasedRLEnvCfg): # Basic settings observations: ObservationsCfg = ObservationsCfg() actions: ActionsCfg = ActionsCfg() - commands: CommandsCfg = CommandsCfg() # MDP settings rewards: RewardsCfg = RewardsCfg() terminations: TerminationsCfg = TerminationsCfg() diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/cabinet/config/franka/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/cabinet/config/franka/__init__.py index 192bcf6d6b..ec861b8cd5 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/cabinet/config/franka/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/cabinet/config/franka/__init__.py @@ -5,7 +5,7 @@ import gymnasium as gym -from . import agents, ik_abs_env_cfg, ik_rel_env_cfg, joint_pos_env_cfg +from . import agents ## # Register Gym environments. @@ -19,7 +19,7 @@ id="Isaac-Open-Drawer-Franka-v0", entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", kwargs={ - "env_cfg_entry_point": joint_pos_env_cfg.FrankaCabinetEnvCfg, + "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:FrankaCabinetEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CabinetPPORunnerCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", @@ -31,7 +31,7 @@ id="Isaac-Open-Drawer-Franka-Play-v0", entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", kwargs={ - "env_cfg_entry_point": joint_pos_env_cfg.FrankaCabinetEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:FrankaCabinetEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CabinetPPORunnerCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", @@ -48,7 +48,7 @@ id="Isaac-Open-Drawer-Franka-IK-Abs-v0", entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", kwargs={ - "env_cfg_entry_point": ik_abs_env_cfg.FrankaCabinetEnvCfg, + "env_cfg_entry_point": f"{__name__}.ik_abs_env_cfg:FrankaCabinetEnvCfg", }, disable_env_checker=True, ) @@ -61,7 +61,7 @@ id="Isaac-Open-Drawer-Franka-IK-Rel-v0", entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", kwargs={ - "env_cfg_entry_point": ik_rel_env_cfg.FrankaCabinetEnvCfg, + "env_cfg_entry_point": f"{__name__}.ik_rel_env_cfg:FrankaCabinetEnvCfg", }, disable_env_checker=True, ) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/cabinet/mdp/observations.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/cabinet/mdp/observations.py index 6c13b989cc..18d7df8347 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/cabinet/mdp/observations.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/cabinet/mdp/observations.py @@ -21,7 +21,7 @@ def rel_ee_object_distance(env: ManagerBasedRLEnv) -> torch.Tensor: ee_tf_data: FrameTransformerData = env.scene["ee_frame"].data object_data: ArticulationData = env.scene["object"].data - return object_data.root_pos_w - ee_tf_data.target_pos_w[..., 0, :] + return object_data.root_link_pos_w - ee_tf_data.target_pos_w[..., 0, :] def rel_ee_drawer_distance(env: ManagerBasedRLEnv) -> torch.Tensor: diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/config/allegro_hand/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/config/allegro_hand/__init__.py index e1d162515d..031a267384 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/config/allegro_hand/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/config/allegro_hand/__init__.py @@ -5,7 +5,7 @@ import gymnasium as gym -from . import agents, allegro_env_cfg +from . import agents ## # Register Gym environments. @@ -20,7 +20,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": allegro_env_cfg.AllegroCubeEnvCfg, + "env_cfg_entry_point": f"{__name__}.allegro_env_cfg:AllegroCubeEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AllegroCubePPORunnerCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", @@ -32,7 +32,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": allegro_env_cfg.AllegroCubeEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.allegro_env_cfg:AllegroCubeEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AllegroCubePPORunnerCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", @@ -48,7 +48,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": allegro_env_cfg.AllegroCubeNoVelObsEnvCfg, + "env_cfg_entry_point": f"{__name__}.allegro_env_cfg:AllegroCubeNoVelObsEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AllegroCubeNoVelObsPPORunnerCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", @@ -60,7 +60,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": allegro_env_cfg.AllegroCubeNoVelObsEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.allegro_env_cfg:AllegroCubeNoVelObsEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AllegroCubeNoVelObsPPORunnerCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/commands/orientation_command.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/commands/orientation_command.py index 153127c45c..167a0b9a43 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/commands/orientation_command.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/commands/orientation_command.py @@ -95,10 +95,10 @@ def _update_metrics(self): # logs data # -- compute the orientation error self.metrics["orientation_error"] = math_utils.quat_error_magnitude( - self.object.data.root_quat_w, self.quat_command_w + self.object.data.root_link_quat_w, self.quat_command_w ) # -- compute the position error - self.metrics["position_error"] = torch.norm(self.object.data.root_pos_w - self.pos_command_w, dim=1) + self.metrics["position_error"] = torch.norm(self.object.data.root_link_pos_w - self.pos_command_w, dim=1) # -- compute the number of consecutive successes successes = self.metrics["orientation_error"] < self.cfg.orientation_success_threshold self.metrics["consecutive_success"] += successes.float() diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/observations.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/observations.py index a7d4335842..f23ebc716a 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/observations.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/observations.py @@ -30,7 +30,7 @@ def goal_quat_diff( # obtain the orientations goal_quat_w = command_term.command[:, 3:7] - asset_quat_w = asset.data.root_quat_w + asset_quat_w = asset.data.root_link_quat_w # compute quaternion difference quat = math_utils.quat_mul(asset_quat_w, math_utils.quat_conjugate(goal_quat_w)) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/rewards.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/rewards.py index 302b85871c..c02f2e301a 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/rewards.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/rewards.py @@ -39,7 +39,7 @@ def success_bonus( # obtain the threshold for the orientation error threshold = command_term.cfg.orientation_success_threshold # calculate the orientation error - dtheta = math_utils.quat_error_magnitude(asset.data.root_quat_w, goal_quat_w) + dtheta = math_utils.quat_error_magnitude(asset.data.root_link_quat_w, goal_quat_w) return dtheta <= threshold @@ -63,7 +63,7 @@ def track_pos_l2( # obtain the goal position goal_pos_e = command_term.command[:, 0:3] # obtain the object position in the environment frame - object_pos_e = asset.data.root_pos_w - env.scene.env_origins + object_pos_e = asset.data.root_link_pos_w - env.scene.env_origins return torch.norm(goal_pos_e - object_pos_e, p=2, dim=-1) @@ -91,6 +91,6 @@ def track_orientation_inv_l2( # obtain the goal orientation goal_quat_w = command_term.command[:, 3:7] # calculate the orientation error - dtheta = math_utils.quat_error_magnitude(asset.data.root_quat_w, goal_quat_w) + dtheta = math_utils.quat_error_magnitude(asset.data.root_link_quat_w, goal_quat_w) return 1.0 / (dtheta + rot_eps) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/terminations.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/terminations.py index a56df2a8b9..4aaea26ad8 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/terminations.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/terminations.py @@ -50,7 +50,7 @@ def object_away_from_goal( asset = env.scene[object_cfg.name] # object pos - asset_pos_e = asset.data.root_pos_w - env.scene.env_origins + asset_pos_e = asset.data.root_link_pos_w - env.scene.env_origins goal_pos_e = command_term.command[:, :3] return torch.norm(asset_pos_e - goal_pos_e, p=2, dim=1) > threshold @@ -78,6 +78,6 @@ def object_away_from_robot( object = env.scene[object_cfg.name] # compute distance - dist = torch.norm(robot.data.root_pos_w - object.data.root_pos_w, dim=1) + dist = torch.norm(robot.data.root_link_pos_w - object.data.root_link_pos_w, dim=1) return dist > threshold diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/config/franka/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/config/franka/__init__.py index 14657e4f92..644ac056e9 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/config/franka/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/config/franka/__init__.py @@ -5,7 +5,7 @@ import gymnasium as gym import os -from . import agents, ik_abs_env_cfg, ik_rel_env_cfg, joint_pos_env_cfg +from . import agents ## # Register Gym environments. @@ -19,10 +19,11 @@ id="Isaac-Lift-Cube-Franka-v0", entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", kwargs={ - "env_cfg_entry_point": joint_pos_env_cfg.FrankaCubeLiftEnvCfg, + "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:FrankaCubeLiftEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:LiftCubePPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", + "sb3_cfg_entry_point": f"{agents.__name__}:sb3_ppo_cfg.yaml", }, disable_env_checker=True, ) @@ -31,10 +32,11 @@ id="Isaac-Lift-Cube-Franka-Play-v0", entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", kwargs={ - "env_cfg_entry_point": joint_pos_env_cfg.FrankaCubeLiftEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:FrankaCubeLiftEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:LiftCubePPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", + "sb3_cfg_entry_point": f"{agents.__name__}:sb3_ppo_cfg.yaml", }, disable_env_checker=True, ) @@ -47,7 +49,7 @@ id="Isaac-Lift-Cube-Franka-IK-Abs-v0", entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", kwargs={ - "env_cfg_entry_point": ik_abs_env_cfg.FrankaCubeLiftEnvCfg, + "env_cfg_entry_point": f"{__name__}.ik_abs_env_cfg:FrankaCubeLiftEnvCfg", }, disable_env_checker=True, ) @@ -56,7 +58,7 @@ id="Isaac-Lift-Teddy-Bear-Franka-IK-Abs-v0", entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", kwargs={ - "env_cfg_entry_point": ik_abs_env_cfg.FrankaTeddyBearLiftEnvCfg, + "env_cfg_entry_point": f"{__name__}.ik_abs_env_cfg:FrankaTeddyBearLiftEnvCfg", }, disable_env_checker=True, ) @@ -69,7 +71,7 @@ id="Isaac-Lift-Cube-Franka-IK-Rel-v0", entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", kwargs={ - "env_cfg_entry_point": ik_rel_env_cfg.FrankaCubeLiftEnvCfg, + "env_cfg_entry_point": f"{__name__}.ik_rel_env_cfg:FrankaCubeLiftEnvCfg", "robomimic_bc_cfg_entry_point": os.path.join(agents.__path__[0], "robomimic/bc.json"), }, disable_env_checker=True, diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/config/franka/agents/sb3_ppo_cfg.yaml b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/config/franka/agents/sb3_ppo_cfg.yaml index 57960b8822..6d6f15781a 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/config/franka/agents/sb3_ppo_cfg.yaml +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/config/franka/agents/sb3_ppo_cfg.yaml @@ -16,7 +16,7 @@ learning_rate: !!float 3e-4 clip_range: 0.2 policy_kwargs: "dict( activation_fn=nn.ELU, - net_arch=[32, 32, dict(pi=[256, 128, 64], vf=[256, 128, 64])] + net_arch=dict(pi=[256, 128, 64], vf=[256, 128, 64]) )" target_kl: 0.01 max_grad_norm: 1.0 diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/observations.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/observations.py index 3a976fa3f5..46468e23f9 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/observations.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/observations.py @@ -24,8 +24,8 @@ def object_position_in_robot_root_frame( """The position of the object in the robot's root frame.""" robot: RigidObject = env.scene[robot_cfg.name] object: RigidObject = env.scene[object_cfg.name] - object_pos_w = object.data.root_pos_w[:, :3] + object_pos_w = object.data.root_link_pos_w[:, :3] object_pos_b, _ = subtract_frame_transforms( - robot.data.root_state_w[:, :3], robot.data.root_state_w[:, 3:7], object_pos_w + robot.data.root_link_state_w[:, :3], robot.data.root_link_state_w[:, 3:7], object_pos_w ) return object_pos_b diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/rewards.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/rewards.py index 334df9ea50..9c48814668 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/rewards.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/rewards.py @@ -22,7 +22,7 @@ def object_is_lifted( ) -> torch.Tensor: """Reward the agent for lifting the object above the minimal height.""" object: RigidObject = env.scene[object_cfg.name] - return torch.where(object.data.root_pos_w[:, 2] > minimal_height, 1.0, 0.0) + return torch.where(object.data.root_link_pos_w[:, 2] > minimal_height, 1.0, 0.0) def object_ee_distance( @@ -36,7 +36,7 @@ def object_ee_distance( object: RigidObject = env.scene[object_cfg.name] ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] # Target object position: (num_envs, 3) - cube_pos_w = object.data.root_pos_w + cube_pos_w = object.data.root_link_pos_w # End-effector position: (num_envs, 3) ee_w = ee_frame.data.target_pos_w[..., 0, :] # Distance of the end-effector to the object: (num_envs,) @@ -60,8 +60,10 @@ def object_goal_distance( command = env.command_manager.get_command(command_name) # compute the desired position in the world frame des_pos_b = command[:, :3] - des_pos_w, _ = combine_frame_transforms(robot.data.root_state_w[:, :3], robot.data.root_state_w[:, 3:7], des_pos_b) + des_pos_w, _ = combine_frame_transforms( + robot.data.root_link_state_w[:, :3], robot.data.root_link_state_w[:, 3:7], des_pos_b + ) # distance of the end-effector to the object: (num_envs,) - distance = torch.norm(des_pos_w - object.data.root_pos_w[:, :3], dim=1) + distance = torch.norm(des_pos_w - object.data.root_link_pos_w[:, :3], dim=1) # rewarded if the object is lifted above the threshold - return (object.data.root_pos_w[:, 2] > minimal_height) * (1 - torch.tanh(distance / std)) + return (object.data.root_link_pos_w[:, 2] > minimal_height) * (1 - torch.tanh(distance / std)) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/terminations.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/terminations.py index 212e192336..bd86490734 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/terminations.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/terminations.py @@ -45,9 +45,11 @@ def object_reached_goal( command = env.command_manager.get_command(command_name) # compute the desired position in the world frame des_pos_b = command[:, :3] - des_pos_w, _ = combine_frame_transforms(robot.data.root_state_w[:, :3], robot.data.root_state_w[:, 3:7], des_pos_b) + des_pos_w, _ = combine_frame_transforms( + robot.data.root_link_state_w[:, :3], robot.data.root_link_state_w[:, 3:7], des_pos_b + ) # distance of the end-effector to the object: (num_envs,) - distance = torch.norm(des_pos_w - object.data.root_pos_w[:, :3], dim=1) + distance = torch.norm(des_pos_w - object.data.root_link_pos_w[:, :3], dim=1) # rewarded if the object is lifted above the threshold return distance < threshold diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/config/franka/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/config/franka/__init__.py index 66146302b2..07fc54526f 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/config/franka/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/config/franka/__init__.py @@ -5,7 +5,7 @@ import gymnasium as gym -from . import agents, ik_abs_env_cfg, ik_rel_env_cfg, joint_pos_env_cfg +from . import agents ## # Register Gym environments. @@ -20,7 +20,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": joint_pos_env_cfg.FrankaReachEnvCfg, + "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:FrankaReachEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaReachPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", @@ -32,7 +32,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": joint_pos_env_cfg.FrankaReachEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:FrankaReachEnvCfg_PLAY", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaReachPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", @@ -48,7 +48,7 @@ id="Isaac-Reach-Franka-IK-Abs-v0", entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", kwargs={ - "env_cfg_entry_point": ik_abs_env_cfg.FrankaReachEnvCfg, + "env_cfg_entry_point": f"{__name__}.ik_abs_env_cfg:FrankaReachEnvCfg", }, disable_env_checker=True, ) @@ -61,7 +61,31 @@ id="Isaac-Reach-Franka-IK-Rel-v0", entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", kwargs={ - "env_cfg_entry_point": ik_rel_env_cfg.FrankaReachEnvCfg, + "env_cfg_entry_point": f"{__name__}.ik_rel_env_cfg:FrankaReachEnvCfg", }, disable_env_checker=True, ) + +## +# Operational Space Control +## + +gym.register( + id="Isaac-Reach-Franka-OSC-v0", + entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.osc_env_cfg:FrankaReachEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaReachPPORunnerCfg", + }, +) + +gym.register( + id="Isaac-Reach-Franka-OSC-Play-v0", + entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.osc_env_cfg:FrankaReachEnvCfg_PLAY", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaReachPPORunnerCfg", + }, +) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/config/franka/osc_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/config/franka/osc_env_cfg.py new file mode 100644 index 0000000000..410e33c467 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/config/franka/osc_env_cfg.py @@ -0,0 +1,71 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from omni.isaac.lab.controllers.operational_space_cfg import OperationalSpaceControllerCfg +from omni.isaac.lab.envs.mdp.actions.actions_cfg import OperationalSpaceControllerActionCfg +from omni.isaac.lab.utils import configclass + +from . import joint_pos_env_cfg + +## +# Pre-defined configs +## +from omni.isaac.lab_assets.franka import FRANKA_PANDA_CFG # isort: skip + + +@configclass +class FrankaReachEnvCfg(joint_pos_env_cfg.FrankaReachEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + + # Set Franka as robot + # We remove stiffness and damping for the shoulder and forearm joints for effort control + self.scene.robot = FRANKA_PANDA_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + self.scene.robot.actuators["panda_shoulder"].stiffness = 0.0 + self.scene.robot.actuators["panda_shoulder"].damping = 0.0 + self.scene.robot.actuators["panda_forearm"].stiffness = 0.0 + self.scene.robot.actuators["panda_forearm"].damping = 0.0 + self.scene.robot.spawn.rigid_props.disable_gravity = True + + # If closed-loop contact force control is desired, contact sensors should be enabled for the robot + # self.scene.robot.spawn.activate_contact_sensors = True + + self.actions.arm_action = OperationalSpaceControllerActionCfg( + asset_name="robot", + joint_names=["panda_joint.*"], + body_name="panda_hand", + # If a task frame different from articulation root/base is desired, a RigidObject, e.g., "task_frame", + # can be added to the scene and its relative path could provided as task_frame_rel_path + # task_frame_rel_path="task_frame", + controller_cfg=OperationalSpaceControllerCfg( + target_types=["pose_abs"], + impedance_mode="variable_kp", + inertial_dynamics_decoupling=True, + partial_inertial_dynamics_decoupling=False, + gravity_compensation=False, + motion_stiffness_task=100.0, + motion_damping_ratio_task=1.0, + motion_stiffness_limits_task=(50.0, 200.0), + ), + position_scale=1.0, + orientation_scale=1.0, + stiffness_scale=100.0, + ) + # Removing these observations as they are not needed for OSC and we want keep the observation space small + self.observations.policy.joint_pos = None + self.observations.policy.joint_vel = None + + +@configclass +class FrankaReachEnvCfg_PLAY(FrankaReachEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + # make a smaller scene for play + self.scene.num_envs = 16 + self.scene.env_spacing = 2.5 + # disable randomization for play + self.observations.policy.enable_corruption = False diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/config/ur_10/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/config/ur_10/__init__.py index ca788c59d9..c2ce4c673d 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/config/ur_10/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/config/ur_10/__init__.py @@ -5,7 +5,7 @@ import gymnasium as gym -from . import agents, joint_pos_env_cfg +from . import agents ## # Register Gym environments. @@ -16,7 +16,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": joint_pos_env_cfg.UR10ReachEnvCfg, + "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:UR10ReachEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UR10ReachPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", @@ -28,7 +28,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": joint_pos_env_cfg.UR10ReachEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:UR10ReachEnvCfg_PLAY", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UR10ReachPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/mdp/rewards.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/mdp/rewards.py index d5ae2a57c5..8f93c7f70a 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/mdp/rewards.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/mdp/rewards.py @@ -28,8 +28,10 @@ def position_command_error(env: ManagerBasedRLEnv, command_name: str, asset_cfg: command = env.command_manager.get_command(command_name) # obtain the desired and current positions des_pos_b = command[:, :3] - des_pos_w, _ = combine_frame_transforms(asset.data.root_state_w[:, :3], asset.data.root_state_w[:, 3:7], des_pos_b) - curr_pos_w = asset.data.body_state_w[:, asset_cfg.body_ids[0], :3] # type: ignore + des_pos_w, _ = combine_frame_transforms( + asset.data.root_link_state_w[:, :3], asset.data.root_link_state_w[:, 3:7], des_pos_b + ) + curr_pos_w = asset.data.body_link_state_w[:, asset_cfg.body_ids[0], :3] # type: ignore return torch.norm(curr_pos_w - des_pos_w, dim=1) @@ -46,8 +48,10 @@ def position_command_error_tanh( command = env.command_manager.get_command(command_name) # obtain the desired and current positions des_pos_b = command[:, :3] - des_pos_w, _ = combine_frame_transforms(asset.data.root_state_w[:, :3], asset.data.root_state_w[:, 3:7], des_pos_b) - curr_pos_w = asset.data.body_state_w[:, asset_cfg.body_ids[0], :3] # type: ignore + des_pos_w, _ = combine_frame_transforms( + asset.data.root_link_state_w[:, :3], asset.data.root_link_state_w[:, 3:7], des_pos_b + ) + curr_pos_w = asset.data.body_link_state_w[:, asset_cfg.body_ids[0], :3] # type: ignore distance = torch.norm(curr_pos_w - des_pos_w, dim=1) return 1 - torch.tanh(distance / std) @@ -64,6 +68,6 @@ def orientation_command_error(env: ManagerBasedRLEnv, command_name: str, asset_c command = env.command_manager.get_command(command_name) # obtain the desired and current orientations des_quat_b = command[:, 3:7] - des_quat_w = quat_mul(asset.data.root_state_w[:, 3:7], des_quat_b) - curr_quat_w = asset.data.body_state_w[:, asset_cfg.body_ids[0], 3:7] # type: ignore + des_quat_w = quat_mul(asset.data.root_link_state_w[:, 3:7], des_quat_b) + curr_quat_w = asset.data.body_link_state_w[:, asset_cfg.body_ids[0], 3:7] # type: ignore return quat_error_magnitude(curr_quat_w, des_quat_w) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/__init__.py new file mode 100644 index 0000000000..a6be077a39 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configurations for the object stack environments.""" + +# We leave this file empty since we don't want to expose any configs in this package directly. +# We still need this file to import the "config" module in the parent package. diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/__init__.py new file mode 100644 index 0000000000..a6be077a39 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configurations for the object stack environments.""" + +# We leave this file empty since we don't want to expose any configs in this package directly. +# We still need this file to import the "config" module in the parent package. diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/__init__.py new file mode 100644 index 0000000000..5a564fc996 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/__init__.py @@ -0,0 +1,64 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +import gymnasium as gym +import os + +from . import ( + agents, + stack_ik_rel_env_cfg, + stack_ik_rel_instance_randomize_env_cfg, + stack_joint_pos_env_cfg, + stack_joint_pos_instance_randomize_env_cfg, +) + +## +# Register Gym environments. +## + +## +# Joint Position Control +## + +gym.register( + id="Isaac-Stack-Cube-Franka-v0", + entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", + kwargs={ + "env_cfg_entry_point": stack_joint_pos_env_cfg.FrankaCubeStackEnvCfg, + }, + disable_env_checker=True, +) + +gym.register( + id="Isaac-Stack-Cube-Instance-Randomize-Franka-v0", + entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", + kwargs={ + "env_cfg_entry_point": stack_joint_pos_instance_randomize_env_cfg.FrankaCubeStackInstanceRandomizeEnvCfg, + }, + disable_env_checker=True, +) + + +## +# Inverse Kinematics - Relative Pose Control +## + +gym.register( + id="Isaac-Stack-Cube-Franka-IK-Rel-v0", + entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", + kwargs={ + "env_cfg_entry_point": stack_ik_rel_env_cfg.FrankaCubeStackEnvCfg, + "robomimic_bc_cfg_entry_point": os.path.join(agents.__path__[0], "robomimic/bc_rnn_low_dim.json"), + }, + disable_env_checker=True, +) + +gym.register( + id="Isaac-Stack-Cube-Instance-Randomize-Franka-IK-Rel-v0", + entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", + kwargs={ + "env_cfg_entry_point": stack_ik_rel_instance_randomize_env_cfg.FrankaCubeStackInstanceRandomizeEnvCfg, + }, + disable_env_checker=True, +) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/agents/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/agents/__init__.py new file mode 100644 index 0000000000..c3ee657052 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/agents/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/agents/robomimic/bc_rnn_low_dim.json b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/agents/robomimic/bc_rnn_low_dim.json new file mode 100644 index 0000000000..24d807ac34 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/agents/robomimic/bc_rnn_low_dim.json @@ -0,0 +1,101 @@ +{ + "algo_name": "bc", + "experiment": { + "name": "bc", + "validate": false, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true + }, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": 100, + "epochs": [], + "on_best_validation": false, + "on_best_rollout_return": false, + "on_best_rollout_success_rate": true + }, + "epoch_every_n_steps": 100, + "env": null, + "additional_envs": null, + "render": false, + "render_video": false, + "rollout": { + "enabled": false + } + }, + "train": { + "data": null, + "num_data_workers": 4, + "hdf5_cache_mode": "all", + "hdf5_use_swmr": true, + "hdf5_normalize_obs": false, + "hdf5_filter_key": null, + "hdf5_validation_filter_key": null, + "seq_length": 10, + "dataset_keys": [ + "actions" + ], + "goal_mode": null, + "cuda": true, + "batch_size": 100, + "num_epochs": 2000, + "seed": 101 + }, + "algo": { + "optim_params": { + "policy": { + "optimizer_type": "adam", + "learning_rate": { + "initial": 0.001, + "decay_factor": 0.1, + "epoch_schedule": [], + "scheduler_type": "multistep" + }, + "regularization": { + "L2": 0.0 + } + } + }, + "loss": { + "l2_weight": 1.0, + "l1_weight": 0.0, + "cos_weight": 0.0 + }, + "actor_layer_dims": [], + "gmm": { + "enabled": true, + "num_modes": 5, + "min_std": 0.0001, + "std_activation": "softplus", + "low_noise_eval": true + }, + "rnn": { + "enabled": true, + "horizon": 10, + "hidden_dim": 400, + "rnn_type": "LSTM", + "num_layers": 2, + "open_loop": false, + "kwargs": { + "bidirectional": false + } + } + }, + "observation": { + "modalities": { + "obs": { + "low_dim": [ + "eef_pos", + "eef_quat", + "gripper_pos", + "object" + ], + "rgb": [], + "depth": [], + "scan": [] + } + } + } +} diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_ik_rel_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_ik_rel_env_cfg.py new file mode 100644 index 0000000000..f95f1c61dc --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_ik_rel_env_cfg.py @@ -0,0 +1,36 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from omni.isaac.lab.controllers.differential_ik_cfg import DifferentialIKControllerCfg +from omni.isaac.lab.envs.mdp.actions.actions_cfg import DifferentialInverseKinematicsActionCfg +from omni.isaac.lab.utils import configclass + +from . import stack_joint_pos_env_cfg + +## +# Pre-defined configs +## +from omni.isaac.lab_assets.franka import FRANKA_PANDA_HIGH_PD_CFG # isort: skip + + +@configclass +class FrankaCubeStackEnvCfg(stack_joint_pos_env_cfg.FrankaCubeStackEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + + # Set Franka as robot + # We switch here to a stiffer PD controller for IK tracking to be better. + self.scene.robot = FRANKA_PANDA_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + # Set actions for the specific robot type (franka) + self.actions.arm_action = DifferentialInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["panda_joint.*"], + body_name="panda_hand", + controller=DifferentialIKControllerCfg(command_type="pose", use_relative_mode=True, ik_method="dls"), + scale=0.5, + body_offset=DifferentialInverseKinematicsActionCfg.OffsetCfg(pos=[0.0, 0.0, 0.107]), + ) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_ik_rel_instance_randomize_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_ik_rel_instance_randomize_env_cfg.py new file mode 100644 index 0000000000..91d1d02730 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_ik_rel_instance_randomize_env_cfg.py @@ -0,0 +1,38 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from omni.isaac.lab.controllers.differential_ik_cfg import DifferentialIKControllerCfg +from omni.isaac.lab.envs.mdp.actions.actions_cfg import DifferentialInverseKinematicsActionCfg +from omni.isaac.lab.utils import configclass + +from . import stack_joint_pos_instance_randomize_env_cfg + +## +# Pre-defined configs +## +from omni.isaac.lab_assets.franka import FRANKA_PANDA_HIGH_PD_CFG # isort: skip + + +@configclass +class FrankaCubeStackInstanceRandomizeEnvCfg( + stack_joint_pos_instance_randomize_env_cfg.FrankaCubeStackInstanceRandomizeEnvCfg +): + def __post_init__(self): + # post init of parent + super().__post_init__() + + # Set Franka as robot + # We switch here to a stiffer PD controller for IK tracking to be better. + self.scene.robot = FRANKA_PANDA_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + # Set actions for the specific robot type (franka) + self.actions.arm_action = DifferentialInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["panda_joint.*"], + body_name="panda_hand", + controller=DifferentialIKControllerCfg(command_type="pose", use_relative_mode=True, ik_method="dls"), + scale=0.5, + body_offset=DifferentialInverseKinematicsActionCfg.OffsetCfg(pos=[0.0, 0.0, 0.107]), + ) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_joint_pos_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_joint_pos_env_cfg.py new file mode 100644 index 0000000000..3b7dc1dc2f --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_joint_pos_env_cfg.py @@ -0,0 +1,180 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.assets import RigidObjectCfg +from omni.isaac.lab.managers import EventTermCfg as EventTerm +from omni.isaac.lab.managers import SceneEntityCfg +from omni.isaac.lab.sensors import CameraCfg, FrameTransformerCfg +from omni.isaac.lab.sensors.frame_transformer.frame_transformer_cfg import OffsetCfg +from omni.isaac.lab.sim.schemas.schemas_cfg import RigidBodyPropertiesCfg +from omni.isaac.lab.sim.spawners.from_files.from_files_cfg import UsdFileCfg +from omni.isaac.lab.utils import configclass +from omni.isaac.lab.utils.assets import ISAAC_NUCLEUS_DIR + +from omni.isaac.lab_tasks.manager_based.manipulation.stack import mdp +from omni.isaac.lab_tasks.manager_based.manipulation.stack.mdp import franka_stack_events +from omni.isaac.lab_tasks.manager_based.manipulation.stack.stack_env_cfg import StackEnvCfg + +## +# Pre-defined configs +## +from omni.isaac.lab.markers.config import FRAME_MARKER_CFG # isort: skip +from omni.isaac.lab_assets.franka import FRANKA_PANDA_CFG # isort: skip + + +@configclass +class EventCfg: + """Configuration for events.""" + + init_franka_arm_pose = EventTerm( + func=franka_stack_events.set_default_joint_pose, + mode="startup", + params={ + "default_pose": [0.0444, -0.1894, -0.1107, -2.5148, 0.0044, 2.3775, 0.6952, 0.0400, 0.0400], + }, + ) + + randomize_franka_joint_state = EventTerm( + func=franka_stack_events.randomize_joint_by_gaussian_offset, + mode="reset", + params={ + "mean": 0.0, + "std": 0.02, + "asset_cfg": SceneEntityCfg("robot"), + }, + ) + + randomize_cube_positions = EventTerm( + func=franka_stack_events.randomize_object_pose, + mode="reset", + params={ + "pose_range": {"x": (0.4, 0.6), "y": (-0.10, 0.10), "z": (0.0203, 0.0203), "yaw": (-1.0, 1, 0)}, + "min_separation": 0.1, + "asset_cfgs": [SceneEntityCfg("cube_1"), SceneEntityCfg("cube_2"), SceneEntityCfg("cube_3")], + }, + ) + + +@configclass +class FrankaCubeStackEnvCfg(StackEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + + # Set events + self.events = EventCfg() + + # Set Franka as robot + self.scene.robot = FRANKA_PANDA_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + # Set actions for the specific robot type (franka) + self.actions.arm_action = mdp.JointPositionActionCfg( + asset_name="robot", joint_names=["panda_joint.*"], scale=0.5, use_default_offset=True + ) + self.actions.gripper_action = mdp.BinaryJointPositionActionCfg( + asset_name="robot", + joint_names=["panda_finger.*"], + open_command_expr={"panda_finger_.*": 0.04}, + close_command_expr={"panda_finger_.*": 0.0}, + ) + + # Rigid body properties of each cube + cube_properties = RigidBodyPropertiesCfg( + solver_position_iteration_count=16, + solver_velocity_iteration_count=1, + max_angular_velocity=1000.0, + max_linear_velocity=1000.0, + max_depenetration_velocity=5.0, + disable_gravity=False, + ) + + # Set each stacking cube deterministically + self.scene.cube_1 = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_1", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.4, 0.0, 0.0203], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/blue_block.usd", + scale=(1.0, 1.0, 1.0), + rigid_props=cube_properties, + ), + ) + self.scene.cube_2 = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_2", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.55, 0.05, 0.0203], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/red_block.usd", + scale=(1.0, 1.0, 1.0), + rigid_props=cube_properties, + ), + ) + self.scene.cube_3 = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_3", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.60, -0.1, 0.0203], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/green_block.usd", + scale=(1.0, 1.0, 1.0), + rigid_props=cube_properties, + ), + ) + + # Set wrist camera + self.scene.wrist_cam = CameraCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_hand/wrist_cam", + update_period=0.0333, + height=84, + width=84, + data_types=["rgb", "distance_to_image_plane"], + spawn=sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 1.0e5) + ), + offset=CameraCfg.OffsetCfg(pos=(0.025, 0.0, 0.0), rot=(0.707, 0.0, 0.0, 0.707), convention="ros"), + ) + + # Set table view camera + self.scene.table_cam = CameraCfg( + prim_path="{ENV_REGEX_NS}/table_cam", + update_period=0.0333, + height=84, + width=84, + data_types=["rgb", "distance_to_image_plane"], + spawn=sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 1.0e5) + ), + offset=CameraCfg.OffsetCfg(pos=(1.0, 0.0, 0.33), rot=(-0.3799, 0.5963, 0.5963, -0.3799), convention="ros"), + ) + + # Listens to the required transforms + marker_cfg = FRAME_MARKER_CFG.copy() + marker_cfg.markers["frame"].scale = (0.1, 0.1, 0.1) + marker_cfg.prim_path = "/Visuals/FrameTransformer" + self.scene.ee_frame = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_link0", + debug_vis=False, + visualizer_cfg=marker_cfg, + target_frames=[ + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_hand", + name="end_effector", + offset=OffsetCfg( + pos=[0.0, 0.0, 0.1034], + ), + ), + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_rightfinger", + name="tool_rightfinger", + offset=OffsetCfg( + pos=(0.0, 0.0, 0.046), + ), + ), + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_leftfinger", + name="tool_leftfinger", + offset=OffsetCfg( + pos=(0.0, 0.0, 0.046), + ), + ), + ], + ) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_joint_pos_instance_randomize_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_joint_pos_instance_randomize_env_cfg.py new file mode 100644 index 0000000000..f9038c8ae9 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/config/franka/stack_joint_pos_instance_randomize_env_cfg.py @@ -0,0 +1,224 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.assets import RigidObjectCfg, RigidObjectCollectionCfg +from omni.isaac.lab.managers import EventTermCfg as EventTerm +from omni.isaac.lab.managers import SceneEntityCfg +from omni.isaac.lab.sensors import CameraCfg, FrameTransformerCfg +from omni.isaac.lab.sensors.frame_transformer.frame_transformer_cfg import OffsetCfg +from omni.isaac.lab.sim.schemas.schemas_cfg import RigidBodyPropertiesCfg +from omni.isaac.lab.sim.spawners.from_files.from_files_cfg import UsdFileCfg +from omni.isaac.lab.utils import configclass +from omni.isaac.lab.utils.assets import ISAAC_NUCLEUS_DIR + +from omni.isaac.lab_tasks.manager_based.manipulation.stack import mdp +from omni.isaac.lab_tasks.manager_based.manipulation.stack.mdp import franka_stack_events +from omni.isaac.lab_tasks.manager_based.manipulation.stack.stack_instance_randomize_env_cfg import ( + StackInstanceRandomizeEnvCfg, +) + +## +# Pre-defined configs +## +from omni.isaac.lab.markers.config import FRAME_MARKER_CFG # isort: skip +from omni.isaac.lab_assets.franka import FRANKA_PANDA_CFG # isort: skip + + +@configclass +class EventCfg: + """Configuration for events.""" + + init_franka_arm_pose = EventTerm( + func=franka_stack_events.set_default_joint_pose, + mode="startup", + params={ + "default_pose": [0.0444, -0.1894, -0.1107, -2.5148, 0.0044, 2.3775, 0.6952, 0.0400, 0.0400], + }, + ) + + randomize_franka_joint_state = EventTerm( + func=franka_stack_events.randomize_joint_by_gaussian_offset, + mode="reset", + params={ + "mean": 0.0, + "std": 0.02, + "asset_cfg": SceneEntityCfg("robot"), + }, + ) + + randomize_cubes_in_focus = EventTerm( + func=franka_stack_events.randomize_rigid_objects_in_focus, + mode="reset", + params={ + "asset_cfgs": [SceneEntityCfg("cube_1"), SceneEntityCfg("cube_2"), SceneEntityCfg("cube_3")], + "out_focus_state": torch.tensor([10.0, 10.0, 10.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), + "pose_range": {"x": (0.4, 0.6), "y": (-0.10, 0.10), "z": (0.0203, 0.0203), "yaw": (-1.0, 1, 0)}, + "min_separation": 0.1, + }, + ) + + +@configclass +class FrankaCubeStackInstanceRandomizeEnvCfg(StackInstanceRandomizeEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + + # Set events + self.events = EventCfg() + + # Set Franka as robot + self.scene.robot = FRANKA_PANDA_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + # Set actions for the specific robot type (franka) + self.actions.arm_action = mdp.JointPositionActionCfg( + asset_name="robot", joint_names=["panda_joint.*"], scale=0.5, use_default_offset=True + ) + self.actions.gripper_action = mdp.BinaryJointPositionActionCfg( + asset_name="robot", + joint_names=["panda_finger.*"], + open_command_expr={"panda_finger_.*": 0.04}, + close_command_expr={"panda_finger_.*": 0.0}, + ) + + # Rigid body properties of each cube + cube_properties = RigidBodyPropertiesCfg( + solver_position_iteration_count=16, + solver_velocity_iteration_count=1, + max_angular_velocity=1000.0, + max_linear_velocity=1000.0, + max_depenetration_velocity=5.0, + disable_gravity=False, + ) + + # Set each stacking cube to be a collection of rigid objects + cube_1_config_dict = { + "blue_cube": RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_1_Blue", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.4, 0.0, 0.0203], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/blue_block.usd", + scale=(1.0, 1.0, 1.0), + rigid_props=cube_properties, + ), + ), + "red_cube": RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_1_Red", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.4, 0.0, 0.0403], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/red_block.usd", + scale=(1.0, 1.0, 1.0), + rigid_props=cube_properties, + ), + ), + } + + cube_2_config_dict = { + "red_cube": RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_2_Red", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.55, 0.05, 0.0203], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/red_block.usd", + scale=(1.0, 1.0, 1.0), + rigid_props=cube_properties, + ), + ), + "yellow_cube": RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_2_Yellow", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.55, 0.05, 0.0403], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/yellow_block.usd", + scale=(1.0, 1.0, 1.0), + rigid_props=cube_properties, + ), + ), + } + + cube_3_config_dict = { + "yellow_cube": RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_3_Yellow", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.60, -0.1, 0.0203], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/yellow_block.usd", + scale=(1.0, 1.0, 1.0), + rigid_props=cube_properties, + ), + ), + "green_cube": RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_2_Green", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.60, -0.1, 0.0403], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/green_block.usd", + scale=(1.0, 1.0, 1.0), + rigid_props=cube_properties, + ), + ), + } + + self.scene.cube_1 = RigidObjectCollectionCfg(rigid_objects=cube_1_config_dict) + self.scene.cube_2 = RigidObjectCollectionCfg(rigid_objects=cube_2_config_dict) + self.scene.cube_3 = RigidObjectCollectionCfg(rigid_objects=cube_3_config_dict) + + # Set wrist camera + self.scene.wrist_cam = CameraCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_hand/wrist_cam", + update_period=0.0333, + height=84, + width=84, + data_types=["rgb", "distance_to_image_plane"], + spawn=sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 1.0e5) + ), + offset=CameraCfg.OffsetCfg(pos=(0.025, 0.0, 0.0), rot=(0.707, 0.0, 0.0, 0.707), convention="ros"), + ) + + # Set table view camera + self.scene.table_cam = CameraCfg( + prim_path="{ENV_REGEX_NS}/table_cam", + update_period=0.0333, + height=84, + width=84, + data_types=["rgb", "distance_to_image_plane"], + spawn=sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 1.0e5) + ), + offset=CameraCfg.OffsetCfg(pos=(1.0, 0.0, 0.33), rot=(-0.3799, 0.5963, 0.5963, -0.3799), convention="ros"), + ) + + # Listens to the required transforms + marker_cfg = FRAME_MARKER_CFG.copy() + marker_cfg.markers["frame"].scale = (0.1, 0.1, 0.1) + marker_cfg.prim_path = "/Visuals/FrameTransformer" + self.scene.ee_frame = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_link0", + debug_vis=False, + visualizer_cfg=marker_cfg, + target_frames=[ + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_hand", + name="end_effector", + offset=OffsetCfg( + pos=[0.0, 0.0, 0.1034], + ), + ), + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_rightfinger", + name="tool_rightfinger", + offset=OffsetCfg( + pos=(0.0, 0.0, 0.046), + ), + ), + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_leftfinger", + name="tool_leftfinger", + offset=OffsetCfg( + pos=(0.0, 0.0, 0.046), + ), + ), + ], + ) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/__init__.py new file mode 100644 index 0000000000..8a6d3a1a71 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2022-2024, 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 lift environments.""" + +from omni.isaac.lab.envs.mdp import * # noqa: F401, F403 + +from .observations import * # noqa: F401, F403 +from .terminations import * # noqa: F401, F403 diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/franka_stack_events.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/franka_stack_events.py new file mode 100644 index 0000000000..8d6aa87351 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/franka_stack_events.py @@ -0,0 +1,186 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + + +from __future__ import annotations + +import math +import random +import torch +from typing import TYPE_CHECKING + +import omni.isaac.lab.utils.math as math_utils +from omni.isaac.lab.assets import Articulation, AssetBase +from omni.isaac.lab.managers import SceneEntityCfg + +if TYPE_CHECKING: + from omni.isaac.lab.envs import ManagerBasedEnv + + +def set_default_joint_pose( + env: ManagerBasedEnv, + env_ids: torch.Tensor, + default_pose: torch.Tensor, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +): + # Set the default pose for robots in all envs + asset = env.scene[asset_cfg.name] + asset.data.default_joint_pos = torch.tensor(default_pose, device=env.device).repeat(env.num_envs, 1) + + +def randomize_joint_by_gaussian_offset( + env: ManagerBasedEnv, + env_ids: torch.Tensor, + mean: float, + std: float, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +): + asset: Articulation = env.scene[asset_cfg.name] + + # Add gaussian noise to joint states + joint_pos = asset.data.default_joint_pos[env_ids].clone() + joint_vel = asset.data.default_joint_vel[env_ids].clone() + joint_pos += math_utils.sample_gaussian(mean, std, joint_pos.shape, joint_pos.device) + + # Clamp joint pos to limits + joint_pos_limits = asset.data.soft_joint_pos_limits[env_ids] + joint_pos = joint_pos.clamp_(joint_pos_limits[..., 0], joint_pos_limits[..., 1]) + + # Don't noise the gripper poses + joint_pos[:, -2:] = asset.data.default_joint_pos[env_ids, -2:] + + # Set into the physics simulation + asset.set_joint_position_target(joint_pos, env_ids=env_ids) + asset.set_joint_velocity_target(joint_vel, env_ids=env_ids) + asset.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids) + + +def randomize_scene_lighting_domelight( + env: ManagerBasedEnv, + env_ids: torch.Tensor, + intensity_range: tuple[float, float], + asset_cfg: SceneEntityCfg = SceneEntityCfg("light"), +): + asset: AssetBase = env.scene[asset_cfg.name] + light_prim = asset.prims[0] + + # Sample new light intensity + new_intensity = random.uniform(intensity_range[0], intensity_range[1]) + + # Set light intensity to light prim + intensity_attr = light_prim.GetAttribute("inputs:intensity") + intensity_attr.Set(new_intensity) + + +def sample_object_poses( + num_objects: int, + min_separation: float = 0.0, + pose_range: dict[str, tuple[float, float]] = {}, + max_sample_tries: int = 5000, +): + range_list = [pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] + pose_list = [] + + for i in range(num_objects): + for j in range(max_sample_tries): + sample = [random.uniform(range[0], range[1]) for range in range_list] + + # Accept pose if it is the first one, or if reached max num tries + if len(pose_list) == 0 or j == max_sample_tries - 1: + pose_list.append(sample) + break + + # Check if pose of object is sufficiently far away from all other objects + separation_check = [math.dist(sample[:3], pose[:3]) > min_separation for pose in pose_list] + if False not in separation_check: + pose_list.append(sample) + break + + return pose_list + + +def randomize_object_pose( + env: ManagerBasedEnv, + env_ids: torch.Tensor, + asset_cfgs: list[SceneEntityCfg], + min_separation: float = 0.0, + pose_range: dict[str, tuple[float, float]] = {}, + max_sample_tries: int = 5000, +): + if env_ids is None: + return + + # Randomize poses in each environment independently + for cur_env in env_ids.tolist(): + pose_list = sample_object_poses( + num_objects=len(asset_cfgs), + min_separation=min_separation, + pose_range=pose_range, + max_sample_tries=max_sample_tries, + ) + + # Randomize pose for each object + for i in range(len(asset_cfgs)): + asset_cfg = asset_cfgs[i] + asset = env.scene[asset_cfg.name] + + # Write pose to simulation + pose_tensor = torch.tensor([pose_list[i]], device=env.device) + positions = pose_tensor[:, 0:3] + env.scene.env_origins[cur_env, 0:3] + orientations = math_utils.quat_from_euler_xyz(pose_tensor[:, 3], pose_tensor[:, 4], pose_tensor[:, 5]) + asset.write_root_link_pose_to_sim( + torch.cat([positions, orientations], dim=-1), env_ids=torch.tensor([cur_env], device=env.device) + ) + asset.write_root_com_velocity_to_sim( + torch.zeros(1, 6, device=env.device), env_ids=torch.tensor([cur_env], device=env.device) + ) + + +def randomize_rigid_objects_in_focus( + env: ManagerBasedEnv, + env_ids: torch.Tensor, + asset_cfgs: list[SceneEntityCfg], + out_focus_state: torch.Tensor, + min_separation: float = 0.0, + pose_range: dict[str, tuple[float, float]] = {}, + max_sample_tries: int = 5000, +): + if env_ids is None: + return + + # List of rigid objects in focus for each env (dim = [num_envs, num_rigid_objects]) + env.rigid_objects_in_focus = [] + + for cur_env in env_ids.tolist(): + # Sample in focus object poses + pose_list = sample_object_poses( + num_objects=len(asset_cfgs), + min_separation=min_separation, + pose_range=pose_range, + max_sample_tries=max_sample_tries, + ) + + selected_ids = [] + for asset_idx in range(len(asset_cfgs)): + asset_cfg = asset_cfgs[asset_idx] + asset = env.scene[asset_cfg.name] + + # Randomly select an object to bring into focus + object_id = random.randint(0, asset.num_objects - 1) + selected_ids.append(object_id) + + # Create object state tensor + object_states = torch.stack([out_focus_state] * asset.num_objects).to(device=env.device) + pose_tensor = torch.tensor([pose_list[asset_idx]], device=env.device) + positions = pose_tensor[:, 0:3] + env.scene.env_origins[cur_env, 0:3] + orientations = math_utils.quat_from_euler_xyz(pose_tensor[:, 3], pose_tensor[:, 4], pose_tensor[:, 5]) + object_states[object_id, 0:3] = positions + object_states[object_id, 3:7] = orientations + + asset.write_object_state_to_sim( + object_state=object_states, env_ids=torch.tensor([cur_env], device=env.device) + ) + + env.rigid_objects_in_focus.append(selected_ids) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/observations.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/observations.py new file mode 100644 index 0000000000..53362934bb --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/observations.py @@ -0,0 +1,326 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from omni.isaac.lab.assets import Articulation, RigidObject, RigidObjectCollection +from omni.isaac.lab.managers import SceneEntityCfg +from omni.isaac.lab.sensors import FrameTransformer + +if TYPE_CHECKING: + from omni.isaac.lab.envs import ManagerBasedRLEnv + + +def cube_positions_in_world_frame( + env: ManagerBasedRLEnv, + cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), + cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), + cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), +) -> torch.Tensor: + """The position of the cubes in the world frame.""" + cube_1: RigidObject = env.scene[cube_1_cfg.name] + cube_2: RigidObject = env.scene[cube_2_cfg.name] + cube_3: RigidObject = env.scene[cube_3_cfg.name] + + return torch.cat((cube_1.data.root_link_pos_w, cube_2.data.root_link_pos_w, cube_3.data.root_link_pos_w), dim=1) + + +def instance_randomize_cube_positions_in_world_frame( + env: ManagerBasedRLEnv, + cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), + cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), + cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), +) -> torch.Tensor: + """The position of the cubes in the world frame.""" + if not hasattr(env, "rigid_objects_in_focus"): + return torch.full((env.num_envs, 9), fill_value=-1) + + cube_1: RigidObjectCollection = env.scene[cube_1_cfg.name] + cube_2: RigidObjectCollection = env.scene[cube_2_cfg.name] + cube_3: RigidObjectCollection = env.scene[cube_3_cfg.name] + + cube_1_pos_w = [] + cube_2_pos_w = [] + cube_3_pos_w = [] + for env_id in range(env.num_envs): + cube_1_pos_w.append(cube_1.data.object_pos_w[env_id, env.rigid_objects_in_focus[env_id][0], :3]) + cube_2_pos_w.append(cube_2.data.object_pos_w[env_id, env.rigid_objects_in_focus[env_id][1], :3]) + cube_3_pos_w.append(cube_3.data.object_pos_w[env_id, env.rigid_objects_in_focus[env_id][2], :3]) + cube_1_pos_w = torch.stack(cube_1_pos_w) + cube_2_pos_w = torch.stack(cube_2_pos_w) + cube_3_pos_w = torch.stack(cube_3_pos_w) + + return torch.cat((cube_1_pos_w, cube_2_pos_w, cube_3_pos_w), dim=1) + + +def cube_orientations_in_world_frame( + env: ManagerBasedRLEnv, + cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), + cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), + cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), +): + """The orientation of the cubes in the world frame.""" + cube_1: RigidObject = env.scene[cube_1_cfg.name] + cube_2: RigidObject = env.scene[cube_2_cfg.name] + cube_3: RigidObject = env.scene[cube_3_cfg.name] + + return torch.cat((cube_1.data.root_link_quat_w, cube_2.data.root_link_quat_w, cube_3.data.root_link_quat_w), dim=1) + + +def instance_randomize_cube_orientations_in_world_frame( + env: ManagerBasedRLEnv, + cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), + cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), + cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), +) -> torch.Tensor: + """The orientation of the cubes in the world frame.""" + if not hasattr(env, "rigid_objects_in_focus"): + return torch.full((env.num_envs, 9), fill_value=-1) + + cube_1: RigidObjectCollection = env.scene[cube_1_cfg.name] + cube_2: RigidObjectCollection = env.scene[cube_2_cfg.name] + cube_3: RigidObjectCollection = env.scene[cube_3_cfg.name] + + cube_1_quat_w = [] + cube_2_quat_w = [] + cube_3_quat_w = [] + for env_id in range(env.num_envs): + cube_1_quat_w.append(cube_1.data.object_quat_w[env_id, env.rigid_objects_in_focus[env_id][0], :4]) + cube_2_quat_w.append(cube_2.data.object_quat_w[env_id, env.rigid_objects_in_focus[env_id][1], :4]) + cube_3_quat_w.append(cube_3.data.object_quat_w[env_id, env.rigid_objects_in_focus[env_id][2], :4]) + cube_1_quat_w = torch.stack(cube_1_quat_w) + cube_2_quat_w = torch.stack(cube_2_quat_w) + cube_3_quat_w = torch.stack(cube_3_quat_w) + + return torch.cat((cube_1_quat_w, cube_2_quat_w, cube_3_quat_w), dim=1) + + +def object_obs( + env: ManagerBasedRLEnv, + cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), + cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), + cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), + ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame"), +): + """ + Object observations (in world frame): + cube_1 pos, + cube_1 quat, + cube_2 pos, + cube_2 quat, + cube_3 pos, + cube_3 quat, + gripper to cube_1, + gripper to cube_2, + gripper to cube_3, + cube_1 to cube_2, + cube_2 to cube_3, + cube_1 to cube_3, + """ + cube_1: RigidObject = env.scene[cube_1_cfg.name] + cube_2: RigidObject = env.scene[cube_2_cfg.name] + cube_3: RigidObject = env.scene[cube_3_cfg.name] + ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] + + cube_1_pos_w = cube_1.data.root_link_pos_w + cube_1_quat_w = cube_1.data.root_link_quat_w + + cube_2_pos_w = cube_2.data.root_link_pos_w + cube_2_quat_w = cube_2.data.root_link_quat_w + + cube_3_pos_w = cube_3.data.root_link_pos_w + cube_3_quat_w = cube_3.data.root_link_quat_w + + ee_pos_w = ee_frame.data.target_pos_w[:, 0, :] + gripper_to_cube_1 = cube_1_pos_w - ee_pos_w + gripper_to_cube_2 = cube_2_pos_w - ee_pos_w + gripper_to_cube_3 = cube_3_pos_w - ee_pos_w + + cube_1_to_2 = cube_1_pos_w - cube_2_pos_w + cube_2_to_3 = cube_2_pos_w - cube_3_pos_w + cube_1_to_3 = cube_1_pos_w - cube_3_pos_w + + return torch.cat( + ( + cube_1_pos_w, + cube_1_quat_w, + cube_2_pos_w, + cube_2_quat_w, + cube_3_pos_w, + cube_3_quat_w, + gripper_to_cube_1, + gripper_to_cube_2, + gripper_to_cube_3, + cube_1_to_2, + cube_2_to_3, + cube_1_to_3, + ), + dim=1, + ) + + +def instance_randomize_object_obs( + env: ManagerBasedRLEnv, + cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), + cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), + cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), + ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame"), +): + """ + Object observations (in world frame): + cube_1 pos, + cube_1 quat, + cube_2 pos, + cube_2 quat, + cube_3 pos, + cube_3 quat, + gripper to cube_1, + gripper to cube_2, + gripper to cube_3, + cube_1 to cube_2, + cube_2 to cube_3, + cube_1 to cube_3, + """ + if not hasattr(env, "rigid_objects_in_focus"): + return torch.full((env.num_envs, 9), fill_value=-1) + + cube_1: RigidObjectCollection = env.scene[cube_1_cfg.name] + cube_2: RigidObjectCollection = env.scene[cube_2_cfg.name] + cube_3: RigidObjectCollection = env.scene[cube_3_cfg.name] + ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] + + cube_1_pos_w = [] + cube_2_pos_w = [] + cube_3_pos_w = [] + cube_1_quat_w = [] + cube_2_quat_w = [] + cube_3_quat_w = [] + for env_id in range(env.num_envs): + cube_1_pos_w.append(cube_1.data.object_pos_w[env_id, env.rigid_objects_in_focus[env_id][0], :3]) + cube_2_pos_w.append(cube_2.data.object_pos_w[env_id, env.rigid_objects_in_focus[env_id][1], :3]) + cube_3_pos_w.append(cube_3.data.object_pos_w[env_id, env.rigid_objects_in_focus[env_id][2], :3]) + cube_1_quat_w.append(cube_1.data.object_quat_w[env_id, env.rigid_objects_in_focus[env_id][0], :4]) + cube_2_quat_w.append(cube_2.data.object_quat_w[env_id, env.rigid_objects_in_focus[env_id][1], :4]) + cube_3_quat_w.append(cube_3.data.object_quat_w[env_id, env.rigid_objects_in_focus[env_id][2], :4]) + cube_1_pos_w = torch.stack(cube_1_pos_w) + cube_2_pos_w = torch.stack(cube_2_pos_w) + cube_3_pos_w = torch.stack(cube_3_pos_w) + cube_1_quat_w = torch.stack(cube_1_quat_w) + cube_2_quat_w = torch.stack(cube_2_quat_w) + cube_3_quat_w = torch.stack(cube_3_quat_w) + + ee_pos_w = ee_frame.data.target_pos_w[:, 0, :] + gripper_to_cube_1 = cube_1_pos_w - ee_pos_w + gripper_to_cube_2 = cube_2_pos_w - ee_pos_w + gripper_to_cube_3 = cube_3_pos_w - ee_pos_w + + cube_1_to_2 = cube_1_pos_w - cube_2_pos_w + cube_2_to_3 = cube_2_pos_w - cube_3_pos_w + cube_1_to_3 = cube_1_pos_w - cube_3_pos_w + + return torch.cat( + ( + cube_1_pos_w, + cube_1_quat_w, + cube_2_pos_w, + cube_2_quat_w, + cube_3_pos_w, + cube_3_quat_w, + gripper_to_cube_1, + gripper_to_cube_2, + gripper_to_cube_3, + cube_1_to_2, + cube_2_to_3, + cube_1_to_3, + ), + dim=1, + ) + + +def ee_frame_pos(env: ManagerBasedRLEnv, ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame")) -> torch.Tensor: + ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] + ee_frame_pos = ee_frame.data.target_pos_w[:, 0, :] + + return ee_frame_pos + + +def ee_frame_quat(env: ManagerBasedRLEnv, ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame")) -> torch.Tensor: + ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] + ee_frame_quat = ee_frame.data.target_quat_w[:, 0, :] + + return ee_frame_quat + + +def gripper_pos(env: ManagerBasedRLEnv, robot_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + robot: Articulation = env.scene[robot_cfg.name] + finger_joint_1 = robot.data.joint_pos[:, -1].clone().unsqueeze(1) + finger_joint_2 = -1 * robot.data.joint_pos[:, -2].clone().unsqueeze(1) + + return torch.cat((finger_joint_1, finger_joint_2), dim=1) + + +def object_grasped( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg, + ee_frame_cfg: SceneEntityCfg, + object_cfg: SceneEntityCfg, + diff_threshold: float = 0.06, + gripper_open_val: torch.tensor = torch.tensor([0.04]), + gripper_threshold: float = 0.005, +) -> torch.Tensor: + """Check if an object is grasped by the specified robot.""" + + robot: Articulation = env.scene[robot_cfg.name] + ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] + object: RigidObject = env.scene[object_cfg.name] + + object_pos = object.data.root_link_pos_w + end_effector_pos = ee_frame.data.target_pos_w[:, 0, :] + pose_diff = torch.linalg.vector_norm(object_pos - end_effector_pos, dim=1) + + grasped = torch.logical_and( + pose_diff < diff_threshold, + torch.abs(robot.data.joint_pos[:, -1] - gripper_open_val.to(env.device)) > gripper_threshold, + ) + grasped = torch.logical_and( + grasped, torch.abs(robot.data.joint_pos[:, -2] - gripper_open_val.to(env.device)) > gripper_threshold + ) + + return grasped + + +def object_stacked( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg, + upper_object_cfg: SceneEntityCfg, + lower_object_cfg: SceneEntityCfg, + xy_threshold: float = 0.05, + height_threshold: float = 0.005, + height_diff: float = 0.0468, + gripper_open_val: torch.tensor = torch.tensor([0.04]), +) -> torch.Tensor: + """Check if an object is stacked by the specified robot.""" + + robot: Articulation = env.scene[robot_cfg.name] + upper_object: RigidObject = env.scene[upper_object_cfg.name] + lower_object: RigidObject = env.scene[lower_object_cfg.name] + + pos_diff = upper_object.data.root_link_pos_w - lower_object.data.root_link_pos_w + height_dist = torch.linalg.vector_norm(pos_diff[:, 2:], dim=1) + xy_dist = torch.linalg.vector_norm(pos_diff[:, :2], dim=1) + + stacked = torch.logical_and(xy_dist < xy_threshold, (height_dist - height_diff) < height_threshold) + + stacked = torch.logical_and( + torch.isclose(robot.data.joint_pos[:, -1], gripper_open_val.to(env.device), atol=1e-4, rtol=1e-4), stacked + ) + stacked = torch.logical_and( + torch.isclose(robot.data.joint_pos[:, -2], gripper_open_val.to(env.device), atol=1e-4, rtol=1e-4), stacked + ) + + return stacked diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/terminations.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/terminations.py new file mode 100644 index 0000000000..b4972006fb --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/terminations.py @@ -0,0 +1,66 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Common functions that can be used to activate certain terminations for the lift task. + +The functions can be passed to the :class:`omni.isaac.lab.managers.TerminationTermCfg` object to enable +the termination introduced by the function. +""" + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from omni.isaac.lab.assets import Articulation, RigidObject +from omni.isaac.lab.managers import SceneEntityCfg + +if TYPE_CHECKING: + from omni.isaac.lab.envs import ManagerBasedRLEnv + + +def cubes_stacked( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), + cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), + cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), + xy_threshold: float = 0.05, + height_threshold: float = 0.005, + height_diff: float = 0.0468, + gripper_open_val: torch.tensor = torch.tensor([0.04]), + atol=0.0001, + rtol=0.0001, +): + robot: Articulation = env.scene[robot_cfg.name] + cube_1: RigidObject = env.scene[cube_1_cfg.name] + cube_2: RigidObject = env.scene[cube_2_cfg.name] + cube_3: RigidObject = env.scene[cube_3_cfg.name] + + pos_diff_c12 = cube_1.data.root_link_pos_w - cube_2.data.root_link_pos_w + pos_diff_c23 = cube_2.data.root_link_pos_w - cube_3.data.root_link_pos_w + + # Compute cube position difference in x-y plane + xy_dist_c12 = torch.norm(pos_diff_c12[:, :2], dim=1) + xy_dist_c23 = torch.norm(pos_diff_c23[:, :2], dim=1) + + # Compute cube height difference + h_dist_c12 = torch.norm(pos_diff_c12[:, 2:], dim=1) + h_dist_c23 = torch.norm(pos_diff_c23[:, 2:], dim=1) + + # Check cube positions + stacked = torch.logical_and(xy_dist_c12 < xy_threshold, xy_dist_c23 < xy_threshold) + stacked = torch.logical_and(h_dist_c12 - height_diff < height_threshold, stacked) + stacked = torch.logical_and(h_dist_c23 - height_diff < height_threshold, stacked) + + # Check gripper positions + stacked = torch.logical_and( + torch.isclose(robot.data.joint_pos[:, -1], gripper_open_val.to(env.device), atol=atol, rtol=rtol), stacked + ) + stacked = torch.logical_and( + torch.isclose(robot.data.joint_pos[:, -2], gripper_open_val.to(env.device), atol=atol, rtol=rtol), stacked + ) + + return stacked diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/stack_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/stack_env_cfg.py new file mode 100644 index 0000000000..6d07861a59 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/stack_env_cfg.py @@ -0,0 +1,205 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import MISSING + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.assets import ArticulationCfg, AssetBaseCfg +from omni.isaac.lab.envs import ManagerBasedRLEnvCfg +from omni.isaac.lab.managers import ObservationGroupCfg as ObsGroup +from omni.isaac.lab.managers import ObservationTermCfg as ObsTerm +from omni.isaac.lab.managers import SceneEntityCfg +from omni.isaac.lab.managers import TerminationTermCfg as DoneTerm +from omni.isaac.lab.scene import InteractiveSceneCfg +from omni.isaac.lab.sensors import CameraCfg +from omni.isaac.lab.sensors.frame_transformer.frame_transformer_cfg import FrameTransformerCfg +from omni.isaac.lab.sim.spawners.from_files.from_files_cfg import GroundPlaneCfg, UsdFileCfg +from omni.isaac.lab.utils import configclass +from omni.isaac.lab.utils.assets import ISAAC_NUCLEUS_DIR + +from . import mdp + + +## +# Scene definition +## +@configclass +class ObjectTableSceneCfg(InteractiveSceneCfg): + """Configuration for the lift scene with a robot and a object. + This is the abstract base implementation, the exact scene is defined in the derived classes + which need to set the target object, robot and end-effector frames + """ + + # robots: will be populated by agent env cfg + robot: ArticulationCfg = MISSING + # end-effector sensor: will be populated by agent env cfg + ee_frame: FrameTransformerCfg = MISSING + + # Cameras + wrist_cam: CameraCfg = MISSING + table_cam: CameraCfg = MISSING + + # Table + table = AssetBaseCfg( + prim_path="{ENV_REGEX_NS}/Table", + init_state=AssetBaseCfg.InitialStateCfg(pos=[0.5, 0, 0], rot=[0.707, 0, 0, 0.707]), + spawn=UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/SeattleLabTable/table_instanceable.usd"), + ) + + # plane + plane = AssetBaseCfg( + prim_path="/World/GroundPlane", + init_state=AssetBaseCfg.InitialStateCfg(pos=[0, 0, -1.05]), + spawn=GroundPlaneCfg(), + ) + + # lights + light = AssetBaseCfg( + prim_path="/World/light", + spawn=sim_utils.DomeLightCfg(color=(0.75, 0.75, 0.75), intensity=3000.0), + ) + + +## +# MDP settings +## +@configclass +class ActionsCfg: + """Action specifications for the MDP.""" + + # will be set by agent env cfg + arm_action: mdp.JointPositionActionCfg = MISSING + gripper_action: mdp.BinaryJointPositionActionCfg = MISSING + + +@configclass +class ObservationsCfg: + """Observation specifications for the MDP.""" + + @configclass + class PolicyCfg(ObsGroup): + """Observations for policy group with state values.""" + + actions = ObsTerm(func=mdp.last_action) + joint_pos = ObsTerm(func=mdp.joint_pos_rel) + joint_vel = ObsTerm(func=mdp.joint_vel_rel) + object = ObsTerm(func=mdp.object_obs) + cube_positions = ObsTerm(func=mdp.cube_positions_in_world_frame) + cube_orientations = ObsTerm(func=mdp.cube_orientations_in_world_frame) + eef_pos = ObsTerm(func=mdp.ee_frame_pos) + eef_quat = ObsTerm(func=mdp.ee_frame_quat) + gripper_pos = ObsTerm(func=mdp.gripper_pos) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = False + + @configclass + class RGBCameraPolicyCfg(ObsGroup): + """Observations for policy group with RGB images.""" + + table_cam = ObsTerm( + func=mdp.image, params={"sensor_cfg": SceneEntityCfg("table_cam"), "data_type": "rgb", "normalize": False} + ) + wrist_cam = ObsTerm( + func=mdp.image, params={"sensor_cfg": SceneEntityCfg("wrist_cam"), "data_type": "rgb", "normalize": False} + ) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = False + + @configclass + class SubtaskCfg(ObsGroup): + """Observations for subtask group.""" + + grasp_1 = ObsTerm( + func=mdp.object_grasped, + params={ + "robot_cfg": SceneEntityCfg("robot"), + "ee_frame_cfg": SceneEntityCfg("ee_frame"), + "object_cfg": SceneEntityCfg("cube_2"), + }, + ) + stack_1 = ObsTerm( + func=mdp.object_stacked, + params={ + "robot_cfg": SceneEntityCfg("robot"), + "upper_object_cfg": SceneEntityCfg("cube_2"), + "lower_object_cfg": SceneEntityCfg("cube_1"), + }, + ) + grasp_2 = ObsTerm( + func=mdp.object_grasped, + params={ + "robot_cfg": SceneEntityCfg("robot"), + "ee_frame_cfg": SceneEntityCfg("ee_frame"), + "object_cfg": SceneEntityCfg("cube_3"), + }, + ) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = False + + # observation groups + policy: PolicyCfg = PolicyCfg() + rgb_camera: RGBCameraPolicyCfg = RGBCameraPolicyCfg() + subtask_terms: SubtaskCfg = SubtaskCfg() + + +@configclass +class TerminationsCfg: + """Termination terms for the MDP.""" + + time_out = DoneTerm(func=mdp.time_out, time_out=True) + + cube_1_dropping = DoneTerm( + func=mdp.root_height_below_minimum, params={"minimum_height": -0.05, "asset_cfg": SceneEntityCfg("cube_1")} + ) + + cube_2_dropping = DoneTerm( + func=mdp.root_height_below_minimum, params={"minimum_height": -0.05, "asset_cfg": SceneEntityCfg("cube_2")} + ) + + cube_3_dropping = DoneTerm( + func=mdp.root_height_below_minimum, params={"minimum_height": -0.05, "asset_cfg": SceneEntityCfg("cube_3")} + ) + + success = DoneTerm(func=mdp.cubes_stacked) + + +@configclass +class StackEnvCfg(ManagerBasedRLEnvCfg): + """Configuration for the stacking environment.""" + + # Scene settings + scene: ObjectTableSceneCfg = ObjectTableSceneCfg(num_envs=4096, env_spacing=2.5, replicate_physics=False) + # Basic settings + observations: ObservationsCfg = ObservationsCfg() + actions: ActionsCfg = ActionsCfg() + # MDP settings + terminations: TerminationsCfg = TerminationsCfg() + + # Unused managers + commands = None + rewards = None + events = None + curriculum = None + + def __post_init__(self): + """Post initialization.""" + # general settings + self.decimation = 5 + self.episode_length_s = 30.0 + # simulation settings + self.sim.dt = 0.01 # 100Hz + self.sim.render_interval = self.decimation + + self.sim.physx.bounce_threshold_velocity = 0.2 + self.sim.physx.bounce_threshold_velocity = 0.01 + self.sim.physx.gpu_found_lost_aggregate_pairs_capacity = 1024 * 1024 * 4 + self.sim.physx.gpu_total_aggregate_pairs_capacity = 16 * 1024 + self.sim.physx.friction_correlation_distance = 0.00625 diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/stack_instance_randomize_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/stack_instance_randomize_env_cfg.py new file mode 100644 index 0000000000..1a5766b918 --- /dev/null +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/stack_instance_randomize_env_cfg.py @@ -0,0 +1,157 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import MISSING + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.assets import ArticulationCfg, AssetBaseCfg +from omni.isaac.lab.envs import ManagerBasedRLEnvCfg +from omni.isaac.lab.managers import ObservationGroupCfg as ObsGroup +from omni.isaac.lab.managers import ObservationTermCfg as ObsTerm +from omni.isaac.lab.managers import SceneEntityCfg +from omni.isaac.lab.managers import TerminationTermCfg as DoneTerm +from omni.isaac.lab.scene import InteractiveSceneCfg +from omni.isaac.lab.sensors import CameraCfg +from omni.isaac.lab.sensors.frame_transformer.frame_transformer_cfg import FrameTransformerCfg +from omni.isaac.lab.sim.spawners.from_files.from_files_cfg import GroundPlaneCfg, UsdFileCfg +from omni.isaac.lab.utils import configclass +from omni.isaac.lab.utils.assets import ISAAC_NUCLEUS_DIR + +from . import mdp + + +## +# Scene definition +## +@configclass +class ObjectTableSceneCfg(InteractiveSceneCfg): + """Configuration for the lift scene with a robot and a object. + This is the abstract base implementation, the exact scene is defined in the derived classes + which need to set the target object, robot and end-effector frames + """ + + # robots: will be populated by agent env cfg + robot: ArticulationCfg = MISSING + # end-effector sensor: will be populated by agent env cfg + ee_frame: FrameTransformerCfg = MISSING + + # Cameras + wrist_cam: CameraCfg = MISSING + table_cam: CameraCfg = MISSING + + # Table + table = AssetBaseCfg( + prim_path="{ENV_REGEX_NS}/Table", + init_state=AssetBaseCfg.InitialStateCfg(pos=[0.5, 0, 0], rot=[0.707, 0, 0, 0.707]), + spawn=UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/SeattleLabTable/table_instanceable.usd"), + ) + + # plane + plane = AssetBaseCfg( + prim_path="/World/GroundPlane", + init_state=AssetBaseCfg.InitialStateCfg(pos=[0, 0, -1.05]), + spawn=GroundPlaneCfg(), + ) + + # lights + light = AssetBaseCfg( + prim_path="/World/light", + spawn=sim_utils.DomeLightCfg(color=(0.75, 0.75, 0.75), intensity=3000.0), + ) + + +## +# MDP settings +## +@configclass +class ActionsCfg: + """Action specifications for the MDP.""" + + # will be set by agent env cfg + arm_action: mdp.JointPositionActionCfg = MISSING + gripper_action: mdp.BinaryJointPositionActionCfg = MISSING + + +@configclass +class ObservationsCfg: + """Observation specifications for the MDP.""" + + @configclass + class PolicyCfg(ObsGroup): + """Observations for policy group with state values.""" + + actions = ObsTerm(func=mdp.last_action) + joint_pos = ObsTerm(func=mdp.joint_pos_rel) + joint_vel = ObsTerm(func=mdp.joint_vel_rel) + object = ObsTerm(func=mdp.instance_randomize_object_obs) + cube_positions = ObsTerm(func=mdp.instance_randomize_cube_positions_in_world_frame) + cube_orientations = ObsTerm(func=mdp.instance_randomize_cube_orientations_in_world_frame) + eef_pos = ObsTerm(func=mdp.ee_frame_pos) + eef_quat = ObsTerm(func=mdp.ee_frame_quat) + gripper_pos = ObsTerm(func=mdp.gripper_pos) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = False + + @configclass + class RGBCameraPolicyCfg(ObsGroup): + """Observations for policy group with RGB images.""" + + table_cam = ObsTerm( + func=mdp.image, params={"sensor_cfg": SceneEntityCfg("table_cam"), "data_type": "rgb", "normalize": False} + ) + wrist_cam = ObsTerm( + func=mdp.image, params={"sensor_cfg": SceneEntityCfg("wrist_cam"), "data_type": "rgb", "normalize": False} + ) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = False + + # observation groups + policy: PolicyCfg = PolicyCfg() + rgb_camera: RGBCameraPolicyCfg = RGBCameraPolicyCfg() + + +@configclass +class TerminationsCfg: + """Termination terms for the MDP.""" + + time_out = DoneTerm(func=mdp.time_out, time_out=True) + + +@configclass +class StackInstanceRandomizeEnvCfg(ManagerBasedRLEnvCfg): + """Configuration for the stacking environment.""" + + # Scene settings + scene: ObjectTableSceneCfg = ObjectTableSceneCfg(num_envs=4096, env_spacing=2.5, replicate_physics=False) + # Basic settings + observations: ObservationsCfg = ObservationsCfg() + actions: ActionsCfg = ActionsCfg() + # MDP settings + terminations: TerminationsCfg = TerminationsCfg() + + # Unused managers + commands = None + rewards = None + events = None + curriculum = None + + def __post_init__(self): + """Post initialization.""" + # general settings + self.decimation = 5 + self.episode_length_s = 30.0 + # simulation settings + self.sim.dt = 0.01 # 100Hz + self.sim.render_interval = self.decimation + + self.sim.physx.bounce_threshold_velocity = 0.2 + self.sim.physx.bounce_threshold_velocity = 0.01 + self.sim.physx.gpu_found_lost_aggregate_pairs_capacity = 1024 * 1024 * 4 + self.sim.physx.gpu_total_aggregate_pairs_capacity = 16 * 1024 + self.sim.physx.friction_correlation_distance = 0.00625 diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/navigation/config/anymal_c/__init__.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/navigation/config/anymal_c/__init__.py index 74ef1ccace..31e7989bc8 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/navigation/config/anymal_c/__init__.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/navigation/config/anymal_c/__init__.py @@ -5,7 +5,7 @@ import gymnasium as gym -from . import agents, navigation_env_cfg +from . import agents ## # Register Gym environments. @@ -16,7 +16,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": navigation_env_cfg.NavigationEnvCfg, + "env_cfg_entry_point": f"{__name__}.navigation_env_cfg:NavigationEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:NavigationEnvPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -27,7 +27,7 @@ entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": navigation_env_cfg.NavigationEnvCfg_PLAY, + "env_cfg_entry_point": f"{__name__}.navigation_env_cfg:NavigationEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:NavigationEnvPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/navigation/config/anymal_c/navigation_env_cfg.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/navigation/config/anymal_c/navigation_env_cfg.py index 6f5f00a025..c13875c545 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/navigation/config/anymal_c/navigation_env_cfg.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/navigation/config/anymal_c/navigation_env_cfg.py @@ -107,13 +107,6 @@ class CommandsCfg: ) -@configclass -class CurriculumCfg: - """Curriculum terms for the MDP.""" - - pass - - @configclass class TerminationsCfg: """Termination terms for the MDP.""" @@ -127,14 +120,16 @@ class TerminationsCfg: @configclass class NavigationEnvCfg(ManagerBasedRLEnvCfg): + """Configuration for the navigation environment.""" + + # environment settings scene: SceneEntityCfg = LOW_LEVEL_ENV_CFG.scene - commands: CommandsCfg = CommandsCfg() actions: ActionsCfg = ActionsCfg() observations: ObservationsCfg = ObservationsCfg() - rewards: RewardsCfg = RewardsCfg() events: EventCfg = EventCfg() - - curriculum: CurriculumCfg = CurriculumCfg() + # mdp settings + commands: CommandsCfg = CommandsCfg() + rewards: RewardsCfg = RewardsCfg() terminations: TerminationsCfg = TerminationsCfg() def __post_init__(self): diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/navigation/mdp/pre_trained_policy_action.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/navigation/mdp/pre_trained_policy_action.py index 7e5af7d313..1936ffa61a 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/navigation/mdp/pre_trained_policy_action.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/navigation/mdp/pre_trained_policy_action.py @@ -128,11 +128,11 @@ def _debug_vis_callback(self, event): return # get marker location # -- base state - base_pos_w = self.robot.data.root_pos_w.clone() + base_pos_w = self.robot.data.root_link_pos_w.clone() base_pos_w[:, 2] += 0.5 # -- resolve the scales and quaternions vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.raw_actions[:, :2]) - vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(self.robot.data.root_lin_vel_b[:, :2]) + vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(self.robot.data.root_com_lin_vel_b[:, :2]) # display markers self.base_vel_goal_visualizer.visualize(base_pos_w, vel_des_arrow_quat, vel_des_arrow_scale) self.base_vel_visualizer.visualize(base_pos_w, vel_arrow_quat, vel_arrow_scale) @@ -153,7 +153,7 @@ def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torc zeros = torch.zeros_like(heading_angle) arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle) # convert everything back from base to world frame - base_quat_w = self.robot.data.root_quat_w + base_quat_w = self.robot.data.root_link_quat_w arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat) return arrow_scale, arrow_quat diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/utils/hydra.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/utils/hydra.py index 1522e9f25e..5d6a8b7c61 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/utils/hydra.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/utils/hydra.py @@ -17,6 +17,7 @@ raise ImportError("Hydra is not installed. Please install it by running 'pip install hydra-core'.") from omni.isaac.lab.envs import DirectRLEnvCfg, ManagerBasedRLEnvCfg +from omni.isaac.lab.envs.utils.spaces import replace_env_cfg_spaces_with_strings, replace_strings_with_env_cfg_spaces from omni.isaac.lab.utils import replace_slices_with_strings, replace_strings_with_slices from omni.isaac.lab_tasks.utils.parse_cfg import load_cfg_from_registry @@ -40,6 +41,9 @@ def register_task_to_hydra( # load the configurations env_cfg = load_cfg_from_registry(task_name, "env_cfg_entry_point") agent_cfg = load_cfg_from_registry(task_name, agent_cfg_entry_point) + # replace gymnasium spaces with strings because OmegaConf does not support them. + # this must be done before converting the env configs to dictionary to avoid internal reinterpretations + replace_env_cfg_spaces_with_strings(env_cfg) # convert the configs to dictionary env_cfg_dict = env_cfg.to_dict() if isinstance(agent_cfg, dict): @@ -83,6 +87,10 @@ def hydra_main(hydra_env_cfg: DictConfig, env_cfg=env_cfg, agent_cfg=agent_cfg): hydra_env_cfg = replace_strings_with_slices(hydra_env_cfg) # update the configs with the Hydra command line arguments env_cfg.from_dict(hydra_env_cfg["env"]) + # replace strings that represent gymnasium spaces because OmegaConf does not support them. + # this must be done after converting the env configs from dictionary to avoid internal reinterpretations + replace_strings_with_env_cfg_spaces(env_cfg) + # get agent configs if isinstance(agent_cfg, dict): agent_cfg = hydra_env_cfg["agent"] else: diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/utils/wrappers/rsl_rl/vecenv_wrapper.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/utils/wrappers/rsl_rl/vecenv_wrapper.py index 0badd08c31..0dedef9ef0 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/utils/wrappers/rsl_rl/vecenv_wrapper.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/utils/wrappers/rsl_rl/vecenv_wrapper.py @@ -70,19 +70,19 @@ def __init__(self, env: ManagerBasedRLEnv | DirectRLEnv): if hasattr(self.unwrapped, "action_manager"): self.num_actions = self.unwrapped.action_manager.total_action_dim else: - self.num_actions = self.unwrapped.num_actions + self.num_actions = gym.spaces.flatdim(self.unwrapped.single_action_space) if hasattr(self.unwrapped, "observation_manager"): self.num_obs = self.unwrapped.observation_manager.group_obs_dim["policy"][0] else: - self.num_obs = self.unwrapped.num_observations + self.num_obs = gym.spaces.flatdim(self.unwrapped.single_observation_space["policy"]) # -- privileged observations if ( hasattr(self.unwrapped, "observation_manager") and "critic" in self.unwrapped.observation_manager.group_obs_dim ): self.num_privileged_obs = self.unwrapped.observation_manager.group_obs_dim["critic"][0] - elif hasattr(self.unwrapped, "num_states"): - self.num_privileged_obs = self.unwrapped.num_states + elif hasattr(self.unwrapped, "num_states") and "critic" in self.unwrapped.single_observation_space: + self.num_privileged_obs = gym.spaces.flatdim(self.unwrapped.single_observation_space["critic"]) else: self.num_privileged_obs = 0 # reset at the start since the RSL-RL runner does not call reset diff --git a/source/extensions/omni.isaac.lab_tasks/test/test_environment_determinism.py b/source/extensions/omni.isaac.lab_tasks/test/test_environment_determinism.py index 3346e8284d..3aa2977935 100644 --- a/source/extensions/omni.isaac.lab_tasks/test/test_environment_determinism.py +++ b/source/extensions/omni.isaac.lab_tasks/test/test_environment_determinism.py @@ -101,13 +101,20 @@ def _obtain_transition_tuples( """Run random actions and obtain transition tuples after fixed number of steps.""" # create a new stage omni.usd.get_context().new_stage() - # parse configuration - env_cfg = parse_env_cfg(task_name, device=device, num_envs=num_envs) - # set seed - env_cfg.seed = 42 - - # create environment - env = gym.make(task_name, cfg=env_cfg) + try: + # parse configuration + env_cfg = parse_env_cfg(task_name, device=device, num_envs=num_envs) + # set seed + env_cfg.seed = 42 + # create environment + env = gym.make(task_name, cfg=env_cfg) + except Exception as e: + if "env" in locals() and hasattr(env, "_is_closed"): + env.close() + else: + if hasattr(e, "obj") and hasattr(e.obj, "_is_closed"): + e.obj.close() + self.fail(f"Failed to set-up the environment for task {task_name}. Error: {e}") # disable control on stop env.unwrapped.sim._app_control_on_stop_handle = None # type: ignore diff --git a/source/extensions/omni.isaac.lab_tasks/test/test_environments.py b/source/extensions/omni.isaac.lab_tasks/test/test_environments.py index 9e92e26156..8cd41fcd86 100644 --- a/source/extensions/omni.isaac.lab_tasks/test/test_environments.py +++ b/source/extensions/omni.isaac.lab_tasks/test/test_environments.py @@ -22,6 +22,7 @@ import omni.usd from omni.isaac.lab.envs import ManagerBasedRLEnvCfg +from omni.isaac.lab.envs.utils.spaces import sample_space import omni.isaac.lab_tasks # noqa: F401 from omni.isaac.lab_tasks.utils.parse_cfg import parse_env_cfg @@ -87,16 +88,24 @@ def _check_random_actions(self, task_name: str, device: str, num_envs: int, num_ """Run random actions and check environments returned signals are valid.""" # create a new stage omni.usd.get_context().new_stage() - # parse configuration - env_cfg: ManagerBasedRLEnvCfg = parse_env_cfg(task_name, device=device, num_envs=num_envs) - - # skip test if the environment is a multi-agent task - if hasattr(env_cfg, "possible_agents"): - print(f"[INFO]: Skipping {task_name} as it is a multi-agent task") - return - - # create environment - env = gym.make(task_name, cfg=env_cfg) + try: + # parse configuration + env_cfg: ManagerBasedRLEnvCfg = parse_env_cfg(task_name, device=device, num_envs=num_envs) + + # skip test if the environment is a multi-agent task + if hasattr(env_cfg, "possible_agents"): + print(f"[INFO]: Skipping {task_name} as it is a multi-agent task") + return + + # create environment + env = gym.make(task_name, cfg=env_cfg) + except Exception as e: + if "env" in locals() and hasattr(env, "_is_closed"): + env.close() + else: + if hasattr(e, "obj") and hasattr(e.obj, "_is_closed"): + e.obj.close() + self.fail(f"Failed to set-up the environment for task {task_name}. Error: {e}") # disable control on stop env.unwrapped.sim._app_control_on_stop_handle = None # type: ignore @@ -108,12 +117,14 @@ def _check_random_actions(self, task_name: str, device: str, num_envs: int, num_ # simulate environment for num_steps steps with torch.inference_mode(): for _ in range(num_steps): - # sample actions from -1 to 1 - actions = 2 * torch.rand(env.action_space.shape, device=env.unwrapped.device) - 1 + # sample actions according to the defined space + actions = sample_space( + env.unwrapped.single_action_space, device=env.unwrapped.device, batch_size=num_envs + ) # apply actions transition = env.step(actions) # check signals - for data in transition: + for data in transition[:-1]: # exclude info self.assertTrue(self._check_valid_tensor(data), msg=f"Invalid data: {data}") # close the environment @@ -131,14 +142,10 @@ def _check_valid_tensor(data: torch.Tensor | dict) -> bool: """ if isinstance(data, torch.Tensor): return not torch.any(torch.isnan(data)) + elif isinstance(data, (tuple, list)): + return all(TestEnvironments._check_valid_tensor(value) for value in data) elif isinstance(data, dict): - valid_tensor = True - for value in data.values(): - if isinstance(value, dict): - valid_tensor &= TestEnvironments._check_valid_tensor(value) - elif isinstance(value, torch.Tensor): - valid_tensor &= not torch.any(torch.isnan(value)) - return valid_tensor + return all(TestEnvironments._check_valid_tensor(value) for value in data.values()) else: raise ValueError(f"Input data of invalid type: {type(data)}.") diff --git a/source/extensions/omni.isaac.lab_tasks/test/test_multi_agent_environments.py b/source/extensions/omni.isaac.lab_tasks/test/test_multi_agent_environments.py index 19fcd88936..441bdb14f7 100644 --- a/source/extensions/omni.isaac.lab_tasks/test/test_multi_agent_environments.py +++ b/source/extensions/omni.isaac.lab_tasks/test/test_multi_agent_environments.py @@ -21,6 +21,7 @@ import omni.usd from omni.isaac.lab.envs import DirectMARLEnv, DirectMARLEnvCfg +from omni.isaac.lab.envs.utils.spaces import sample_space import omni.isaac.lab_tasks # noqa: F401 from omni.isaac.lab_tasks.utils.parse_cfg import parse_env_cfg @@ -83,19 +84,28 @@ def _check_random_actions(self, task_name: str, device: str, num_envs: int, num_ """Run random actions and check environments return valid signals.""" # create a new stage omni.usd.get_context().new_stage() - # parse configuration - env_cfg: DirectMARLEnvCfg = parse_env_cfg(task_name, device=device, num_envs=num_envs) + try: + # parse configuration + env_cfg: DirectMARLEnvCfg = parse_env_cfg(task_name, device=device, num_envs=num_envs) + + # skip test if the environment is not a multi-agent task + if not hasattr(env_cfg, "possible_agents"): + print(f"[INFO]: Skipping {task_name} as it is not a multi-agent task") + return + + # create environment + env: DirectMARLEnv = gym.make(task_name, cfg=env_cfg) + except Exception as e: + if "env" in locals() and hasattr(env, "_is_closed"): + env.close() + else: + if hasattr(e, "obj") and hasattr(e.obj, "_is_closed"): + e.obj.close() + self.fail(f"Failed to set-up the environment for task {task_name}. Error: {e}") - # skip test if the environment is not a multi-agent task - if not hasattr(env_cfg, "possible_agents"): - print(f"[INFO]: Skipping {task_name} as it is not a multi-agent task") - return - - # create environment - env: DirectMARLEnv = gym.make(task_name, cfg=env_cfg) # this flag is necessary to prevent a bug where the simulation gets stuck randomly when running the # test on many environments. - env.sim.set_setting("/physics/cooking/ujitsoCollisionCooking", False) + env.unwrapped.sim.set_setting("/physics/cooking/ujitsoCollisionCooking", False) # reset environment obs, _ = env.reset() @@ -104,9 +114,11 @@ def _check_random_actions(self, task_name: str, device: str, num_envs: int, num_ # simulate environment for num_steps steps with torch.inference_mode(): for _ in range(num_steps): - # sample actions from -1 to 1 + # sample actions according to the defined space actions = { - agent: 2 * torch.rand(env.action_space(agent).shape, device=env.unwrapped.device) - 1 + agent: sample_space( + env.unwrapped.action_spaces[agent], device=env.unwrapped.device, batch_size=num_envs + ) for agent in env.unwrapped.possible_agents } # apply actions @@ -131,14 +143,10 @@ def _check_valid_tensor(data: torch.Tensor | dict) -> bool: """ if isinstance(data, torch.Tensor): return not torch.any(torch.isnan(data)) + elif isinstance(data, (tuple, list)): + return all(TestEnvironments._check_valid_tensor(value) for value in data) elif isinstance(data, dict): - valid_tensor = True - for value in data.values(): - if isinstance(value, dict): - valid_tensor &= TestEnvironments._check_valid_tensor(value) - elif isinstance(value, torch.Tensor): - valid_tensor &= not torch.any(torch.isnan(value)) - return valid_tensor + return all(TestEnvironments._check_valid_tensor(value) for value in data.values()) else: raise ValueError(f"Input data of invalid type: {type(data)}.") diff --git a/source/extensions/omni.isaac.lab_tasks/test/wrappers/test_rl_games_wrapper.py b/source/extensions/omni.isaac.lab_tasks/test/wrappers/test_rl_games_wrapper.py index e97ff5929e..0cbf01ea75 100644 --- a/source/extensions/omni.isaac.lab_tasks/test/wrappers/test_rl_games_wrapper.py +++ b/source/extensions/omni.isaac.lab_tasks/test/wrappers/test_rl_games_wrapper.py @@ -20,6 +20,8 @@ import omni.usd +from omni.isaac.lab.envs import DirectMARLEnv, multi_agent_to_single_agent + import omni.isaac.lab_tasks # noqa: F401 from omni.isaac.lab_tasks.utils.parse_cfg import load_cfg_from_registry, parse_env_cfg from omni.isaac.lab_tasks.utils.wrappers.rl_games import RlGamesVecEnvWrapper @@ -55,13 +57,24 @@ def test_random_actions(self): print(f">>> Running test for environment: {task_name}") # create a new stage omni.usd.get_context().new_stage() - # parse configuration - env_cfg = parse_env_cfg(task_name, device=self.device, num_envs=self.num_envs) - agent_cfg = load_cfg_from_registry(task_name, "rl_games_cfg_entry_point") # noqa: F841 - # create environment - env = gym.make(task_name, cfg=env_cfg) - # wrap environment - env = RlGamesVecEnvWrapper(env, "cuda:0", 100, 100) + try: + # parse configuration + env_cfg = parse_env_cfg(task_name, device=self.device, num_envs=self.num_envs) + agent_cfg = load_cfg_from_registry(task_name, "rl_games_cfg_entry_point") # noqa: F841 + # create environment + env = gym.make(task_name, cfg=env_cfg) + # convert to single-agent instance if required by the RL algorithm + if isinstance(env.unwrapped, DirectMARLEnv): + env = multi_agent_to_single_agent(env) + # wrap environment + env = RlGamesVecEnvWrapper(env, "cuda:0", 100, 100) + except Exception as e: + if "env" in locals() and hasattr(env, "_is_closed"): + env.close() + else: + if hasattr(e, "obj") and hasattr(e.obj, "_is_closed"): + e.obj.close() + self.fail(f"Failed to set-up the environment for task {task_name}. Error: {e}") # reset environment obs = env.reset() diff --git a/source/extensions/omni.isaac.lab_tasks/test/wrappers/test_rsl_rl_wrapper.py b/source/extensions/omni.isaac.lab_tasks/test/wrappers/test_rsl_rl_wrapper.py index 154171e4cd..c4e7c797f2 100644 --- a/source/extensions/omni.isaac.lab_tasks/test/wrappers/test_rsl_rl_wrapper.py +++ b/source/extensions/omni.isaac.lab_tasks/test/wrappers/test_rsl_rl_wrapper.py @@ -20,6 +20,8 @@ import omni.usd +from omni.isaac.lab.envs import DirectMARLEnv, multi_agent_to_single_agent + import omni.isaac.lab_tasks # noqa: F401 from omni.isaac.lab_tasks.utils.parse_cfg import load_cfg_from_registry, parse_env_cfg from omni.isaac.lab_tasks.utils.wrappers.rsl_rl import RslRlVecEnvWrapper @@ -55,13 +57,24 @@ def test_random_actions(self): print(f">>> Running test for environment: {task_name}") # create a new stage omni.usd.get_context().new_stage() - # parse configuration - env_cfg = parse_env_cfg(task_name, device=self.device, num_envs=self.num_envs) - agent_cfg = load_cfg_from_registry(task_name, "rsl_rl_cfg_entry_point") # noqa: F841 - # create environment - env = gym.make(task_name, cfg=env_cfg) - # wrap environment - env = RslRlVecEnvWrapper(env) + try: + # parse configuration + env_cfg = parse_env_cfg(task_name, device=self.device, num_envs=self.num_envs) + agent_cfg = load_cfg_from_registry(task_name, "rsl_rl_cfg_entry_point") # noqa: F841 + # create environment + env = gym.make(task_name, cfg=env_cfg) + # convert to single-agent instance if required by the RL algorithm + if isinstance(env.unwrapped, DirectMARLEnv): + env = multi_agent_to_single_agent(env) + # wrap environment + env = RslRlVecEnvWrapper(env) + except Exception as e: + if "env" in locals() and hasattr(env, "_is_closed"): + env.close() + else: + if hasattr(e, "obj") and hasattr(e.obj, "_is_closed"): + e.obj.close() + self.fail(f"Failed to set-up the environment for task {task_name}. Error: {e}") # reset environment obs, extras = env.reset() @@ -69,9 +82,9 @@ def test_random_actions(self): self.assertTrue(self._check_valid_tensor(obs)) self.assertTrue(self._check_valid_tensor(extras)) - # simulate environment for 1000 steps + # simulate environment for 100 steps with torch.inference_mode(): - for _ in range(1000): + for _ in range(100): # sample actions from -1 to 1 actions = 2 * torch.rand(env.action_space.shape, device=env.unwrapped.device) - 1 # apply actions diff --git a/source/extensions/omni.isaac.lab_tasks/test/wrappers/test_sb3_wrapper.py b/source/extensions/omni.isaac.lab_tasks/test/wrappers/test_sb3_wrapper.py index afee3a8f04..d0a9aa507c 100644 --- a/source/extensions/omni.isaac.lab_tasks/test/wrappers/test_sb3_wrapper.py +++ b/source/extensions/omni.isaac.lab_tasks/test/wrappers/test_sb3_wrapper.py @@ -21,6 +21,8 @@ import omni.usd +from omni.isaac.lab.envs import DirectMARLEnv, multi_agent_to_single_agent + import omni.isaac.lab_tasks # noqa: F401 from omni.isaac.lab_tasks.utils.parse_cfg import load_cfg_from_registry, parse_env_cfg from omni.isaac.lab_tasks.utils.wrappers.sb3 import Sb3VecEnvWrapper @@ -56,22 +58,33 @@ def test_random_actions(self): print(f">>> Running test for environment: {task_name}") # create a new stage omni.usd.get_context().new_stage() - # parse configuration - env_cfg = parse_env_cfg(task_name, device=self.device, num_envs=self.num_envs) - agent_cfg = load_cfg_from_registry(task_name, "sb3_cfg_entry_point") # noqa: F841 - # create environment - env = gym.make(task_name, cfg=env_cfg) - # wrap environment - env = Sb3VecEnvWrapper(env) + try: + # parse configuration + env_cfg = parse_env_cfg(task_name, device=self.device, num_envs=self.num_envs) + agent_cfg = load_cfg_from_registry(task_name, "sb3_cfg_entry_point") # noqa: F841 + # create environment + env = gym.make(task_name, cfg=env_cfg) + # convert to single-agent instance if required by the RL algorithm + if isinstance(env.unwrapped, DirectMARLEnv): + env = multi_agent_to_single_agent(env) + # wrap environment + env = Sb3VecEnvWrapper(env) + except Exception as e: + if "env" in locals() and hasattr(env, "_is_closed"): + env.close() + else: + if hasattr(e, "obj") and hasattr(e.obj, "_is_closed"): + e.obj.close() + self.fail(f"Failed to set-up the environment for task {task_name}. Error: {e}") # reset environment obs = env.reset() # check signal self.assertTrue(self._check_valid_array(obs)) - # simulate environment for 1000 steps + # simulate environment for 100 steps with torch.inference_mode(): - for _ in range(1000): + for _ in range(100): # sample actions from -1 to 1 actions = 2 * np.random.rand(env.num_envs, *env.action_space.shape) - 1 # apply actions diff --git a/source/extensions/omni.isaac.lab_tasks/test/wrappers/test_skrl_wrapper.py b/source/extensions/omni.isaac.lab_tasks/test/wrappers/test_skrl_wrapper.py index 90ae6eb3aa..58085c1671 100644 --- a/source/extensions/omni.isaac.lab_tasks/test/wrappers/test_skrl_wrapper.py +++ b/source/extensions/omni.isaac.lab_tasks/test/wrappers/test_skrl_wrapper.py @@ -57,25 +57,32 @@ def test_random_actions(self): print(f">>> Running test for environment: {task_name}") # create a new stage omni.usd.get_context().new_stage() - # parse configuration - env_cfg = parse_env_cfg(task_name, device=self.device, num_envs=self.num_envs) - agent_cfg = load_cfg_from_registry(task_name, "skrl_cfg_entry_point") # noqa: F841 - # create environment - env = gym.make(task_name, cfg=env_cfg) - if isinstance(env.unwrapped, DirectMARLEnv): - env = multi_agent_to_single_agent(env) - # wrap environment - env = SkrlVecEnvWrapper(env) - + try: + # parse configuration + env_cfg = parse_env_cfg(task_name, device=self.device, num_envs=self.num_envs) + agent_cfg = load_cfg_from_registry(task_name, "skrl_cfg_entry_point") # noqa: F841 + # create environment + env = gym.make(task_name, cfg=env_cfg) + if isinstance(env.unwrapped, DirectMARLEnv): + env = multi_agent_to_single_agent(env) + # wrap environment + env = SkrlVecEnvWrapper(env) + except Exception as e: + if "env" in locals() and hasattr(env, "_is_closed"): + env.close() + else: + if hasattr(e, "obj") and hasattr(e.obj, "_is_closed"): + e.obj.close() + self.fail(f"Failed to set-up the environment for task {task_name}. Error: {e}") # reset environment obs, extras = env.reset() # check signal self.assertTrue(self._check_valid_tensor(obs)) self.assertTrue(self._check_valid_tensor(extras)) - # simulate environment for 1000 steps + # simulate environment for 100 steps with torch.inference_mode(): - for _ in range(10): + for _ in range(100): # sample actions from -1 to 1 actions = ( 2 * torch.rand(self.num_envs, *env.action_space.shape, device=env.unwrapped.device) - 1 diff --git a/source/standalone/benchmarks/benchmark_load_robot.py b/source/standalone/benchmarks/benchmark_load_robot.py index 3ac0345e37..0704397244 100644 --- a/source/standalone/benchmarks/benchmark_load_robot.py +++ b/source/standalone/benchmarks/benchmark_load_robot.py @@ -115,7 +115,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world root_state = robot.data.default_root_state.clone() root_state[:, :3] += scene.env_origins - robot.write_root_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) # set joint positions with some noise joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() joint_pos += torch.rand_like(joint_pos) * 0.1 diff --git a/source/standalone/demos/arms.py b/source/standalone/demos/arms.py index d50034d4d9..d68cde9f6c 100644 --- a/source/standalone/demos/arms.py +++ b/source/standalone/demos/arms.py @@ -179,7 +179,8 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula # root state root_state = robot.data.default_root_state.clone() root_state[:, :3] += origins[index] - robot.write_root_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) # set joint positions joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() robot.write_joint_state_to_sim(joint_pos, joint_vel) diff --git a/source/standalone/demos/bipeds.py b/source/standalone/demos/bipeds.py index 961af39e19..0ec291be74 100644 --- a/source/standalone/demos/bipeds.py +++ b/source/standalone/demos/bipeds.py @@ -97,7 +97,8 @@ def main(): robot.write_joint_state_to_sim(joint_pos, joint_vel) root_state = robot.data.default_root_state.clone() root_state[:, :3] += origins[index] - robot.write_root_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) robot.reset() # reset command print(">>>>>>>> Reset!") diff --git a/source/standalone/demos/cameras.py b/source/standalone/demos/cameras.py index 6c5583d55a..204d8436df 100644 --- a/source/standalone/demos/cameras.py +++ b/source/standalone/demos/cameras.py @@ -187,7 +187,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world root_state = scene["robot"].data.default_root_state.clone() root_state[:, :3] += scene.env_origins - scene["robot"].write_root_state_to_sim(root_state) + scene["robot"].write_root_link_pose_to_sim(root_state[:, :7]) + scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:]) # set joint positions with some noise joint_pos, joint_vel = ( scene["robot"].data.default_joint_pos.clone(), diff --git a/source/standalone/demos/hands.py b/source/standalone/demos/hands.py index e939d63cf5..9d67c3665e 100644 --- a/source/standalone/demos/hands.py +++ b/source/standalone/demos/hands.py @@ -113,7 +113,8 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula # root state root_state = robot.data.default_root_state.clone() root_state[:, :3] += origins[index] - robot.write_root_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) # joint state joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() robot.write_joint_state_to_sim(joint_pos, joint_vel) diff --git a/source/standalone/demos/multi_asset.py b/source/standalone/demos/multi_asset.py index 6363999949..e640e976d2 100644 --- a/source/standalone/demos/multi_asset.py +++ b/source/standalone/demos/multi_asset.py @@ -23,7 +23,7 @@ # add argparse arguments parser = argparse.ArgumentParser(description="Demo on spawning different objects in multiple environments.") -parser.add_argument("--num_envs", type=int, default=1024, help="Number of environments to spawn.") +parser.add_argument("--num_envs", type=int, default=512, help="Number of environments to spawn.") # append AppLauncher cli args AppLauncher.add_app_launcher_args(parser) # parse the arguments @@ -41,7 +41,15 @@ from pxr import Gf, Sdf import omni.isaac.lab.sim as sim_utils -from omni.isaac.lab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg +from omni.isaac.lab.assets import ( + Articulation, + ArticulationCfg, + AssetBaseCfg, + RigidObject, + RigidObjectCfg, + RigidObjectCollection, + RigidObjectCollectionCfg, +) from omni.isaac.lab.scene import InteractiveScene, InteractiveSceneCfg from omni.isaac.lab.sim import SimulationContext from omni.isaac.lab.utils import Timer, configclass @@ -124,6 +132,52 @@ class MultiObjectSceneCfg(InteractiveSceneCfg): init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 2.0)), ) + # object collection + object_collection: RigidObjectCollectionCfg = RigidObjectCollectionCfg( + rigid_objects={ + "object_A": RigidObjectCfg( + prim_path="/World/envs/env_.*/Object_A", + spawn=sim_utils.SphereCfg( + radius=0.1, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), metallic=0.2), + rigid_props=sim_utils.RigidBodyPropertiesCfg( + solver_position_iteration_count=4, solver_velocity_iteration_count=0 + ), + mass_props=sim_utils.MassPropertiesCfg(mass=1.0), + collision_props=sim_utils.CollisionPropertiesCfg(), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, -0.5, 2.0)), + ), + "object_B": RigidObjectCfg( + prim_path="/World/envs/env_.*/Object_B", + spawn=sim_utils.CuboidCfg( + size=(0.1, 0.1, 0.1), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), metallic=0.2), + rigid_props=sim_utils.RigidBodyPropertiesCfg( + solver_position_iteration_count=4, solver_velocity_iteration_count=0 + ), + mass_props=sim_utils.MassPropertiesCfg(mass=1.0), + collision_props=sim_utils.CollisionPropertiesCfg(), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.5, 2.0)), + ), + "object_C": RigidObjectCfg( + prim_path="/World/envs/env_.*/Object_C", + spawn=sim_utils.ConeCfg( + radius=0.1, + height=0.3, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), metallic=0.2), + rigid_props=sim_utils.RigidBodyPropertiesCfg( + solver_position_iteration_count=4, solver_velocity_iteration_count=0 + ), + mass_props=sim_utils.MassPropertiesCfg(mass=1.0), + collision_props=sim_utils.CollisionPropertiesCfg(), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.5, 0.0, 2.0)), + ), + } + ) + # articulation robot: ArticulationCfg = ArticulationCfg( prim_path="/World/envs/env_.*/Robot", @@ -170,27 +224,35 @@ def run_simulator(sim: SimulationContext, scene: InteractiveScene): """Runs the simulation loop.""" # Extract scene entities # note: we only do this here for readability. - rigid_object = scene["object"] - robot = scene["robot"] + rigid_object: RigidObject = scene["object"] + rigid_object_collection: RigidObjectCollection = scene["object_collection"] + robot: Articulation = scene["robot"] # Define simulation stepping sim_dt = sim.get_physics_dt() count = 0 # Simulation loop while simulation_app.is_running(): # Reset - if count % 500 == 0: + if count % 250 == 0: # reset counter count = 0 # reset the scene entities # object root_state = rigid_object.data.default_root_state.clone() root_state[:, :3] += scene.env_origins - rigid_object.write_root_state_to_sim(root_state) + rigid_object.write_root_link_pose_to_sim(root_state[:, :7]) + rigid_object.write_root_com_velocity_to_sim(root_state[:, 7:]) + # object collection + object_state = rigid_object_collection.data.default_object_state.clone() + object_state[..., :3] += scene.env_origins.unsqueeze(1) + rigid_object_collection.write_object_link_pose_to_sim(object_state[..., :7]) + rigid_object_collection.write_object_com_velocity_to_sim(object_state[..., 7:]) # robot # -- root state root_state = robot.data.default_root_state.clone() root_state[:, :3] += scene.env_origins - robot.write_root_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) # -- joint state joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() robot.write_joint_state_to_sim(joint_pos, joint_vel) diff --git a/source/standalone/demos/quadcopter.py b/source/standalone/demos/quadcopter.py index 2110ae1219..88e00745fc 100644 --- a/source/standalone/demos/quadcopter.py +++ b/source/standalone/demos/quadcopter.py @@ -91,8 +91,8 @@ def main(): # reset dof state joint_pos, joint_vel = robot.data.default_joint_pos, robot.data.default_joint_vel robot.write_joint_state_to_sim(joint_pos, joint_vel) - robot.write_root_pose_to_sim(robot.data.default_root_state[:, :7]) - robot.write_root_velocity_to_sim(robot.data.default_root_state[:, 7:]) + robot.write_root_link_pose_to_sim(robot.data.default_root_state[:, :7]) + robot.write_root_com_velocity_to_sim(robot.data.default_root_state[:, 7:]) robot.reset() # reset command print(">>>>>>>> Reset!") diff --git a/source/standalone/demos/quadrupeds.py b/source/standalone/demos/quadrupeds.py index 3122a4741c..8cabf4e92b 100644 --- a/source/standalone/demos/quadrupeds.py +++ b/source/standalone/demos/quadrupeds.py @@ -142,7 +142,8 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula # root state root_state = robot.data.default_root_state.clone() root_state[:, :3] += origins[index] - robot.write_root_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) # joint state joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() robot.write_joint_state_to_sim(joint_pos, joint_vel) diff --git a/source/standalone/demos/sensors/contact_sensor.py b/source/standalone/demos/sensors/contact_sensor.py new file mode 100644 index 0000000000..1559bdd701 --- /dev/null +++ b/source/standalone/demos/sensors/contact_sensor.py @@ -0,0 +1,177 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Launch Isaac Sim Simulator first.""" + +import argparse + +from omni.isaac.lab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Tutorial on adding sensors on a robot.") +parser.add_argument("--num_envs", type=int, default=1, help="Number of environments to spawn.") +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli = parser.parse_args() + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import torch + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.assets import AssetBaseCfg, RigidObjectCfg +from omni.isaac.lab.scene import InteractiveScene, InteractiveSceneCfg +from omni.isaac.lab.sensors import ContactSensorCfg +from omni.isaac.lab.utils import configclass + +## +# Pre-defined configs +## +from omni.isaac.lab_assets.anymal import ANYMAL_C_CFG # isort: skip + + +@configclass +class ContactSensorSceneCfg(InteractiveSceneCfg): + """Design the scene with sensors on the robot.""" + + # ground plane + ground = AssetBaseCfg(prim_path="/World/defaultGroundPlane", spawn=sim_utils.GroundPlaneCfg()) + + # lights + dome_light = AssetBaseCfg( + prim_path="/World/Light", spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75)) + ) + + # robot + robot = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + # Rigid Object + cube = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube", + spawn=sim_utils.CuboidCfg( + size=(0.5, 0.5, 0.1), + rigid_props=sim_utils.RigidBodyPropertiesCfg(), + mass_props=sim_utils.MassPropertiesCfg(mass=100.0), + collision_props=sim_utils.CollisionPropertiesCfg(), + physics_material=sim_utils.RigidBodyMaterialCfg(static_friction=1.0), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 1.0, 0.0), metallic=0.2), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.5, 0.5, 0.05)), + ) + + contact_forces_LF = ContactSensorCfg( + prim_path="{ENV_REGEX_NS}/Robot/LF_FOOT", + update_period=0.0, + history_length=6, + debug_vis=True, + filter_prim_paths_expr=["{ENV_REGEX_NS}/Cube"], + ) + + contact_forces_RF = ContactSensorCfg( + prim_path="{ENV_REGEX_NS}/Robot/RF_FOOT", + update_period=0.0, + history_length=6, + debug_vis=True, + filter_prim_paths_expr=["{ENV_REGEX_NS}/Cube"], + ) + + contact_forces_H = ContactSensorCfg( + prim_path="{ENV_REGEX_NS}/Robot/.*H_FOOT", + update_period=0.0, + history_length=6, + debug_vis=True, + ) + + +def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): + """Run the simulator.""" + # Define simulation stepping + sim_dt = sim.get_physics_dt() + sim_time = 0.0 + count = 0 + + # Simulate physics + while simulation_app.is_running(): + + if count % 500 == 0: + # reset counter + count = 0 + # reset the scene entities + # root state + # we offset the root state by the origin since the states are written in simulation world frame + # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world + root_state = scene["robot"].data.default_root_state.clone() + root_state[:, :3] += scene.env_origins + scene["robot"].write_root_link_pose_to_sim(root_state[:, :7]) + scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:]) + # set joint positions with some noise + joint_pos, joint_vel = ( + scene["robot"].data.default_joint_pos.clone(), + scene["robot"].data.default_joint_vel.clone(), + ) + joint_pos += torch.rand_like(joint_pos) * 0.1 + scene["robot"].write_joint_state_to_sim(joint_pos, joint_vel) + # clear internal buffers + scene.reset() + print("[INFO]: Resetting robot state...") + # Apply default actions to the robot + # -- generate actions/commands + targets = scene["robot"].data.default_joint_pos + # -- apply action to the robot + scene["robot"].set_joint_position_target(targets) + # -- write data to sim + scene.write_data_to_sim() + # perform step + sim.step() + # update sim-time + sim_time += sim_dt + count += 1 + # update buffers + scene.update(sim_dt) + + # print information from the sensors + print("-------------------------------") + print(scene["contact_forces_LF"]) + print("Received force matrix of: ", scene["contact_forces_LF"].data.force_matrix_w) + print("Received contact force of: ", scene["contact_forces_LF"].data.net_forces_w) + print("-------------------------------") + print(scene["contact_forces_RF"]) + print("Received force matrix of: ", scene["contact_forces_RF"].data.force_matrix_w) + print("Received contact force of: ", scene["contact_forces_RF"].data.net_forces_w) + print("-------------------------------") + print(scene["contact_forces_H"]) + print("Received force matrix of: ", scene["contact_forces_H"].data.force_matrix_w) + print("Received contact force of: ", scene["contact_forces_H"].data.net_forces_w) + + +def main(): + """Main function.""" + + # Initialize the simulation context + sim_cfg = sim_utils.SimulationCfg(dt=0.005, device=args_cli.device) + sim = sim_utils.SimulationContext(sim_cfg) + # Set main camera + sim.set_camera_view(eye=[3.5, 3.5, 3.5], target=[0.0, 0.0, 0.0]) + # design scene + scene_cfg = ContactSensorSceneCfg(num_envs=args_cli.num_envs, env_spacing=2.0) + scene = InteractiveScene(scene_cfg) + # Play the simulator + sim.reset() + # Now we are ready! + print("[INFO]: Setup complete...") + # Run the simulator + run_simulator(sim, scene) + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/source/standalone/demos/sensors/frame_transformer_sensor.py b/source/standalone/demos/sensors/frame_transformer_sensor.py new file mode 100644 index 0000000000..6d009a59d0 --- /dev/null +++ b/source/standalone/demos/sensors/frame_transformer_sensor.py @@ -0,0 +1,171 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import argparse + +from omni.isaac.lab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Tutorial on adding sensors on a robot.") +parser.add_argument("--num_envs", type=int, default=1, help="Number of environments to spawn.") +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli = parser.parse_args() + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import torch + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.assets import AssetBaseCfg, RigidObjectCfg +from omni.isaac.lab.scene import InteractiveScene, InteractiveSceneCfg +from omni.isaac.lab.sensors import FrameTransformerCfg +from omni.isaac.lab.utils import configclass + +## +# Pre-defined configs +## +from omni.isaac.lab_assets.anymal import ANYMAL_C_CFG # isort: skip + + +@configclass +class FrameTransformerSensorSceneCfg(InteractiveSceneCfg): + """Design the scene with sensors on the robot.""" + + # ground plane + ground = AssetBaseCfg(prim_path="/World/defaultGroundPlane", spawn=sim_utils.GroundPlaneCfg()) + + # lights + dome_light = AssetBaseCfg( + prim_path="/World/Light", spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75)) + ) + + # robot + robot = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + # Rigid Object + cube = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube", + spawn=sim_utils.CuboidCfg( + size=(1, 1, 1), + rigid_props=sim_utils.RigidBodyPropertiesCfg(), + mass_props=sim_utils.MassPropertiesCfg(mass=100.0), + collision_props=sim_utils.CollisionPropertiesCfg(), + physics_material=sim_utils.RigidBodyMaterialCfg(static_friction=1.0), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 1.0, 0.0), metallic=0.2), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(5, 0, 0.5)), + ) + + specific_transforms = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/base", + target_frames=[ + FrameTransformerCfg.FrameCfg(prim_path="{ENV_REGEX_NS}/Robot/LF_FOOT"), + FrameTransformerCfg.FrameCfg(prim_path="{ENV_REGEX_NS}/Robot/RF_FOOT"), + ], + debug_vis=True, + ) + + cube_transform = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/base", + target_frames=[FrameTransformerCfg.FrameCfg(prim_path="{ENV_REGEX_NS}/Cube")], + debug_vis=False, + ) + + robot_transforms = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/base", + target_frames=[FrameTransformerCfg.FrameCfg(prim_path="{ENV_REGEX_NS}/Robot/.*")], + debug_vis=False, + ) + + +def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): + """Run the simulator.""" + # Define simulation stepping + sim_dt = sim.get_physics_dt() + sim_time = 0.0 + count = 0 + + # Simulate physics + while simulation_app.is_running(): + + if count % 500 == 0: + # reset counter + count = 0 + # reset the scene entities + # root state + # we offset the root state by the origin since the states are written in simulation world frame + # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world + root_state = scene["robot"].data.default_root_state.clone() + root_state[:, :3] += scene.env_origins + scene["robot"].write_root_link_pose_to_sim(root_state[:, :7]) + scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:]) + # set joint positions with some noise + joint_pos, joint_vel = ( + scene["robot"].data.default_joint_pos.clone(), + scene["robot"].data.default_joint_vel.clone(), + ) + joint_pos += torch.rand_like(joint_pos) * 0.1 + scene["robot"].write_joint_state_to_sim(joint_pos, joint_vel) + # clear internal buffers + scene.reset() + print("[INFO]: Resetting robot state...") + # Apply default actions to the robot + # -- generate actions/commands + targets = scene["robot"].data.default_joint_pos + # -- apply action to the robot + scene["robot"].set_joint_position_target(targets) + # -- write data to sim + scene.write_data_to_sim() + # perform step + sim.step() + # update sim-time + sim_time += sim_dt + count += 1 + # update buffers + scene.update(sim_dt) + + # print information from the sensors + print("-------------------------------") + print(scene["specific_transforms"]) + print("relative transforms:", scene["specific_transforms"].data.target_pos_source) + print("relative orientations:", scene["specific_transforms"].data.target_quat_source) + print("-------------------------------") + print(scene["cube_transform"]) + print("relative transform:", scene["cube_transform"].data.target_pos_source) + print("-------------------------------") + print(scene["robot_transforms"]) + print("relative transforms:", scene["robot_transforms"].data.target_pos_source) + + +def main(): + """Main function.""" + + # Initialize the simulation context + sim_cfg = sim_utils.SimulationCfg(dt=0.005, device=args_cli.device) + sim = sim_utils.SimulationContext(sim_cfg) + # Set main camera + sim.set_camera_view(eye=[3.5, 3.5, 3.5], target=[0.0, 0.0, 0.0]) + # design scene + scene_cfg = FrameTransformerSensorSceneCfg(num_envs=args_cli.num_envs, env_spacing=2.0) + scene = InteractiveScene(scene_cfg) + # Play the simulator + sim.reset() + # Now we are ready! + print("[INFO]: Setup complete...") + # Run the simulator + run_simulator(sim, scene) + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/source/standalone/demos/sensors/raycaster_sensor.py b/source/standalone/demos/sensors/raycaster_sensor.py new file mode 100644 index 0000000000..846ee8a41c --- /dev/null +++ b/source/standalone/demos/sensors/raycaster_sensor.py @@ -0,0 +1,162 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import argparse +import numpy as np + +from omni.isaac.lab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Tutorial on adding sensors on a robot.") +parser.add_argument("--num_envs", type=int, default=1, help="Number of environments to spawn.") +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli = parser.parse_args() + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import torch + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.assets import AssetBaseCfg +from omni.isaac.lab.scene import InteractiveScene, InteractiveSceneCfg +from omni.isaac.lab.sensors.ray_caster import RayCasterCfg, patterns +from omni.isaac.lab.utils import configclass +from omni.isaac.lab.utils.assets import ISAAC_NUCLEUS_DIR + +## +# Pre-defined configs +## +from omni.isaac.lab_assets.anymal import ANYMAL_C_CFG # isort: skip + + +@configclass +class RaycasterSensorSceneCfg(InteractiveSceneCfg): + """Design the scene with sensors on the robot.""" + + # ground plane + ground = AssetBaseCfg( + prim_path="/World/Ground", + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Environments/Terrains/rough_plane.usd", + scale=(1, 1, 1), + ), + ) + + # lights + dome_light = AssetBaseCfg( + prim_path="/World/Light", spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75)) + ) + + # robot + robot = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + ray_caster = RayCasterCfg( + prim_path="{ENV_REGEX_NS}/Robot/base/lidar_cage", + update_period=1 / 60, + offset=RayCasterCfg.OffsetCfg(pos=(0, 0, 0.5)), + mesh_prim_paths=["/World/Ground"], + attach_yaw_only=True, + pattern_cfg=patterns.LidarPatternCfg( + channels=100, vertical_fov_range=[-90, 90], horizontal_fov_range=[-90, 90], horizontal_res=1.0 + ), + debug_vis=not args_cli.headless, + ) + + +def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): + """Run the simulator.""" + # Define simulation stepping + sim_dt = sim.get_physics_dt() + sim_time = 0.0 + count = 0 + + triggered = True + countdown = 42 + + # Simulate physics + while simulation_app.is_running(): + + if count % 500 == 0: + # reset counter + count = 0 + # reset the scene entities + # root state + # we offset the root state by the origin since the states are written in simulation world frame + # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world + root_state = scene["robot"].data.default_root_state.clone() + root_state[:, :3] += scene.env_origins + scene["robot"].write_root_link_pose_to_sim(root_state[:, :7]) + scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:]) + # set joint positions with some noise + joint_pos, joint_vel = ( + scene["robot"].data.default_joint_pos.clone(), + scene["robot"].data.default_joint_vel.clone(), + ) + joint_pos += torch.rand_like(joint_pos) * 0.1 + scene["robot"].write_joint_state_to_sim(joint_pos, joint_vel) + # clear internal buffers + scene.reset() + print("[INFO]: Resetting robot state...") + # Apply default actions to the robot + # -- generate actions/commands + targets = scene["robot"].data.default_joint_pos + # -- apply action to the robot + scene["robot"].set_joint_position_target(targets) + # -- write data to sim + scene.write_data_to_sim() + # perform step + sim.step() + # update sim-time + sim_time += sim_dt + count += 1 + # update buffers + scene.update(sim_dt) + + # print information from the sensors + print("-------------------------------") + print(scene["ray_caster"]) + print("Ray cast hit results: ", scene["ray_caster"].data.ray_hits_w) + + if not triggered: + if countdown > 0: + countdown -= 1 + continue + data = scene["ray_caster"].data.ray_hits_w.cpu().numpy() + np.save("cast_data.npy", data) + triggered = True + else: + continue + + +def main(): + """Main function.""" + + # Initialize the simulation context + sim_cfg = sim_utils.SimulationCfg(dt=0.005, device=args_cli.device) + sim = sim_utils.SimulationContext(sim_cfg) + # Set main camera + sim.set_camera_view(eye=[3.5, 3.5, 3.5], target=[0.0, 0.0, 0.0]) + # design scene + scene_cfg = RaycasterSensorSceneCfg(num_envs=args_cli.num_envs, env_spacing=2.0) + scene = InteractiveScene(scene_cfg) + # Play the simulator + sim.reset() + # Now we are ready! + print("[INFO]: Setup complete...") + # Run the simulator + run_simulator(sim, scene) + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/source/standalone/environments/state_machine/lift_cube_sm.py b/source/standalone/environments/state_machine/lift_cube_sm.py index bd14dcaf0d..503f5d8c0c 100644 --- a/source/standalone/environments/state_machine/lift_cube_sm.py +++ b/source/standalone/environments/state_machine/lift_cube_sm.py @@ -81,6 +81,11 @@ class PickSmWaitTime: LIFT_OBJECT = wp.constant(1.0) +@wp.func +def distance_below_threshold(current_pos: wp.vec3, desired_pos: wp.vec3, threshold: float) -> bool: + return wp.length(current_pos - desired_pos) < threshold + + @wp.kernel def infer_state_machine( dt: wp.array(dtype=float), @@ -92,6 +97,7 @@ def infer_state_machine( des_ee_pose: wp.array(dtype=wp.transform), gripper_state: wp.array(dtype=float), offset: wp.array(dtype=wp.transform), + position_threshold: float, ): # retrieve thread id tid = wp.tid() @@ -109,21 +115,28 @@ def infer_state_machine( elif state == PickSmState.APPROACH_ABOVE_OBJECT: des_ee_pose[tid] = wp.transform_multiply(offset[tid], object_pose[tid]) gripper_state[tid] = GripperState.OPEN - # TODO: error between current and desired ee pose below threshold - # wait for a while - if sm_wait_time[tid] >= PickSmWaitTime.APPROACH_OBJECT: - # move to next state and reset wait time - sm_state[tid] = PickSmState.APPROACH_OBJECT - sm_wait_time[tid] = 0.0 + if distance_below_threshold( + wp.transform_get_translation(ee_pose[tid]), + wp.transform_get_translation(des_ee_pose[tid]), + position_threshold, + ): + # wait for a while + if sm_wait_time[tid] >= PickSmWaitTime.APPROACH_OBJECT: + # move to next state and reset wait time + sm_state[tid] = PickSmState.APPROACH_OBJECT + sm_wait_time[tid] = 0.0 elif state == PickSmState.APPROACH_OBJECT: des_ee_pose[tid] = object_pose[tid] gripper_state[tid] = GripperState.OPEN - # TODO: error between current and desired ee pose below threshold - # wait for a while - if sm_wait_time[tid] >= PickSmWaitTime.APPROACH_OBJECT: - # move to next state and reset wait time - sm_state[tid] = PickSmState.GRASP_OBJECT - sm_wait_time[tid] = 0.0 + if distance_below_threshold( + wp.transform_get_translation(ee_pose[tid]), + wp.transform_get_translation(des_ee_pose[tid]), + position_threshold, + ): + if sm_wait_time[tid] >= PickSmWaitTime.APPROACH_OBJECT: + # move to next state and reset wait time + sm_state[tid] = PickSmState.GRASP_OBJECT + sm_wait_time[tid] = 0.0 elif state == PickSmState.GRASP_OBJECT: des_ee_pose[tid] = object_pose[tid] gripper_state[tid] = GripperState.CLOSE @@ -135,12 +148,16 @@ def infer_state_machine( elif state == PickSmState.LIFT_OBJECT: des_ee_pose[tid] = des_object_pose[tid] gripper_state[tid] = GripperState.CLOSE - # TODO: error between current and desired ee pose below threshold - # wait for a while - if sm_wait_time[tid] >= PickSmWaitTime.LIFT_OBJECT: - # move to next state and reset wait time - sm_state[tid] = PickSmState.LIFT_OBJECT - sm_wait_time[tid] = 0.0 + if distance_below_threshold( + wp.transform_get_translation(ee_pose[tid]), + wp.transform_get_translation(des_ee_pose[tid]), + position_threshold, + ): + # wait for a while + if sm_wait_time[tid] >= PickSmWaitTime.LIFT_OBJECT: + # move to next state and reset wait time + sm_state[tid] = PickSmState.LIFT_OBJECT + sm_wait_time[tid] = 0.0 # increment wait time sm_wait_time[tid] = sm_wait_time[tid] + dt[tid] @@ -160,7 +177,7 @@ class PickAndLiftSm: 5. LIFT_OBJECT: The robot lifts the object to the desired pose. This is the final state. """ - def __init__(self, dt: float, num_envs: int, device: torch.device | str = "cpu"): + def __init__(self, dt: float, num_envs: int, device: torch.device | str = "cpu", position_threshold=0.01): """Initialize the state machine. Args: @@ -172,6 +189,7 @@ def __init__(self, dt: float, num_envs: int, device: torch.device | str = "cpu") self.dt = float(dt) self.num_envs = num_envs self.device = device + self.position_threshold = position_threshold # initialize state machine self.sm_dt = torch.full((self.num_envs,), self.dt, device=self.device) self.sm_state = torch.full((self.num_envs,), 0, dtype=torch.int32, device=self.device) @@ -201,7 +219,7 @@ def reset_idx(self, env_ids: Sequence[int] = None): self.sm_state[env_ids] = 0 self.sm_wait_time[env_ids] = 0.0 - def compute(self, ee_pose: torch.Tensor, object_pose: torch.Tensor, des_object_pose: torch.Tensor): + def compute(self, ee_pose: torch.Tensor, object_pose: torch.Tensor, des_object_pose: torch.Tensor) -> torch.Tensor: """Compute the desired state of the robot's end-effector and the gripper.""" # convert all transformations from (w, x, y, z) to (x, y, z, w) ee_pose = ee_pose[:, [0, 1, 2, 4, 5, 6, 3]] @@ -227,6 +245,7 @@ def compute(self, ee_pose: torch.Tensor, object_pose: torch.Tensor, des_object_p self.des_ee_pose_wp, self.des_gripper_state_wp, self.offset_wp, + self.position_threshold, ], device=self.device, ) @@ -257,7 +276,9 @@ def main(): desired_orientation = torch.zeros((env.unwrapped.num_envs, 4), device=env.unwrapped.device) desired_orientation[:, 1] = 1.0 # create state machine - pick_sm = PickAndLiftSm(env_cfg.sim.dt * env_cfg.decimation, env.unwrapped.num_envs, env.unwrapped.device) + pick_sm = PickAndLiftSm( + env_cfg.sim.dt * env_cfg.decimation, env.unwrapped.num_envs, env.unwrapped.device, position_threshold=0.01 + ) while simulation_app.is_running(): # run everything in inference mode @@ -272,7 +293,7 @@ def main(): tcp_rest_orientation = ee_frame_sensor.data.target_quat_w[..., 0, :].clone() # -- object frame object_data: RigidObjectData = env.unwrapped.scene["object"].data - object_position = object_data.root_pos_w - env.unwrapped.scene.env_origins + object_position = object_data.root_link_pos_w - env.unwrapped.scene.env_origins # -- target object frame desired_position = env.unwrapped.command_manager.get_command("object_pose")[..., :3] diff --git a/source/standalone/environments/state_machine/lift_teddy_bear.py b/source/standalone/environments/state_machine/lift_teddy_bear.py index 896aa614bc..016476066a 100644 --- a/source/standalone/environments/state_machine/lift_teddy_bear.py +++ b/source/standalone/environments/state_machine/lift_teddy_bear.py @@ -80,6 +80,11 @@ class PickSmWaitTime: OPEN_GRIPPER = wp.constant(0.0) +@wp.func +def distance_below_threshold(current_pos: wp.vec3, desired_pos: wp.vec3, threshold: float) -> bool: + return wp.length(current_pos - desired_pos) < threshold + + @wp.kernel def infer_state_machine( dt: wp.array(dtype=float), @@ -91,6 +96,7 @@ def infer_state_machine( des_ee_pose: wp.array(dtype=wp.transform), gripper_state: wp.array(dtype=float), offset: wp.array(dtype=wp.transform), + position_threshold: float, ): # retrieve thread id tid = wp.tid() @@ -108,21 +114,29 @@ def infer_state_machine( elif state == PickSmState.APPROACH_ABOVE_OBJECT: des_ee_pose[tid] = wp.transform_multiply(offset[tid], object_pose[tid]) gripper_state[tid] = GripperState.OPEN - # TODO: error between current and desired ee pose below threshold - # wait for a while - if sm_wait_time[tid] >= PickSmWaitTime.APPROACH_OBJECT: - # move to next state and reset wait time - sm_state[tid] = PickSmState.APPROACH_OBJECT - sm_wait_time[tid] = 0.0 + if distance_below_threshold( + wp.transform_get_translation(ee_pose[tid]), + wp.transform_get_translation(des_ee_pose[tid]), + position_threshold, + ): + # wait for a while + if sm_wait_time[tid] >= PickSmWaitTime.APPROACH_OBJECT: + # move to next state and reset wait time + sm_state[tid] = PickSmState.APPROACH_OBJECT + sm_wait_time[tid] = 0.0 elif state == PickSmState.APPROACH_OBJECT: des_ee_pose[tid] = object_pose[tid] gripper_state[tid] = GripperState.OPEN - # TODO: error between current and desired ee pose below threshold - # wait for a while - if sm_wait_time[tid] >= PickSmWaitTime.APPROACH_OBJECT: - # move to next state and reset wait time - sm_state[tid] = PickSmState.GRASP_OBJECT - sm_wait_time[tid] = 0.0 + if distance_below_threshold( + wp.transform_get_translation(ee_pose[tid]), + wp.transform_get_translation(des_ee_pose[tid]), + position_threshold, + ): + # wait for a while + if sm_wait_time[tid] >= PickSmWaitTime.APPROACH_OBJECT: + # move to next state and reset wait time + sm_state[tid] = PickSmState.GRASP_OBJECT + sm_wait_time[tid] = 0.0 elif state == PickSmState.GRASP_OBJECT: des_ee_pose[tid] = object_pose[tid] gripper_state[tid] = GripperState.CLOSE @@ -134,12 +148,16 @@ def infer_state_machine( elif state == PickSmState.LIFT_OBJECT: des_ee_pose[tid] = des_object_pose[tid] gripper_state[tid] = GripperState.CLOSE - # TODO: error between current and desired ee pose below threshold - # wait for a while - if sm_wait_time[tid] >= PickSmWaitTime.LIFT_OBJECT: - # move to next state and reset wait time - sm_state[tid] = PickSmState.OPEN_GRIPPER - sm_wait_time[tid] = 0.0 + if distance_below_threshold( + wp.transform_get_translation(ee_pose[tid]), + wp.transform_get_translation(des_ee_pose[tid]), + position_threshold, + ): + # wait for a while + if sm_wait_time[tid] >= PickSmWaitTime.LIFT_OBJECT: + # move to next state and reset wait time + sm_state[tid] = PickSmState.OPEN_GRIPPER + sm_wait_time[tid] = 0.0 elif state == PickSmState.OPEN_GRIPPER: # des_ee_pose[tid] = object_pose[tid] gripper_state[tid] = GripperState.OPEN @@ -167,7 +185,7 @@ class PickAndLiftSm: 5. LIFT_OBJECT: The robot lifts the object to the desired pose. This is the final state. """ - def __init__(self, dt: float, num_envs: int, device: torch.device | str = "cpu"): + def __init__(self, dt: float, num_envs: int, device: torch.device | str = "cpu", position_threshold=0.01): """Initialize the state machine. Args: @@ -179,6 +197,7 @@ def __init__(self, dt: float, num_envs: int, device: torch.device | str = "cpu") self.dt = float(dt) self.num_envs = num_envs self.device = device + self.position_threshold = position_threshold # initialize state machine self.sm_dt = torch.full((self.num_envs,), self.dt, device=self.device) self.sm_state = torch.full((self.num_envs,), 0, dtype=torch.int32, device=self.device) @@ -234,6 +253,7 @@ def compute(self, ee_pose: torch.Tensor, object_pose: torch.Tensor, des_object_p self.des_ee_pose_wp, self.des_gripper_state_wp, self.offset_wp, + self.position_threshold, ], device=self.device, ) diff --git a/source/standalone/environments/state_machine/open_cabinet_sm.py b/source/standalone/environments/state_machine/open_cabinet_sm.py index ad40653fca..624defead2 100644 --- a/source/standalone/environments/state_machine/open_cabinet_sm.py +++ b/source/standalone/environments/state_machine/open_cabinet_sm.py @@ -83,6 +83,11 @@ class OpenDrawerSmWaitTime: RELEASE_HANDLE = wp.constant(0.2) +@wp.func +def distance_below_threshold(current_pos: wp.vec3, desired_pos: wp.vec3, threshold: float) -> bool: + return wp.length(current_pos - desired_pos) < threshold + + @wp.kernel def infer_state_machine( dt: wp.array(dtype=float), @@ -95,6 +100,7 @@ def infer_state_machine( handle_approach_offset: wp.array(dtype=wp.transform), handle_grasp_offset: wp.array(dtype=wp.transform), drawer_opening_rate: wp.array(dtype=wp.transform), + position_threshold: float, ): # retrieve thread id tid = wp.tid() @@ -112,21 +118,29 @@ def infer_state_machine( elif state == OpenDrawerSmState.APPROACH_INFRONT_HANDLE: des_ee_pose[tid] = wp.transform_multiply(handle_approach_offset[tid], handle_pose[tid]) gripper_state[tid] = GripperState.OPEN - # TODO: error between current and desired ee pose below threshold - # wait for a while - if sm_wait_time[tid] >= OpenDrawerSmWaitTime.APPROACH_INFRONT_HANDLE: - # move to next state and reset wait time - sm_state[tid] = OpenDrawerSmState.APPROACH_HANDLE - sm_wait_time[tid] = 0.0 + if distance_below_threshold( + wp.transform_get_translation(ee_pose[tid]), + wp.transform_get_translation(des_ee_pose[tid]), + position_threshold, + ): + # wait for a while + if sm_wait_time[tid] >= OpenDrawerSmWaitTime.APPROACH_INFRONT_HANDLE: + # move to next state and reset wait time + sm_state[tid] = OpenDrawerSmState.APPROACH_HANDLE + sm_wait_time[tid] = 0.0 elif state == OpenDrawerSmState.APPROACH_HANDLE: des_ee_pose[tid] = handle_pose[tid] gripper_state[tid] = GripperState.OPEN - # TODO: error between current and desired ee pose below threshold - # wait for a while - if sm_wait_time[tid] >= OpenDrawerSmWaitTime.APPROACH_HANDLE: - # move to next state and reset wait time - sm_state[tid] = OpenDrawerSmState.GRASP_HANDLE - sm_wait_time[tid] = 0.0 + if distance_below_threshold( + wp.transform_get_translation(ee_pose[tid]), + wp.transform_get_translation(des_ee_pose[tid]), + position_threshold, + ): + # wait for a while + if sm_wait_time[tid] >= OpenDrawerSmWaitTime.APPROACH_HANDLE: + # move to next state and reset wait time + sm_state[tid] = OpenDrawerSmState.GRASP_HANDLE + sm_wait_time[tid] = 0.0 elif state == OpenDrawerSmState.GRASP_HANDLE: des_ee_pose[tid] = wp.transform_multiply(handle_grasp_offset[tid], handle_pose[tid]) gripper_state[tid] = GripperState.CLOSE @@ -170,7 +184,7 @@ class OpenDrawerSm: 5. RELEASE_HANDLE: The robot releases the handle of the drawer. This is the final state. """ - def __init__(self, dt: float, num_envs: int, device: torch.device | str = "cpu"): + def __init__(self, dt: float, num_envs: int, device: torch.device | str = "cpu", position_threshold=0.01): """Initialize the state machine. Args: @@ -182,6 +196,7 @@ def __init__(self, dt: float, num_envs: int, device: torch.device | str = "cpu") self.dt = float(dt) self.num_envs = num_envs self.device = device + self.position_threshold = position_threshold # initialize state machine self.sm_dt = torch.full((self.num_envs,), self.dt, device=self.device) self.sm_state = torch.full((self.num_envs,), 0, dtype=torch.int32, device=self.device) @@ -248,6 +263,7 @@ def compute(self, ee_pose: torch.Tensor, handle_pose: torch.Tensor): self.handle_approach_offset_wp, self.handle_grasp_offset_wp, self.drawer_opening_rate_wp, + self.position_threshold, ], device=self.device, ) diff --git a/source/standalone/tools/convert_mesh.py b/source/standalone/tools/convert_mesh.py index 5cebf53d7f..5a88061067 100644 --- a/source/standalone/tools/convert_mesh.py +++ b/source/standalone/tools/convert_mesh.py @@ -106,10 +106,6 @@ def main(): if not os.path.isabs(dest_path): dest_path = os.path.abspath(dest_path) - print(dest_path) - print(os.path.dirname(dest_path)) - print(os.path.basename(dest_path)) - # Mass properties if args_cli.mass is not None: mass_props = schemas_cfg.MassPropertiesCfg(mass=args_cli.mass) diff --git a/source/standalone/tools/record_demos.py b/source/standalone/tools/record_demos.py new file mode 100644 index 0000000000..8c8e9bc208 --- /dev/null +++ b/source/standalone/tools/record_demos.py @@ -0,0 +1,187 @@ +# Copyright (c) 2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +""" +Script to record demonstrations with Isaac Lab environments using human teleoperation. + +This script allows users to record demonstrations operated by human teleoperation for a specified task. +The recorded demonstrations are stored as episodes in a hdf5 file. Users can specify the task, teleoperation +device, dataset directory, and environment stepping rate through command-line arguments. + +required arguments: + --task Name of the task. + +optional arguments: + -h, --help Show this help message and exit + --teleop_device Device for interacting with environment. (default: keyboard) + --dataset_file File path to export recorded demos. (default: "./datasets/dataset.hdf5") + --step_hz Environment stepping rate in Hz. (default: 30) +""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse + +from omni.isaac.lab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Record demonstrations for Isaac Lab environments.") +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +parser.add_argument("--teleop_device", type=str, default="keyboard", help="Device for interacting with environment.") +parser.add_argument( + "--dataset_file", type=str, default="./datasets/dataset.hdf5", help="File path to export recorded demos." +) +parser.add_argument("--step_hz", type=int, default=30, help="Environment stepping rate in Hz.") +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli = parser.parse_args() + +# launch the simulator +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import contextlib +import gymnasium as gym +import os +import time +import torch + +from omni.isaac.lab.devices import Se3Keyboard, Se3SpaceMouse +from omni.isaac.lab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg + +import omni.isaac.lab_tasks # noqa: F401 +from omni.isaac.lab_tasks.utils.parse_cfg import parse_env_cfg + + +class RateLimiter: + """Convenience class for enforcing rates in loops.""" + + def __init__(self, hz): + """ + Args: + hz (int): frequency to enforce + """ + self.hz = hz + self.last_time = time.time() + self.sleep_duration = 1.0 / hz + self.render_period = min(0.033, self.sleep_duration) + + def sleep(self, env): + """Attempt to sleep at the specified rate in hz.""" + next_wakeup_time = self.last_time + self.sleep_duration + while time.time() < next_wakeup_time: + time.sleep(self.render_period) + env.unwrapped.sim.render() + + self.last_time = self.last_time + self.sleep_duration + + # detect time jumping forwards (e.g. loop is too slow) + if self.last_time < time.time(): + while self.last_time < time.time(): + self.last_time += self.sleep_duration + + +def pre_process_actions(delta_pose: torch.Tensor, gripper_command: bool) -> torch.Tensor: + """Pre-process actions for the environment.""" + # compute actions based on environment + if "Reach" in args_cli.task: + # note: reach is the only one that uses a different action space + # compute actions + return delta_pose + else: + # resolve gripper command + gripper_vel = torch.zeros((delta_pose.shape[0], 1), dtype=torch.float, device=delta_pose.device) + gripper_vel[:] = -1 if gripper_command else 1 + # compute actions + return torch.concat([delta_pose, gripper_vel], dim=1) + + +def main(): + """Collect demonstrations from the environment using teleop interfaces.""" + + rate_limiter = RateLimiter(args_cli.step_hz) + + # get directory path and file name (without extension) from cli arguments + output_dir = os.path.dirname(args_cli.dataset_file) + output_file_name = os.path.splitext(os.path.basename(args_cli.dataset_file))[0] + + # create directory if it does not exist + if not os.path.exists(output_dir): + os.makedirs(output_dir) + + # parse configuration + env_cfg = parse_env_cfg(args_cli.task, device=args_cli.device, num_envs=1) + env_cfg.env_name = args_cli.task + + # modify configuration such that the environment runs indefinitely + # until goal is reached + env_cfg.terminations.time_out = None + + env_cfg.observations.policy.concatenate_terms = False + + env_cfg.recorders: ActionStateRecorderManagerCfg = ActionStateRecorderManagerCfg() + env_cfg.recorders.dataset_export_dir_path = output_dir + env_cfg.recorders.dataset_filename = output_file_name + + # create environment + env = gym.make(args_cli.task, cfg=env_cfg) + + # create controller + if args_cli.teleop_device.lower() == "keyboard": + teleop_interface = Se3Keyboard(pos_sensitivity=0.2, rot_sensitivity=0.5) + elif args_cli.teleop_device.lower() == "spacemouse": + teleop_interface = Se3SpaceMouse(pos_sensitivity=0.2, rot_sensitivity=0.5) + else: + raise ValueError(f"Invalid device interface '{args_cli.teleop_device}'. Supported: 'keyboard', 'spacemouse'.") + + # add teleoperation key for reset current recording instance + should_reset_recording_instance = False + + def reset_recording_instance(): + nonlocal should_reset_recording_instance + should_reset_recording_instance = True + + teleop_interface.add_callback("R", reset_recording_instance) + print(teleop_interface) + + # reset before starting + env.reset() + teleop_interface.reset() + + # simulate environment -- run everything in inference mode + with contextlib.suppress(KeyboardInterrupt) and torch.inference_mode(): + while True: + # get keyboard command + delta_pose, gripper_command = teleop_interface.advance() + # convert to torch + delta_pose = torch.tensor(delta_pose, dtype=torch.float, device=env.device).repeat(env.num_envs, 1) + # compute actions based on environment + actions = pre_process_actions(delta_pose, gripper_command) + + # perform action on environment + env.step(actions) + + if should_reset_recording_instance: + env.unwrapped.recorder_manager.reset() + env.reset() + should_reset_recording_instance = False + + # check that simulation is stopped or not + if env.unwrapped.sim.is_stopped(): + break + + rate_limiter.sleep(env.unwrapped) + + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/source/standalone/tools/replay_demos.py b/source/standalone/tools/replay_demos.py new file mode 100644 index 0000000000..1c0a9b948c --- /dev/null +++ b/source/standalone/tools/replay_demos.py @@ -0,0 +1,229 @@ +# Copyright (c) 2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Script to replay demonstrations with Isaac Lab environments.""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse + +from omni.isaac.lab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Replay demonstrations in Isaac Lab environments.") +parser.add_argument("--num_envs", type=int, default=1, help="Number of environments to replay episodes.") +parser.add_argument("--task", type=str, default=None, help="Force to use the specified task.") +parser.add_argument( + "--select_episodes", + type=int, + nargs="+", + default=[], + help="A list of episode indices to be replayed. Keep empty to replay all in the dataset file.", +) +parser.add_argument("--dataset_file", type=str, default="datasets/dataset.hdf5", help="Dataset file to be replayed.") +parser.add_argument( + "--validate_states", + action="store_true", + default=False, + help=( + "Validate if the states, if available, match between loaded from datasets and replayed. Only valid if" + " --num_envs is 1." + ), +) + +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli = parser.parse_args() +# args_cli.headless = True + +# launch the simulator +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import contextlib +import gymnasium as gym +import os +import torch + +from omni.isaac.lab.devices import Se3Keyboard +from omni.isaac.lab.utils.datasets import EpisodeData, HDF5DatasetFileHandler + +import omni.isaac.lab_tasks # noqa: F401 +from omni.isaac.lab_tasks.utils.parse_cfg import parse_env_cfg + +is_paused = False + + +def play_cb(): + global is_paused + is_paused = False + + +def pause_cb(): + global is_paused + is_paused = True + + +def compare_states(state_from_dataset, runtime_state, runtime_env_index) -> (bool, str): + """Compare states from dataset and runtime. + + Args: + state_from_dataset: State from dataset. + runtime_state: State from runtime. + runtime_env_index: Index of the environment in the runtime states to be compared. + + Returns: + bool: True if states match, False otherwise. + str: Log message if states don't match. + """ + states_matched = True + output_log = "" + for asset_type in ["articulation", "rigid_object"]: + for asset_name in runtime_state[asset_type].keys(): + for state_name in runtime_state[asset_type][asset_name].keys(): + runtime_asset_state = runtime_state[asset_type][asset_name][state_name][runtime_env_index] + dataset_asset_state = state_from_dataset[asset_type][asset_name][state_name] + if len(dataset_asset_state) != len(runtime_asset_state): + raise ValueError(f"State shape of {state_name} for asset {asset_name} don't match") + for i in range(len(dataset_asset_state)): + if abs(dataset_asset_state[i] - runtime_asset_state[i]) > 0.01: + states_matched = False + output_log += f'\tState ["{asset_type}"]["{asset_name}"]["{state_name}"][{i}] don\'t match\r\n' + output_log += f"\t Dataset:\t{dataset_asset_state[i]}\r\n" + output_log += f"\t Runtime: \t{runtime_asset_state[i]}\r\n" + return states_matched, output_log + + +def main(): + """Replay episodes loaded from a file.""" + global is_paused + + # Load dataset + if not os.path.exists(args_cli.dataset_file): + raise FileNotFoundError(f"The dataset file {args_cli.dataset_file} does not exist.") + dataset_file_handler = HDF5DatasetFileHandler() + dataset_file_handler.open(args_cli.dataset_file) + env_name = dataset_file_handler.get_env_name() + episode_count = dataset_file_handler.get_num_episodes() + + if episode_count == 0: + print("No episodes found in the dataset.") + exit() + + episode_indices_to_replay = args_cli.select_episodes + if len(episode_indices_to_replay) == 0: + episode_indices_to_replay = list(range(episode_count)) + + if args_cli.task is not None: + env_name = args_cli.task + if env_name is None: + raise ValueError("Task/env name was not specified nor found in the dataset.") + + num_envs = args_cli.num_envs + + env_cfg = parse_env_cfg(env_name, device=args_cli.device, num_envs=num_envs) + + # Disable all recorders and terminations + env_cfg.recorders = {} + env_cfg.terminations = {} + + # create environment from loaded config + env = gym.make(env_name, cfg=env_cfg) + + teleop_interface = Se3Keyboard(pos_sensitivity=0.1, rot_sensitivity=0.1) + teleop_interface.add_callback("N", play_cb) + teleop_interface.add_callback("B", pause_cb) + print(teleop_interface) + + # Determine if state validation should be conducted + state_validation_enabled = False + if args_cli.validate_states and num_envs == 1: + state_validation_enabled = True + elif args_cli.validate_states and num_envs > 1: + print("Warning: State validation is only supported with a single environment. Skipping state validation.") + + # reset before starting + env.reset() + teleop_interface.reset() + + # simulate environment -- run everything in inference mode + episode_names = list(dataset_file_handler.get_episode_names()) + replayed_episode_count = 0 + with contextlib.suppress(KeyboardInterrupt) and torch.inference_mode(): + while simulation_app.is_running() and not simulation_app.is_exiting(): + env_episode_data_map = {index: EpisodeData() for index in range(num_envs)} + first_loop = True + has_next_action = True + while has_next_action: + # initialize actions with zeros so those without next action will not move + actions = torch.zeros(env.unwrapped.action_space.shape) + has_next_action = False + for env_id in range(num_envs): + env_next_action = env_episode_data_map[env_id].get_next_action() + if env_next_action is None: + next_episode_index = None + while episode_indices_to_replay: + next_episode_index = episode_indices_to_replay.pop(0) + if next_episode_index < episode_count: + break + next_episode_index = None + + if next_episode_index is not None: + replayed_episode_count += 1 + print(f"{replayed_episode_count :4}: Loading #{next_episode_index} episode to env_{env_id}") + episode_data = dataset_file_handler.load_episode( + episode_names[next_episode_index], env.unwrapped.device + ) + env_episode_data_map[env_id] = episode_data + # Set initial state for the new episode + initial_state = episode_data.get_initial_state() + env.unwrapped.reset_to( + initial_state, torch.tensor([env_id], device=env.unwrapped.device), is_relative=True + ) + # Get the first action for the new episode + env_next_action = env_episode_data_map[env_id].get_next_action() + has_next_action = True + else: + continue + else: + has_next_action = True + actions[env_id] = env_next_action + if first_loop: + first_loop = False + else: + while is_paused: + env.unwrapped.sim.render() + continue + env.step(actions) + + if state_validation_enabled: + state_from_dataset = env_episode_data_map[0].get_next_state() + if state_from_dataset is not None: + print( + f"Validating states at action-index: {env_episode_data_map[0].next_state_index - 1 :4}", + end="", + ) + current_runtime_state = env.unwrapped.scene.get_state(is_relative=True) + states_matched, comparison_log = compare_states(state_from_dataset, current_runtime_state, 0) + if states_matched: + print("\t- matched.") + else: + print("\t- mismatched.") + print(comparison_log) + break + # Close environment after replay in complete + plural_trailing_s = "s" if replayed_episode_count > 1 else "" + print(f"Finished replaying {replayed_episode_count} episode{plural_trailing_s}.") + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/source/standalone/tutorials/00_sim/log_time.py b/source/standalone/tutorials/00_sim/log_time.py index e55c9f2a62..f47d96222a 100644 --- a/source/standalone/tutorials/00_sim/log_time.py +++ b/source/standalone/tutorials/00_sim/log_time.py @@ -40,12 +40,14 @@ def main(): """Main function.""" # Specify that the logs must be in logs/docker_tutorial - log_dir_path = os.path.join("logs", "docker_tutorial") + log_dir_path = os.path.join("logs") + if not os.path.isdir(log_dir_path): + os.mkdir(log_dir_path) # In the container, the absolute path will be # /workspace/isaaclab/logs/docker_tutorial, because # all python execution is done through /workspace/isaaclab/isaaclab.sh # and the calling process' path will be /workspace/isaaclab - log_dir_path = os.path.abspath(log_dir_path) + log_dir_path = os.path.abspath(os.path.join(log_dir_path, "docker_tutorial")) if not os.path.isdir(log_dir_path): os.mkdir(log_dir_path) print(f"[INFO] Logging experiment to directory: {log_dir_path}") diff --git a/source/standalone/tutorials/01_assets/run_articulation.py b/source/standalone/tutorials/01_assets/run_articulation.py index d06874ac40..bcfa409046 100644 --- a/source/standalone/tutorials/01_assets/run_articulation.py +++ b/source/standalone/tutorials/01_assets/run_articulation.py @@ -94,7 +94,8 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world root_state = robot.data.default_root_state.clone() root_state[:, :3] += origins - robot.write_root_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) # set joint positions with some noise joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() joint_pos += torch.rand_like(joint_pos) * 0.1 diff --git a/source/standalone/tutorials/01_assets/run_deformable_object.py b/source/standalone/tutorials/01_assets/run_deformable_object.py index eccbfab06e..8a6177d1d4 100644 --- a/source/standalone/tutorials/01_assets/run_deformable_object.py +++ b/source/standalone/tutorials/01_assets/run_deformable_object.py @@ -109,7 +109,7 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Deformab # write nodal state to simulation cube_object.write_nodal_state_to_sim(nodal_state) - # write kinematic target to nodal state and free all vertices + # Write the nodal state to the kinematic target and free all vertices nodal_kinematic_target[..., :3] = nodal_state[..., :3] nodal_kinematic_target[..., 3] = 1.0 cube_object.write_nodal_kinematic_target_to_sim(nodal_kinematic_target) diff --git a/source/standalone/tutorials/01_assets/run_rigid_object.py b/source/standalone/tutorials/01_assets/run_rigid_object.py index cede00b2c1..2f5cd327ce 100644 --- a/source/standalone/tutorials/01_assets/run_rigid_object.py +++ b/source/standalone/tutorials/01_assets/run_rigid_object.py @@ -103,7 +103,8 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, RigidObj radius=0.1, h_range=(0.25, 0.5), size=cone_object.num_instances, device=cone_object.device ) # write root state to simulation - cone_object.write_root_state_to_sim(root_state) + cone_object.write_root_link_pose_to_sim(root_state[:, :7]) + cone_object.write_root_com_velocity_to_sim(root_state[:, 7:]) # reset buffers cone_object.reset() print("----------------------------------------") @@ -119,7 +120,7 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, RigidObj cone_object.update(sim_dt) # print the root position if count % 50 == 0: - print(f"Root position (in world): {cone_object.data.root_state_w[:, :3]}") + print(f"Root position (in world): {cone_object.data.root_link_state_w[:, :3]}") def main(): diff --git a/source/standalone/tutorials/02_scene/create_scene.py b/source/standalone/tutorials/02_scene/create_scene.py index 36096d35d5..549d820508 100644 --- a/source/standalone/tutorials/02_scene/create_scene.py +++ b/source/standalone/tutorials/02_scene/create_scene.py @@ -83,7 +83,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world root_state = robot.data.default_root_state.clone() root_state[:, :3] += scene.env_origins - robot.write_root_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) # set joint positions with some noise joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() joint_pos += torch.rand_like(joint_pos) * 0.1 diff --git a/source/standalone/tutorials/03_envs/create_cube_base_env.py b/source/standalone/tutorials/03_envs/create_cube_base_env.py index 92a47ed6cf..ee6643c396 100644 --- a/source/standalone/tutorials/03_envs/create_cube_base_env.py +++ b/source/standalone/tutorials/03_envs/create_cube_base_env.py @@ -120,11 +120,11 @@ def process_actions(self, actions: torch.Tensor): def apply_actions(self): # implement a PD controller to track the target position - pos_error = self._processed_actions - (self._asset.data.root_pos_w - self._env.scene.env_origins) - vel_error = -self._asset.data.root_lin_vel_w + pos_error = self._processed_actions - (self._asset.data.root_link_pos_w - self._env.scene.env_origins) + vel_error = -self._asset.data.root_com_lin_vel_w # set velocity targets self._vel_command[:, :3] = self.p_gain * pos_error + self.d_gain * vel_error - self._asset.write_root_velocity_to_sim(self._vel_command) + self._asset.write_root_com_velocity_to_sim(self._vel_command) @configclass @@ -149,7 +149,7 @@ def base_position(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg) -> torch.Tens """Root linear velocity in the asset's root frame.""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return asset.data.root_pos_w - env.scene.env_origins + return asset.data.root_link_pos_w - env.scene.env_origins ## diff --git a/source/standalone/tutorials/04_sensors/add_sensors_on_robot.py b/source/standalone/tutorials/04_sensors/add_sensors_on_robot.py index 8ba29c2921..8c367a1230 100644 --- a/source/standalone/tutorials/04_sensors/add_sensors_on_robot.py +++ b/source/standalone/tutorials/04_sensors/add_sensors_on_robot.py @@ -113,7 +113,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world root_state = scene["robot"].data.default_root_state.clone() root_state[:, :3] += scene.env_origins - scene["robot"].write_root_state_to_sim(root_state) + scene["robot"].write_root_link_pose_to_sim(root_state[:, :7]) + scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:]) # set joint positions with some noise joint_pos, joint_vel = ( scene["robot"].data.default_joint_pos.clone(), diff --git a/source/standalone/tutorials/04_sensors/run_ray_caster.py b/source/standalone/tutorials/04_sensors/run_ray_caster.py index 921eebad53..c56c2f9ec7 100644 --- a/source/standalone/tutorials/04_sensors/run_ray_caster.py +++ b/source/standalone/tutorials/04_sensors/run_ray_caster.py @@ -109,7 +109,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene_entities: dict): # Reset the scene if step_count % 250 == 0: # reset the balls - balls.write_root_state_to_sim(ball_default_state) + balls.write_root_link_pose_to_sim(ball_default_state[:, :7]) + balls.write_root_com_velocity_to_sim(ball_default_state[:, 7:]) # reset the sensor ray_caster.reset() # reset the counter diff --git a/source/standalone/tutorials/04_sensors/run_ray_caster_camera.py b/source/standalone/tutorials/04_sensors/run_ray_caster_camera.py index a813bd74ea..22958ec3a3 100644 --- a/source/standalone/tutorials/04_sensors/run_ray_caster_camera.py +++ b/source/standalone/tutorials/04_sensors/run_ray_caster_camera.py @@ -129,14 +129,9 @@ def run_simulator(sim: sim_utils.SimulationContext, scene_entities: dict): # Extract camera data camera_index = 0 # note: BasicWriter only supports saving data in numpy format, so we need to convert the data to numpy. - if sim.backend == "torch": - # tensordict allows easy indexing of tensors in the dictionary - single_cam_data = convert_dict_to_backend(camera.data.output[camera_index], backend="numpy") - else: - # for numpy, we need to manually index the data - single_cam_data = dict() - for key, value in camera.data.output.items(): - single_cam_data[key] = value[camera_index] + single_cam_data = convert_dict_to_backend( + {k: v[camera_index] for k, v in camera.data.output.items()}, backend="numpy" + ) # Extract the other information single_cam_info = camera.data.info[camera_index] diff --git a/source/standalone/tutorials/04_sensors/run_usd_camera.py b/source/standalone/tutorials/04_sensors/run_usd_camera.py index 129a76f566..22d0871441 100644 --- a/source/standalone/tutorials/04_sensors/run_usd_camera.py +++ b/source/standalone/tutorials/04_sensors/run_usd_camera.py @@ -228,8 +228,9 @@ def run_simulator(sim: sim_utils.SimulationContext, scene_entities: dict): if args_cli.save: # Save images from camera at camera_index # note: BasicWriter only supports saving data in numpy format, so we need to convert the data to numpy. - # tensordict allows easy indexing of tensors in the dictionary - single_cam_data = convert_dict_to_backend(camera.data.output[camera_index], backend="numpy") + single_cam_data = convert_dict_to_backend( + {k: v[camera_index] for k, v in camera.data.output.items()}, backend="numpy" + ) # Extract the other information single_cam_info = camera.data.info[camera_index] diff --git a/source/standalone/tutorials/05_controllers/run_diff_ik.py b/source/standalone/tutorials/05_controllers/run_diff_ik.py index 51c2102070..68483f2902 100644 --- a/source/standalone/tutorials/05_controllers/run_diff_ik.py +++ b/source/standalone/tutorials/05_controllers/run_diff_ik.py @@ -160,8 +160,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): else: # obtain quantities from simulation jacobian = robot.root_physx_view.get_jacobians()[:, ee_jacobi_idx, :, robot_entity_cfg.joint_ids] - ee_pose_w = robot.data.body_state_w[:, robot_entity_cfg.body_ids[0], 0:7] - root_pose_w = robot.data.root_state_w[:, 0:7] + ee_pose_w = robot.data.body_link_state_w[:, robot_entity_cfg.body_ids[0], 0:7] + root_pose_w = robot.data.root_link_state_w[:, 0:7] joint_pos = robot.data.joint_pos[:, robot_entity_cfg.joint_ids] # compute frame in root frame ee_pos_b, ee_quat_b = subtract_frame_transforms( @@ -181,7 +181,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): scene.update(sim_dt) # obtain quantities from simulation - ee_pose_w = robot.data.body_state_w[:, robot_entity_cfg.body_ids[0], 0:7] + ee_pose_w = robot.data.body_link_state_w[:, robot_entity_cfg.body_ids[0], 0:7] # update marker positions ee_marker.visualize(ee_pose_w[:, 0:3], ee_pose_w[:, 3:7]) goal_marker.visualize(ik_commands[:, 0:3] + scene.env_origins, ik_commands[:, 3:7]) diff --git a/source/standalone/tutorials/05_controllers/run_osc.py b/source/standalone/tutorials/05_controllers/run_osc.py new file mode 100644 index 0000000000..252b93ce01 --- /dev/null +++ b/source/standalone/tutorials/05_controllers/run_osc.py @@ -0,0 +1,442 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +""" +This script demonstrates how to use the operational space controller (OSC) with the simulator. + +The OSC controller can be configured in different modes. It uses the dynamical quantities such as Jacobians and +mass matricescomputed by PhysX. + +.. code-block:: bash + + # Usage + ./isaaclab.sh -p source/standalone/tutorials/05_controllers/run_osc.py + +""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse + +from omni.isaac.lab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Tutorial on using the operational space controller.") +parser.add_argument("--num_envs", type=int, default=128, help="Number of environments to spawn.") +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli = parser.parse_args() + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import torch + +import omni.isaac.lab.sim as sim_utils +from omni.isaac.lab.assets import Articulation, AssetBaseCfg +from omni.isaac.lab.controllers import OperationalSpaceController, OperationalSpaceControllerCfg +from omni.isaac.lab.markers import VisualizationMarkers +from omni.isaac.lab.markers.config import FRAME_MARKER_CFG +from omni.isaac.lab.scene import InteractiveScene, InteractiveSceneCfg +from omni.isaac.lab.sensors import ContactSensorCfg +from omni.isaac.lab.utils import configclass +from omni.isaac.lab.utils.math import ( + combine_frame_transforms, + matrix_from_quat, + quat_inv, + quat_rotate_inverse, + subtract_frame_transforms, +) + +## +# Pre-defined configs +## +from omni.isaac.lab_assets import FRANKA_PANDA_HIGH_PD_CFG # isort:skip + + +@configclass +class SceneCfg(InteractiveSceneCfg): + """Configuration for a simple scene with a tilted wall.""" + + # ground plane + ground = AssetBaseCfg( + prim_path="/World/defaultGroundPlane", + spawn=sim_utils.GroundPlaneCfg(), + ) + + # lights + dome_light = AssetBaseCfg( + prim_path="/World/Light", spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75)) + ) + + # Tilted wall + tilted_wall = AssetBaseCfg( + prim_path="{ENV_REGEX_NS}/TiltedWall", + spawn=sim_utils.CuboidCfg( + size=(2.0, 1.5, 0.01), + collision_props=sim_utils.CollisionPropertiesCfg(), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), opacity=0.1), + rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True), + activate_contact_sensors=True, + ), + init_state=AssetBaseCfg.InitialStateCfg( + pos=(0.6 + 0.085, 0.0, 0.3), rot=(0.9238795325, 0.0, -0.3826834324, 0.0) + ), + ) + + contact_forces = ContactSensorCfg( + prim_path="/World/envs/env_.*/TiltedWall", + update_period=0.0, + history_length=2, + debug_vis=False, + ) + + robot = FRANKA_PANDA_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + robot.actuators["panda_shoulder"].stiffness = 0.0 + robot.actuators["panda_shoulder"].damping = 0.0 + robot.actuators["panda_forearm"].stiffness = 0.0 + robot.actuators["panda_forearm"].damping = 0.0 + robot.spawn.rigid_props.disable_gravity = True + + +def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): + """Runs the simulation loop. + + Args: + sim: (SimulationContext) Simulation context. + scene: (InteractiveScene) Interactive scene. + """ + + # Extract scene entities for readability. + robot = scene["robot"] + contact_forces = scene["contact_forces"] + + # Obtain indices for the end-effector and arm joints + ee_frame_name = "panda_leftfinger" + arm_joint_names = ["panda_joint.*"] + ee_frame_idx = robot.find_bodies(ee_frame_name)[0][0] + arm_joint_ids = robot.find_joints(arm_joint_names)[0] + + # Create the OSC + osc_cfg = OperationalSpaceControllerCfg( + target_types=["pose_abs", "wrench_abs"], + impedance_mode="variable_kp", + inertial_dynamics_decoupling=True, + partial_inertial_dynamics_decoupling=False, + gravity_compensation=False, + motion_damping_ratio_task=1.0, + contact_wrench_stiffness_task=[0.0, 0.0, 0.1, 0.0, 0.0, 0.0], + motion_control_axes_task=[1, 1, 0, 1, 1, 1], + contact_wrench_control_axes_task=[0, 0, 1, 0, 0, 0], + ) + osc = OperationalSpaceController(osc_cfg, num_envs=scene.num_envs, device=sim.device) + + # Markers + frame_marker_cfg = FRAME_MARKER_CFG.copy() + frame_marker_cfg.markers["frame"].scale = (0.1, 0.1, 0.1) + ee_marker = VisualizationMarkers(frame_marker_cfg.replace(prim_path="/Visuals/ee_current")) + goal_marker = VisualizationMarkers(frame_marker_cfg.replace(prim_path="/Visuals/ee_goal")) + + # Define targets for the arm + ee_goal_pose_set_tilted_b = torch.tensor( + [ + [0.6, 0.15, 0.3, 0.0, 0.92387953, 0.0, 0.38268343], + [0.6, -0.3, 0.3, 0.0, 0.92387953, 0.0, 0.38268343], + [0.8, 0.0, 0.5, 0.0, 0.92387953, 0.0, 0.38268343], + ], + device=sim.device, + ) + ee_goal_wrench_set_tilted_task = torch.tensor( + [ + [0.0, 0.0, 10.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 10.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 10.0, 0.0, 0.0, 0.0], + ], + device=sim.device, + ) + kp_set_task = torch.tensor( + [ + [360.0, 360.0, 360.0, 360.0, 360.0, 360.0], + [420.0, 420.0, 420.0, 420.0, 420.0, 420.0], + [320.0, 320.0, 320.0, 320.0, 320.0, 320.0], + ], + device=sim.device, + ) + ee_target_set = torch.cat([ee_goal_pose_set_tilted_b, ee_goal_wrench_set_tilted_task, kp_set_task], dim=-1) + + # Define simulation stepping + sim_dt = sim.get_physics_dt() + + # Update existing buffers + # Note: We need to update buffers before the first step for the controller. + robot.update(dt=sim_dt) + + # get the updated states + jacobian_b, mass_matrix, gravity, ee_pose_b, ee_vel_b, root_pose_w, ee_pose_w, ee_force_b = update_states( + sim, scene, robot, ee_frame_idx, arm_joint_ids, contact_forces + ) + + # Track the given target command + current_goal_idx = 0 # Current goal index for the arm + command = torch.zeros( + scene.num_envs, osc.action_dim, device=sim.device + ) # Generic target command, which can be pose, position, force, etc. + ee_target_pose_b = torch.zeros(scene.num_envs, 7, device=sim.device) # Target pose in the body frame + ee_target_pose_w = torch.zeros(scene.num_envs, 7, device=sim.device) # Target pose in the world frame (for marker) + + # Set joint efforts to zero + zero_joint_efforts = torch.zeros(scene.num_envs, robot.num_joints, device=sim.device) + joint_efforts = torch.zeros(scene.num_envs, len(arm_joint_ids), device=sim.device) + + count = 0 + # Simulation loop + while simulation_app.is_running(): + # reset every 500 steps + if count % 500 == 0: + # reset joint state to default + default_joint_pos = robot.data.default_joint_pos.clone() + default_joint_vel = robot.data.default_joint_vel.clone() + robot.write_joint_state_to_sim(default_joint_pos, default_joint_vel) + robot.set_joint_effort_target(zero_joint_efforts) # Set zero torques in the initial step + robot.write_data_to_sim() + robot.reset() + # reset contact sensor + contact_forces.reset() + # reset target pose + robot.update(sim_dt) + _, _, _, ee_pose_b, _, _, _, _ = update_states( + sim, scene, robot, ee_frame_idx, arm_joint_ids, contact_forces + ) # at reset, the jacobians are not updated to the latest state + command, ee_target_pose_b, ee_target_pose_w, current_goal_idx = update_target( + sim, scene, osc, root_pose_w, ee_target_set, current_goal_idx + ) + # set the osc command + osc.reset() + command, task_frame_pose_b = convert_to_task_frame(osc, command=command, ee_target_pose_b=ee_target_pose_b) + osc.set_command(command=command, current_ee_pose_b=ee_pose_b, current_task_frame_pose_b=task_frame_pose_b) + else: + # get the updated states + jacobian_b, mass_matrix, gravity, ee_pose_b, ee_vel_b, root_pose_w, ee_pose_w, ee_force_b = update_states( + sim, scene, robot, ee_frame_idx, arm_joint_ids, contact_forces + ) + # compute the joint commands + joint_efforts = osc.compute( + jacobian_b=jacobian_b, + current_ee_pose_b=ee_pose_b, + current_ee_vel_b=ee_vel_b, + current_ee_force_b=ee_force_b, + mass_matrix=mass_matrix, + gravity=gravity, + ) + # apply actions + robot.set_joint_effort_target(joint_efforts, joint_ids=arm_joint_ids) + robot.write_data_to_sim() + + # update marker positions + ee_marker.visualize(ee_pose_w[:, 0:3], ee_pose_w[:, 3:7]) + goal_marker.visualize(ee_target_pose_w[:, 0:3], ee_target_pose_w[:, 3:7]) + + # perform step + sim.step(render=True) + # update robot buffers + robot.update(sim_dt) + # update buffers + scene.update(sim_dt) + # update sim-time + count += 1 + + +# Update robot states +def update_states( + sim: sim_utils.SimulationContext, + scene: InteractiveScene, + robot: Articulation, + ee_frame_idx: int, + arm_joint_ids: list[int], + contact_forces, +): + """Update the robot states. + + Args: + sim: (SimulationContext) Simulation context. + scene: (InteractiveScene) Interactive scene. + robot: (Articulation) Robot articulation. + ee_frame_idx: (int) End-effector frame index. + arm_joint_ids: (list[int]) Arm joint indices. + contact_forces: (ContactSensor) Contact sensor. + + Returns: + jacobian_b (torch.tensor): Jacobian in the body frame. + mass_matrix (torch.tensor): Mass matrix. + gravity (torch.tensor): Gravity vector. + ee_pose_b (torch.tensor): End-effector pose in the body frame. + ee_vel_b (torch.tensor): End-effector velocity in the body frame. + root_pose_w (torch.tensor): Root pose in the world frame. + ee_pose_w (torch.tensor): End-effector pose in the world frame. + ee_force_b (torch.tensor): End-effector force in the body frame. + + Raises: + ValueError: Undefined target_type. + """ + # obtain dynamics related quantities from simulation + ee_jacobi_idx = ee_frame_idx - 1 + jacobian_w = robot.root_physx_view.get_jacobians()[:, ee_jacobi_idx, :, arm_joint_ids] + mass_matrix = robot.root_physx_view.get_mass_matrices()[:, arm_joint_ids, :][:, :, arm_joint_ids] + gravity = robot.root_physx_view.get_generalized_gravity_forces()[:, arm_joint_ids] + # Convert the Jacobian from world to root frame + jacobian_b = jacobian_w.clone() + root_rot_matrix = matrix_from_quat(quat_inv(robot.data.root_link_quat_w)) + jacobian_b[:, :3, :] = torch.bmm(root_rot_matrix, jacobian_b[:, :3, :]) + jacobian_b[:, 3:, :] = torch.bmm(root_rot_matrix, jacobian_b[:, 3:, :]) + + # Compute current pose of the end-effector + root_pos_w = robot.data.root_link_pos_w + root_quat_w = robot.data.root_link_quat_w + ee_pos_w = robot.data.body_link_pos_w[:, ee_frame_idx] + ee_quat_w = robot.data.body_link_quat_w[:, ee_frame_idx] + ee_pos_b, ee_quat_b = subtract_frame_transforms(root_pos_w, root_quat_w, ee_pos_w, ee_quat_w) + root_pose_w = torch.cat([root_pos_w, root_quat_w], dim=-1) + ee_pose_w = torch.cat([ee_pos_w, ee_quat_w], dim=-1) + ee_pose_b = torch.cat([ee_pos_b, ee_quat_b], dim=-1) + + # Compute the current velocity of the end-effector + ee_vel_w = robot.data.body_com_vel_w[:, ee_frame_idx, :] # Extract end-effector velocity in the world frame + root_vel_w = robot.data.root_com_vel_w # Extract root velocity in the world frame + relative_vel_w = ee_vel_w - root_vel_w # Compute the relative velocity in the world frame + ee_lin_vel_b = quat_rotate_inverse(robot.data.root_link_quat_w, relative_vel_w[:, 0:3]) # From world to root frame + ee_ang_vel_b = quat_rotate_inverse(robot.data.root_link_quat_w, relative_vel_w[:, 3:6]) + ee_vel_b = torch.cat([ee_lin_vel_b, ee_ang_vel_b], dim=-1) + + # Calculate the contact force + ee_force_w = torch.zeros(scene.num_envs, 3, device=sim.device) + sim_dt = sim.get_physics_dt() + contact_forces.update(sim_dt) # update contact sensor + # Calculate the contact force by averaging over last four time steps (i.e., to smoothen) and + # taking the max of three surfaces as only one should be the contact of interest + ee_force_w, _ = torch.max(torch.mean(contact_forces.data.net_forces_w_history, dim=1), dim=1) + + # This is a simplification, only for the sake of testing. + ee_force_b = ee_force_w + + return jacobian_b, mass_matrix, gravity, ee_pose_b, ee_vel_b, root_pose_w, ee_pose_w, ee_force_b + + +# Update the target commands +def update_target( + sim: sim_utils.SimulationContext, + scene: InteractiveScene, + osc: OperationalSpaceController, + root_pose_w: torch.tensor, + ee_target_set: torch.tensor, + current_goal_idx: int, +): + """Update the targets for the operational space controller. + + Args: + sim: (SimulationContext) Simulation context. + scene: (InteractiveScene) Interactive scene. + osc: (OperationalSpaceController) Operational space controller. + root_pose_w: (torch.tensor) Root pose in the world frame. + ee_target_set: (torch.tensor) End-effector target set. + current_goal_idx: (int) Current goal index. + + Returns: + command (torch.tensor): Updated target command. + ee_target_pose_b (torch.tensor): Updated target pose in the body frame. + ee_target_pose_w (torch.tensor): Updated target pose in the world frame. + next_goal_idx (int): Next goal index. + + Raises: + ValueError: Undefined target_type. + """ + + # update the ee desired command + command = torch.zeros(scene.num_envs, osc.action_dim, device=sim.device) + command[:] = ee_target_set[current_goal_idx] + + # update the ee desired pose + ee_target_pose_b = torch.zeros(scene.num_envs, 7, device=sim.device) + for target_type in osc.cfg.target_types: + if target_type == "pose_abs": + ee_target_pose_b[:] = command[:, :7] + elif target_type == "wrench_abs": + pass # ee_target_pose_b could stay at the root frame for force control, what matters is ee_target_b + else: + raise ValueError("Undefined target_type within update_target().") + + # update the target desired pose in world frame (for marker) + ee_target_pos_w, ee_target_quat_w = combine_frame_transforms( + root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_target_pose_b[:, 0:3], ee_target_pose_b[:, 3:7] + ) + ee_target_pose_w = torch.cat([ee_target_pos_w, ee_target_quat_w], dim=-1) + + next_goal_idx = (current_goal_idx + 1) % len(ee_target_set) + + return command, ee_target_pose_b, ee_target_pose_w, next_goal_idx + + +# Convert the target commands to the task frame +def convert_to_task_frame(osc: OperationalSpaceController, command: torch.tensor, ee_target_pose_b: torch.tensor): + """Converts the target commands to the task frame. + + Args: + osc: OperationalSpaceController object. + command: Command to be converted. + ee_target_pose_b: Target pose in the body frame. + + Returns: + command (torch.tensor): Target command in the task frame. + task_frame_pose_b (torch.tensor): Target pose in the task frame. + + Raises: + ValueError: Undefined target_type. + """ + command = command.clone() + task_frame_pose_b = ee_target_pose_b.clone() + + cmd_idx = 0 + for target_type in osc.cfg.target_types: + if target_type == "pose_abs": + command[:, :3], command[:, 3:7] = subtract_frame_transforms( + task_frame_pose_b[:, :3], task_frame_pose_b[:, 3:], command[:, :3], command[:, 3:7] + ) + cmd_idx += 7 + elif target_type == "wrench_abs": + # These are already defined in target frame for ee_goal_wrench_set_tilted_task (since it is + # easier), so not transforming + cmd_idx += 6 + else: + raise ValueError("Undefined target_type within _convert_to_task_frame().") + + return command, task_frame_pose_b + + +def main(): + """Main function.""" + # Load kit helper + sim_cfg = sim_utils.SimulationCfg(dt=0.01, device=args_cli.device) + sim = sim_utils.SimulationContext(sim_cfg) + # Set main camera + sim.set_camera_view([2.5, 2.5, 2.5], [0.0, 0.0, 0.0]) + # Design scene + scene_cfg = SceneCfg(num_envs=args_cli.num_envs, env_spacing=2.0) + scene = InteractiveScene(scene_cfg) + # Play the simulator + sim.reset() + # Now we are ready! + print("[INFO]: Setup complete...") + # Run the simulator + run_simulator(sim, scene) + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/source/standalone/workflows/ray/cluster_configs/Dockerfile b/source/standalone/workflows/ray/cluster_configs/Dockerfile new file mode 100644 index 0000000000..48dfe56f7f --- /dev/null +++ b/source/standalone/workflows/ray/cluster_configs/Dockerfile @@ -0,0 +1,19 @@ +FROM isaac-lab-base:latest + +# Set NVIDIA paths +ENV PATH="/usr/local/nvidia/bin:$PATH" +ENV LD_LIBRARY_PATH="/usr/local/nvidia/lib64" + +# Link NVIDIA binaries +RUN ln -sf /usr/local/nvidia/bin/nvidia* /usr/bin + +# Install Ray and configure it +RUN /workspace/isaaclab/_isaac_sim/python.sh -m pip install "ray[default, tune]"==2.31.0 && \ +sed -i "1i $(echo "#!/workspace/isaaclab/_isaac_sim/python.sh")" \ +/isaac-sim/kit/python/bin/ray && ln -s /isaac-sim/kit/python/bin/ray /usr/local/bin/ray + +# Install tuning dependencies +RUN /workspace/isaaclab/_isaac_sim/python.sh -m pip install optuna bayesian-optimization + +# Install MLflow for logging +RUN /workspace/isaaclab/_isaac_sim/python.sh -m pip install mlflow diff --git a/source/standalone/workflows/ray/cluster_configs/google_cloud/kuberay.yaml.jinja b/source/standalone/workflows/ray/cluster_configs/google_cloud/kuberay.yaml.jinja new file mode 100644 index 0000000000..2d43075445 --- /dev/null +++ b/source/standalone/workflows/ray/cluster_configs/google_cloud/kuberay.yaml.jinja @@ -0,0 +1,204 @@ +# Jinja is used for templating here as full helm setup is excessive for application +apiVersion: ray.io/v1alpha1 +kind: RayCluster +metadata: + name: {{ name }} + namespace: {{ namespace }} +spec: + rayVersion: "2.8.0" + enableInTreeAutoscaling: true + autoscalerOptions: + upscalingMode: Default + idleTimeoutSeconds: 120 + imagePullPolicy: Always + securityContext: {} + envFrom: [] + + headGroupSpec: + rayStartParams: + block: "true" + dashboard-host: 0.0.0.0 + dashboard-port: "8265" + node-ip-address: "0.0.0.0" + port: "6379" + include-dashboard: "true" + ray-debugger-external: "true" + object-manager-port: "8076" + num-gpus: "0" + num-cpus: "0" # prevent scheduling jobs to the head node - workers only + headService: + apiVersion: v1 + kind: Service + metadata: + name: head + spec: + type: LoadBalancer + template: + metadata: + labels: + app.kubernetes.io/instance: tuner + app.kubernetes.io/name: kuberay + cloud.google.com/gke-ray-node-type: head + spec: + serviceAccountName: {{ service_account_name }} + affinity: {} + securityContext: + fsGroup: 100 + containers: + - env: + image: {{ image }} + imagePullPolicy: Always + name: head + resources: + limits: + cpu: "{{ num_head_cpu }}" + memory: {{ head_ram_gb }}G + nvidia.com/gpu: "0" + requests: + cpu: "{{ num_head_cpu }}" + memory: {{ head_ram_gb }}G + nvidia.com/gpu: "0" + securityContext: {} + volumeMounts: + - mountPath: /tmp/ray + name: ray-logs + command: ["/bin/bash", "-c", "ray start --head --port=6379 --object-manager-port=8076 --dashboard-host=0.0.0.0 --dashboard-port=8265 --include-dashboard=true && tail -f /dev/null"] + - image: fluent/fluent-bit:1.9.6 + name: fluentbit + resources: + limits: + cpu: 100m + memory: 128Mi + requests: + cpu: 100m + memory: 128Mi + volumeMounts: + - mountPath: /tmp/ray + name: ray-logs + imagePullSecrets: [] + nodeSelector: + iam.gke.io/gke-metadata-server-enabled: "true" + volumes: + - configMap: + name: fluentbit-config + name: fluentbit-config + - name: ray-logs + emptyDir: {} + + workerGroupSpecs: + {% for it in range(gpu_per_worker|length) %} + - groupName: "{{ worker_accelerator[it] }}x{{ gpu_per_worker[it] }}-cpu-{{ cpu_per_worker[it] }}-ram-gb-{{ ram_gb_per_worker[it] }}" + replicas: {{ num_workers[it] }} + maxReplicas: {{ num_workers[it] }} + minReplicas: {{ num_workers[it] }} + rayStartParams: + block: "true" + ray-debugger-external: "true" + replicas: "{{num_workers[it]}}" + template: + metadata: + annotations: {} + labels: + app.kubernetes.io/instance: tuner + app.kubernetes.io/name: kuberay + cloud.google.com/gke-ray-node-type: worker + spec: + serviceAccountName: {{ service_account_name }} + affinity: {} + securityContext: + fsGroup: 100 + containers: + - env: + - name: NVIDIA_VISIBLE_DEVICES + value: "all" + - name: NVIDIA_DRIVER_CAPABILITIES + value: "compute,utility" + + image: {{ image }} + imagePullPolicy: Always + name: ray-worker + resources: + limits: + cpu: "{{ cpu_per_worker[it] }}" + memory: {{ ram_gb_per_worker[it] }}G + nvidia.com/gpu: "{{ gpu_per_worker[it] }}" + requests: + cpu: "{{ cpu_per_worker[it] }}" + memory: {{ ram_gb_per_worker[it] }}G + nvidia.com/gpu: "{{ gpu_per_worker[it] }}" + securityContext: {} + volumeMounts: + - mountPath: /tmp/ray + name: ray-logs + command: ["/bin/bash", "-c", "ray start --address=head.{{ namespace }}.svc.cluster.local:6379 && tail -f /dev/null"] + - image: fluent/fluent-bit:1.9.6 + name: fluentbit + resources: + limits: + cpu: 100m + memory: 128Mi + requests: + cpu: 100m + memory: 128Mi + volumeMounts: + - mountPath: /tmp/ray + name: ray-logs + + imagePullSecrets: [] + nodeSelector: + cloud.google.com/gke-accelerator: {{ worker_accelerator[it] }} + iam.gke.io/gke-metadata-server-enabled: "true" + tolerations: + - key: "nvidia.com/gpu" + operator: "Exists" + effect: "NoSchedule" + volumes: + - configMap: + name: fluentbit-config + name: fluentbit-config + - name: ray-logs + emptyDir: {} + {% endfor %} + +--- +# ML Flow Server - for fetching logs +apiVersion: apps/v1 +kind: Deployment +metadata: + name: {{name}}-mlflow + namespace: {{ namespace }} +spec: + replicas: 1 + selector: + matchLabels: + app: mlflow + template: + metadata: + labels: + app: mlflow + spec: + containers: + - name: mlflow + image: ghcr.io/mlflow/mlflow:v2.9.2 + ports: + - containerPort: 5000 + command: ["mlflow"] + args: + - server + - --host=0.0.0.0 + - --port=5000 + - --backend-store-uri=sqlite:///mlflow.db +--- +# ML Flow Service (for port forwarding, kubectl port-forward service/{name}-mlflow 5000:5000) +apiVersion: v1 +kind: Service +metadata: + name: {{name}}-mlflow + namespace: {{ namespace }} +spec: + selector: + app: mlflow + ports: + - port: 5000 + targetPort: 5000 + type: ClusterIP diff --git a/source/standalone/workflows/ray/grok_cluster_with_kubectl.py b/source/standalone/workflows/ray/grok_cluster_with_kubectl.py new file mode 100644 index 0000000000..a7e9b3debf --- /dev/null +++ b/source/standalone/workflows/ray/grok_cluster_with_kubectl.py @@ -0,0 +1,274 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import argparse +import os +import re +import subprocess +import threading +import time +from concurrent.futures import ThreadPoolExecutor, as_completed + +""" +This script requires that kubectl is installed and KubeRay was used to create the cluster. + +Creates a config file containing ``name: address: http://:`` on +a new line for each cluster, and also fetches the MLFlow URI. + +Usage: + +.. code-block:: bash + + ./isaaclab.sh -p source/standalone/workflows/ray/grok_cluster_with_kubectl.py + # For options, supply -h arg +""" + + +def get_namespace() -> str: + """Get the current Kubernetes namespace from the context, fallback to default if not set""" + try: + namespace = ( + subprocess.check_output(["kubectl", "config", "view", "--minify", "--output", "jsonpath={..namespace}"]) + .decode() + .strip() + ) + if not namespace: + namespace = "default" + except subprocess.CalledProcessError: + namespace = "default" + return namespace + + +def get_pods(namespace: str = "default") -> list[tuple]: + """Get a list of all of the pods in the namespace""" + cmd = ["kubectl", "get", "pods", "-n", namespace, "--no-headers"] + output = subprocess.check_output(cmd).decode() + pods = [] + for line in output.strip().split("\n"): + fields = line.split() + pod_name = fields[0] + status = fields[2] + pods.append((pod_name, status)) + return pods + + +def get_clusters(pods: list, cluster_name_prefix: str) -> set: + """ + Get unique cluster name(s). Works for one or more clusters, based off of the number of head nodes. + Excludes MLflow deployments. + """ + clusters = set() + for pod_name, _ in pods: + # Skip MLflow pods + if "-mlflow" in pod_name: + continue + + match = re.match(r"(" + re.escape(cluster_name_prefix) + r"[-\w]+)", pod_name) + if match: + # Get base name without head/worker suffix + base_name = match.group(1).split("-head")[0].split("-worker")[0] + clusters.add(base_name) + return sorted(clusters) + + +def get_mlflow_info(namespace: str = None, cluster_prefix: str = "isaacray") -> str: + """ + Get MLflow service information if it exists in the namespace with the given prefix. + Only works for a single cluster instance. + Args: + namespace: Kubernetes namespace + cluster_prefix: Base cluster name (without -head/-worker suffixes) + Returns: + MLflow service URL + """ + # Strip any -head or -worker suffixes to get base name + if namespace is None: + namespace = get_namespace() + pods = get_pods(namespace=namespace) + clusters = get_clusters(pods=pods, cluster_name_prefix=cluster_prefix) + if len(clusters) > 1: + raise ValueError("More than one cluster matches prefix, could not automatically determine mlflow info.") + + base_name = cluster_prefix.split("-head")[0].split("-worker")[0] + mlflow_name = f"{base_name}-mlflow" + + cmd = ["kubectl", "get", "svc", mlflow_name, "-n", namespace, "--no-headers"] + try: + output = subprocess.check_output(cmd).decode() + fields = output.strip().split() + + # Get cluster IP + cluster_ip = fields[2] + port = "5000" # Default MLflow port + + return f"http://{cluster_ip}:{port}" + except subprocess.CalledProcessError as e: + raise ValueError(f"Could not grok MLflow: {e}") # Fixed f-string + + +def check_clusters_running(pods: list, clusters: set) -> bool: + """ + Check that all of the pods in all provided clusters are running. + + Args: + pods (list): A list of tuples where each tuple contains the pod name and its status. + clusters (set): A set of cluster names to check. + + Returns: + bool: True if all pods in any of the clusters are running, False otherwise. + """ + clusters_running = False + for cluster in clusters: + cluster_pods = [p for p in pods if p[0].startswith(cluster)] + total_pods = len(cluster_pods) + running_pods = len([p for p in cluster_pods if p[1] == "Running"]) + if running_pods == total_pods and running_pods > 0: + clusters_running = True + break + return clusters_running + + +def get_ray_address(head_pod: str, namespace: str = "default", ray_head_name: str = "head") -> str: + """ + Given a cluster head pod, check its logs, which should include the ray address which can accept job requests. + + Args: + head_pod (str): The name of the head pod. + namespace (str, optional): The Kubernetes namespace. Defaults to "default". + ray_head_name (str, optional): The name of the ray head container. Defaults to "head". + + Returns: + str: The ray address if found, None otherwise. + + Raises: + ValueError: If the logs cannot be retrieved or the ray address is not found. + """ + cmd = ["kubectl", "logs", head_pod, "-c", ray_head_name, "-n", namespace] + try: + output = subprocess.check_output(cmd).decode() + except subprocess.CalledProcessError as e: + raise ValueError( + f"Could not enter head container with cmd {cmd}: {e}Perhaps try a different namespace or ray head name." + ) + match = re.search(r"RAY_ADDRESS='([^']+)'", output) + if match: + return match.group(1) + else: + return None + + +def process_cluster(cluster_info: dict, ray_head_name: str = "head") -> str: + """ + For each cluster, check that it is running, and get the Ray head address that will accept jobs. + + Args: + cluster_info (dict): A dictionary containing cluster information with keys 'cluster', 'pods', and 'namespace'. + ray_head_name (str, optional): The name of the ray head container. Defaults to "head". + + Returns: + str: A string containing the cluster name and its Ray head address, or an error message if the head pod or Ray address is not found. + """ + cluster, pods, namespace = cluster_info + head_pod = None + for pod_name, status in pods: + if pod_name.startswith(cluster + "-head"): + head_pod = pod_name + break + if not head_pod: + return f"Error: Could not find head pod for cluster {cluster}\n" + + # Get RAY_ADDRESS and status + ray_address = get_ray_address(head_pod, namespace=namespace, ray_head_name=ray_head_name) + if not ray_address: + return f"Error: Could not find RAY_ADDRESS for cluster {cluster}\n" + + # Return only cluster and ray address + return f"name: {cluster} address: {ray_address}\n" + + +def main(): + # Parse command-line arguments + parser = argparse.ArgumentParser(description="Process Ray clusters and save their specifications.") + parser.add_argument("--prefix", default="isaacray", help="The prefix for the cluster names.") + parser.add_argument("--output", default="~/.cluster_config", help="The file to save cluster specifications.") + parser.add_argument("--ray_head_name", default="head", help="The metadata name for the ray head container") + parser.add_argument( + "--namespace", help="Kubernetes namespace to use. If not provided, will detect from current context." + ) + args = parser.parse_args() + + # Get namespace from args or detect it + current_namespace = args.namespace if args.namespace else get_namespace() + print(f"Using namespace: {current_namespace}") + + cluster_name_prefix = args.prefix + cluster_spec_file = os.path.expanduser(args.output) + + # Get all pods + pods = get_pods(namespace=current_namespace) + + # Get clusters + clusters = get_clusters(pods, cluster_name_prefix) + if not clusters: + print(f"No clusters found with prefix {cluster_name_prefix}") + return + + # Wait for clusters to be running + while True: + pods = get_pods(namespace=current_namespace) + if check_clusters_running(pods, clusters): + break + print("Waiting for all clusters to spin up...") + time.sleep(5) + + print("Checking for MLflow:") + # Check MLflow status for each cluster + for cluster in clusters: + try: + mlflow_address = get_mlflow_info(current_namespace, cluster) + print(f"MLflow address for {cluster}: {mlflow_address}") + except ValueError as e: + print(f"ML Flow not located: {e}") + print() + + # Prepare cluster info for parallel processing + cluster_infos = [] + for cluster in clusters: + cluster_pods = [p for p in pods if p[0].startswith(cluster)] + cluster_infos.append((cluster, cluster_pods, current_namespace)) + + # Use ThreadPoolExecutor to process clusters in parallel + results = [] + results_lock = threading.Lock() + + with ThreadPoolExecutor() as executor: + future_to_cluster = { + executor.submit(process_cluster, info, args.ray_head_name): info[0] for info in cluster_infos + } + for future in as_completed(future_to_cluster): + cluster_name = future_to_cluster[future] + try: + result = future.result() + with results_lock: + results.append(result) + except Exception as exc: + print(f"{cluster_name} generated an exception: {exc}") + + # Sort results alphabetically by cluster name + results.sort() + + # Write sorted results to the output file (Ray info only) + with open(cluster_spec_file, "w") as f: + for result in results: + f.write(result) + + print(f"Cluster spec information saved to {cluster_spec_file}") + # Display the contents of the config file + with open(cluster_spec_file) as f: + print(f.read()) + + +if __name__ == "__main__": + main() diff --git a/source/standalone/workflows/ray/hyperparameter_tuning/vision_cartpole_cfg.py b/source/standalone/workflows/ray/hyperparameter_tuning/vision_cartpole_cfg.py new file mode 100644 index 0000000000..d0abc75f0c --- /dev/null +++ b/source/standalone/workflows/ray/hyperparameter_tuning/vision_cartpole_cfg.py @@ -0,0 +1,49 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +import pathlib +import sys + +# Allow for import of items from the ray workflow. +CUR_DIR = pathlib.Path(__file__).parent +UTIL_DIR = CUR_DIR.parent +sys.path.extend([str(UTIL_DIR), str(CUR_DIR)]) +import util +import vision_cfg +from ray import tune + + +class CartpoleRGBNoTuneJobCfg(vision_cfg.CameraJobCfg): + def __init__(self, cfg: dict = {}): + cfg = util.populate_isaac_ray_cfg_args(cfg) + cfg["runner_args"]["--task"] = tune.choice(["Isaac-Cartpole-RGB-v0"]) + super().__init__(cfg, vary_env_count=False, vary_cnn=False, vary_mlp=False) + + +class CartpoleRGBCNNOnlyJobCfg(vision_cfg.CameraJobCfg): + def __init__(self, cfg: dict = {}): + cfg = util.populate_isaac_ray_cfg_args(cfg) + cfg["runner_args"]["--task"] = tune.choice(["Isaac-Cartpole-RGB-v0"]) + super().__init__(cfg, vary_env_count=False, vary_cnn=True, vary_mlp=False) + + +class CartpoleRGBJobCfg(vision_cfg.CameraJobCfg): + def __init__(self, cfg: dict = {}): + cfg = util.populate_isaac_ray_cfg_args(cfg) + cfg["runner_args"]["--task"] = tune.choice(["Isaac-Cartpole-RGB-v0"]) + super().__init__(cfg, vary_env_count=True, vary_cnn=True, vary_mlp=True) + + +class CartpoleResNetJobCfg(vision_cfg.ResNetCameraJob): + def __init__(self, cfg: dict = {}): + cfg = util.populate_isaac_ray_cfg_args(cfg) + cfg["runner_args"]["--task"] = tune.choice(["Isaac-Cartpole-RGB-ResNet18-v0"]) + super().__init__(cfg) + + +class CartpoleTheiaJobCfg(vision_cfg.TheiaCameraJob): + def __init__(self, cfg: dict = {}): + cfg = util.populate_isaac_ray_cfg_args(cfg) + cfg["runner_args"]["--task"] = tune.choice(["Isaac-Cartpole-RGB-TheiaTiny-v0"]) + super().__init__(cfg) diff --git a/source/standalone/workflows/ray/hyperparameter_tuning/vision_cfg.py b/source/standalone/workflows/ray/hyperparameter_tuning/vision_cfg.py new file mode 100644 index 0000000000..be23e48cf0 --- /dev/null +++ b/source/standalone/workflows/ray/hyperparameter_tuning/vision_cfg.py @@ -0,0 +1,151 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import pathlib +import sys + +# Allow for import of items from the ray workflow. +UTIL_DIR = pathlib.Path(__file__).parent.parent.parent +sys.path.append(str(UTIL_DIR)) +import tuner +import util +from ray import tune + + +class CameraJobCfg(tuner.JobCfg): + """In order to be compatible with :meth: invoke_tuning_run, and + :class:IsaacLabTuneTrainable , configurations should + be in a similar format to this class. This class can vary env count/horizon length, + CNN structure, and MLP structure. Broad possible ranges are set, the specific values + that work can be found via tuning. Tuning results can inform better ranges for a second tuning run. + These ranges were selected for demonstration purposes. Best ranges are run/task specific.""" + + @staticmethod + def _get_batch_size_divisors(batch_size: int, min_size: int = 128) -> list[int]: + """Get valid batch divisors to combine with num_envs and horizon length""" + divisors = [i for i in range(min_size, batch_size + 1) if batch_size % i == 0] + return divisors if divisors else [min_size] + + def __init__(self, cfg={}, vary_env_count: bool = False, vary_cnn: bool = False, vary_mlp: bool = False): + cfg = util.populate_isaac_ray_cfg_args(cfg) + + # Basic configuration + cfg["runner_args"]["headless_singleton"] = "--headless" + cfg["runner_args"]["enable_cameras_singleton"] = "--enable_cameras" + cfg["hydra_args"]["agent.params.config.max_epochs"] = 200 + + if vary_env_count: # Vary the env count, and horizon length, and select a compatible mini-batch size + # Check from 512 to 8196 envs in powers of 2 + # check horizon lengths of 8 to 256 + # More envs should be better, but different batch sizes can improve gradient estimation + env_counts = [2**x for x in range(9, 13)] + horizon_lengths = [2**x for x in range(3, 8)] + + selected_env_count = tune.choice(env_counts) + selected_horizon = tune.choice(horizon_lengths) + + cfg["runner_args"]["--num_envs"] = selected_env_count + cfg["hydra_args"]["agent.params.config.horizon_length"] = selected_horizon + + def get_valid_batch_size(config): + num_envs = config["runner_args"]["--num_envs"] + horizon_length = config["hydra_args"]["agent.params.config.horizon_length"] + total_batch = horizon_length * num_envs + divisors = self._get_batch_size_divisors(total_batch) + return divisors[0] + + cfg["hydra_args"]["agent.params.config.minibatch_size"] = tune.sample_from(get_valid_batch_size) + + if vary_cnn: # Vary the depth, and size of the layers in the CNN part of the agent + # Also varies kernel size, and stride. + num_layers = tune.randint(2, 3) + cfg["hydra_args"]["agent.params.network.cnn.type"] = "conv2d" + cfg["hydra_args"]["agent.params.network.cnn.activation"] = tune.choice(["relu", "elu"]) + cfg["hydra_args"]["agent.params.network.cnn.initializer"] = "{name:default}" + cfg["hydra_args"]["agent.params.network.cnn.regularizer"] = "{name:None}" + + def get_cnn_layers(_): + layers = [] + size = 64 # Initial input size + + for _ in range(num_layers.sample()): + # Get valid kernel sizes for current size + valid_kernels = [k for k in [3, 4, 6, 8, 10, 12] if k <= size] + if not valid_kernels: + break + + kernel = int(tune.choice([str(k) for k in valid_kernels]).sample()) + stride = int(tune.choice(["1", "2", "3", "4"]).sample()) + padding = int(tune.choice(["0", "1"]).sample()) + + # Calculate next size + next_size = ((size + 2 * padding - kernel) // stride) + 1 + if next_size <= 0: + break + + layers.append({ + "filters": tune.randint(16, 32).sample(), + "kernel_size": str(kernel), + "strides": str(stride), + "padding": str(padding), + }) + size = next_size + + return layers + + cfg["hydra_args"]["agent.params.network.cnn.convs"] = tune.sample_from(get_cnn_layers) + + if vary_mlp: # Vary the MLP structure; neurons (units) per layer, number of layers, + + max_num_layers = 6 + max_neurons_per_layer = 128 + if "env.observations.policy.image.params.model_name" in cfg["hydra_args"]: + # By decreasing MLP size when using pretrained helps prevent out of memory on L4 + max_num_layers = 3 + max_neurons_per_layer = 32 + if "agent.params.network.cnn.convs" in cfg["hydra_args"]: + # decrease MLP size to prevent running out of memory on L4 + max_num_layers = 2 + max_neurons_per_layer = 32 + + num_layers = tune.randint(1, max_num_layers) + + def get_mlp_layers(_): + return [tune.randint(4, max_neurons_per_layer).sample() for _ in range(num_layers.sample())] + + cfg["hydra_args"]["agent.params.network.mlp.units"] = tune.sample_from(get_mlp_layers) + cfg["hydra_args"]["agent.params.network.mlp.initializer.name"] = tune.choice(["default"]).sample() + cfg["hydra_args"]["agent.params.network.mlp.activation"] = tune.choice( + ["relu", "tanh", "sigmoid", "elu"] + ).sample() + + super().__init__(cfg) + + +class ResNetCameraJob(CameraJobCfg): + """Try different ResNet sizes.""" + + def __init__(self, cfg: dict = {}): + cfg = util.populate_isaac_ray_cfg_args(cfg) + cfg["hydra_args"]["env.observations.policy.image.params.model_name"] = tune.choice( + ["resnet18", "resnet34", "resnet50", "resnet101"] + ) + super().__init__(cfg, vary_env_count=True, vary_cnn=False, vary_mlp=True) + + +class TheiaCameraJob(CameraJobCfg): + """Try different Theia sizes.""" + + def __init__(self, cfg: dict = {}): + cfg = util.populate_isaac_ray_cfg_args(cfg) + cfg["hydra_args"]["env.observations.policy.image.params.model_name"] = tune.choice([ + "theia-tiny-patch16-224-cddsv", + "theia-tiny-patch16-224-cdiv", + "theia-small-patch16-224-cdiv", + "theia-base-patch16-224-cdiv", + "theia-small-patch16-224-cddsv", + "theia-base-patch16-224-cddsv", + ]) + super().__init__(cfg, vary_env_count=True, vary_cnn=False, vary_mlp=True) diff --git a/source/standalone/workflows/ray/launch.py b/source/standalone/workflows/ray/launch.py new file mode 100644 index 0000000000..416d1f1f84 --- /dev/null +++ b/source/standalone/workflows/ray/launch.py @@ -0,0 +1,180 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import argparse +import pathlib +import subprocess +import yaml + +from jinja2 import Environment, FileSystemLoader +from kubernetes import config + +import source.standalone.workflows.ray.util as util + +"""This script helps create one or more KubeRay clusters. + +Usage: + +.. code-block:: bash + # If the head node is stuck on container creating, make sure to create a secret + ./isaaclab.sh -p source/standalone/workflows/ray/launch.py -h + + # Examples + + # The following creates 8 GPUx1 nvidia l4 workers + ./isaaclab.sh -p source/standalone/workflows/ray/launch.py --cluster_host google_cloud \ + --namespace --image \ + --num_workers 8 --num_clusters 1 --worker_accelerator nvidia-l4 --gpu_per_worker 1 + + # The following creates 1 GPUx1 nvidia l4 worker, 2 GPUx2 nvidia-tesla-t4 workers, + # and 2 GPUx4 nvidia-tesla-t4 GPU workers + ./isaaclab.sh -p source/standalone/workflows/ray/launch.py --cluster_host google_cloud \ + --namespace --image \ + --num_workers 1 2 --num_clusters 1 \ + --worker_accelerator nvidia-l4 nvidia-tesla-t4 --gpu_per_worker 1 2 4 +""" +RAY_DIR = pathlib.Path(__file__).parent + + +def apply_manifest(args: argparse.Namespace) -> None: + """Provided a Jinja templated ray.io/v1alpha1 file, + populate the arguments and create the cluster. Additionally, create + kubernetes containers for resources separated by '---' from the rest + of the file. + + Args: + args: Possible arguments concerning cluster parameters. + """ + # Load Kubernetes configuration + config.load_kube_config() + + # Set up Jinja2 environment for loading templates + templates_dir = RAY_DIR / "cluster_configs" / args.cluster_host + file_loader = FileSystemLoader(str(templates_dir)) + jinja_env = Environment(loader=file_loader, keep_trailing_newline=True) + + # Define template filename + template_file = "kuberay.yaml.jinja" + + # Convert args namespace to a dictionary + template_params = vars(args) + + # Load and render the template + template = jinja_env.get_template(template_file) + file_contents = template.render(template_params) + + # Parse all YAML documents in the rendered template + all_yamls = [] + for doc in yaml.safe_load_all(file_contents): + all_yamls.append(doc) + + # Convert back to YAML string, preserving multiple documents + cleaned_yaml_string = "" + for i, doc in enumerate(all_yamls): + if i > 0: + cleaned_yaml_string += "\n---\n" + cleaned_yaml_string += yaml.dump(doc) + + # Apply the Kubernetes manifest using kubectl + try: + subprocess.run(["kubectl", "apply", "-f", "-"], input=cleaned_yaml_string, text=True, check=True) + except subprocess.CalledProcessError as e: + exit(f"An error occurred while running `kubectl`: {e}") + + +def parse_args() -> argparse.Namespace: + """ + Parse command-line arguments for Kubernetes deployment script. + + Returns: + argparse.Namespace: Parsed command-line arguments. + """ + arg_parser = argparse.ArgumentParser( + description="Script to apply manifests to create Kubernetes objects for Ray clusters.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + + arg_parser.add_argument( + "--cluster_host", + type=str, + default="google_cloud", + choices=["google_cloud"], + help=( + "In the cluster_configs directory, the name of the folder where a tune.yaml.jinja" + "file exists defining the KubeRay config. Currently only google_cloud is supported." + ), + ) + + arg_parser.add_argument( + "--name", + type=str, + required=False, + default="isaacray", + help="Name of the Kubernetes deployment.", + ) + + arg_parser.add_argument( + "--namespace", + type=str, + required=False, + default="default", + help="Kubernetes namespace to deploy the Ray cluster.", + ) + + arg_parser.add_argument( + "--service_acount_name", type=str, required=False, default="default", help="The service account name to use." + ) + + arg_parser.add_argument( + "--image", + type=str, + required=True, + help="Docker image for the Ray cluster pods.", + ) + + arg_parser.add_argument( + "--worker_accelerator", + nargs="+", + type=str, + default=["nvidia-l4"], + help="GPU accelerator name. Supply more than one for heterogeneous resources.", + ) + + arg_parser = util.add_resource_arguments(arg_parser, cluster_create_defaults=True) + + arg_parser.add_argument( + "--num_clusters", + type=int, + default=1, + help="How many Ray Clusters to create.", + ) + arg_parser.add_argument( + "--num_head_cpu", + type=float, # to be able to schedule partial CPU heads + default=8, + help="The number of CPUs to give the Ray head.", + ) + + arg_parser.add_argument("--head_ram_gb", type=int, default=8, help="How many gigs of ram to give the Ray head") + args = arg_parser.parse_args() + return util.fill_in_missing_resources(args, cluster_creation_flag=True) + + +def main(): + args = parse_args() + + if "head" in args.name: + raise ValueError("For compatibility with other scripts, do not include head in the name") + if args.num_clusters == 1: + apply_manifest(args) + else: + default_name = args.name + for i in range(args.num_clusters): + args.name = default_name + "-" + str(i) + apply_manifest(args) + + +if __name__ == "__main__": + main() diff --git a/source/standalone/workflows/ray/mlflow_to_local_tensorboard.py b/source/standalone/workflows/ray/mlflow_to_local_tensorboard.py new file mode 100644 index 0000000000..652055ccdf --- /dev/null +++ b/source/standalone/workflows/ray/mlflow_to_local_tensorboard.py @@ -0,0 +1,149 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import argparse +import logging +import multiprocessing as mp +import os +import sys +from concurrent.futures import ProcessPoolExecutor, as_completed +from torch.utils.tensorboard import SummaryWriter + +import mlflow +from mlflow.tracking import MlflowClient + + +def setup_logging(level=logging.INFO): + logging.basicConfig(level=level, format="%(asctime)s - %(levelname)s - %(message)s") + + +def get_existing_runs(download_dir: str) -> set[str]: + """Get set of run IDs that have already been downloaded.""" + existing_runs = set() + tensorboard_dir = os.path.join(download_dir, "tensorboard") + if os.path.exists(tensorboard_dir): + for entry in os.listdir(tensorboard_dir): + if entry.startswith("run_"): + existing_runs.add(entry[4:]) + return existing_runs + + +def process_run(args): + """Convert MLflow run to TensorBoard format.""" + run_id, download_dir, tracking_uri = args + + try: + # Set up MLflow client + mlflow.set_tracking_uri(tracking_uri) + client = MlflowClient() + run = client.get_run(run_id) + + # Create TensorBoard writer + tensorboard_log_dir = os.path.join(download_dir, "tensorboard", f"run_{run_id}") + writer = SummaryWriter(log_dir=tensorboard_log_dir) + + # Log parameters + for key, value in run.data.params.items(): + writer.add_text(f"params/{key}", str(value)) + + # Log metrics with history + for key in run.data.metrics.keys(): + history = client.get_metric_history(run_id, key) + for m in history: + writer.add_scalar(f"metrics/{key}", m.value, m.step) + + # Log tags + for key, value in run.data.tags.items(): + writer.add_text(f"tags/{key}", str(value)) + + writer.close() + return run_id, True + except Exception: + return run_id, False + + +def download_experiment_tensorboard_logs(uri: str, experiment_name: str, download_dir: str) -> None: + """Download MLflow experiment logs and convert to TensorBoard format.""" + logger = logging.getLogger(__name__) + + try: + # Set up MLflow + mlflow.set_tracking_uri(uri) + logger.info(f"Connected to MLflow tracking server at {uri}") + + # Get experiment + experiment = mlflow.get_experiment_by_name(experiment_name) + if experiment is None: + raise ValueError(f"Experiment '{experiment_name}' not found at URI '{uri}'.") + + # Get all runs + runs = mlflow.search_runs([experiment.experiment_id]) + logger.info(f"Found {len(runs)} total runs in experiment '{experiment_name}'") + + # Check existing runs + existing_runs = get_existing_runs(download_dir) + logger.info(f"Found {len(existing_runs)} existing runs in {download_dir}") + + # Create directory structure + os.makedirs(os.path.join(download_dir, "tensorboard"), exist_ok=True) + + # Process new runs + new_run_ids = [run.run_id for _, run in runs.iterrows() if run.run_id not in existing_runs] + + if not new_run_ids: + logger.info("No new runs to process") + return + + logger.info(f"Processing {len(new_run_ids)} new runs...") + + # Process runs in parallel + num_processes = min(mp.cpu_count(), len(new_run_ids)) + processed = 0 + + with ProcessPoolExecutor(max_workers=num_processes) as executor: + future_to_run = { + executor.submit(process_run, (run_id, download_dir, uri)): run_id for run_id in new_run_ids + } + + for future in as_completed(future_to_run): + run_id = future_to_run[future] + try: + run_id, success = future.result() + processed += 1 + if success: + logger.info(f"[{processed}/{len(new_run_ids)}] Successfully processed run {run_id}") + else: + logger.error(f"[{processed}/{len(new_run_ids)}] Failed to process run {run_id}") + except Exception as e: + logger.error(f"Error processing run {run_id}: {e}") + + logger.info(f"\nAll data saved to {download_dir}/tensorboard") + + except Exception as e: + logger.error(f"Error during download: {e}") + raise + + +def main(): + parser = argparse.ArgumentParser(description="Download MLflow experiment logs for TensorBoard visualization.") + parser.add_argument("--uri", required=True, help="The MLflow tracking URI (e.g., http://localhost:5000)") + parser.add_argument("--experiment-name", required=True, help="Name of the experiment to download") + parser.add_argument("--download-dir", required=True, help="Directory to save TensorBoard logs") + parser.add_argument("--debug", action="store_true", help="Enable debug logging") + + args = parser.parse_args() + setup_logging(level=logging.DEBUG if args.debug else logging.INFO) + + try: + download_experiment_tensorboard_logs(args.uri, args.experiment_name, args.download_dir) + print("\nSuccess! To view the logs, run:") + print(f"tensorboard --logdir {os.path.join(args.download_dir, 'tensorboard')}") + except Exception as e: + logging.error(f"Failed to download experiment logs: {e}") + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/source/standalone/workflows/ray/submit_job.py b/source/standalone/workflows/ray/submit_job.py new file mode 100644 index 0000000000..492c952216 --- /dev/null +++ b/source/standalone/workflows/ray/submit_job.py @@ -0,0 +1,148 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import argparse +import os +import time +from concurrent.futures import ThreadPoolExecutor + +from ray import job_submission + +""" +This script submits aggregate job(s) to cluster(s) described in a +config file containing ``name: address: http://:`` on +a new line for each cluster. For KubeRay clusters, this file +can be automatically created with :file:`grok_cluster_with_kubectl.py` + +Aggregate job(s) are matched with cluster(s) via the following relation: +cluster_line_index_submitted_to = job_index % total_cluster_count + +Aggregate jobs are separated by the * delimiter. The ``--aggregate_jobs`` argument must be +the last argument supplied to the script. + +An aggregate job could be a :file:`../tuner.py` tuning job, which automatically +creates several individual jobs when started on a cluster. Alternatively, an aggregate job +could be a :file:'../wrap_resources.py` resource-wrapped job, +which may contain several individual sub-jobs separated by +the + delimiter. + +If there are more aggregate jobs than cluster(s), aggregate jobs will be submitted +as clusters become available via the defined relation above. If there are less aggregate job(s) +than clusters, some clusters will not receive aggregate job(s). The maximum number of +aggregate jobs that can be run simultaneously is equal to the number of workers created by +default by a ThreadPoolExecutor on the machine submitting jobs due to fetching the log output after +jobs finish, which is unlikely to constrain overall-job submission. + +Usage: + +.. code-block:: bash + + # Example; submitting a tuning job + ./isaaclab.sh -p source/standalone/workflows/ray/submit_job.py \ + --aggregate_jobs /workspace/isaaclab/source/standalone/workflows/ray/tuner.py \ + --cfg_file hyperparameter_tuning/vision_cartpole_cfg.py \ + --cfg_class CartpoleRGBNoTuneJobCfg --mlflow_uri + + # Example: Submitting resource wrapped job + ./isaaclab.sh -p source/standalone/workflows/ray/submit_job.py --aggregate_jobs wrap_resources.py --sub_jobs ./isaaclab.sh -p source/standalone/workflows/rl_games/train.py --task Isaac-Cartpole-v0 --headless+./isaaclab.sh -p source/standalone/workflows/rl_games/train.py --task Isaac-Cartpole-RGB-Camera-Direct-v0 --headless --enable_cameras agent.params.config.max_epochs=150 + + # For all command line arguments + ./isaaclab.sh -p source/standalone/workflows/ray/submit_job.py -h +""" +script_directory = os.path.dirname(os.path.abspath(__file__)) +CONFIG = {"working_dir": script_directory, "executable": "/workspace/isaaclab/isaaclab.sh -p"} + + +def read_cluster_spec(fn: str | None = None) -> list[dict]: + if fn is None: + cluster_spec_path = os.path.expanduser("~/.cluster_config") + else: + cluster_spec_path = os.path.expanduser(fn) + + if not os.path.exists(cluster_spec_path): + raise FileNotFoundError(f"Cluster spec file not found at {cluster_spec_path}") + + clusters = [] + with open(cluster_spec_path) as f: + for line in f: + parts = line.strip().split(" ") + http_address = parts[3] + cluster_info = {"name": parts[1], "address": http_address} + print(f"[INFO] Setting {cluster_info['name']}") # with {cluster_info['num_gpu']} GPUs.") + clusters.append(cluster_info) + + return clusters + + +def submit_job(cluster: dict, job_command: str) -> None: + """ + Submits a job to a single cluster, prints the final result and Ray dashboard URL at the end. + """ + address = cluster["address"] + cluster_name = cluster["name"] + print(f"[INFO]: Submitting job to cluster '{cluster_name}' at {address}") # with {num_gpus} GPUs.") + client = job_submission.JobSubmissionClient(address) + runtime_env = {"working_dir": CONFIG["working_dir"], "executable": CONFIG["executable"]} + print(f"[INFO]: Checking contents of the directory: {CONFIG['working_dir']}") + try: + dir_contents = os.listdir(CONFIG["working_dir"]) + print(f"[INFO]: Directory contents: {dir_contents}") + except Exception as e: + print(f"[INFO]: Failed to list directory contents: {str(e)}") + entrypoint = f"{CONFIG['executable']} {job_command}" + print(f"[INFO]: Attempting entrypoint {entrypoint=} in cluster {cluster}") + job_id = client.submit_job(entrypoint=entrypoint, runtime_env=runtime_env) + status = client.get_job_status(job_id) + while status in [job_submission.JobStatus.PENDING, job_submission.JobStatus.RUNNING]: + time.sleep(5) + status = client.get_job_status(job_id) + + final_logs = client.get_job_logs(job_id) + print("----------------------------------------------------") + print(f"[INFO]: Cluster {cluster_name} Logs: \n") + print(final_logs) + print("----------------------------------------------------") + + +def submit_jobs_to_clusters(jobs: list[str], clusters: list[dict]) -> None: + """ + Submit all jobs to their respective clusters, cycling through clusters if there are more jobs than clusters. + """ + if not clusters: + raise ValueError("No clusters available for job submission.") + + if len(jobs) < len(clusters): + print("[INFO]: Less jobs than clusters, some clusters will not receive jobs") + elif len(jobs) == len(clusters): + print("[INFO]: Exactly one job per cluster") + else: + print("[INFO]: More jobs than clusters, jobs submitted as clusters become available.") + with ThreadPoolExecutor() as executor: + for idx, job_command in enumerate(jobs): + # Cycle through clusters using modulus to wrap around if there are more jobs than clusters + cluster = clusters[idx % len(clusters)] + executor.submit(submit_job, cluster, job_command) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Submit multiple GPU jobs to multiple Ray clusters.") + parser.add_argument("--config_file", default="~/.cluster_config", help="The cluster config path.") + parser.add_argument( + "--aggregate_jobs", + type=str, + nargs=argparse.REMAINDER, + help="This should be last argument. The aggregate jobs to submit separated by the * delimiter.", + ) + args = parser.parse_args() + if args.aggregate_jobs is not None: + jobs = " ".join(args.aggregate_jobs) + formatted_jobs = jobs.split("*") + if len(formatted_jobs) > 1: + print("Warning; Split jobs by cluster with the * delimiter") + else: + formatted_jobs = [] + print(f"[INFO]: Isaac Ray Wrapper received jobs {formatted_jobs=}") + clusters = read_cluster_spec(args.config_file) + submit_jobs_to_clusters(formatted_jobs, clusters) diff --git a/source/standalone/workflows/ray/tuner.py b/source/standalone/workflows/ray/tuner.py new file mode 100644 index 0000000000..f9de457f8d --- /dev/null +++ b/source/standalone/workflows/ray/tuner.py @@ -0,0 +1,343 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +import argparse +import importlib.util +import os +import sys +from time import sleep + +import ray +import util +from ray import air, tune +from ray.tune.search.optuna import OptunaSearch +from ray.tune.search.repeater import Repeater + +""" +This script breaks down an aggregate tuning job, as defined by a hyperparameter sweep configuration, +into individual jobs (shell commands) to run on the GPU-enabled nodes of the cluster. +By default, (unless combined as a sub-job in a resource-wrapped aggregate job), one worker is created +for each GPU-enabled node in the cluster for each individual job. + +Each hyperparameter sweep configuration should include the workflow, +runner arguments, and hydra arguments to vary. + +This assumes that all workers in a cluster are homogeneous. For heterogeneous workloads, +create several heterogeneous clusters (with homogeneous nodes in each cluster), +then submit several overall-cluster jobs with :file:`../submit_job.py`. +KubeRay clusters on Google GKE can be created with :file:`../launch.py` + +To report tune metrics on clusters, a running MLFlow server with a known URI that the cluster has +access to is required. For KubeRay clusters configured with :file:`../launch.py`, this is included +automatically, and can be easily found with with :file:`grok_cluster_with_kubectl.py` + +Usage: + +.. code-block:: bash + + ./isaaclab.sh -p source/standalone/workflows/ray/tuner.py -h + + # Examples + # Local (not within a docker container, when within a local docker container, do not supply run_mode argument) + ./isaaclab.sh -p source/standalone/workflows/ray/tuner.py --run_mode local \ + --cfg_file source/standalone/workflows/ray/hyperparameter_tuning/vision_cartpole_cfg.py \ + --cfg_class CartpoleRGBNoTuneJobCfg + # Local docker: start the ray server and run above command in the same running container without run_mode arg + # Remote (run grok cluster or create config file mentioned in :file:`submit_job.py`) + ./isaaclab.sh -p source/standalone/workflows/ray/submit_job.py \ + --aggregate_jobs tuner.py \ + --cfg_file hyperparameter_tuning/vision_cartpole_cfg.py \ + --cfg_class CartpoleRGBNoTuneJobCfg --mlflow_uri + +""" + +DOCKER_PREFIX = "/workspace/isaaclab/" +BASE_DIR = os.path.expanduser("~") +PYTHON_EXEC = "./isaaclab.sh -p" +WORKFLOW = "source/standalone/workflows/rl_games/train.py" +NUM_WORKERS_PER_NODE = 1 # needed for local parallelism + + +class IsaacLabTuneTrainable(tune.Trainable): + """The Isaac Lab Ray Tune Trainable. + This class uses the standalone workflows to start jobs, along with the hydra integration. + This class achieves Ray-based logging through reading the tensorboard logs from + the standalone workflows. This depends on a config generated in the format of + :class:`JobCfg` + """ + + def setup(self, config: dict) -> None: + """Get the invocation command, return quick for easy scheduling.""" + self.data = None + self.invoke_cmd = util.get_invocation_command_from_cfg(cfg=config, python_cmd=PYTHON_EXEC, workflow=WORKFLOW) + print(f"[INFO]: Recovered invocation with {self.invoke_cmd}") + self.experiment = None + + def reset_config(self, new_config): + """Allow environments to be re-used by fetching a new invocation command""" + self.setup(new_config) + return True + + def step(self) -> dict: + if self.experiment is None: # start experiment + # When including this as first step instead of setup, experiments get scheduled faster + # Don't want to block the scheduler while the experiment spins up + print(f"[INFO]: Invoking experiment as first step with {self.invoke_cmd}...") + experiment = util.execute_job( + self.invoke_cmd, + identifier_string="", + extract_experiment=True, + persistent_dir=BASE_DIR, + ) + self.experiment = experiment + print(f"[INFO]: Tuner recovered experiment info {experiment}") + self.proc = experiment["proc"] + self.experiment_name = experiment["experiment_name"] + self.isaac_logdir = experiment["logdir"] + self.tensorboard_logdir = self.isaac_logdir + f"/{self.experiment_name}/summaries" + self.done = False + + if self.proc is None: + raise ValueError("Could not start trial.") + + if self.proc.poll() is not None: # process finished, signal finish + self.data["done"] = True + print("[INFO]: Process finished, returning...") + else: # wait until the logs are ready or fresh + data = util.load_tensorboard_logs(self.tensorboard_logdir) + + while data is None: + data = util.load_tensorboard_logs(self.tensorboard_logdir) + sleep(2) # Lazy report metrics to avoid performance overhead + + if self.data is not None: + while util._dicts_equal(data, self.data): + data = util.load_tensorboard_logs(self.tensorboard_logdir) + sleep(2) # Lazy report metrics to avoid performance overhead + + self.data = data + self.data["done"] = False + return self.data + + def default_resource_request(self): + """How many resources each trainable uses. Assumes homogeneous resources across gpu nodes, + and that each trainable is meant for one node, where it uses all available resources.""" + resources = util.get_gpu_node_resources(one_node_only=True) + if NUM_WORKERS_PER_NODE != 1: + print("[WARNING]: Splitting node into more than one worker") + return tune.PlacementGroupFactory( + [{"CPU": resources["CPU"] / NUM_WORKERS_PER_NODE, "GPU": resources["GPU"] / NUM_WORKERS_PER_NODE}], + strategy="STRICT_PACK", + ) + + +def invoke_tuning_run(cfg: dict, args: argparse.Namespace) -> None: + """Invoke an Isaac-Ray tuning run. + + Log either to a local directory or to MLFlow. + Args: + cfg: Configuration dictionary extracted from job setup + args: Command-line arguments related to tuning. + """ + # Allow for early exit + os.environ["TUNE_DISABLE_STRICT_METRIC_CHECKING"] = "1" + + print("[WARNING]: Not saving checkpoints, just running experiment...") + print("[INFO]: Model parameters and metrics will be preserved.") + print("[WARNING]: For homogeneous cluster resources only...") + # Get available resources + resources = util.get_gpu_node_resources() + print(f"[INFO]: Available resources {resources}") + + if not ray.is_initialized(): + ray.init( + address=args.ray_address, + log_to_driver=True, + num_gpus=len(resources), + ) + + print(f"[INFO]: Using config {cfg}") + + # Configure the search algorithm and the repeater + searcher = OptunaSearch( + metric=args.metric, + mode=args.mode, + ) + repeat_search = Repeater(searcher, repeat=args.repeat_run_count) + + if args.run_mode == "local": # Standard config, to file + run_config = air.RunConfig( + storage_path="/tmp/ray", + name=f"IsaacRay-{args.cfg_class}-tune", + verbose=1, + checkpoint_config=air.CheckpointConfig( + checkpoint_frequency=0, # Disable periodic checkpointing + checkpoint_at_end=False, # Disable final checkpoint + ), + ) + + elif args.run_mode == "remote": # MLFlow, to MLFlow server + mlflow_callback = MLflowLoggerCallback( + tracking_uri=args.mlflow_uri, + experiment_name=f"IsaacRay-{args.cfg_class}-tune", + save_artifact=False, + tags={"run_mode": "remote", "cfg_class": args.cfg_class}, + ) + + run_config = ray.train.RunConfig( + name="mlflow", + storage_path="/tmp/ray", + callbacks=[mlflow_callback], + checkpoint_config=ray.train.CheckpointConfig(checkpoint_frequency=0, checkpoint_at_end=False), + ) + else: + raise ValueError("Unrecognized run mode.") + + # Configure the tuning job + tuner = tune.Tuner( + IsaacLabTuneTrainable, + param_space=cfg, + tune_config=tune.TuneConfig( + search_alg=repeat_search, + num_samples=args.num_samples, + reuse_actors=True, + ), + run_config=run_config, + ) + + # Execute the tuning + tuner.fit() + + # Save results to mounted volume + if args.run_mode == "local": + print("[DONE!]: Check results with tensorboard dashboard") + else: + print("[DONE!]: Check results with MLFlow dashboard") + + +class JobCfg: + """To be compatible with :meth: invoke_tuning_run and :class:IsaacLabTuneTrainable, + at a minimum, the tune job should inherit from this class.""" + + def __init__(self, cfg): + assert "runner_args" in cfg, "No runner arguments specified." + assert "--task" in cfg["runner_args"], "No task specified." + assert "hydra_args" in cfg, "No hypeparameters specified." + self.cfg = cfg + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Tune Isaac Lab hyperparameters.") + parser.add_argument("--ray_address", type=str, default="auto", help="the Ray address.") + parser.add_argument( + "--cfg_file", + type=str, + default="hyperparameter_tuning/vision_cartpole_cfg.py", + required=False, + help="The relative filepath where a hyperparameter sweep is defined", + ) + parser.add_argument( + "--cfg_class", + type=str, + default="CartpoleRGBNoTuneJobCfg", + required=False, + help="Name of the hyperparameter sweep class to use", + ) + parser.add_argument( + "--run_mode", + choices=["local", "remote"], + default="remote", + help=( + "Set to local to use ./isaaclab.sh -p python, set to " + "remote to use /workspace/isaaclab/isaaclab.sh -p python" + ), + ) + parser.add_argument( + "--workflow", + default=None, # populated with RL Games + help="The absolute path of the workflow to use for the experiment. By default, RL Games is used.", + ) + parser.add_argument( + "--mlflow_uri", + type=str, + default=None, + required=False, + help="The MLFlow Uri.", + ) + parser.add_argument( + "--num_workers_per_node", + type=int, + default=1, + help="Number of workers to run on each GPU node. Only supply for parallelism on multi-gpu nodes", + ) + + parser.add_argument("--metric", type=str, default="rewards/time", help="What metric to tune for.") + + parser.add_argument( + "--mode", + choices=["max", "min"], + default="max", + help="What to optimize the metric to while tuning", + ) + parser.add_argument( + "--num_samples", + type=int, + default=100, + help="How many hyperparameter runs to try total.", + ) + parser.add_argument( + "--repeat_run_count", + type=int, + default=3, + help="How many times to repeat each hyperparameter config.", + ) + + args = parser.parse_args() + NUM_WORKERS_PER_NODE = args.num_workers_per_node + print(f"[INFO]: Using {NUM_WORKERS_PER_NODE} workers per node.") + if args.run_mode == "remote": + BASE_DIR = DOCKER_PREFIX # ensure logs are dumped to persistent location + PYTHON_EXEC = DOCKER_PREFIX + PYTHON_EXEC[2:] + if args.workflow is None: + WORKFLOW = DOCKER_PREFIX + WORKFLOW + else: + WORKFLOW = args.workflow + print(f"[INFO]: Using remote mode {PYTHON_EXEC=} {WORKFLOW=}") + + if args.mlflow_uri is not None: + import mlflow + + mlflow.set_tracking_uri(args.mlflow_uri) + from ray.air.integrations.mlflow import MLflowLoggerCallback + else: + raise ValueError("Please provide a result MLFLow URI server.") + else: # local + PYTHON_EXEC = os.getcwd() + "/" + PYTHON_EXEC[2:] + if args.workflow is None: + WORKFLOW = os.getcwd() + "/" + WORKFLOW + else: + WORKFLOW = args.workflow + BASE_DIR = os.getcwd() + print(f"[INFO]: Using local mode {PYTHON_EXEC=} {WORKFLOW=}") + file_path = args.cfg_file + class_name = args.cfg_class + print(f"[INFO]: Attempting to use sweep config from {file_path=} {class_name=}") + module_name = os.path.splitext(os.path.basename(file_path))[0] + + spec = importlib.util.spec_from_file_location(module_name, file_path) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + print(f"[INFO]: Successfully imported {module_name} from {file_path}") + if hasattr(module, class_name): + ClassToInstantiate = getattr(module, class_name) + print(f"[INFO]: Found correct class {ClassToInstantiate}") + instance = ClassToInstantiate() + print(f"[INFO]: Successfully instantiated class '{class_name}' from {file_path}") + cfg = instance.cfg + print(f"[INFO]: Grabbed the following hyperparameter sweep config: \n {cfg}") + invoke_tuning_run(cfg, args) + + else: + raise AttributeError(f"[ERROR]:Class '{class_name}' not found in {file_path}") diff --git a/source/standalone/workflows/ray/util.py b/source/standalone/workflows/ray/util.py new file mode 100644 index 0000000000..4432468c3b --- /dev/null +++ b/source/standalone/workflows/ray/util.py @@ -0,0 +1,409 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +import argparse +import os +import re +import subprocess +from datetime import datetime +from math import isclose + +import ray +from tensorboard.backend.event_processing.event_accumulator import EventAccumulator + + +def load_tensorboard_logs(directory: str) -> dict: + """From a tensorboard directory, get the latest scalar values. + + Args: + directory: The directory of the tensorboard logging. + + Returns: + The latest available scalar values. + """ + # Initialize the event accumulator with a size guidance for only the latest entry + size_guidance = {"scalars": 1} # Load only the latest entry for scalars + event_acc = EventAccumulator(directory, size_guidance=size_guidance) + event_acc.Reload() # Load all data from the directory + + # Extract the latest scalars logged + latest_scalars = {} + for tag in event_acc.Tags()["scalars"]: + events = event_acc.Scalars(tag) + if events: # Check if there is at least one entry + latest_event = events[-1] # Get the latest entry + latest_scalars[tag] = latest_event.value + return latest_scalars + + +def get_invocation_command_from_cfg( + cfg: dict, + python_cmd: str = "/workspace/isaaclab/isaaclab.sh -p", + workflow: str = "source/standalone/workflows/rl_games/train.py", +) -> str: + """Generate command with proper Hydra arguments""" + runner_args = [] + hydra_args = [] + + def process_args(args, target_list, is_hydra=False): + for key, value in args.items(): + if not is_hydra: + if key.endswith("_singleton"): + target_list.append(value) + elif key.startswith("--"): + target_list.append(f"{key} {value}") # Space instead of = for runner args + else: + target_list.append(f"{value}") + else: + if isinstance(value, list): + # Check the type of the first item to determine formatting + if value and isinstance(value[0], dict): + # Handle list of dictionaries (e.g., CNN convs) + formatted_items = [f"{{{','.join(f'{k}:{v}' for k, v in item.items())}}}" for item in value] + else: + # Handle list of primitives (e.g., MLP units) + formatted_items = [str(x) for x in value] + target_list.append(f"'{key}=[{','.join(formatted_items)}]'") + elif isinstance(value, str) and ("{" in value or "}" in value): + target_list.append(f"'{key}={value}'") + else: + target_list.append(f"{key}={value}") + + print(f"[INFO]: Starting workflow {workflow}") + process_args(cfg["runner_args"], runner_args) + print(f"[INFO]: Retrieved workflow runner args: {runner_args}") + process_args(cfg["hydra_args"], hydra_args, is_hydra=True) + print(f"[INFO]: Retrieved hydra args: {hydra_args}") + + invoke_cmd = f"{python_cmd} {workflow} " + invoke_cmd += " ".join(runner_args) + " " + " ".join(hydra_args) + return invoke_cmd + + +@ray.remote +def remote_execute_job( + job_cmd: str, identifier_string: str, test_mode: bool = False, extract_experiment: bool = False +) -> str | dict: + """This method has an identical signature to :meth:`execute_job`, with the ray remote decorator""" + return execute_job( + job_cmd=job_cmd, identifier_string=identifier_string, test_mode=test_mode, extract_experiment=extract_experiment + ) + + +def execute_job( + job_cmd: str, + identifier_string: str = "job 0", + test_mode: bool = False, + extract_experiment: bool = False, + persistent_dir: str | None = None, + log_all_output: bool = False, +) -> str | dict: + """Issue a job (shell command). + + Args: + job_cmd: The shell command to run. + identifier_string: What prefix to add to make logs easier to differentiate + across clusters or jobs. Defaults to "job 0". + test_mode: When true, only run 'nvidia-smi'. Defaults to False. + extract_experiment: When true, search for experiment details from a training run. Defaults to False. + persistent_dir: When supplied, change to run the directory in a persistent + directory. Can be used to avoid losing logs in the /tmp directory. Defaults to None. + log_all_output: When true, print all output to the console. Defaults to False. + Raises: + ValueError: If the job is unable to start, or throws an error. Most likely to happen + due to running out of memory. + + Returns: + Relevant information from the job + """ + start_time = datetime.now().strftime("%H:%M:%S.%f") + result_details = [f"{identifier_string}: ---------------------------------\n"] + result_details.append(f"{identifier_string}:[INFO]: Invocation {job_cmd} \n") + node_id = ray.get_runtime_context().get_node_id() + result_details.append(f"{identifier_string}:[INFO]: Ray Node ID: {node_id} \n") + + if test_mode: + import torch + + try: + result = subprocess.run( + ["nvidia-smi", "--query-gpu=name,memory.free,serial", "--format=csv,noheader,nounits"], + capture_output=True, + check=True, + text=True, + ) + output = result.stdout.strip().split("\n") + for gpu_info in output: + name, memory_free, serial = gpu_info.split(", ") + result_details.append( + f"{identifier_string}[INFO]: Name: {name}|Memory Available: {memory_free} MB|Serial Number" + f" {serial} \n" + ) + + # Get GPU count from PyTorch + num_gpus_detected = torch.cuda.device_count() + result_details.append(f"{identifier_string}[INFO]: Detected GPUs from PyTorch: {num_gpus_detected} \n") + + # Check CUDA_VISIBLE_DEVICES and count the number of visible GPUs + cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES") + if cuda_visible_devices: + visible_devices_count = len(cuda_visible_devices.split(",")) + result_details.append( + f"{identifier_string}[INFO]: GPUs visible via CUDA_VISIBLE_DEVICES: {visible_devices_count} \n" + ) + else: + visible_devices_count = len(output) # All GPUs visible if CUDA_VISIBLE_DEVICES is not set + result_details.append( + f"{identifier_string}[INFO]: CUDA_VISIBLE_DEVICES not set; all GPUs visible" + f" ({visible_devices_count}) \n" + ) + + # If PyTorch GPU count disagrees with nvidia-smi, reset CUDA_VISIBLE_DEVICES and rerun detection + if num_gpus_detected != len(output): + result_details.append( + f"{identifier_string}[WARNING]: PyTorch and nvidia-smi disagree on GPU count! Re-running with all" + " GPUs visible. \n" + ) + result_details.append(f"{identifier_string}[INFO]: This shows that GPU resources were isolated.\n") + os.environ["CUDA_VISIBLE_DEVICES"] = ",".join([str(i) for i in range(len(output))]) + num_gpus_detected_after_reset = torch.cuda.device_count() + result_details.append( + f"{identifier_string}[INFO]: After setting CUDA_VISIBLE_DEVICES, PyTorch detects" + f" {num_gpus_detected_after_reset} GPUs \n" + ) + + except subprocess.CalledProcessError as e: + print(f"Error calling nvidia-smi: {e.stderr}") + result_details.append({"error": "Failed to retrieve GPU information"}) + else: + if persistent_dir: + og_dir = os.getcwd() + os.chdir(persistent_dir) + process = subprocess.Popen( + job_cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1 + ) + if persistent_dir: + os.chdir(og_dir) + experiment_name = None + logdir = None + experiment_info_pattern = re.compile("Exact experiment name requested from command line: (.+)") + logdir_pattern = re.compile(r"\[INFO\] Logging experiment in directory: (.+)$") + err_pattern = re.compile("There was an error (.+)$") + with process.stdout as stdout: + for line in iter(stdout.readline, ""): + line = line.strip() + result_details.append(f"{identifier_string}: {line} \n") + if log_all_output: + print(f"{identifier_string}: {line}") + + if extract_experiment: + exp_match = experiment_info_pattern.search(line) + log_match = logdir_pattern.search(line) + err_match = err_pattern.search(line) + if err_match: + raise ValueError(f"Encountered an error during trial run. {' '.join(result_details)}") + + if exp_match: + experiment_name = exp_match.group(1) + if log_match: + logdir = log_match.group(1) + + if experiment_name and logdir: + result = { + "experiment_name": experiment_name, + "logdir": logdir, + "proc": process, + "result": " ".join(result_details), + } + return result + + with process.stderr as stderr: + for line in iter(stderr.readline, ""): + line = line.strip() + result_details.append(f"{identifier_string}: {line}") + print(f"{identifier_string}: {line}") + + process.wait() # Wait for the subprocess to finish naturally if not exited early + + now = datetime.now().strftime("%H:%M:%S.%f") + completion_info = f"\n[INFO]: {identifier_string}: Job Started at {start_time}, completed at {now}\n" + print(completion_info) + result_details.append(completion_info) + return " ".join(result_details) + + +def get_gpu_node_resources( + total_resources: bool = False, + one_node_only: bool = False, + include_gb_ram: bool = False, + include_id: bool = False, + ray_address: str = "auto", +) -> list[dict] | dict: + """Get information about available GPU node resources. + + Args: + total_resources: When true, return total available resources. Defaults to False. + one_node_only: When true, return resources for a single node. Defaults to False. + include_gb_ram: Set to true to convert MB to GB in result + include_id: Set to true to include node ID + ray_address: The ray address to connect to. + + Returns: + Resource information for all nodes, sorted by descending GPU count, then descending CPU + count, then descending RAM capacity, and finally by node ID in ascending order if available, + or simply the resource for a single node if requested. + """ + if not ray.is_initialized(): + ray.init(address=ray_address) + + nodes = ray.nodes() + node_resources = [] + total_cpus = 0 + total_gpus = 0 + total_memory = 0 # in bytes + + for node in nodes: + if node["Alive"] and "GPU" in node["Resources"]: + node_id = node["NodeID"] + resources = node["Resources"] + cpus = resources.get("CPU", 0) + gpus = resources.get("GPU", 0) + memory = resources.get("memory", 0) + node_resources.append({"CPU": cpus, "GPU": gpus, "memory": memory}) + + if include_id: + node_resources[-1]["id"] = node_id + if include_gb_ram: + node_resources[-1]["ram_gb"] = memory / 1024**3 + + total_cpus += cpus + total_gpus += gpus + total_memory += memory + node_resources = sorted(node_resources, key=lambda x: (-x["GPU"], -x["CPU"], -x["memory"], x.get("id", ""))) + + if total_resources: + # Return summed total resources + return {"CPU": total_cpus, "GPU": total_gpus, "memory": total_memory} + + if one_node_only and node_resources: + return node_resources[0] + + return node_resources + + +def add_resource_arguments( + arg_parser: argparse.ArgumentParser, + defaults: list | None = None, + cluster_create_defaults: bool = False, +) -> argparse.ArgumentParser: + """Add resource arguments to a cluster; this is shared across both + wrapping resources and launching clusters. + + Args: + arg_parser: the argparser to add the arguments to. This argparser is mutated. + defaults: The default values for GPUs, CPUs, RAM, and Num Workers + cluster_create_defaults: Set to true to populate reasonable defaults for creating clusters. + Returns: + The argparser with the standard resource arguments. + """ + if defaults is None: + if cluster_create_defaults: + defaults = [[1], [8], [16], [1]] + else: + defaults = [None, None, None, [1]] + arg_parser.add_argument( + "--gpu_per_worker", + nargs="+", + type=int, + default=defaults[0], + help="Number of GPUs per worker node. Supply more than one for heterogeneous resources", + ) + arg_parser.add_argument( + "--cpu_per_worker", + nargs="+", + type=int, + default=defaults[1], + help="Number of CPUs per worker node. Supply more than one for heterogeneous resources", + ) + arg_parser.add_argument( + "--ram_gb_per_worker", + nargs="+", + type=int, + default=defaults[2], + help="RAM in GB per worker node. Supply more than one for heterogeneous resources.", + ) + arg_parser.add_argument( + "--num_workers", + nargs="+", + type=int, + default=defaults[3], + help="Number of desired workers. Supply more than one for heterogeneous resources.", + ) + return arg_parser + + +def fill_in_missing_resources( + args: argparse.Namespace, resources: dict | None = None, cluster_creation_flag: bool = False, policy: callable = max +): + """Normalize the lengths of resource lists based on the longest list provided.""" + print("[INFO]: Filling in missing command line arguments with best guess...") + if resources is None: + resources = { + "gpu_per_worker": args.gpu_per_worker, + "cpu_per_worker": args.cpu_per_worker, + "ram_gb_per_worker": args.ram_gb_per_worker, + "num_workers": args.num_workers, + } + if cluster_creation_flag: + cluster_creation_resources = {"worker_accelerator": args.worker_accelerator} + resources.update(cluster_creation_resources) + + # Calculate the maximum length of any list + max_length = max(len(v) for v in resources.values()) + print("[INFO]: Resource list lengths:") + for key, value in resources.items(): + print(f"[INFO] {key}: {len(value)} values {value}") + + # Extend each list to match the maximum length using the maximum value in each list + for key, value in resources.items(): + potential_value = getattr(args, key) + if potential_value is not None: + max_value = policy(policy(value), policy(potential_value)) + else: + max_value = policy(value) + extension_length = max_length - len(value) + if extension_length > 0: # Only extend if the current list is shorter than max_length + print(f"\n[WARNING]: Resource '{key}' needs extension:") + print(f"[INFO] Current length: {len(value)}") + print(f"[INFO] Target length: {max_length}") + print(f"[INFO] Filling in {extension_length} missing values with {max_value}") + print(f"[INFO] To avoid auto-filling, provide {extension_length} more {key} value(s)") + value.extend([max_value] * extension_length) + setattr(args, key, value) + resources[key] = value + print(f"[INFO] Final {key} values: {getattr(args, key)}") + print("[INFO]: Done filling in command line arguments...\n\n") + return args + + +def populate_isaac_ray_cfg_args(cfg: dict = {}) -> dict: + """Small utility method to create empty fields if needed for a configuration.""" + if "runner_args" not in cfg: + cfg["runner_args"] = {} + if "hydra_args" not in cfg: + cfg["hydra_args"] = {} + return cfg + + +def _dicts_equal(d1: dict, d2: dict, tol=1e-9) -> bool: + """Check if two dicts are equal; helps ensure only new logs are returned.""" + if d1.keys() != d2.keys(): + return False + for key in d1: + if isinstance(d1[key], float) and isinstance(d2[key], float): + if not isclose(d1[key], d2[key], abs_tol=tol): + return False + elif d1[key] != d2[key]: + return False + return True diff --git a/source/standalone/workflows/ray/wrap_resources.py b/source/standalone/workflows/ray/wrap_resources.py new file mode 100644 index 0000000000..d1b2cf547e --- /dev/null +++ b/source/standalone/workflows/ray/wrap_resources.py @@ -0,0 +1,151 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import argparse + +import ray +from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy + +import source.standalone.workflows.ray.util as util + +""" +This script dispatches sub-job(s) (either individual jobs or tuning aggregate jobs) +to worker(s) on GPU-enabled node(s) of a specific cluster as part of an resource-wrapped aggregate +job. If no desired compute resources for each sub-job are specified, +this script creates one worker per available node for each node with GPU(s) in the cluster. +If the desired resources for each sub-job is specified, +the maximum number of workers possible with the desired resources are created for each node +with GPU(s) in the cluster. It is also possible to split available node resources for each node +into the desired number of workers with the ``--num_workers`` flag, to be able to easily +parallelize sub-jobs on multi-GPU nodes. Due to Isaac Lab requiring a GPU, +this ignores all CPU only nodes such as loggers. + +Sub-jobs are matched with node(s) in a cluster via the following relation: +sorted_nodes = Node sorted by descending GPUs, then descending CPUs, then descending RAM, then node ID +node_submitted_to = sorted_nodes[job_index % total_node_count] + +To check the ordering of sorted nodes, supply the ``--test`` argument and run the script. + +Sub-jobs are separated by the + delimiter. The ``--sub_jobs`` argument must be the last +argument supplied to the script. + +If there is more than one available worker, and more than one sub-job, +sub-jobs will be executed in parallel. If there are more sub-jobs than workers, sub-jobs will +be dispatched to workers as they become available. There is no limit on the number +of sub-jobs that can be near-simultaneously submitted. + +This script is meant to be executed on a Ray cluster head node as an aggregate cluster job. +To submit aggregate cluster jobs such as this script to one or more remote clusters, +see :file:`../submit_isaac_ray_job.py`. + +KubeRay clusters on Google GKE can be created with :file:`../launch.py` + +Usage: + +.. code-block:: bash + # **Ensure that sub-jobs are separated by the ``+`` delimiter.** + # Generic Templates----------------------------------- + ./isaaclab.sh -p source/standalone/workflows/ray/wrap_resources.py -h + # No resource isolation; no parallelization: + ./isaaclab.sh -p source/standalone/workflows/ray/wrap_resources.py + --sub_jobs ++ + # Automatic Resource Isolation; Example A: needed for parallelization + ./isaaclab.sh -p source/standalone/workflows/ray/wrap_resources.py \ + --num_workers \ + --sub_jobs + + # Manual Resource Isolation; Example B: needed for parallelization + ./isaaclab.sh -p source/standalone/workflows/ray/wrap_resources.py --num_cpu_per_worker \ + --gpu_per_worker --ram_gb_per_worker --sub_jobs + + # Manual Resource Isolation; Example C: Needed for parallelization, for heterogeneous workloads + ./isaaclab.sh -p source/standalone/workflows/ray/wrap_resources.py --num_cpu_per_worker \ + --gpu_per_worker --ram_gb_per_worker --sub_jobs + + # to see all arguments + ./isaaclab.sh -p source/standalone/workflows/ray/wrap_resources.py -h +""" + + +def wrap_resources_to_jobs(jobs: list[str], args: argparse.Namespace) -> None: + """ + Provided a list of jobs, dispatch jobs to one worker per available node, + unless otherwise specified by resource constraints. + + Args: + jobs: bash commands to execute on a Ray cluster + args: The arguments for resource allocation + + """ + if not ray.is_initialized(): + ray.init(address=args.ray_address, log_to_driver=True) + job_results = [] + gpu_node_resources = util.get_gpu_node_resources(include_id=True, include_gb_ram=True) + + if any([args.gpu_per_worker, args.cpu_per_worker, args.ram_gb_per_worker]) and args.num_workers: + raise ValueError("Either specify only num_workers or only granular resources(GPU,CPU,RAM_GB).") + + num_nodes = len(gpu_node_resources) + # Populate arguments + formatted_node_resources = { + "gpu_per_worker": [gpu_node_resources[i]["GPU"] for i in range(num_nodes)], + "cpu_per_worker": [gpu_node_resources[i]["CPU"] for i in range(num_nodes)], + "ram_gb_per_worker": [gpu_node_resources[i]["ram_gb"] for i in range(num_nodes)], + "num_workers": args.num_workers, # By default, 1 worker por node + } + args = util.fill_in_missing_resources(args, resources=formatted_node_resources, policy=min) + print(f"[INFO]: Number of GPU nodes found: {num_nodes}") + if args.test: + jobs = ["nvidia-smi"] * num_nodes + for i, job in enumerate(jobs): + gpu_node = gpu_node_resources[i % num_nodes] + print(f"[INFO]: Submitting job {i + 1} of {len(jobs)} with job '{job}' to node {gpu_node}") + print( + f"[INFO]: Resource parameters: GPU: {args.gpu_per_worker[i]}" + f" CPU: {args.cpu_per_worker[i]} RAM {args.ram_gb_per_worker[i]}" + ) + print(f"[INFO] For the node parameters, creating {args.num_workers[i]} workers") + num_gpus = args.gpu_per_worker[i] / args.num_workers[i] + num_cpus = args.cpu_per_worker[i] / args.num_workers[i] + memory = (args.ram_gb_per_worker[i] * 1024**3) / args.num_workers[i] + print(f"[INFO]: Requesting {num_gpus=} {num_cpus=} {memory=} id={gpu_node['id']}") + job = util.remote_execute_job.options( + num_gpus=num_gpus, + num_cpus=num_cpus, + memory=memory, + scheduling_strategy=NodeAffinitySchedulingStrategy(gpu_node["id"], soft=False), + ).remote(job, f"Job {i}", args.test) + job_results.append(job) + + results = ray.get(job_results) + for i, result in enumerate(results): + print(f"[INFO]: Job {i} result: {result}") + print("[INFO]: All jobs completed.") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Submit multiple jobs with optional GPU testing.") + parser = util.add_resource_arguments(arg_parser=parser) + parser.add_argument("--ray_address", type=str, default="auto", help="the Ray address.") + parser.add_argument( + "--test", + action="store_true", + help=( + "Run nvidia-smi test instead of the arbitrary job," + "can use as a sanity check prior to any jobs to check " + "that GPU resources are correctly isolated." + ), + ) + parser.add_argument( + "--sub_jobs", + type=str, + nargs=argparse.REMAINDER, + help="This should be last wrapper argument. Jobs separated by the + delimiter to run on a cluster.", + ) + args = parser.parse_args() + if args.sub_jobs is not None: + jobs = " ".join(args.sub_jobs) + formatted_jobs = jobs.split("+") + else: + formatted_jobs = [] + print(f"[INFO]: Isaac Ray Wrapper received jobs {formatted_jobs=}") + wrap_resources_to_jobs(jobs=formatted_jobs, args=args) diff --git a/source/standalone/workflows/rl_games/play.py b/source/standalone/workflows/rl_games/play.py index d16e46b2ef..7aa1456fe0 100644 --- a/source/standalone/workflows/rl_games/play.py +++ b/source/standalone/workflows/rl_games/play.py @@ -155,7 +155,7 @@ def main(): # convert obs to agent format obs = agent.obs_to_torch(obs) # agent stepping - actions = agent.get_action(obs, is_deterministic=True) + actions = agent.get_action(obs, is_deterministic=agent.is_deterministic) # env stepping obs, _, dones, _ = env.step(actions) diff --git a/tools/per_test_timeouts.py b/tools/per_test_timeouts.py index e98595372f..340abde85c 100644 --- a/tools/per_test_timeouts.py +++ b/tools/per_test_timeouts.py @@ -9,7 +9,8 @@ """ PER_TEST_TIMEOUTS = { "test_articulation.py": 200, - "test_environments.py": 1200, # This test runs through all the environments for 100 steps each + "test_deformable_object.py": 200, + "test_environments.py": 1500, # This test runs through all the environments for 100 steps each "test_environment_determinism.py": 200, # This test runs through many the environments for 100 steps each "test_env_rendering_logic.py": 300, "test_camera.py": 500, @@ -17,4 +18,5 @@ "test_rsl_rl_wrapper.py": 200, "test_sb3_wrapper.py": 200, "test_skrl_wrapper.py": 200, + "test_operational_space.py": 200, } diff --git a/tools/run_all_tests.py b/tools/run_all_tests.py index 733af27c02..bfed59e343 100644 --- a/tools/run_all_tests.py +++ b/tools/run_all_tests.py @@ -23,6 +23,7 @@ import argparse import logging import os +import re import subprocess import sys import time @@ -118,57 +119,9 @@ def test_all( # Set up logger logging.basicConfig(level=logging.INFO, format="%(message)s", handlers=logging_handlers) - # Discover all tests under current directory - all_test_paths = [str(path) for path in Path(test_dir).resolve().rglob("*test_*.py")] - skipped_test_paths = [] - test_paths = [] - # Check that all tests to skip are actually in the tests - for test_to_skip in tests_to_skip: - for test_path in all_test_paths: - if test_to_skip in test_path: - break - else: - raise ValueError(f"Test to skip '{test_to_skip}' not found in tests.") - - # Filter tests by extension - if extension is not None: - all_tests_in_selected_extension = [] - - for test_path in all_test_paths: - # Extract extension name from test path - extension_name = test_path[test_path.find("extensions") :].split("/")[1] - - # Skip tests that are not in the selected extension - if extension_name != extension: - continue - - all_tests_in_selected_extension.append(test_path) - - all_test_paths = all_tests_in_selected_extension - - # Remove tests to skip from the list of tests to run - if len(tests_to_skip) != 0: - for test_path in all_test_paths: - if any([test_to_skip in test_path for test_to_skip in tests_to_skip]): - skipped_test_paths.append(test_path) - else: - test_paths.append(test_path) - else: - test_paths = all_test_paths - - # Sort test paths so they're always in the same order - all_test_paths.sort() - test_paths.sort() - skipped_test_paths.sort() - - # Initialize all tests to have the same timeout - test_timeouts = {test_path: timeout for test_path in all_test_paths} - - # Overwrite timeouts for specific tests - for test_path_with_timeout, test_timeout in per_test_timeouts.items(): - for test_path in all_test_paths: - if test_path_with_timeout in test_path: - test_timeouts[test_path] = test_timeout + all_test_paths, test_paths, skipped_test_paths, test_timeouts = extract_tests_and_timeouts( + test_dir, extension, tests_to_skip, timeout, per_test_timeouts + ) # Print tests to be run logging.info("\n" + "=" * 60 + "\n") @@ -213,36 +166,31 @@ def test_all( except Exception as e: logging.error(f"Unexpected exception {e}. Please report this issue on the repository.") result = "FAILED" - stdout = str(e) - stderr = str(e) + stdout = None + stderr = None else: - # Should only get here if the process ran successfully, e.g. no exceptions were raised - # but we still check the returncode just in case - result = "PASSED" if completed_process.returncode == 0 else "FAILED" + result = "COMPLETED" stdout = completed_process.stdout stderr = completed_process.stderr after = time.time() time_elapsed = after - before - # Decode stdout and stderr and write to file and print to console if desired - if stdout is not None: - if isinstance(stdout, str): - stdout_str = stdout - else: - stdout_str = stdout.decode("utf-8") - else: - stdout_str = "" - if stderr is not None: - if isinstance(stderr, str): - stderr_str = stderr + + # Decode stdout and stderr + stdout = stdout.decode("utf-8") if stdout is not None else "" + stderr = stderr.decode("utf-8") if stderr is not None else "" + + if result == "COMPLETED": + # Check for success message in the output + success_pattern = r"Ran \d+ tests? in [\d.]+s\s+OK" + if re.search(success_pattern, stdout) or re.search(success_pattern, stderr): + result = "PASSED" else: - stderr_str = stderr.decode("utf-8") - else: - stderr_str = "" + result = "FAILED" # Write to log file - logging.info(stdout_str) - logging.info(stderr_str) + logging.info(stdout) + logging.info(stderr) logging.info(f"[INFO] Time elapsed: {time_elapsed:.2f} s") logging.info(f"[INFO] Result '{test_path}': {result}") # Collect results @@ -307,8 +255,89 @@ def test_all( return num_failing + num_timing_out == 0 +def extract_tests_and_timeouts( + test_dir: str, + extension: str | None = None, + tests_to_skip: list[str] = [], + timeout: float = DEFAULT_TIMEOUT, + per_test_timeouts: dict[str, float] = {}, +) -> tuple[list[str], list[str], list[str], dict[str, float]]: + """Extract all tests under the given directory or extension and their respective timeouts. + + Args: + test_dir: Path to the directory containing the tests. + extension: Run tests only for the given extension. Defaults to None, which means all extensions' + tests will be run. + tests_to_skip: List of tests to skip. + timeout: Timeout for each test in seconds. Defaults to DEFAULT_TIMEOUT. + per_test_timeouts: A dictionary of tests and their timeouts in seconds. Any tests not listed here will use the + timeout specified by `timeout`. Defaults to an empty dictionary. + + Returns: + A tuple containing the paths of all tests, tests to run, tests to skip, and their respective timeouts. + + Raises: + ValueError: If any test to skip is not found under the given `test_dir`. + """ + + # Discover all tests under current directory + all_test_paths = [str(path) for path in Path(test_dir).resolve().rglob("*test_*.py")] + skipped_test_paths = [] + test_paths = [] + # Check that all tests to skip are actually in the tests + for test_to_skip in tests_to_skip: + for test_path in all_test_paths: + if test_to_skip in test_path: + break + else: + raise ValueError(f"Test to skip '{test_to_skip}' not found in tests.") + + # Filter tests by extension + if extension is not None: + all_tests_in_selected_extension = [] + + for test_path in all_test_paths: + # Extract extension name from test path + extension_name = test_path[test_path.find("extensions") :].split("/")[1] + + # Skip tests that are not in the selected extension + if extension_name != extension: + continue + + all_tests_in_selected_extension.append(test_path) + + all_test_paths = all_tests_in_selected_extension + + # Remove tests to skip from the list of tests to run + if len(tests_to_skip) != 0: + for test_path in all_test_paths: + if any([test_to_skip in test_path for test_to_skip in tests_to_skip]): + skipped_test_paths.append(test_path) + else: + test_paths.append(test_path) + else: + test_paths = all_test_paths + + # Sort test paths so they're always in the same order + all_test_paths.sort() + test_paths.sort() + skipped_test_paths.sort() + + # Initialize all tests to have the same timeout + test_timeouts = {test_path: timeout for test_path in all_test_paths} + + # Overwrite timeouts for specific tests + for test_path_with_timeout, test_timeout in per_test_timeouts.items(): + for test_path in all_test_paths: + if test_path_with_timeout in test_path: + test_timeouts[test_path] = test_timeout + + return all_test_paths, test_paths, skipped_test_paths, test_timeouts + + def warm_start_app(): """Warm start the app to compile shaders before running the tests.""" + print("[INFO] Warm starting the simulation app before running tests.") before = time.time() # headless experience