diff --git a/.flake8 b/.flake8 new file mode 100644 index 00000000..bf8ef702 --- /dev/null +++ b/.flake8 @@ -0,0 +1,23 @@ +[flake8] +show-source=True +statistics=True +per-file-ignores=*/__init__.py:F401 +# E402: Module level import not at top of file +# E501: Line too long +# W503: Line break before binary operator +# W605: Invalid escape sequence +# E203: Whitespace before ':' -> conflicts with black +# D401: First line should be in imperative mood +# R504: Unnecessary variable assignment before return statement. +# R505: Unnecessary elif after return statement +# SIM102: Use a single if-statement instead of nested if-statements +# SIM117: Merge with statements for context managers that have same scope. +ignore=E402,E501,W503,W605,E203,D401,R504,R505,SIM102,SIM117 +max-line-length = 120 +max-complexity = 30 +exclude=_*,.vscode,.git,docs/** +# docstrings +docstring-convention=google +# annotations +suppress-none-returning=True +allow-star-arg-any=True diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 00000000..827ccf18 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,18 @@ +# copied and modified from https://github.com/NVIDIA/warp/blob/main/.gitattributes +* text=auto +*.sh text eol=LF + +# copied from https://github.com/isaac-orbit/isaaclab_assets/blob/main/.gitattributes +*.usd filter=lfs diff=lfs merge=lfs -text +*.dae filter=lfs diff=lfs merge=lfs -text +*.mtl filter=lfs diff=lfs merge=lfs -text +*.obj filter=lfs diff=lfs merge=lfs -text +*.gif filter=lfs diff=lfs merge=lfs -text +*.png filter=lfs diff=lfs merge=lfs -text +*.jpg filter=lfs diff=lfs merge=lfs -text +*.psd filter=lfs diff=lfs merge=lfs -text +*.mp4 filter=lfs diff=lfs merge=lfs -text +*.usda filter=lfs diff=lfs merge=lfs -text +*.hdr filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.jit filter=lfs diff=lfs merge=lfs -text diff --git a/.github/workflows/docs.yaml b/.github/workflows/docs.yaml new file mode 100644 index 00000000..881d2a6d --- /dev/null +++ b/.github/workflows/docs.yaml @@ -0,0 +1,81 @@ +name: Build & deploy docs + +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: "UW-Lab/UWLab" + BRANCH_REF: "refs/heads/main" + 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: 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: Check branch docs building + working-directory: ./docs + if: needs.check-secrets.outputs.trigger-deploy != 'true' + run: make current-docs + + - 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: + name: Deploy Docs + runs-on: ubuntu-latest + needs: [check-secrets, build-docs] + if: needs.check-secrets.outputs.trigger-deploy == 'true' + + steps: + - 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/.github/workflows/pre-commit.yaml b/.github/workflows/pre-commit.yaml new file mode 100644 index 00000000..d34fcde4 --- /dev/null +++ b/.github/workflows/pre-commit.yaml @@ -0,0 +1,14 @@ +name: Run linters using pre-commit + +on: + pull_request: + push: + branches: [main] + +jobs: + pre-commit: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - uses: actions/setup-python@v3 + - uses: pre-commit/action@v3.0.0 diff --git a/.gitignore b/.gitignore new file mode 100644 index 00000000..e0b18488 --- /dev/null +++ b/.gitignore @@ -0,0 +1,68 @@ +# C++ +**/cmake-build*/ +**/build*/ +**/*.so +**/*.log* + +# Omniverse +**/*.dmp +**/.thumbs + +# No USD files allowed in the repo +**/*.usd +**/*.usda +**/*.usdc +**/*.usdz + +# Python +.DS_Store +**/*.egg-info/ +**/__pycache__/ +**/.pytest_cache/ +**/*.pyc +**/*.pb + +# Docker/Singularity +**/*.sif +docker/exports/ +docker/.container.yaml + +# IDE +**/.idea/ +**/.vscode/ +# Don't ignore the top-level .vscode directory as it is +# used to configure VS Code settings + +code_log.md + +# Outputs +**/output/* +**/outputs/* +**/videos/* +**/wandb/* +**/.neptune/* +docker/artifacts/ +*.tmp + +# Doc Outputs +**/docs/_build/* +**/generated/* + +# Isaac-Sim packman +_isaac_sim* +_repo +_build +.lastformat + +# RL-Games +**/runs/* +**/logs/* +**/recordings/* +**/datasets/ + +# node generated files +**/node_modules/* +**/dist/* + +# credentials +**/credentials/* diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 00000000..3d355d02 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,49 @@ +repos: + - repo: https://github.com/python/black + rev: 23.10.1 + hooks: + - id: black + args: ["--line-length", "120", "--preview"] + - repo: https://github.com/pycqa/flake8 + rev: 6.1.0 + hooks: + - id: flake8 + additional_dependencies: [flake8-simplify, flake8-return] + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.5.0 + hooks: + - id: trailing-whitespace + - id: check-symlinks + - id: destroyed-symlinks + - id: check-yaml + - id: check-merge-conflict + - id: check-case-conflict + - id: check-executables-have-shebangs + - id: check-toml + - id: end-of-file-fixer + - id: check-shebang-scripts-are-executable + - id: detect-private-key + - id: debug-statements + - repo: https://github.com/pycqa/isort + rev: 5.12.0 + hooks: + - id: isort + name: isort (python) + args: ["--profile", "black", "--filter-files"] + - repo: https://github.com/asottile/pyupgrade + rev: v3.15.0 + hooks: + - id: pyupgrade + args: ["--py37-plus"] + - repo: https://github.com/codespell-project/codespell + rev: v2.2.6 + hooks: + - id: codespell + additional_dependencies: + - tomli + - repo: https://github.com/pre-commit/pygrep-hooks + rev: v1.10.0 + hooks: + - id: rst-backticks + - id: rst-directive-colons + - id: rst-inline-touching-normal diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md new file mode 100644 index 00000000..73998c5d --- /dev/null +++ b/CONTRIBUTORS.md @@ -0,0 +1,27 @@ +# UW Lab Developers and Contributors + +This is the official list of UW Lab Project developers and contributors. + +To see the full list of contributors, please check the revision history in the source control. + +Guidelines for modifications: + +* Please keep the **lists sorted alphabetically**. +* Names should be added to this file as: *individual names* or *organizations*. +* E-mail addresses are tracked elsewhere to avoid spam. + +## Developers + +* NVIDIA Corporation & Affiliates +* Shanghai Jiao Tong University +* University of Washington + + +--- + +* Feng Yu +* Mateo Guaman Castro +* Patrick Yin +* Quanquan Peng +* Rosario Scalise +* Zhengyu Zhang diff --git a/LICENCE b/LICENCE new file mode 100644 index 00000000..3e687eaa --- /dev/null +++ b/LICENCE @@ -0,0 +1,30 @@ +Copyright (c) 2022-2025, The UW Lab Project Developers. + +All rights reserved. + +SPDX-License-Identifier: BSD-3-Clause + +Redistribution and use in source and binary forms, with or without modification, +are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its contributors + may be used to endorse or promote products derived from this software without + specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md index 48d112ad..dfca2c67 100644 --- a/README.md +++ b/README.md @@ -1 +1,73 @@ -# UWLab +![Isaac Lab](docs/source/_static/uwlab.png) + +--- + +# UW Lab + +[![IsaacSim](https://img.shields.io/badge/IsaacSim-4.5.0-silver.svg)](https://docs.isaacsim.omniverse.nvidia.com/latest/index.html) +[![Python](https://img.shields.io/badge/python-3.10-blue.svg)](https://docs.python.org/3/whatsnew/3.10.html) +[![Linux platform](https://img.shields.io/badge/platform-linux--64-orange.svg)](https://releases.ubuntu.com/20.04/) +[![Windows platform](https://img.shields.io/badge/platform-windows--64-orange.svg)](https://www.microsoft.com/en-us/) +[![pre-commit](https://img.shields.io/github/actions/workflow/status/isaac-sim/IsaacLab/pre-commit.yaml?logo=pre-commit&logoColor=white&label=pre-commit&color=brightgreen)](https://github.com/isaac-sim/IsaacLab/actions/workflows/pre-commit.yaml) +[![docs status](https://img.shields.io/github/actions/workflow/status/isaac-sim/IsaacLab/docs.yaml?label=docs&color=brightgreen)](https://github.com/isaac-sim/IsaacLab/actions/workflows/docs.yaml) +[![License](https://img.shields.io/badge/license-BSD--3-yellow.svg)](https://opensource.org/licenses/BSD-3-Clause) +[![License](https://img.shields.io/badge/license-Apache--2.0-yellow.svg)](https://opensource.org/license/apache-2-0) + +## Overview + +**UW Lab**, open source projects built upon the robust foundation established by Isaac Lab, organized by UW(University of Washington) Robotic Students, aims to creates accelerated, unified, organized research in the field of robotics simulation in Isaac Ecosystem. This repo is designed and structured to reuse toolkit of ongoing IsaacLab development, track IsaacLab version, and extend the components from novelty and engineers from UW. + +## Key Features + +In addition to what IsaacLab provides, UW Lab brings: + +- **Environments**: Cleaned Implementation of reputable environments in Manager-Based format +- **Sim to Real**: Providing robots and configuration that has been tested in UW Robotic Lab amd deliver the Simulation Setup that can directly transfer to reals + + +## 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) +- [Available environments](https://isaac-sim.github.io/IsaacLab/main/source/overview/environments.html) + + +## Contributing to UW Lab + +Please refer to Isaac Lab +[contribution guideline](https://isaac-sim.github.io/IsaacLab/main/source/refs/contributing.html). + + +## Troubleshooting + +Please for bug and troubleshooting [submit an issue](https://github.com/UW-Lab/UWLab/issues). + +For issues related to Isaac Sim, we recommend checking its [documentation](https://docs.omniverse.nvidia.com/app_isaacsim/app_isaacsim/overview.html) +or opening a question on its [forums](https://forums.developer.nvidia.com/c/agx-autonomous-machines/isaac/67). + +## Support + +* Please use GitHub [Discussions](https://github.com/UW-Lab/UWLab/discussions) for discussing ideas, asking questions, and requests for new features. +* Github [Issues](https://github.com/UW-Lab/UWLab/issues) should only be used to track executable pieces of work with a definite scope and a clear deliverable. These can be fixing bugs, documentation issues, new features, or general updates. + +## License + +UW Lab is released under [BSD-3 License](LICENSE). The Isaac Lab framework is released under [BSD-3 License](LICENSE). + +## Acknowledgement + +UW Lab development initialed from the Isaac Lab, and closely track the development of Isaac Lab. As gratitude we appreciate if you cite Isaac Lab in academic publications: + +``` +@article{mittal2023orbit, + author={Mittal, Mayank and Yu, Calvin and Yu, Qinxi and Liu, Jingzhou and Rudin, Nikita and Hoeller, David and Yuan, Jia Lin and Singh, Ritvik and Guo, Yunrong and Mazhar, Hammad and Mandlekar, Ajay and Babich, Buck and State, Gavriel and Hutter, Marco and Garg, Animesh}, + journal={IEEE Robotics and Automation Letters}, + title={Orbit: A Unified Simulation Framework for Interactive Robot Learning Environments}, + year={2023}, + volume={8}, + number={6}, + pages={3740-3747}, + doi={10.1109/LRA.2023.3270034} +} +``` diff --git a/VERSION b/VERSION new file mode 100644 index 00000000..ac39a106 --- /dev/null +++ b/VERSION @@ -0,0 +1 @@ +0.9.0 diff --git a/docs/Makefile b/docs/Makefile new file mode 100644 index 00000000..ce33dad5 --- /dev/null +++ b/docs/Makefile @@ -0,0 +1,18 @@ +# Minimal makefile for Sphinx documentation +# + +# You can set these variables from the command line, and also +# from the environment for the first two. +SPHINXOPTS ?= +SPHINXBUILD ?= sphinx-build +SOURCEDIR = . +BUILDDIR = _build + +.PHONY: multi-docs +multi-docs: + @sphinx-multiversion "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) + @cp _redirect/index.html $(BUILDDIR)/index.html + +.PHONY: current-docs +current-docs: + @$(SPHINXBUILD) "$(SOURCEDIR)" "$(BUILDDIR)/current" $(SPHINXOPTS) diff --git a/docs/README.md b/docs/README.md new file mode 100644 index 00000000..69a77a48 --- /dev/null +++ b/docs/README.md @@ -0,0 +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 and generating our documentation. + +> **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. + +## Current-Version Documentation + +This section describes how to build the documentation for the current version of the project. + +
+Linux + +```bash +# 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 +xdg-open _build/current/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 00000000..5208597e --- /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 00000000..eb67be60 --- /dev/null +++ b/docs/_templates/versioning.html @@ -0,0 +1,21 @@ +{% if versions %} + +{% endif %} diff --git a/docs/conf.py b/docs/conf.py new file mode 100644 index 00000000..8932e479 --- /dev/null +++ b/docs/conf.py @@ -0,0 +1,291 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +# Configuration file for the Sphinx documentation builder. +# +# This file only contains a selection of the most common options. For a full +# list see the documentation: +# https://www.sphinx-doc.org/en/master/usage/configuration.html + +# -- Path setup -------------------------------------------------------------- + +# If extensions (or modules to document with autodoc) are in another directory, +# add these directories to sys.path here. If the directory is relative to the +# documentation root, use os.path.abspath to make it absolute, like shown here. +# +import os +import sys + +sys.path.insert(0, os.path.abspath("../source/uwlab")) +sys.path.insert(0, os.path.abspath("../source/uwlab/uwlab")) +sys.path.insert(0, os.path.abspath("../source/uwlab_tasks")) +sys.path.insert(0, os.path.abspath("../source/uwlab_tasks/uwlab_tasks")) +sys.path.insert(0, os.path.abspath("../source/uwlab_apps")) +sys.path.insert(0, os.path.abspath("../source/uwlab_apps/uwlab_apps")) + +# -- Project information ----------------------------------------------------- + +project = "UW Lab" +copyright = "2022-2025, The UW Lab Project Developers." +author = "The UW Lab Project Developers." + +# Read version from the package +with open(os.path.join(os.path.dirname(__file__), "..", "VERSION")) as f: + full_version = f.read().strip() + version = ".".join(full_version.split(".")[:3]) + +# -- General configuration --------------------------------------------------- + +# Add any Sphinx extension module names here, as strings. They can be +# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom +# ones. +extensions = [ + "autodocsumm", + "myst_parser", + "sphinx.ext.napoleon", + "sphinxemoji.sphinxemoji", + "sphinx.ext.autodoc", + "sphinx.ext.autosummary", + "sphinx.ext.githubpages", + "sphinx.ext.intersphinx", + "sphinx.ext.mathjax", + "sphinx.ext.todo", + "sphinx.ext.viewcode", + "sphinxcontrib.bibtex", + "sphinxcontrib.icon", + "sphinx_copybutton", + "sphinx_design", + "sphinx_tabs.tabs", # backwards compatibility for building docs on v1.0.0 + "sphinx_multiversion", +] + +# mathjax hacks +mathjax3_config = { + "tex": { + "inlineMath": [["\\(", "\\)"]], + "displayMath": [["\\[", "\\]"]], + }, +} + +# panels hacks +panels_add_bootstrap_css = False +panels_add_fontawesome_css = True + +# supported file extensions for source files +source_suffix = { + ".rst": "restructuredtext", + ".md": "markdown", +} + +# make sure we don't have any unknown references +# TODO: Enable this by default once we have fixed all the warnings +# nitpicky = True + +# put type hints inside the signature instead of the description (easier to maintain) +autodoc_typehints = "signature" +# autodoc_typehints_format = "fully-qualified" +# document class *and* __init__ methods +autoclass_content = "class" # +# separate class docstring from __init__ docstring +autodoc_class_signature = "separated" +# sort members by source order +autodoc_member_order = "bysource" +# inherit docstrings from base classes +autodoc_inherit_docstrings = True +# BibTeX configuration +bibtex_bibfiles = ["source/_static/refs.bib"] +# generate autosummary even if no references +autosummary_generate = True +autosummary_generate_overwrite = False +# default autodoc settings +autodoc_default_options = { + "autosummary": True, +} + +# generate links to the documentation of objects in external projects +intersphinx_mapping = { + "python": ("https://docs.python.org/3", None), + "numpy": ("https://numpy.org/doc/stable/", None), + "torch": ("https://pytorch.org/docs/stable/", None), + "isaac": ("https://docs.omniverse.nvidia.com/py/isaacsim", None), + "gymnasium": ("https://gymnasium.farama.org/", None), + "warp": ("https://nvidia.github.io/warp/", None), +} + +# Add any paths that contain templates here, relative to this directory. +templates_path = [] + +# 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", "_redirect", "_templates", "Thumbs.db", ".DS_Store", "README.md", "licenses/*"] + +# Mock out modules that are not available on RTD +autodoc_mock_imports = [ + "torch", + "numpy", + "matplotlib", + "scipy", + "carb", + "warp", + "pxr", + "isaacsim", + "omni", + "pyrealsense2", + "cv2", + "lz4", + "omni.kit", + "omni.log", + "omni.usd", + "omni.client", + "omni.physx", + "omni.physics", + "pxr.PhysxSchema", + "pxr.PhysicsSchemaTools", + "omni.replicator", + "omni.isaac.core", + "omni.isaac.kit", + "omni.isaac.cloner", + "omni.isaac.urdf", + "omni.isaac.version", + "omni.isaac.motion_generation", + "omni.isaac.ui", + "isaacsim", + "isaacsim.core.api", + "isaacsim.core.cloner", + "isaacsim.core.version", + "isaacsim.robot_motion.motion_generation", + "isaacsim.gui.components", + "isaacsim.asset.importer.urdf", + "isaacsim.asset.importer.mjcf", + "omni.syntheticdata", + "omni.timeline", + "omni.ui", + "gym", + "skrl", + "stable_baselines3", + "rsl_rl", + "rl_games", + "ray", + "h5py", + "hid", + "prettytable", + "tqdm", + "tensordict", + "trimesh", + "toml", + "isaaclab", + "isaaclab_assets", + "isaaclab_tasks", + "isaaclab_rl", +] + +# List of zero or more Sphinx-specific warning categories to be squelched (i.e., +# suppressed, ignored). +suppress_warnings = [ + # Generally speaking, we do want Sphinx to inform + # us about cross-referencing failures. Remove this entirely after Sphinx + # resolves this open issue: + # https://github.com/sphinx-doc/sphinx/issues/4961 + # Squelch mostly ignorable warnings resembling: + # WARNING: more than one target found for cross-reference 'TypeHint': + # beartype.door._doorcls.TypeHint, beartype.door.TypeHint + # + # Sphinx currently emits *MANY* of these warnings against our + # documentation. All of these warnings appear to be ignorable. Although we + # could explicitly squelch *SOME* of these warnings by canonicalizing + # relative to absolute references in docstrings, Sphinx emits still others + # of these warnings when parsing PEP-compliant type hints via static + # analysis. Since those hints are actual hints that *CANNOT* by definition + # by canonicalized, our only recourse is to squelch warnings altogether. + "ref.python", +] + +# -- Internationalization ---------------------------------------------------- + +# specifying the natural language populates some key tags +language = "en" + +# -- Options for HTML output ------------------------------------------------- + +import sphinx_book_theme + +html_title = "UW Lab Documentation" +html_theme_path = [sphinx_book_theme.get_html_theme_path()] +html_theme = "sphinx_book_theme" +html_favicon = "source/_static/favicon.ico" +html_show_copyright = True +html_show_sphinx = False +html_last_updated_fmt = "" # to reveal the build date in the pages meta + +# Add any paths that contain custom static files (such as style sheets) here, +# relative to this directory. They are copied after the builtin static files, +# so a file named "default.css" will overwrite the builtin "default.css". +html_static_path = ["source/_static/css"] +html_css_files = ["custom.css"] + +html_theme_options = { + "collapse_navigation": True, + "repository_url": "https://github.com/UW-Lab/UWLab", + "use_repository_button": True, + "use_issues_button": True, + "use_edit_page_button": True, + "show_toc_level": 1, + "use_sidenotes": True, + "logo": { + "text": "UW Lab Documentation", + "image_light": "source/_static/UW-logo-black.png", + "image_dark": "source/_static/UW-logo-white.png", + }, + "icon_links": [ + { + "name": "UWLab", + "url": "https://github.com/UW-Lab/UWLab", + "icon": "fa-brands fa-github-alt", + "type": "fontawesome", + }, + { + "name": "IsaacLab", + "url": "https://github.com/isaac-sim/IsaacLab", + "icon": "fa-brands fa-square-github", + "type": "fontawesome", + }, + { + "name": "Isaac Sim", + "url": "https://developer.nvidia.com/isaac-sim", + "icon": "https://img.shields.io/badge/IsaacSim-4.5.0-silver.svg", + "type": "url", + }, + ], + "icon_links_label": "Quick Links", +} + +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 ------------------------------------------------- + + +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", "validate", "__post_init__"] + if name in exclusions: + return True + return None + + +def setup(app): + app.connect("autodoc-skip-member", skip_member) diff --git a/docs/index.rst b/docs/index.rst new file mode 100644 index 00000000..53a839f6 --- /dev/null +++ b/docs/index.rst @@ -0,0 +1,98 @@ +Overview +======== + +.. figure:: source/_static/cover.png + :width: 100% + :alt: Cover + +Preamble +================== + +**UW Lab** build upon the solid foundation laid by ``Isaac Lab`` / ``NVIDIA Isaac Sim``, +expanding its framework to embrace a broader spectrum of algorithms, robots, and environments. While adhering to the +principles of modularity, agility, openness, and battery-included as Isaac Lab. + +In the short term, our mission is to unify and facilitate the research efforts of our colleagues within a single, +cohesive framework. Looking ahead, UW Lab envisions a future where AI, robotics, and the boundaries between +reality and digital world seamlessly converge, offering profound insights into the interaction and development of +intelligent systems. + +We recognize that this is a long and evolving journey, which is why we place immense value on the journey of +development, prioritizing principled, flexible, and extensible structures over mere results. At UW Lab, we are +committed to crafting a value where the process is as significant as the outcome, fostering innovation that resonates +deeply with our vision. + + +License +======= + +The UW Lab framework is open-sourced under the BSD-3-Clause license. +Please refer to :ref:`license` for more details. + + +Acknowledgement +=============== +UW Lab development initiated from the `Orbit `_ framework. +please cite orbit in academic publications as well in honor of the original authors.: + +.. code:: bibtex + + @article{mittal2023orbit, + author={Mittal, Mayank and Yu, Calvin and Yu, Qinxi and Liu, Jingzhou and Rudin, Nikita and Hoeller, David and Yuan, Jia Lin and Singh, Ritvik and Guo, Yunrong and Mazhar, Hammad and Mandlekar, Ajay and Babich, Buck and State, Gavriel and Hutter, Marco and Garg, Animesh}, + journal={IEEE Robotics and Automation Letters}, + title={Orbit: A Unified Simulation Framework for Interactive Robot Learning Environments}, + year={2023}, + volume={8}, + number={6}, + pages={3740-3747}, + doi={10.1109/LRA.2023.3270034} + } + + +Table of Contents +================= + +.. toctree:: + :maxdepth: 2 + :caption: Getting Started + + source/setup/installation/local_installation + +.. toctree:: + :maxdepth: 1 + :caption: Publications + :titlesonly: + + source/publications/pg1 + +.. toctree:: + :maxdepth: 3 + :caption: Overview + :titlesonly: + + source/overview/isaac_environments + source/overview/uw_environments + +.. toctree:: + :maxdepth: 1 + :caption: References + + source/refs/license + +.. toctree:: + :hidden: + :caption: Project Links + + UW Lab + Isaac Lab + NVIDIA Isaac Sim + NVIDIA PhysX + +Indices and tables +================== + +* :ref:`genindex` +* :ref:`modindex` +* :ref:`search` + +.. _NVIDIA Isaac Sim: https://docs.omniverse.nvidia.com/isaacsim/latest/index.html diff --git a/docs/licenses/assets/anymal_b-license.txt b/docs/licenses/assets/anymal_b-license.txt new file mode 100644 index 00000000..5d62a455 --- /dev/null +++ b/docs/licenses/assets/anymal_b-license.txt @@ -0,0 +1,29 @@ +Copyright 2019, ANYbotics AG. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions +are met: + +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in + the documentation and/or other materials provided with the + distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/docs/licenses/assets/anymal_c-license.txt b/docs/licenses/assets/anymal_c-license.txt new file mode 100644 index 00000000..4defe407 --- /dev/null +++ b/docs/licenses/assets/anymal_c-license.txt @@ -0,0 +1,29 @@ +Copyright 2020, ANYbotics AG. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions +are met: + +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in + the documentation and/or other materials provided with the + distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/docs/licenses/assets/franka-license.txt b/docs/licenses/assets/franka-license.txt new file mode 100644 index 00000000..d9a10c0d --- /dev/null +++ b/docs/licenses/assets/franka-license.txt @@ -0,0 +1,176 @@ + 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 diff --git a/docs/licenses/assets/kinova-license.txt b/docs/licenses/assets/kinova-license.txt new file mode 100644 index 00000000..9b1705a2 --- /dev/null +++ b/docs/licenses/assets/kinova-license.txt @@ -0,0 +1,95 @@ +Copyright (c) 2017, Kinova Robotics inc. + +All rights reserved. + +Redistribution and use in source and binary forms, with or without modification, +are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + * Neither the name of the copyright holder nor the names of its contributors + may be used to endorse or promote products derived from this software + without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +============================================================================== + +Copyright (c) 2018, Kinova inc. + +All rights reserved. + +Redistribution and use in source and binary forms, with or without modification, +are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + * Neither the name of the copyright holder nor the names of its contributors + may be used to endorse or promote products derived from this software + without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +____________________________________________________________________ + + +Protocol Buffer license + +Copyright 2008 Google Inc. All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + + * Redistributions of source code must retain the above copyright +notice, this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above +copyright notice, this list of conditions and the following disclaimer +in the documentation and/or other materials provided with the +distribution. + * Neither the name of Google Inc. nor the names of its +contributors may be used to endorse or promote products derived from +this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +Code generated by the Protocol Buffer compiler is owned by the owner +of the input file used when generating it. This code is not +standalone and requires a support library to be linked with it. This +support library is itself covered by the above license. diff --git a/docs/licenses/assets/robotiq-license.txt b/docs/licenses/assets/robotiq-license.txt new file mode 100644 index 00000000..78c0f735 --- /dev/null +++ b/docs/licenses/assets/robotiq-license.txt @@ -0,0 +1,23 @@ +Copyright (c) 2013, ROS-Industrial +All rights reserved. + +Redistribution and use in source and binary forms, with or without modification, +are permitted provided that the following conditions are met: + + Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + + Redistributions in binary form must reproduce the above copyright notice, this + list of conditions and the following disclaimer in the documentation and/or + other materials provided with the distribution. + +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/assets/unitree-license.txt b/docs/licenses/assets/unitree-license.txt new file mode 100644 index 00000000..e472bbd9 --- /dev/null +++ b/docs/licenses/assets/unitree-license.txt @@ -0,0 +1,29 @@ +BSD 3-Clause License + +Copyright (c) 2016-2022 HangZhou YuShu TECHNOLOGY CO.,LTD. ("Unitree Robotics") +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +* Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/docs/licenses/dependencies/black-license.txt b/docs/licenses/dependencies/black-license.txt new file mode 100644 index 00000000..7a9b891f --- /dev/null +++ b/docs/licenses/dependencies/black-license.txt @@ -0,0 +1,21 @@ +The MIT License (MIT) + +Copyright (c) 2018 Ɓukasz Langa + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/docs/licenses/dependencies/codespell-license.txt b/docs/licenses/dependencies/codespell-license.txt new file mode 100644 index 00000000..d159169d --- /dev/null +++ b/docs/licenses/dependencies/codespell-license.txt @@ -0,0 +1,339 @@ + GNU GENERAL PUBLIC LICENSE + Version 2, June 1991 + + Copyright (C) 1989, 1991 Free Software Foundation, Inc., + 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA + Everyone is permitted to copy and distribute verbatim copies + of this license document, but changing it is not allowed. + + Preamble + + The licenses for most software are designed to take away your +freedom to share and change it. By contrast, the GNU General Public +License is intended to guarantee your freedom to share and change free +software--to make sure the software is free for all its users. This +General Public License applies to most of the Free Software +Foundation's software and to any other program whose authors commit to +using it. (Some other Free Software Foundation software is covered by +the GNU Lesser General Public License instead.) You can apply it to +your programs, too. + + When we speak of free software, we are referring to freedom, not +price. Our General Public Licenses are designed to make sure that you +have the freedom to distribute copies of free software (and charge for +this service if you wish), that you receive source code or can get it +if you want it, that you can change the software or use pieces of it +in new free programs; and that you know you can do these things. + + To protect your rights, we need to make restrictions that forbid +anyone to deny you these rights or to ask you to surrender the rights. +These restrictions translate to certain responsibilities for you if you +distribute copies of the software, or if you modify it. + + For example, if you distribute copies of such a program, whether +gratis or for a fee, you must give the recipients all the rights that +you have. You must make sure that they, too, receive or can get the +source code. And you must show them these terms so they know their +rights. + + We protect your rights with two steps: (1) copyright the software, and +(2) offer you this license which gives you legal permission to copy, +distribute and/or modify the software. + + Also, for each author's protection and ours, we want to make certain +that everyone understands that there is no warranty for this free +software. If the software is modified by someone else and passed on, we +want its recipients to know that what they have is not the original, so +that any problems introduced by others will not reflect on the original +authors' reputations. + + Finally, any free program is threatened constantly by software +patents. We wish to avoid the danger that redistributors of a free +program will individually obtain patent licenses, in effect making the +program proprietary. To prevent this, we have made it clear that any +patent must be licensed for everyone's free use or not licensed at all. + + The precise terms and conditions for copying, distribution and +modification follow. + + GNU GENERAL PUBLIC LICENSE + TERMS AND CONDITIONS FOR COPYING, DISTRIBUTION AND MODIFICATION + + 0. This License applies to any program or other work which contains +a notice placed by the copyright holder saying it may be distributed +under the terms of this General Public License. The "Program", below, +refers to any such program or work, and a "work based on the Program" +means either the Program or any derivative work under copyright law: +that is to say, a work containing the Program or a portion of it, +either verbatim or with modifications and/or translated into another +language. (Hereinafter, translation is included without limitation in +the term "modification".) Each licensee is addressed as "you". + +Activities other than copying, distribution and modification are not +covered by this License; they are outside its scope. The act of +running the Program is not restricted, and the output from the Program +is covered only if its contents constitute a work based on the +Program (independent of having been made by running the Program). +Whether that is true depends on what the Program does. + + 1. You may copy and distribute verbatim copies of the Program's +source code as you receive it, in any medium, provided that you +conspicuously and appropriately publish on each copy an appropriate +copyright notice and disclaimer of warranty; keep intact all the +notices that refer to this License and to the absence of any warranty; +and give any other recipients of the Program a copy of this License +along with the Program. + +You may charge a fee for the physical act of transferring a copy, and +you may at your option offer warranty protection in exchange for a fee. + + 2. You may modify your copy or copies of the Program or any portion +of it, thus forming a work based on the Program, and copy and +distribute such modifications or work under the terms of Section 1 +above, provided that you also meet all of these conditions: + + a) You must cause the modified files to carry prominent notices + stating that you changed the files and the date of any change. + + b) You must cause any work that you distribute or publish, that in + whole or in part contains or is derived from the Program or any + part thereof, to be licensed as a whole at no charge to all third + parties under the terms of this License. + + c) If the modified program normally reads commands interactively + when run, you must cause it, when started running for such + interactive use in the most ordinary way, to print or display an + announcement including an appropriate copyright notice and a + notice that there is no warranty (or else, saying that you provide + a warranty) and that users may redistribute the program under + these conditions, and telling the user how to view a copy of this + License. (Exception: if the Program itself is interactive but + does not normally print such an announcement, your work based on + the Program is not required to print an announcement.) + +These requirements apply to the modified work as a whole. If +identifiable sections of that work are not derived from the Program, +and can be reasonably considered independent and separate works in +themselves, then this License, and its terms, do not apply to those +sections when you distribute them as separate works. But when you +distribute the same sections as part of a whole which is a work based +on the Program, the distribution of the whole must be on the terms of +this License, whose permissions for other licensees extend to the +entire whole, and thus to each and every part regardless of who wrote it. + +Thus, it is not the intent of this section to claim rights or contest +your rights to work written entirely by you; rather, the intent is to +exercise the right to control the distribution of derivative or +collective works based on the Program. + +In addition, mere aggregation of another work not based on the Program +with the Program (or with a work based on the Program) on a volume of +a storage or distribution medium does not bring the other work under +the scope of this License. + + 3. You may copy and distribute the Program (or a work based on it, +under Section 2) in object code or executable form under the terms of +Sections 1 and 2 above provided that you also do one of the following: + + a) Accompany it with the complete corresponding machine-readable + source code, which must be distributed under the terms of Sections + 1 and 2 above on a medium customarily used for software interchange; or, + + b) Accompany it with a written offer, valid for at least three + years, to give any third party, for a charge no more than your + cost of physically performing source distribution, a complete + machine-readable copy of the corresponding source code, to be + distributed under the terms of Sections 1 and 2 above on a medium + customarily used for software interchange; or, + + c) Accompany it with the information you received as to the offer + to distribute corresponding source code. (This alternative is + allowed only for noncommercial distribution and only if you + received the program in object code or executable form with such + an offer, in accord with Subsection b above.) + +The source code for a work means the preferred form of the work for +making modifications to it. For an executable work, complete source +code means all the source code for all modules it contains, plus any +associated interface definition files, plus the scripts used to +control compilation and installation of the executable. However, as a +special exception, the source code distributed need not include +anything that is normally distributed (in either source or binary +form) with the major components (compiler, kernel, and so on) of the +operating system on which the executable runs, unless that component +itself accompanies the executable. + +If distribution of executable or object code is made by offering +access to copy from a designated place, then offering equivalent +access to copy the source code from the same place counts as +distribution of the source code, even though third parties are not +compelled to copy the source along with the object code. + + 4. You may not copy, modify, sublicense, or distribute the Program +except as expressly provided under this License. Any attempt +otherwise to copy, modify, sublicense or distribute the Program is +void, and will automatically terminate your rights under this License. +However, parties who have received copies, or rights, from you under +this License will not have their licenses terminated so long as such +parties remain in full compliance. + + 5. You are not required to accept this License, since you have not +signed it. However, nothing else grants you permission to modify or +distribute the Program or its derivative works. These actions are +prohibited by law if you do not accept this License. Therefore, by +modifying or distributing the Program (or any work based on the +Program), you indicate your acceptance of this License to do so, and +all its terms and conditions for copying, distributing or modifying +the Program or works based on it. + + 6. Each time you redistribute the Program (or any work based on the +Program), the recipient automatically receives a license from the +original licensor to copy, distribute or modify the Program subject to +these terms and conditions. You may not impose any further +restrictions on the recipients' exercise of the rights granted herein. +You are not responsible for enforcing compliance by third parties to +this License. + + 7. If, as a consequence of a court judgment or allegation of patent +infringement or for any other reason (not limited to patent issues), +conditions are imposed on you (whether by court order, agreement or +otherwise) that contradict the conditions of this License, they do not +excuse you from the conditions of this License. If you cannot +distribute so as to satisfy simultaneously your obligations under this +License and any other pertinent obligations, then as a consequence you +may not distribute the Program at all. For example, if a patent +license would not permit royalty-free redistribution of the Program by +all those who receive copies directly or indirectly through you, then +the only way you could satisfy both it and this License would be to +refrain entirely from distribution of the Program. + +If any portion of this section is held invalid or unenforceable under +any particular circumstance, the balance of the section is intended to +apply and the section as a whole is intended to apply in other +circumstances. + +It is not the purpose of this section to induce you to infringe any +patents or other property right claims or to contest validity of any +such claims; this section has the sole purpose of protecting the +integrity of the free software distribution system, which is +implemented by public license practices. Many people have made +generous contributions to the wide range of software distributed +through that system in reliance on consistent application of that +system; it is up to the author/donor to decide if he or she is willing +to distribute software through any other system and a licensee cannot +impose that choice. + +This section is intended to make thoroughly clear what is believed to +be a consequence of the rest of this License. + + 8. If the distribution and/or use of the Program is restricted in +certain countries either by patents or by copyrighted interfaces, the +original copyright holder who places the Program under this License +may add an explicit geographical distribution limitation excluding +those countries, so that distribution is permitted only in or among +countries not thus excluded. In such case, this License incorporates +the limitation as if written in the body of this License. + + 9. The Free Software Foundation may publish revised and/or new versions +of the General Public License from time to time. Such new versions will +be similar in spirit to the present version, but may differ in detail to +address new problems or concerns. + +Each version is given a distinguishing version number. If the Program +specifies a version number of this License which applies to it and "any +later version", you have the option of following the terms and conditions +either of that version or of any later version published by the Free +Software Foundation. If the Program does not specify a version number of +this License, you may choose any version ever published by the Free Software +Foundation. + + 10. If you wish to incorporate parts of the Program into other free +programs whose distribution conditions are different, write to the author +to ask for permission. For software which is copyrighted by the Free +Software Foundation, write to the Free Software Foundation; we sometimes +make exceptions for this. Our decision will be guided by the two goals +of preserving the free status of all derivatives of our free software and +of promoting the sharing and reuse of software generally. + + NO WARRANTY + + 11. BECAUSE THE PROGRAM IS LICENSED FREE OF CHARGE, THERE IS NO WARRANTY +FOR THE PROGRAM, TO THE EXTENT PERMITTED BY APPLICABLE LAW. EXCEPT WHEN +OTHERWISE STATED IN WRITING THE COPYRIGHT HOLDERS AND/OR OTHER PARTIES +PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESSED +OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF +MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE. THE ENTIRE RISK AS +TO THE QUALITY AND PERFORMANCE OF THE PROGRAM IS WITH YOU. SHOULD THE +PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF ALL NECESSARY SERVICING, +REPAIR OR CORRECTION. + + 12. IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING +WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MAY MODIFY AND/OR +REDISTRIBUTE THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, +INCLUDING ANY GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING +OUT OF THE USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED +TO LOSS OF DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY +YOU OR THIRD PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER +PROGRAMS), EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE +POSSIBILITY OF SUCH DAMAGES. + + END OF TERMS AND CONDITIONS + + How to Apply These Terms to Your New Programs + + If you develop a new program, and you want it to be of the greatest +possible use to the public, the best way to achieve this is to make it +free software which everyone can redistribute and change under these terms. + + To do so, attach the following notices to the program. It is safest +to attach them to the start of each source file to most effectively +convey the exclusion of warranty; and each file should have at least +the "copyright" line and a pointer to where the full notice is found. + + + Copyright (C) + + This program is free software; you can redistribute it and/or modify + it under the terms of the GNU General Public License as published by + the Free Software Foundation; either version 2 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU General Public License for more details. + + You should have received a copy of the GNU General Public License along + with this program; if not, write to the Free Software Foundation, Inc., + 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. + +Also add information on how to contact you by electronic and paper mail. + +If the program is interactive, make it output a short notice like this +when it starts in an interactive mode: + + Gnomovision version 69, Copyright (C) year name of author + Gnomovision comes with ABSOLUTELY NO WARRANTY; for details type `show w'. + This is free software, and you are welcome to redistribute it + under certain conditions; type `show c' for details. + +The hypothetical commands `show w' and `show c' should show the appropriate +parts of the General Public License. Of course, the commands you use may +be called something other than `show w' and `show c'; they could even be +mouse-clicks or menu items--whatever suits your program. + +You should also get your employer (if you work as a programmer) or your +school, if any, to sign a "copyright disclaimer" for the program, if +necessary. Here is a sample; alter the names: + + Yoyodyne, Inc., hereby disclaims all copyright interest in the program + `Gnomovision' (which makes passes at compilers) written by James Hacker. + + , 1 April 1989 + Ty Coon, President of Vice + +This General Public License does not permit incorporating your program into +proprietary programs. If your program is a subroutine library, you may +consider it more useful to permit linking proprietary applications with the +library. If this is what you want to do, use the GNU Lesser General +Public License instead of this License. diff --git a/docs/licenses/dependencies/flake8-license.txt b/docs/licenses/dependencies/flake8-license.txt new file mode 100644 index 00000000..e5e3d6f9 --- /dev/null +++ b/docs/licenses/dependencies/flake8-license.txt @@ -0,0 +1,22 @@ +== Flake8 License (MIT) == + +Copyright (C) 2011-2013 Tarek Ziade +Copyright (C) 2012-2016 Ian Cordasco + +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/gym-license.txt b/docs/licenses/dependencies/gym-license.txt new file mode 100644 index 00000000..979a5ce5 --- /dev/null +++ b/docs/licenses/dependencies/gym-license.txt @@ -0,0 +1,34 @@ +The MIT License + +Copyright (c) 2016 OpenAI (https://openai.com) + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. + +# Mujoco models +This work is derived from [MuJuCo models](http://www.mujoco.org/forum/index.php?resources/) used under the following license: +``` +This file is part of MuJoCo. +Copyright 2009-2015 Roboti LLC. +Mujoco :: Advanced physics simulation engine +Source : www.roboti.us +Version : 1.31 +Released : 23Apr16 +Author :: Vikash Kumar +Contacts : kumar@roboti.us +``` diff --git a/docs/licenses/dependencies/h5py-license.txt b/docs/licenses/dependencies/h5py-license.txt new file mode 100644 index 00000000..28ca5627 --- /dev/null +++ b/docs/licenses/dependencies/h5py-license.txt @@ -0,0 +1,30 @@ +Copyright (c) 2008 Andrew Collette and contributors +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the + distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/docs/licenses/dependencies/hidapi-license.txt b/docs/licenses/dependencies/hidapi-license.txt new file mode 100644 index 00000000..538cdf95 --- /dev/null +++ b/docs/licenses/dependencies/hidapi-license.txt @@ -0,0 +1,26 @@ +Copyright (c) 2010, Alan Ott, Signal 11 Software +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + * Neither the name of Signal 11 Software nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +POSSIBILITY OF SUCH DAMAGE. diff --git a/docs/licenses/dependencies/isaacsim-license.txt b/docs/licenses/dependencies/isaacsim-license.txt new file mode 100644 index 00000000..80cff4a4 --- /dev/null +++ b/docs/licenses/dependencies/isaacsim-license.txt @@ -0,0 +1,13 @@ +Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. + +Note: Licenses for assets such as Robots and Props used within these environments can be found inside their respective folders on the Nucleus server where they are hosted. + +For more information: https://docs.omniverse.nvidia.com/app_isaacsim/common/NVIDIA_Omniverse_License_Agreement.html + +For sub-dependencies of Isaac Sim: https://docs.omniverse.nvidia.com/app_isaacsim/common/licenses.html diff --git a/docs/licenses/dependencies/isort-license.txt b/docs/licenses/dependencies/isort-license.txt new file mode 100644 index 00000000..b5083a50 --- /dev/null +++ b/docs/licenses/dependencies/isort-license.txt @@ -0,0 +1,21 @@ +The MIT License (MIT) + +Copyright (c) 2013 Timothy Edmund Crosley + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. diff --git a/docs/licenses/dependencies/legged-gym-license.txt b/docs/licenses/dependencies/legged-gym-license.txt new file mode 100644 index 00000000..92e75a2e --- /dev/null +++ b/docs/licenses/dependencies/legged-gym-license.txt @@ -0,0 +1,31 @@ +Copyright (c) 2021, ETH Zurich, Nikita Rudin +Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES +All rights reserved. + +Redistribution and use in source and binary forms, with or without modification, +are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its contributors + may be used to endorse or promote products derived from this software without + specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +See licenses/assets for license information for assets included in this repository. +See licenses/dependencies for license information of dependencies of this package. diff --git a/docs/licenses/dependencies/matplotlib-license.txt b/docs/licenses/dependencies/matplotlib-license.txt new file mode 100644 index 00000000..0752df48 --- /dev/null +++ b/docs/licenses/dependencies/matplotlib-license.txt @@ -0,0 +1,99 @@ +License agreement for matplotlib versions 1.3.0 and later +========================================================= + +1. This LICENSE AGREEMENT is between the Matplotlib Development Team +("MDT"), and the Individual or Organization ("Licensee") accessing and +otherwise using matplotlib software in source or binary form and its +associated documentation. + +2. Subject to the terms and conditions of this License Agreement, MDT +hereby grants Licensee a nonexclusive, royalty-free, world-wide license +to reproduce, analyze, test, perform and/or display publicly, prepare +derivative works, distribute, and otherwise use matplotlib +alone or in any derivative version, provided, however, that MDT's +License Agreement and MDT's notice of copyright, i.e., "Copyright (c) +2012- Matplotlib Development Team; All Rights Reserved" are retained in +matplotlib alone or in any derivative version prepared by +Licensee. + +3. In the event Licensee prepares a derivative work that is based on or +incorporates matplotlib or any part thereof, and wants to +make the derivative work available to others as provided herein, then +Licensee hereby agrees to include in any such work a brief summary of +the changes made to matplotlib . + +4. MDT is making matplotlib available to Licensee on an "AS +IS" basis. MDT MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR +IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, MDT MAKES NO AND +DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS +FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF MATPLOTLIB +WILL NOT INFRINGE ANY THIRD PARTY RIGHTS. + +5. MDT SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF MATPLOTLIB + FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR +LOSS AS A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING +MATPLOTLIB , OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF +THE POSSIBILITY THEREOF. + +6. This License Agreement will automatically terminate upon a material +breach of its terms and conditions. + +7. Nothing in this License Agreement shall be deemed to create any +relationship of agency, partnership, or joint venture between MDT and +Licensee. This License Agreement does not grant permission to use MDT +trademarks or trade name in a trademark sense to endorse or promote +products or services of Licensee, or any third party. + +8. By copying, installing or otherwise using matplotlib , +Licensee agrees to be bound by the terms and conditions of this License +Agreement. + +License agreement for matplotlib versions prior to 1.3.0 +======================================================== + +1. This LICENSE AGREEMENT is between John D. Hunter ("JDH"), and the +Individual or Organization ("Licensee") accessing and otherwise using +matplotlib software in source or binary form and its associated +documentation. + +2. Subject to the terms and conditions of this License Agreement, JDH +hereby grants Licensee a nonexclusive, royalty-free, world-wide license +to reproduce, analyze, test, perform and/or display publicly, prepare +derivative works, distribute, and otherwise use matplotlib +alone or in any derivative version, provided, however, that JDH's +License Agreement and JDH's notice of copyright, i.e., "Copyright (c) +2002-2011 John D. Hunter; All Rights Reserved" are retained in +matplotlib alone or in any derivative version prepared by +Licensee. + +3. In the event Licensee prepares a derivative work that is based on or +incorporates matplotlib or any part thereof, and wants to +make the derivative work available to others as provided herein, then +Licensee hereby agrees to include in any such work a brief summary of +the changes made to matplotlib. + +4. JDH is making matplotlib available to Licensee on an "AS +IS" basis. JDH MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR +IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, JDH MAKES NO AND +DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS +FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF MATPLOTLIB +WILL NOT INFRINGE ANY THIRD PARTY RIGHTS. + +5. JDH SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF MATPLOTLIB + FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR +LOSS AS A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING +MATPLOTLIB , OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF +THE POSSIBILITY THEREOF. + +6. This License Agreement will automatically terminate upon a material +breach of its terms and conditions. + +7. Nothing in this License Agreement shall be deemed to create any +relationship of agency, partnership, or joint venture between JDH and +Licensee. This License Agreement does not grant permission to use JDH +trademarks or trade name in a trademark sense to endorse or promote +products or services of Licensee, or any third party. + +8. By copying, installing or otherwise using matplotlib, +Licensee agrees to be bound by the terms and conditions of this License +Agreement. diff --git a/docs/licenses/dependencies/numpy-license.txt b/docs/licenses/dependencies/numpy-license.txt new file mode 100644 index 00000000..e1fb614d --- /dev/null +++ b/docs/licenses/dependencies/numpy-license.txt @@ -0,0 +1,30 @@ +Copyright (c) 2005-2022, NumPy Developers. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of the NumPy Developers nor the names of any + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/docs/licenses/dependencies/onnx-license.txt b/docs/licenses/dependencies/onnx-license.txt new file mode 100644 index 00000000..d6456956 --- /dev/null +++ b/docs/licenses/dependencies/onnx-license.txt @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/docs/licenses/dependencies/open3d-license.txt b/docs/licenses/dependencies/open3d-license.txt new file mode 100644 index 00000000..79d16287 --- /dev/null +++ b/docs/licenses/dependencies/open3d-license.txt @@ -0,0 +1,22 @@ +The MIT License (MIT) + +Open3D: www.open3d.org +Copyright (c) 2018-2021 www.open3d.org + +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/pillow-license.txt b/docs/licenses/dependencies/pillow-license.txt new file mode 100644 index 00000000..40aabc32 --- /dev/null +++ b/docs/licenses/dependencies/pillow-license.txt @@ -0,0 +1,30 @@ +The Python Imaging Library (PIL) is + + Copyright © 1997-2011 by Secret Labs AB + Copyright © 1995-2011 by Fredrik Lundh + +Pillow is the friendly PIL fork. It is + + Copyright © 2010-2022 by Alex Clark and contributors + +Like PIL, Pillow is licensed under the open source HPND License: + +By obtaining, using, and/or copying this software and/or its associated +documentation, you agree that you have read, understood, and will comply +with the following terms and conditions: + +Permission to use, copy, modify, and distribute this software and its +associated documentation for any purpose and without fee is hereby granted, +provided that the above copyright notice appears in all copies, and that +both that copyright notice and this permission notice appear in supporting +documentation, and that the name of Secret Labs AB or the author not be +used in advertising or publicity pertaining to distribution of the software +without specific, written prior permission. + +SECRET LABS AB AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS +SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. +IN NO EVENT SHALL SECRET LABS AB OR THE AUTHOR BE LIABLE FOR ANY SPECIAL, +INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE +OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +PERFORMANCE OF THIS SOFTWARE. diff --git a/docs/licenses/dependencies/pre-commit-hooks-license.txt b/docs/licenses/dependencies/pre-commit-hooks-license.txt new file mode 100644 index 00000000..4a071fc5 --- /dev/null +++ b/docs/licenses/dependencies/pre-commit-hooks-license.txt @@ -0,0 +1,19 @@ +Copyright (c) 2014 pre-commit dev team: Anthony Sottile, Ken Struys + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. diff --git a/docs/licenses/dependencies/pre-commit-license.txt b/docs/licenses/dependencies/pre-commit-license.txt new file mode 100644 index 00000000..4a071fc5 --- /dev/null +++ b/docs/licenses/dependencies/pre-commit-license.txt @@ -0,0 +1,19 @@ +Copyright (c) 2014 pre-commit dev team: Anthony Sottile, Ken Struys + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. diff --git a/docs/licenses/dependencies/protobuf-license.txt b/docs/licenses/dependencies/protobuf-license.txt new file mode 100644 index 00000000..19b305b0 --- /dev/null +++ b/docs/licenses/dependencies/protobuf-license.txt @@ -0,0 +1,32 @@ +Copyright 2008 Google Inc. All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + + * Redistributions of source code must retain the above copyright +notice, this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above +copyright notice, this list of conditions and the following disclaimer +in the documentation and/or other materials provided with the +distribution. + * Neither the name of Google Inc. nor the names of its +contributors may be used to endorse or promote products derived from +this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +Code generated by the Protocol Buffer compiler is owned by the owner +of the input file used when generating it. This code is not +standalone and requires a support library to be linked with it. This +support library is itself covered by the above license. diff --git a/docs/licenses/dependencies/pygrep-hooks-license.txt b/docs/licenses/dependencies/pygrep-hooks-license.txt new file mode 100644 index 00000000..b7af5ef2 --- /dev/null +++ b/docs/licenses/dependencies/pygrep-hooks-license.txt @@ -0,0 +1,19 @@ +Copyright (c) 2018 Anthony Sottile + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. diff --git a/docs/licenses/dependencies/pyright-license.txt b/docs/licenses/dependencies/pyright-license.txt new file mode 100644 index 00000000..dd44ffc0 --- /dev/null +++ b/docs/licenses/dependencies/pyright-license.txt @@ -0,0 +1,47 @@ +MIT License + +Copyright (c) 2021 Robert Craigie + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. + + +=============================================================================== + +MIT License + +Pyright - A static type checker for the Python language +Copyright (c) Microsoft Corporation. All rights reserved. + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE diff --git a/docs/licenses/dependencies/pytorch3d-license.txt b/docs/licenses/dependencies/pytorch3d-license.txt new file mode 100644 index 00000000..c55382ff --- /dev/null +++ b/docs/licenses/dependencies/pytorch3d-license.txt @@ -0,0 +1,30 @@ +BSD License + +For PyTorch3D software + +Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved. + +Redistribution and use in source and binary forms, with or without modification, +are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + + * Neither the name Meta nor the names of its contributors may be used to + endorse or promote products derived from this software without specific + prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON +ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/docs/licenses/dependencies/pyupgrade-license.txt b/docs/licenses/dependencies/pyupgrade-license.txt new file mode 100644 index 00000000..522fbe20 --- /dev/null +++ b/docs/licenses/dependencies/pyupgrade-license.txt @@ -0,0 +1,19 @@ +Copyright (c) 2017 Anthony Sottile + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. diff --git a/docs/licenses/dependencies/ray-license.txt b/docs/licenses/dependencies/ray-license.txt new file mode 100644 index 00000000..b813b34b --- /dev/null +++ b/docs/licenses/dependencies/ray-license.txt @@ -0,0 +1,450 @@ + 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. + +-------------------------------------------------------------------------------- + +Code in python/ray/rllib/{evolution_strategies, dqn} adapted from +https://github.com/openai (MIT License) + +Copyright (c) 2016 OpenAI (http://openai.com) + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. + +-------------------------------------------------------------------------------- + +Code in python/ray/rllib/impala/vtrace.py from +https://github.com/deepmind/scalable_agent + +Copyright 2018 Google LLC + +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 + + https://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. + +-------------------------------------------------------------------------------- +Code in python/ray/rllib/ars is adapted from https://github.com/modestyachts/ARS + +Copyright (c) 2018, ARS contributors (Horia Mania, Aurelia Guy, Benjamin Recht) +All rights reserved. + +Redistribution and use of ARS 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. + +------------------ +Code in python/ray/_private/prometheus_exporter.py is adapted from https://github.com/census-instrumentation/opencensus-python/blob/master/contrib/opencensus-ext-prometheus/opencensus/ext/prometheus/stats_exporter/__init__.py + +# Copyright 2018, OpenCensus Authors +# +# 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. + +-------------------------------------------------------------------------------- +Code in python/ray/tests/modin/test_modin and +python/ray/tests/modin/modin_test_utils adapted from: +- http://github.com/modin-project/modin/master/modin/pandas/test/test_general.py +- http://github.com/modin-project/modin/master/modin/pandas/test/utils.py + +Copyright (c) 2018-2020 Modin Developers. + +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 + + https://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. + +-------------------------------------------------------------------------------- +Code in src/ray/util/logging.h is adapted from +https://github.com/google/glog/blob/master/src/glog/logging.h.in + +Copyright (c) 2008, Google Inc. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + + * Redistributions of source code must retain the above copyright +notice, this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above +copyright notice, this list of conditions and the following disclaimer +in the documentation and/or other materials provided with the +distribution. + * Neither the name of Google Inc. nor the names of its +contributors may be used to endorse or promote products derived from +this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +-------------------------------------------------------------------------------- +Code in python/ray/_private/runtime_env/conda_utils.py is adapted from +https://github.com/mlflow/mlflow/blob/master/mlflow/utils/conda.py + +Copyright (c) 2018, Databricks, Inc. +All rights reserved. + +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. + +-------------------------------------------------------------------------------- +Code in python/ray/_private/runtime_env/_clonevirtualenv.py is adapted from https://github.com/edwardgeorge/virtualenv-clone/blob/master/clonevirtualenv.py + +Copyright (c) 2011, Edward George, based on code contained within the +virtualenv project. + +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. + +-------------------------------------------------------------------------------- +Code in python/ray/_private/async_compat.py is adapted from +https://github.com/python-trio/async_generator/blob/master/async_generator/_util.py + +Copyright (c) 2022, Nathaniel J. Smith + +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. + +--------------------------------------------------------------------------------------------------------------- +Code in python/ray/_private/thirdparty/tabulate/tabulate.py is adapted from https://github.com/astanin/python-tabulate/blob/4892c6e9a79638c7897ccea68b602040da9cc7a7/tabulate.py + +Copyright (c) 2011-2020 Sergey Astanin and contributors + +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. + +-------------------------------------------------------------------------------- +Code in python/ray/_private/thirdparty/dacite is adapted from https://github.com/konradhalas/dacite/blob/master/dacite + +Copyright (c) 2018 Konrad HaƂas + +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/rl_games-license.txt b/docs/licenses/dependencies/rl_games-license.txt new file mode 100644 index 00000000..313ca229 --- /dev/null +++ b/docs/licenses/dependencies/rl_games-license.txt @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2019 Denys88 + +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/robomimc-license.txt b/docs/licenses/dependencies/robomimc-license.txt new file mode 100644 index 00000000..934eaa87 --- /dev/null +++ b/docs/licenses/dependencies/robomimc-license.txt @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2021 Stanford Vision and Learning Lab + +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/rsl_rl-license.txt b/docs/licenses/dependencies/rsl_rl-license.txt new file mode 100644 index 00000000..5e09cdaf --- /dev/null +++ b/docs/licenses/dependencies/rsl_rl-license.txt @@ -0,0 +1,141 @@ +Copyright (c) 2021, ETH Zurich, Nikita Rudin +Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES +All rights reserved. + +Redistribution and use in source and binary forms, with or without modification, +are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its contributors + may be used to endorse or promote products derived from this software without + specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +See licenses/dependencies for license information of dependencies of this package. + +=============================================================================== + +Copyright (c) 2005-2021, NumPy Developers. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of the NumPy Developers nor the names of any + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +=============================================================================== + +From PyTorch: + +Copyright (c) 2016- Facebook, Inc (Adam Paszke) +Copyright (c) 2014- Facebook, Inc (Soumith Chintala) +Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert) +Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu) +Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu) +Copyright (c) 2011-2013 NYU (Clement Farabet) +Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston) +Copyright (c) 2006 Idiap Research Institute (Samy Bengio) +Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz) + +=============================================================================== + +From Caffe2: + +Copyright (c) 2016-present, Facebook Inc. All rights reserved. + +All contributions by Facebook: +Copyright (c) 2016 Facebook Inc. + +All contributions by Google: +Copyright (c) 2015 Google Inc. +All rights reserved. + +All contributions by Yangqing Jia: +Copyright (c) 2015 Yangqing Jia +All rights reserved. + +All contributions by Kakao Brain: +Copyright 2019-2020 Kakao Brain + +All contributions from Caffe: +Copyright(c) 2013, 2014, 2015, the respective contributors +All rights reserved. + +All other contributions: +Copyright(c) 2015, 2016 the respective contributors +All rights reserved. + +Caffe2 uses a copyright model similar to Caffe: each contributor holds +copyright over their contributions to Caffe2. The project versioning records +all such contribution and copyright details. If a contributor wants to further +mark their specific copyright on a particular contribution, they should +indicate their copyright solely in the commit message of the change when it is +committed. + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + +3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America + and IDIAP Research Institute nor the names of its contributors may be + used to endorse or promote products derived from this software without + specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +POSSIBILITY OF SUCH DAMAGE. diff --git a/docs/licenses/dependencies/skrl-license.txt b/docs/licenses/dependencies/skrl-license.txt new file mode 100644 index 00000000..198089ec --- /dev/null +++ b/docs/licenses/dependencies/skrl-license.txt @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2021 Toni-SM + +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 00000000..172d6b3f --- /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/tensorboard-license.txt b/docs/licenses/dependencies/tensorboard-license.txt new file mode 100644 index 00000000..15ae4214 --- /dev/null +++ b/docs/licenses/dependencies/tensorboard-license.txt @@ -0,0 +1,203 @@ +Copyright 2017 The TensorFlow Authors. 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 2017, The TensorFlow Authors. + + 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/licenses/dependencies/tensordict-license.txt b/docs/licenses/dependencies/tensordict-license.txt new file mode 100644 index 00000000..426ba7a3 --- /dev/null +++ b/docs/licenses/dependencies/tensordict-license.txt @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) Meta Platforms, Inc. and affiliates. + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + SOFTWARE. diff --git a/docs/licenses/dependencies/torch-license.txt b/docs/licenses/dependencies/torch-license.txt new file mode 100644 index 00000000..69570b3a --- /dev/null +++ b/docs/licenses/dependencies/torch-license.txt @@ -0,0 +1,27 @@ +Copyright (c) 2016, Soumith Chintala, Ronan Collobert, Koray Kavukcuoglu, Clement Farabet +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +* Neither the name of distro nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/docs/licenses/dependencies/walk-these-ways-license.txt b/docs/licenses/dependencies/walk-these-ways-license.txt new file mode 100644 index 00000000..04211683 --- /dev/null +++ b/docs/licenses/dependencies/walk-these-ways-license.txt @@ -0,0 +1,25 @@ +MIT License + +Copyright (c) 2022 MIT Improbable AI Lab + +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. + + +See licenses/legged_gym and licenses/rsl_rl for additional license information for some files in this package. +Files associated with these additional licenses indicate so in the header. diff --git a/docs/licenses/dependencies/warp-license.txt b/docs/licenses/dependencies/warp-license.txt new file mode 100644 index 00000000..9fd8e1dc --- /dev/null +++ b/docs/licenses/dependencies/warp-license.txt @@ -0,0 +1,36 @@ +# NVIDIA Source Code License for Warp + +## 1. Definitions + +“Licensor” means any person or entity that distributes its Work. +“Software” means the original work of authorship made available under this License. +“Work” means the Software and any additions to or derivative works of the Software that are made available under this License. +The terms “reproduce,” “reproduction,” “derivative works,” and “distribution” have the meaning as provided under U.S. copyright law; provided, however, that 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. +Works, including the Software, are “made available” under this License by including in or with the Work either (a) a copyright notice referencing the applicability of this License to the Work, or (b) a copy of this License. + +## 2. License Grant + +2.1 Copyright Grant. Subject to the terms and conditions of this License, each Licensor grants to you a perpetual, worldwide, non-exclusive, royalty-free, copyright license to reproduce, prepare derivative works of, publicly display, publicly perform, sublicense and distribute its Work and any resulting derivative works in any form. + +## 3. Limitations + +3.1 Redistribution. You may reproduce or distribute the Work only if (a) you do so under this License, (b) you include a complete copy of this License with your distribution, and (c) you retain without modification any copyright, patent, trademark, or attribution notices that are present in the Work. + +3.2 Derivative Works. You may specify that additional or different terms apply to the use, reproduction, and distribution of your derivative works of the Work (“Your Terms”) only if (a) Your Terms provide that the use limitation in Section 3.3 applies to your derivative works, and (b) you identify the specific derivative works that are subject to Your Terms. Notwithstanding Your Terms, this License (including the redistribution requirements in Section 3.1) will continue to apply to the Work itself. + +3.3 Use Limitation. The Work and any derivative works thereof only may be used or intended for use non-commercially. Notwithstanding the foregoing, NVIDIA and its affiliates may use the Work and any derivative works commercially. As used herein, “non-commercially” means for research or evaluation purposes only. + +3.4 Patent Claims. If you bring or threaten to bring a patent claim against any Licensor (including any claim, cross-claim or counterclaim in a lawsuit) to enforce any patents that you allege are infringed by any Work, then your rights under this License from such Licensor (including the grant in Section 2.1) will terminate immediately. + +3.5 Trademarks. This License does not grant any rights to use any Licensor’s or its affiliates’ names, logos, or trademarks, except as necessary to reproduce the notices described in this License. + +3.6 Termination. If you violate any term of this License, then your rights under this License (including the grant in Section 2.1) will terminate immediately. + +## 4. Disclaimer of Warranty. + +THE WORK IS PROVIDED “AS IS” WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER THIS LICENSE. + +## 5. Limitation of Liability. + +EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK (INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION, LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER COMM ERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. diff --git a/docs/make.bat b/docs/make.bat new file mode 100644 index 00000000..29fcc2e7 --- /dev/null +++ b/docs/make.bat @@ -0,0 +1,64 @@ +@ECHO OFF + +pushd %~dp0 + +REM Command file to build Sphinx documentation + +set SOURCEDIR=. +set BUILDDIR=_build + +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 +) + +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 +) + +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 new file mode 100644 index 00000000..13b2bfe9 --- /dev/null +++ b/docs/requirements.txt @@ -0,0 +1,18 @@ +# for building the docs +sphinx-book-theme==1.0.1 +myst-parser +sphinxcontrib-bibtex==2.5.0 +autodocsumm +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 +matplotlib +warp-lang +# learning +gymnasium diff --git a/docs/source/_static/UW-logo-black.png b/docs/source/_static/UW-logo-black.png new file mode 100644 index 00000000..ab6bda5e --- /dev/null +++ b/docs/source/_static/UW-logo-black.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2abffa0c3d827e270c3619329e838d4c3623f9a4861aac22fd26a12989101889 +size 160785 diff --git a/docs/source/_static/UW-logo-white.png b/docs/source/_static/UW-logo-white.png new file mode 100644 index 00000000..91b7abe8 --- /dev/null +++ b/docs/source/_static/UW-logo-white.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1bdfe1ae00549a107442f4c82ef8a91fb298c4ca59d8fd09a973cfdf617ffd35 +size 143163 diff --git a/docs/source/_static/actuator-group/actuator-dark.svg b/docs/source/_static/actuator-group/actuator-dark.svg new file mode 100644 index 00000000..b9e06823 --- /dev/null +++ b/docs/source/_static/actuator-group/actuator-dark.svg @@ -0,0 +1,522 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + image/svg+xml + + + + + + + + DC Motor + Actuator Net(MLP/LSTM) + + Gripper + + + Arm + Base + Mimic Group + + + + open/close (1) + joint position(6) + joint position(12) + joint torque(12) + joint torque(6) + joint velocity(6) + + Simulation + + + + Actions + + + + + + + + + + + diff --git a/docs/source/_static/actuator-group/actuator-light.svg b/docs/source/_static/actuator-group/actuator-light.svg new file mode 100644 index 00000000..214b5a7f --- /dev/null +++ b/docs/source/_static/actuator-group/actuator-light.svg @@ -0,0 +1,10214 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + image/svg+xml + + + + + + + + DC Motor + Actuator Net(MLP/LSTM) + + Gripper + + + Arm + Base + Mimic Group + + + + open/close (1) + joint position(6) + joint position(12) + joint torque(12) + joint torque(6) + joint velocity(6) + + Simulation + + + + Actions + + + + + + + + + + + diff --git a/docs/source/_static/cover.png b/docs/source/_static/cover.png new file mode 100644 index 00000000..56a2ae36 --- /dev/null +++ b/docs/source/_static/cover.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e1f95e26072f91e8d98fab031b2e5445b179a2c39652f552c96dbb803d30cfad +size 1154561 diff --git a/docs/source/_static/css/custom.css b/docs/source/_static/css/custom.css new file mode 100644 index 00000000..3a76e124 --- /dev/null +++ b/docs/source/_static/css/custom.css @@ -0,0 +1,104 @@ +/* + * For reference: https://pydata-sphinx-theme.readthedocs.io/en/v0.9.0/user_guide/customizing.html + * For colors: https://clrs.cc/ + */ + @import url('https://fonts.googleapis.com/css2?family=Josefin+Sans:wght@400;700&display=swap'); + + body { + font-family: 'Josefin Sans', sans-serif; + } + + /* Apply the font to headings */ + h1, h2, h3, h4, h5, h6 { + font-family: 'Josefin Sans', sans-serif; + font-weight: 100; /* Adjust weight for headings if desired */ + } + +/* Style footnotes or smaller text if needed */ +.note { + font-size: 0.80em; /* Smaller font size */ + padding: 8px; /* Less padding around the note */ + margin-top: 10px; + border: 1px solid #ddd; /* Lighter border for subtlety */ + background-color: #f9f9f9; /* Softer background color */ +} + +/* anything related to the light theme */ +html[data-theme="light"] { + --pst-color-primary: #76B900; + --pst-color-secondary: #5b8e03; + --pst-color-secondary-highlight: #5b8e03; + --pst-color-inline-code-links: #76B900; + --pst-color-info: var(--pst-color-primary); + --pst-color-info-highlight: var(--pst-color-primary); + --pst-color-info-bg: #daedb9; + --pst-color-attention: #ffc107; + --pst-color-text-base: #323232; + --pst-color-text-muted: #646464; + --pst-color-shadow: #d8d8d8; + --pst-color-border: #c9c9c9; + --pst-color-inline-code: #76B900; + --pst-color-target: #fbe54e; + --pst-color-background: #fff; + --pst-color-on-background: #fff; + --pst-color-surface: #f5f5f5; + --pst-color-on-surface: #e1e1e1; + --pst-color-link: var(--pst-color-primary); + --pst-color-link-hover: #789841; + --pst-color-table-row-hover-bg: #daedb9; + --pst-color-accent: var(--pst-color-primary); +} + +/* anything related to the dark theme */ +html[data-theme="dark"] { + --pst-color-primary: #76B900; + --pst-color-secondary: #c2f26f; + --pst-color-secondary-highlight: #c2f26f; + --pst-color-inline-code-links: #b6e664; + --pst-color-info: var(--pst-color-primary); + --pst-color-info-highlight: var(--pst-color-primary); + --pst-color-info-bg: #3a550b; + --pst-color-attention: #dca90f; + --pst-color-text-base: #cecece; + --pst-color-text-muted: #a6a6a6; + --pst-color-shadow: #212121; + --pst-color-border: silver; + --pst-color-inline-code: #76B900; + --pst-color-target: #472700; + --pst-color-background: #121212; + --pst-color-on-background: #1e1e1e; + --pst-color-surface: #212121; + --pst-color-on-surface: #373737; + --pst-color-link: var(--pst-color-primary); + --pst-color-link-hover: #aee354; + --pst-color-table-row-hover-bg: #3a550b; + --pst-color-accent: var(--pst-color-primary); +} + +a { + text-decoration: none !important; +} + +/* for the announcement link */ +.bd-header-announcement a, +.bd-header-version-warning a { + color: #7FDBFF; +} + +/* for the search box in the navbar */ +.form-control { + border-radius: 0 !important; + border: none !important; + outline: none !important; +} + +/* reduce padding for logo */ +.navbar-brand { + padding-top: 0.0rem !important; + padding-bottom: 0.0rem !important; +} + +.navbar-icon-links { + padding-top: 0.0rem !important; + padding-bottom: 0.0rem !important; +} diff --git a/docs/source/_static/favicon.ico b/docs/source/_static/favicon.ico new file mode 100644 index 00000000..7b39f8af Binary files /dev/null and b/docs/source/_static/favicon.ico differ diff --git a/docs/source/_static/multi-gpu-rl/a3c-dark.svg b/docs/source/_static/multi-gpu-rl/a3c-dark.svg new file mode 100644 index 00000000..48cdb849 --- /dev/null +++ b/docs/source/_static/multi-gpu-rl/a3c-dark.svg @@ -0,0 +1,364 @@ + + + + + + image/svg+xml + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/docs/source/_static/multi-gpu-rl/a3c-light.svg b/docs/source/_static/multi-gpu-rl/a3c-light.svg new file mode 100644 index 00000000..d24bfc20 --- /dev/null +++ b/docs/source/_static/multi-gpu-rl/a3c-light.svg @@ -0,0 +1,385 @@ + + + + + + image/svg+xml + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/docs/source/_static/publications/pg1/pg1.png b/docs/source/_static/publications/pg1/pg1.png new file mode 100644 index 00000000..fe344ad7 --- /dev/null +++ b/docs/source/_static/publications/pg1/pg1.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7040c0cf3fb55fa90ae9a140c66ea94322fddc57e90cec8f6fc1cb9d36d7ae85 +size 634980 diff --git a/docs/source/_static/refs.bib b/docs/source/_static/refs.bib new file mode 100644 index 00000000..e7f82d17 --- /dev/null +++ b/docs/source/_static/refs.bib @@ -0,0 +1,141 @@ +@inproceedings{rudin2022learning, + title={Learning to walk in minutes using massively parallel deep reinforcement learning}, + author={Rudin, Nikita and Hoeller, David and Reist, Philipp and Hutter, Marco}, + booktitle={Conference on Robot Learning}, + pages={91--100}, + year={2022}, + organization={PMLR} +} + +@article{hwangbo2019learning, + title={Learning agile and dynamic motor skills for legged robots}, + author={Hwangbo, Jemin and Lee, Joonho and Dosovitskiy, Alexey and Bellicoso, Dario and Tsounis, Vassilios and Koltun, Vladlen and Hutter, Marco}, + journal={Science Robotics}, + volume={4}, + number={26}, + pages={eaau5872}, + year={2019}, + publisher={American Association for the Advancement of Science} +} + +@article{khatib1987osc, + author={Khatib, O.}, + journal={IEEE Journal on Robotics and Automation}, + title={A unified approach for motion and force control of robot manipulators: The operational space formulation}, + year={1987}, + volume={3}, + number={1}, + pages={43-53}, + doi={10.1109/JRA.1987.1087068} +} + +@book{siciliano2009force, + title={Force control}, + author={Siciliano, Bruno and Sciavicco, Lorenzo and Villani, Luigi and Oriolo, Giuseppe}, + year={2009}, + publisher={Springer} +} + +@article{cheng2021rmpflow, + author={Cheng, Ching-An and Mukadam, Mustafa and Issac, Jan and Birchfield, Stan and Fox, Dieter and Boots, Byron and Ratliff, Nathan}, + journal={IEEE Transactions on Automation Science and Engineering}, + title={RMPflow: A Geometric Framework for Generation of Multitask Motion Policies}, + year={2021}, + volume={18}, + number={3}, + pages={968-987}, + doi={10.1109/TASE.2021.3053422} +} + +@article{buss2004ik, + author = {Buss, Samuel}, + year = {2004}, + pages = {}, + title = {Introduction to inverse kinematics with Jacobian transpose, pseudoinverse and damped least squares methods}, + volume = {17}, + journal={IEEE Transactions in Robotics and Automation}, +} + +@article{sucan2012ompl, + Author = {Ioan A. {\c{S}}ucan and Mark Moll and Lydia E. Kavraki}, + Doi = {10.1109/MRA.2012.2205651}, + Journal = {{IEEE} Robotics \& Automation Magazine}, + Month = {December}, + Number = {4}, + Pages = {72--82}, + Title = {The {O}pen {M}otion {P}lanning {L}ibrary}, + Note = {\url{https://ompl.kavrakilab.org}}, + Volume = {19}, + Year = {2012} +} + +@article{mittal2021articulated, + title={Articulated object interaction in unknown scenes with whole-body mobile manipulation}, + author={Mittal, Mayank and Hoeller, David and Farshidian, Farbod and Hutter, Marco and Garg, Animesh}, + journal={IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)}, + year={2022} +} + +@INPROCEEDINGS{rudin2022advanced, + author={Rudin, Nikita and Hoeller, David and Bjelonic, Marko and Hutter, Marco}, + booktitle={2022 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)}, + title={Advanced Skills by Learning Locomotion and Local Navigation End-to-End}, + year={2022}, + volume={}, + number={}, + pages={2497-2503}, + doi={10.1109/IROS47612.2022.9981198} +} + +@ARTICLE{frankhauser2018probabilistic, + author={Fankhauser, PĂ©ter and Bloesch, Michael and Hutter, Marco}, + journal={IEEE Robotics and Automation Letters}, + title={Probabilistic Terrain Mapping for Mobile Robots With Uncertain Localization}, + year={2018}, + volume={3}, + number={4}, + pages={3019-3026}, + doi={10.1109/LRA.2018.2849506} +} + +@article{makoviychuk2021isaac, + title={Isaac gym: High performance gpu-based physics simulation for robot learning}, + author={Makoviychuk, Viktor and Wawrzyniak, Lukasz and Guo, Yunrong and Lu, Michelle and Storey, Kier and Macklin, Miles and Hoeller, David and Rudin, Nikita and Allshire, Arthur and Handa, Ankur and others}, + journal={arXiv preprint arXiv:2108.10470}, + year={2021} +} + + +@article{handa2022dextreme, + title={DeXtreme: Transfer of Agile In-hand Manipulation from Simulation to Reality}, + author={Handa, Ankur and Allshire, Arthur and Makoviychuk, Viktor and Petrenko, Aleksei and Singh, Ritvik and Liu, Jingzhou and Makoviichuk, Denys and Van Wyk, Karl and Zhurkevich, Alexander and Sundaralingam, Balakumar and others}, + journal={arXiv preprint arXiv:2210.13702}, + year={2022} +} + +@article{narang2022factory, + title={Factory: Fast contact for robotic assembly}, + author={Narang, Yashraj and Storey, Kier and Akinola, Iretiayo and Macklin, Miles and Reist, Philipp and Wawrzyniak, Lukasz and Guo, Yunrong and Moravanszky, Adam and State, Gavriel and Lu, Michelle and others}, + journal={arXiv preprint arXiv:2205.03532}, + year={2022} +} + +@inproceedings{allshire2022transferring, + title={Transferring dexterous manipulation from gpu simulation to a remote real-world trifinger}, + author={Allshire, Arthur and MittaI, Mayank and Lodaya, Varun and Makoviychuk, Viktor and Makoviichuk, Denys and Widmaier, Felix and W{\"u}thrich, Manuel and Bauer, Stefan and Handa, Ankur and Garg, Animesh}, + booktitle={2022 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)}, + pages={11802--11809}, + year={2022}, + organization={IEEE} +} + +@article{mittal2023orbit, + author={Mittal, Mayank and Yu, Calvin and Yu, Qinxi and Liu, Jingzhou and Rudin, Nikita and Hoeller, David and Yuan, Jia Lin and Singh, Ritvik and Guo, Yunrong and Mazhar, Hammad and Mandlekar, Ajay and Babich, Buck and State, Gavriel and Hutter, Marco and Garg, Animesh}, + journal={IEEE Robotics and Automation Letters}, + title={Orbit: A Unified Simulation Framework for Interactive Robot Learning Environments}, + year={2023}, + volume={8}, + number={6}, + pages={3740-3747}, + doi={10.1109/LRA.2023.3270034} +} diff --git a/docs/source/_static/tasks.jpg b/docs/source/_static/tasks.jpg new file mode 100644 index 00000000..780bc7c3 --- /dev/null +++ b/docs/source/_static/tasks.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e5e3c8a8881c532722d69951760fd471b70e69cc1c972793ee565f18133b00ab +size 464725 diff --git a/docs/source/_static/tasks/classic/ant.jpg b/docs/source/_static/tasks/classic/ant.jpg new file mode 100644 index 00000000..3497047b --- /dev/null +++ b/docs/source/_static/tasks/classic/ant.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:43ac3daa07efc112314c9544063897b6c9bfcd7643dba7eb40dfd2b47dbd0f55 +size 234491 diff --git a/docs/source/_static/tasks/classic/cart_double_pendulum.jpg b/docs/source/_static/tasks/classic/cart_double_pendulum.jpg new file mode 100644 index 00000000..5ebc25ec --- /dev/null +++ b/docs/source/_static/tasks/classic/cart_double_pendulum.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4dea89acbef59318dfa099ba48c7408282e18251e88935c023884f635cf915ea +size 219224 diff --git a/docs/source/_static/tasks/classic/cartpole.jpg b/docs/source/_static/tasks/classic/cartpole.jpg new file mode 100644 index 00000000..7ed42994 --- /dev/null +++ b/docs/source/_static/tasks/classic/cartpole.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5ae2cfc71b0c18c41388f5437750daa64c0722c4be869dcf2144c7b747b28eac +size 278750 diff --git a/docs/source/_static/tasks/classic/humanoid.jpg b/docs/source/_static/tasks/classic/humanoid.jpg new file mode 100644 index 00000000..6c74125f --- /dev/null +++ b/docs/source/_static/tasks/classic/humanoid.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1cad36b1be44f6daab6de8fa8d7155ea758b5c5d6f28430c771e49027b9732ce +size 220946 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 00000000..ff15d745 --- /dev/null +++ b/docs/source/_static/tasks/factory/gear_mesh.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5fec8395a26bfab8db6e6d3d2996b868ecc9774433f354fbd02b12aa73b5d4c3 +size 47279 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 00000000..45811da4 --- /dev/null +++ b/docs/source/_static/tasks/factory/nut_thread.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:236d040dd57e82be7db94b719ed593b51bf8f7443a677881d72d21fa7fde5b7a +size 44543 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 00000000..86ee1ba9 --- /dev/null +++ b/docs/source/_static/tasks/factory/peg_insert.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9376e7c43b9417063380d4252fa4de5cb492e22f70c6d99aea50aa081dc1f228 +size 43938 diff --git a/docs/source/_static/tasks/future_plans/deformable/drop.jpg b/docs/source/_static/tasks/future_plans/deformable/drop.jpg new file mode 100644 index 00000000..a2c67876 --- /dev/null +++ b/docs/source/_static/tasks/future_plans/deformable/drop.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:49dbbdfe1ab8aa707924010042e5e054717f9912dfaa8fa9fe6c533058ac0180 +size 25903 diff --git a/docs/source/_static/tasks/future_plans/deformable/flag.jpg b/docs/source/_static/tasks/future_plans/deformable/flag.jpg new file mode 100644 index 00000000..1848fadc --- /dev/null +++ b/docs/source/_static/tasks/future_plans/deformable/flag.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:179ee3a947c8dfc57c3e66270525a8d822833d71b4162efa87b8d5042b3e6e7d +size 17833 diff --git a/docs/source/_static/tasks/future_plans/deformable/fluid_pour.jpg b/docs/source/_static/tasks/future_plans/deformable/fluid_pour.jpg new file mode 100644 index 00000000..a27f3e43 --- /dev/null +++ b/docs/source/_static/tasks/future_plans/deformable/fluid_pour.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c26fd00b3bf13d51f7b77c84b4eb6c3094b489c652f122c8bc0cad0e274bc851 +size 61788 diff --git a/docs/source/_static/tasks/future_plans/deformable/fluid_transport.jpg b/docs/source/_static/tasks/future_plans/deformable/fluid_transport.jpg new file mode 100644 index 00000000..d19a7499 --- /dev/null +++ b/docs/source/_static/tasks/future_plans/deformable/fluid_transport.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0a846a7293c531f12cc750b8afc50b93b7a6e2141afb841e805812dbc51a0b54 +size 89228 diff --git a/docs/source/_static/tasks/future_plans/deformable/pick.jpg b/docs/source/_static/tasks/future_plans/deformable/pick.jpg new file mode 100644 index 00000000..30ce39e2 --- /dev/null +++ b/docs/source/_static/tasks/future_plans/deformable/pick.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b23eb958bff30418969c14ee9be231d358cf62e29bc94394778225d765c9864d +size 37247 diff --git a/docs/source/_static/tasks/future_plans/deformable/place.jpg b/docs/source/_static/tasks/future_plans/deformable/place.jpg new file mode 100644 index 00000000..11e38782 --- /dev/null +++ b/docs/source/_static/tasks/future_plans/deformable/place.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2efc4970b668c97da4dde611f902499131f9a1614cd29fa98f7ef6885ecffd5e +size 31492 diff --git a/docs/source/_static/tasks/future_plans/deformable/rope.jpg b/docs/source/_static/tasks/future_plans/deformable/rope.jpg new file mode 100644 index 00000000..6047cf71 --- /dev/null +++ b/docs/source/_static/tasks/future_plans/deformable/rope.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:219d7084f1ad6b92265a251eb037cdbc5b8363dbfe7fbe62bb053cdf2f97b084 +size 34371 diff --git a/docs/source/_static/tasks/future_plans/deformable/shirt-basket.jpg b/docs/source/_static/tasks/future_plans/deformable/shirt-basket.jpg new file mode 100644 index 00000000..535b02cb --- /dev/null +++ b/docs/source/_static/tasks/future_plans/deformable/shirt-basket.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:772b1347bd8bafe74237c2f4197b905d208c6f81b6afbce83753670ecd2071df +size 27122 diff --git a/docs/source/_static/tasks/future_plans/deformable/shirt.jpg b/docs/source/_static/tasks/future_plans/deformable/shirt.jpg new file mode 100644 index 00000000..28f4dd71 --- /dev/null +++ b/docs/source/_static/tasks/future_plans/deformable/shirt.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b1c6526fa6695afce262bd87a31a5443dcf7554b255fb5d3db8e392bdfde4dc3 +size 40687 diff --git a/docs/source/_static/tasks/future_plans/deformable/stacking.jpg b/docs/source/_static/tasks/future_plans/deformable/stacking.jpg new file mode 100644 index 00000000..45f3464b --- /dev/null +++ b/docs/source/_static/tasks/future_plans/deformable/stacking.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f6ea4825c0eb146caf27d13221f987e2914bbfd07eb9d6847cfc619f8307eeeb +size 24790 diff --git a/docs/source/_static/tasks/future_plans/deformable/sweater.jpg b/docs/source/_static/tasks/future_plans/deformable/sweater.jpg new file mode 100644 index 00000000..a7bf2c47 --- /dev/null +++ b/docs/source/_static/tasks/future_plans/deformable/sweater.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d6d4fc136077e2a3ae05347c72e6a26523993c9236424a3cc7d1542b624d24bf +size 36202 diff --git a/docs/source/_static/tasks/future_plans/deformable/tower_of_hanoi.jpg b/docs/source/_static/tasks/future_plans/deformable/tower_of_hanoi.jpg new file mode 100644 index 00000000..68bae27c --- /dev/null +++ b/docs/source/_static/tasks/future_plans/deformable/tower_of_hanoi.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f264a0dc761103ae922c525a8e5014ce594825e9d54ab726a49c17f689dd7794 +size 28248 diff --git a/docs/source/_static/tasks/future_plans/deformable/vest.jpg b/docs/source/_static/tasks/future_plans/deformable/vest.jpg new file mode 100644 index 00000000..93e91e0a --- /dev/null +++ b/docs/source/_static/tasks/future_plans/deformable/vest.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8d04ea51607d6ebb63b9b1f7dea1955bca2d036629b8978b2fd596207e253a21 +size 44279 diff --git a/docs/source/_static/tasks/future_plans/rigid/beat-the-buzz.jpg b/docs/source/_static/tasks/future_plans/rigid/beat-the-buzz.jpg new file mode 100644 index 00000000..1cb8259a --- /dev/null +++ b/docs/source/_static/tasks/future_plans/rigid/beat-the-buzz.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1e1d05a79a20853a74dbad460846905cf2f1d254935583b5d7800bd4916d4921 +size 24401 diff --git a/docs/source/_static/tasks/future_plans/rigid/cabinet.jpg b/docs/source/_static/tasks/future_plans/rigid/cabinet.jpg new file mode 100644 index 00000000..d903bf0e --- /dev/null +++ b/docs/source/_static/tasks/future_plans/rigid/cabinet.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ff357f2b5dc1471cb4343d320f90abe1b40b2d8c00c8080c60a5c2365932410b +size 21502 diff --git a/docs/source/_static/tasks/future_plans/rigid/hockey.jpg b/docs/source/_static/tasks/future_plans/rigid/hockey.jpg new file mode 100644 index 00000000..ec5f2279 --- /dev/null +++ b/docs/source/_static/tasks/future_plans/rigid/hockey.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0a55efa4340be1b3f4657300e5650fbb3b8c14160dc9b7c7d273b96d6318cf17 +size 32399 diff --git a/docs/source/_static/tasks/future_plans/rigid/jenga.jpg b/docs/source/_static/tasks/future_plans/rigid/jenga.jpg new file mode 100644 index 00000000..fb408e2a --- /dev/null +++ b/docs/source/_static/tasks/future_plans/rigid/jenga.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:667f1c02b099193f2c41eda41501d71e1f12513999b0cb1beee0885bfae3f461 +size 38419 diff --git a/docs/source/_static/tasks/future_plans/rigid/lift.jpg b/docs/source/_static/tasks/future_plans/rigid/lift.jpg new file mode 100644 index 00000000..0c29dd2b --- /dev/null +++ b/docs/source/_static/tasks/future_plans/rigid/lift.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:106603401f1c9c413fc07d001063081c4f649693a03721fc340e0f30c7d07e90 +size 59761 diff --git a/docs/source/_static/tasks/future_plans/rigid/locomotion.jpg b/docs/source/_static/tasks/future_plans/rigid/locomotion.jpg new file mode 100644 index 00000000..fb3ef199 --- /dev/null +++ b/docs/source/_static/tasks/future_plans/rigid/locomotion.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c1f17ef11bb0a83f1065c3446d941a7bc3c5a9940c55ce0709d9553762197354 +size 90707 diff --git a/docs/source/_static/tasks/future_plans/rigid/mobile_cabinet.jpg b/docs/source/_static/tasks/future_plans/rigid/mobile_cabinet.jpg new file mode 100644 index 00000000..27dd6486 --- /dev/null +++ b/docs/source/_static/tasks/future_plans/rigid/mobile_cabinet.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:95279089168f2c6368654cd4437092e91ed0900c6d56de804dabbbcdde7662ee +size 61198 diff --git a/docs/source/_static/tasks/future_plans/rigid/mobile_reach.jpg b/docs/source/_static/tasks/future_plans/rigid/mobile_reach.jpg new file mode 100644 index 00000000..40ad56a5 --- /dev/null +++ b/docs/source/_static/tasks/future_plans/rigid/mobile_reach.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:58fc672d64064f3ff8315aa3f8a977f581f1072637dd6e414faa62674db0c1b2 +size 27568 diff --git a/docs/source/_static/tasks/future_plans/rigid/nut-bolt.jpg b/docs/source/_static/tasks/future_plans/rigid/nut-bolt.jpg new file mode 100644 index 00000000..5f1b1000 --- /dev/null +++ b/docs/source/_static/tasks/future_plans/rigid/nut-bolt.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:dea7f255657161fe3598661d14be7311050fdfa7b552688a46531021114222fa +size 29808 diff --git a/docs/source/_static/tasks/future_plans/rigid/peg-in-hole.jpg b/docs/source/_static/tasks/future_plans/rigid/peg-in-hole.jpg new file mode 100644 index 00000000..90ac266d --- /dev/null +++ b/docs/source/_static/tasks/future_plans/rigid/peg-in-hole.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:962ab4439aff83fd846ba081297f525ab58064428ae1ccdf8796207a7c23b59a +size 23574 diff --git a/docs/source/_static/tasks/future_plans/rigid/pyramid.jpg b/docs/source/_static/tasks/future_plans/rigid/pyramid.jpg new file mode 100644 index 00000000..6ea4c4fa --- /dev/null +++ b/docs/source/_static/tasks/future_plans/rigid/pyramid.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e840225634627d1c1934b80f82958ad07c566e546582b9c2d6a28b796a1028fd +size 36821 diff --git a/docs/source/_static/tasks/future_plans/rigid/reach.jpg b/docs/source/_static/tasks/future_plans/rigid/reach.jpg new file mode 100644 index 00000000..599a7da4 --- /dev/null +++ b/docs/source/_static/tasks/future_plans/rigid/reach.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cd6d05f3fd5f728a7672309c9624e80894442c1fb1f55d6288f7f4dbbf21de5c +size 38995 diff --git a/docs/source/_static/tasks/future_plans/rigid/shadow.jpg b/docs/source/_static/tasks/future_plans/rigid/shadow.jpg new file mode 100644 index 00000000..49eaeabd --- /dev/null +++ b/docs/source/_static/tasks/future_plans/rigid/shadow.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:173cb408d269c023dc9cc42b4b51b5bdbfb351f690cfe4f002411c266ea70f78 +size 73939 diff --git a/docs/source/_static/tasks/locomotion/a1_flat.jpg b/docs/source/_static/tasks/locomotion/a1_flat.jpg new file mode 100644 index 00000000..5c6bb0b4 --- /dev/null +++ b/docs/source/_static/tasks/locomotion/a1_flat.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:368bd1dc3699093ed6478a9e98e86cdd64cf87d7e543bf4cf5f625c04e2a3bc8 +size 155080 diff --git a/docs/source/_static/tasks/locomotion/a1_rough.jpg b/docs/source/_static/tasks/locomotion/a1_rough.jpg new file mode 100644 index 00000000..4e1641e5 --- /dev/null +++ b/docs/source/_static/tasks/locomotion/a1_rough.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:40f088b5c00042302934a9a6cb4f68f1eecb77e4397fd084b32934e5bbf9b371 +size 244321 diff --git a/docs/source/_static/tasks/locomotion/anymal_b_flat.jpg b/docs/source/_static/tasks/locomotion/anymal_b_flat.jpg new file mode 100644 index 00000000..730a6730 --- /dev/null +++ b/docs/source/_static/tasks/locomotion/anymal_b_flat.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:04d0da1516731d518ea77c92a4f8effdb904fc97afa2f770cde15cd20d3b34ca +size 169317 diff --git a/docs/source/_static/tasks/locomotion/anymal_b_rough.jpg b/docs/source/_static/tasks/locomotion/anymal_b_rough.jpg new file mode 100644 index 00000000..555b1af7 --- /dev/null +++ b/docs/source/_static/tasks/locomotion/anymal_b_rough.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3af47c26bd252b49969a58f250c780fd9a010abaa8d30d6ee4cae7bfe2e443d2 +size 254756 diff --git a/docs/source/_static/tasks/locomotion/anymal_c_flat.jpg b/docs/source/_static/tasks/locomotion/anymal_c_flat.jpg new file mode 100644 index 00000000..0d480f79 --- /dev/null +++ b/docs/source/_static/tasks/locomotion/anymal_c_flat.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6ff12b5c7a75975e5475a7a42451487b45d7894d98946aa603ecb96bfc316097 +size 220053 diff --git a/docs/source/_static/tasks/locomotion/anymal_c_rough.jpg b/docs/source/_static/tasks/locomotion/anymal_c_rough.jpg new file mode 100644 index 00000000..32bacdb6 --- /dev/null +++ b/docs/source/_static/tasks/locomotion/anymal_c_rough.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:894e19a454ef6f0718de4abb5289b46c2267a9c26c3c652cb4a502da9b2bc387 +size 310484 diff --git a/docs/source/_static/tasks/locomotion/anymal_d_flat.jpg b/docs/source/_static/tasks/locomotion/anymal_d_flat.jpg new file mode 100644 index 00000000..1d6f595d --- /dev/null +++ b/docs/source/_static/tasks/locomotion/anymal_d_flat.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9df159248847a1230f6694acaf9ca09021b060d65bb3780f6532aaeacedd72e6 +size 245789 diff --git a/docs/source/_static/tasks/locomotion/anymal_d_rough.jpg b/docs/source/_static/tasks/locomotion/anymal_d_rough.jpg new file mode 100644 index 00000000..50bb8686 --- /dev/null +++ b/docs/source/_static/tasks/locomotion/anymal_d_rough.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d113feb51428cfdc86583ed85adb7d60507be30ad25c27f73c2a87d5d534a611 +size 312429 diff --git a/docs/source/_static/tasks/locomotion/g1_flat.jpg b/docs/source/_static/tasks/locomotion/g1_flat.jpg new file mode 100644 index 00000000..624391a0 --- /dev/null +++ b/docs/source/_static/tasks/locomotion/g1_flat.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9b934872a68d4d0b7a548d7056e4185751dffa13e8e7dd1d3c14fd2089c09df0 +size 1163142 diff --git a/docs/source/_static/tasks/locomotion/g1_rough.jpg b/docs/source/_static/tasks/locomotion/g1_rough.jpg new file mode 100644 index 00000000..71c716a3 --- /dev/null +++ b/docs/source/_static/tasks/locomotion/g1_rough.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:69ee9c49e1d84b41cb78527e0a538517bbc9240f692bb5a286b6f078e5728283 +size 1169796 diff --git a/docs/source/_static/tasks/locomotion/go1_flat.jpg b/docs/source/_static/tasks/locomotion/go1_flat.jpg new file mode 100644 index 00000000..3ab9caf5 --- /dev/null +++ b/docs/source/_static/tasks/locomotion/go1_flat.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:621b878e694ec444b01acdd61ba0d150b6f6fa5b86cc502e8ad45d784a1c83b1 +size 144390 diff --git a/docs/source/_static/tasks/locomotion/go1_rough.jpg b/docs/source/_static/tasks/locomotion/go1_rough.jpg new file mode 100644 index 00000000..b5ed7ac5 --- /dev/null +++ b/docs/source/_static/tasks/locomotion/go1_rough.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:15fee781b73a76e01cdb74d837d843e22f26207dac906183553cf11723659de1 +size 216384 diff --git a/docs/source/_static/tasks/locomotion/go2_flat.jpg b/docs/source/_static/tasks/locomotion/go2_flat.jpg new file mode 100644 index 00000000..256bd303 --- /dev/null +++ b/docs/source/_static/tasks/locomotion/go2_flat.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ab93bf91249bbd900d0fbda5a1f4221f01177e64e4ede79fee075c52e5729395 +size 177792 diff --git a/docs/source/_static/tasks/locomotion/go2_rough.jpg b/docs/source/_static/tasks/locomotion/go2_rough.jpg new file mode 100644 index 00000000..7e242aff --- /dev/null +++ b/docs/source/_static/tasks/locomotion/go2_rough.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a167898a57bfc1cfcbefe5ba40b75a6d20cf2d82402ba49439d2c6dac666f793 +size 233147 diff --git a/docs/source/_static/tasks/locomotion/h1_flat.jpg b/docs/source/_static/tasks/locomotion/h1_flat.jpg new file mode 100644 index 00000000..2578a9a9 --- /dev/null +++ b/docs/source/_static/tasks/locomotion/h1_flat.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fef313f7c89e39fc41b02d3d07773b568194b31eb2f012928bc251b3ab889f1c +size 106872 diff --git a/docs/source/_static/tasks/locomotion/h1_rough.jpg b/docs/source/_static/tasks/locomotion/h1_rough.jpg new file mode 100644 index 00000000..f52e9649 --- /dev/null +++ b/docs/source/_static/tasks/locomotion/h1_rough.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7930c748ebea16f91aa61cd17fdff36bda134bb7ec5d4aab6feb737ac4fb209c +size 169799 diff --git a/docs/source/_static/tasks/locomotion/spot_flat.jpg b/docs/source/_static/tasks/locomotion/spot_flat.jpg new file mode 100644 index 00000000..97ff55dd --- /dev/null +++ b/docs/source/_static/tasks/locomotion/spot_flat.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fc249939c5ccae8dd7475e065ee5b6e2751f3ff5739fcaf40face94c96a7c7d7 +size 87832 diff --git a/docs/source/_static/tasks/locomotion/spot_gap.jpg b/docs/source/_static/tasks/locomotion/spot_gap.jpg new file mode 100644 index 00000000..7bcbb29e --- /dev/null +++ b/docs/source/_static/tasks/locomotion/spot_gap.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ce3a265b342a44c55bbe82b82b5380c7ddca42737b91f05af551ab932f85e682 +size 80681 diff --git a/docs/source/_static/tasks/locomotion/spot_obstacle.jpg b/docs/source/_static/tasks/locomotion/spot_obstacle.jpg new file mode 100644 index 00000000..05d56ac3 --- /dev/null +++ b/docs/source/_static/tasks/locomotion/spot_obstacle.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7f3589a45ab3f9ad1322aaf16f51dfe459f73c0af8caa62f3a49ee3a4b8a750d +size 140945 diff --git a/docs/source/_static/tasks/locomotion/spot_pit.jpg b/docs/source/_static/tasks/locomotion/spot_pit.jpg new file mode 100644 index 00000000..f809c076 --- /dev/null +++ b/docs/source/_static/tasks/locomotion/spot_pit.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b147c5ae16f5cffc2dd710178ac5e4dc85fd88c10a64c4c5e7ab06dc093d6957 +size 99166 diff --git a/docs/source/_static/tasks/locomotion/spot_slope.jpg b/docs/source/_static/tasks/locomotion/spot_slope.jpg new file mode 100644 index 00000000..07fea596 --- /dev/null +++ b/docs/source/_static/tasks/locomotion/spot_slope.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0d6815182aa3fbbd647b0654c1fef8f96bb82ec0bc28e8c68bb439903ec86c70 +size 119710 diff --git a/docs/source/_static/tasks/locomotion/spot_stepping_stone.jpg b/docs/source/_static/tasks/locomotion/spot_stepping_stone.jpg new file mode 100644 index 00000000..68083acd --- /dev/null +++ b/docs/source/_static/tasks/locomotion/spot_stepping_stone.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:845c70e73709008ff89842a9ec811cbbfa32281be49510746f33c9022c51e05c +size 113151 diff --git a/docs/source/_static/tasks/manipulation/allegro_cube.jpg b/docs/source/_static/tasks/manipulation/allegro_cube.jpg new file mode 100644 index 00000000..b75b6ae9 --- /dev/null +++ b/docs/source/_static/tasks/manipulation/allegro_cube.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9aa6fbb937a4732ac06baf77c0610cad4bd5a7ba1ef91f3639651560636979ae +size 103465 diff --git a/docs/source/_static/tasks/manipulation/factory_ext/gear_mesh_ext.jpg b/docs/source/_static/tasks/manipulation/factory_ext/gear_mesh_ext.jpg new file mode 100644 index 00000000..3b44c189 --- /dev/null +++ b/docs/source/_static/tasks/manipulation/factory_ext/gear_mesh_ext.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:21fa36e9aeacc4691008866e7e909ac8ae7cf331f139d744647a43063b827292 +size 69853 diff --git a/docs/source/_static/tasks/manipulation/factory_ext/nut_thread_ext.jpg b/docs/source/_static/tasks/manipulation/factory_ext/nut_thread_ext.jpg new file mode 100644 index 00000000..90432bd8 --- /dev/null +++ b/docs/source/_static/tasks/manipulation/factory_ext/nut_thread_ext.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:770c738c2a3603d5f40960c1ff3a8caa58bec5cd606da5a6483322d593761ddd +size 67872 diff --git a/docs/source/_static/tasks/manipulation/factory_ext/peg_insert_ext.jpg b/docs/source/_static/tasks/manipulation/factory_ext/peg_insert_ext.jpg new file mode 100644 index 00000000..ed7c24f4 --- /dev/null +++ b/docs/source/_static/tasks/manipulation/factory_ext/peg_insert_ext.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:72a1b299211c9b8624842a3f596b22a1f79db01177d30684d3ea3c174e7d634c +size 71767 diff --git a/docs/source/_static/tasks/manipulation/franka_lift.jpg b/docs/source/_static/tasks/manipulation/franka_lift.jpg new file mode 100644 index 00000000..7fa4e618 --- /dev/null +++ b/docs/source/_static/tasks/manipulation/franka_lift.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6d0d93ba33e8e851a6139ad82cb6970b39236f75d1bf1c471aad16323e4a2909 +size 193446 diff --git a/docs/source/_static/tasks/manipulation/franka_open_drawer.jpg b/docs/source/_static/tasks/manipulation/franka_open_drawer.jpg new file mode 100644 index 00000000..94ce61af --- /dev/null +++ b/docs/source/_static/tasks/manipulation/franka_open_drawer.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:23e49861bd120df01f2b3496e522850c1456697e9ba678b4cc6b373ef5c25a56 +size 124090 diff --git a/docs/source/_static/tasks/manipulation/franka_reach.jpg b/docs/source/_static/tasks/manipulation/franka_reach.jpg new file mode 100644 index 00000000..27216fa7 --- /dev/null +++ b/docs/source/_static/tasks/manipulation/franka_reach.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a2dc7ce3445fc8e1a6f8fd462c5f7b61acd55131fe3907585f7c6d7eea8a6745 +size 198124 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 00000000..845051d0 --- /dev/null +++ b/docs/source/_static/tasks/manipulation/franka_stack.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a22a0cf57db9af77cffe320e545263818899dcbec35f6563ca36c25556fff23a +size 131565 diff --git a/docs/source/_static/tasks/manipulation/shadow_cube.jpg b/docs/source/_static/tasks/manipulation/shadow_cube.jpg new file mode 100644 index 00000000..9123f86b --- /dev/null +++ b/docs/source/_static/tasks/manipulation/shadow_cube.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:90620afca4e77b0206c98cf21dab4c8397ecb0fcff10a0185f5a494d5c1e1fc8 +size 126215 diff --git a/docs/source/_static/tasks/manipulation/shadow_hand_over.jpg b/docs/source/_static/tasks/manipulation/shadow_hand_over.jpg new file mode 100644 index 00000000..b15ca673 --- /dev/null +++ b/docs/source/_static/tasks/manipulation/shadow_hand_over.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:75760e0603a11354c806d8e13ec5f9d93a981ab08a6a90a76eca70231d958331 +size 144111 diff --git a/docs/source/_static/tasks/manipulation/tycho_track_goal.jpg b/docs/source/_static/tasks/manipulation/tycho_track_goal.jpg new file mode 100644 index 00000000..5736a3d8 --- /dev/null +++ b/docs/source/_static/tasks/manipulation/tycho_track_goal.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:50622595a7a3b50c5d24c03edba6d2f91d051fb877eeb5ad1480d55f6ba1d0dc +size 160398 diff --git a/docs/source/_static/tasks/manipulation/ur10_reach.jpg b/docs/source/_static/tasks/manipulation/ur10_reach.jpg new file mode 100644 index 00000000..faa0e662 --- /dev/null +++ b/docs/source/_static/tasks/manipulation/ur10_reach.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4a2c652db6613bd1676e8a0706166c2cc8cf367692387719f94d0329051daf98 +size 201139 diff --git a/docs/source/_static/tasks/manipulation/ur5_track_goal.jpg b/docs/source/_static/tasks/manipulation/ur5_track_goal.jpg new file mode 100644 index 00000000..cd6bd2b6 --- /dev/null +++ b/docs/source/_static/tasks/manipulation/ur5_track_goal.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:89fcc92bc0ac84a297e4b700f510368c86c746d309b1f477ab6a901fe16a3dab +size 222444 diff --git a/docs/source/_static/tasks/manipulation/xarm_leap_track_goal.jpg b/docs/source/_static/tasks/manipulation/xarm_leap_track_goal.jpg new file mode 100644 index 00000000..c24354cd --- /dev/null +++ b/docs/source/_static/tasks/manipulation/xarm_leap_track_goal.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:515014e830ef4aaab8e3cdf47ea1380cd6926a45fe3f4cbf3bc599768ace0b7d +size 186002 diff --git a/docs/source/_static/tasks/navigation/anymal_c_nav.jpg b/docs/source/_static/tasks/navigation/anymal_c_nav.jpg new file mode 100644 index 00000000..9cc7cd52 --- /dev/null +++ b/docs/source/_static/tasks/navigation/anymal_c_nav.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7ad045aeb1f7146e7c4ccb925116d7fe09933ad6409e2c17678c29fd497637fc +size 142532 diff --git a/docs/source/_static/tasks/others/humanoid_amp.jpg b/docs/source/_static/tasks/others/humanoid_amp.jpg new file mode 100644 index 00000000..5503dfd9 --- /dev/null +++ b/docs/source/_static/tasks/others/humanoid_amp.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ccc036923ba73ee8fb005674ea520ec9e83f6678b4e34ab9c7251cb13ed8749c +size 121642 diff --git a/docs/source/_static/tasks/others/quadcopter.jpg b/docs/source/_static/tasks/others/quadcopter.jpg new file mode 100644 index 00000000..0bc7dbd7 --- /dev/null +++ b/docs/source/_static/tasks/others/quadcopter.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8ade4726bed5e43b60f1305835bbb4a790a153aa4cbe3f59ae4b1cd954fa4a49 +size 58039 diff --git a/docs/source/_static/terrains/height_field/discrete_obstacles_terrain.jpg b/docs/source/_static/terrains/height_field/discrete_obstacles_terrain.jpg new file mode 100644 index 00000000..43947d4c --- /dev/null +++ b/docs/source/_static/terrains/height_field/discrete_obstacles_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4927163189666c35680684d50ce202d895489d4dbc23d9969adc398d769a1232 +size 91428 diff --git a/docs/source/_static/terrains/height_field/discrete_obstacles_terrain_choice.jpg b/docs/source/_static/terrains/height_field/discrete_obstacles_terrain_choice.jpg new file mode 100644 index 00000000..d7d43da5 --- /dev/null +++ b/docs/source/_static/terrains/height_field/discrete_obstacles_terrain_choice.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0e582967fb981a074ac2fcd1bada3f8f542f8fc3eea2fa5b59b86b15cd791450 +size 33975 diff --git a/docs/source/_static/terrains/height_field/discrete_obstacles_terrain_fixed.jpg b/docs/source/_static/terrains/height_field/discrete_obstacles_terrain_fixed.jpg new file mode 100644 index 00000000..7f54eabf --- /dev/null +++ b/docs/source/_static/terrains/height_field/discrete_obstacles_terrain_fixed.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fd9a77be2a92dc714b580b2f90bd72916028acd0ecacb51286de815bfd1d4d67 +size 32427 diff --git a/docs/source/_static/terrains/height_field/inverted_pyramid_sloped_terrain.jpg b/docs/source/_static/terrains/height_field/inverted_pyramid_sloped_terrain.jpg new file mode 100644 index 00000000..a08de8e7 --- /dev/null +++ b/docs/source/_static/terrains/height_field/inverted_pyramid_sloped_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8c07615709339b8c01610ec3bd1eb0f53338df0bb39442d2b16443d204b8fc1a +size 148994 diff --git a/docs/source/_static/terrains/height_field/inverted_pyramid_stairs_terrain.jpg b/docs/source/_static/terrains/height_field/inverted_pyramid_stairs_terrain.jpg new file mode 100644 index 00000000..cd86e924 --- /dev/null +++ b/docs/source/_static/terrains/height_field/inverted_pyramid_stairs_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:826e96d56ebf28b27bd9cde1822fc79a98c5650d1e3a88d44ecc86570c967d38 +size 47349 diff --git a/docs/source/_static/terrains/height_field/pyramid_sloped_terrain.jpg b/docs/source/_static/terrains/height_field/pyramid_sloped_terrain.jpg new file mode 100644 index 00000000..a8feb347 --- /dev/null +++ b/docs/source/_static/terrains/height_field/pyramid_sloped_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:080d9fc8a3d54e14d5f370c982dd913b98b5b93c035314b6333f0bae4583556e +size 150200 diff --git a/docs/source/_static/terrains/height_field/pyramid_stairs_terrain.jpg b/docs/source/_static/terrains/height_field/pyramid_stairs_terrain.jpg new file mode 100644 index 00000000..a9f53ee1 --- /dev/null +++ b/docs/source/_static/terrains/height_field/pyramid_stairs_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8c4a038949988eab9887fdd078421acc743c245d9b4d78a2934f5c431b3998ff +size 29207 diff --git a/docs/source/_static/terrains/height_field/random_uniform_terrain.jpg b/docs/source/_static/terrains/height_field/random_uniform_terrain.jpg new file mode 100644 index 00000000..47c194ac --- /dev/null +++ b/docs/source/_static/terrains/height_field/random_uniform_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8023cf72f1addf993ccf4037c061ac3fbe8c2938e52870c8d7cec3ca12c5fd15 +size 333513 diff --git a/docs/source/_static/terrains/height_field/stepping_stones_terrain.jpg b/docs/source/_static/terrains/height_field/stepping_stones_terrain.jpg new file mode 100644 index 00000000..b276619d --- /dev/null +++ b/docs/source/_static/terrains/height_field/stepping_stones_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8ec37e7717d8c92e6e39988edfa2a934e7aace62c72a92ce794bff1f71e1baa8 +size 248578 diff --git a/docs/source/_static/terrains/height_field/wave_terrain.jpg b/docs/source/_static/terrains/height_field/wave_terrain.jpg new file mode 100644 index 00000000..841e3206 --- /dev/null +++ b/docs/source/_static/terrains/height_field/wave_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:24656af96f83e10872cc6112b57f40e47e7e6d4f4f983bc728a74c922dac9139 +size 282380 diff --git a/docs/source/_static/terrains/trimesh/box_terrain.jpg b/docs/source/_static/terrains/trimesh/box_terrain.jpg new file mode 100644 index 00000000..c4af3eb2 --- /dev/null +++ b/docs/source/_static/terrains/trimesh/box_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f233bebb033a074a71720b853e64bbc863c57ded47b69d18064059cf4304f06f +size 2956 diff --git a/docs/source/_static/terrains/trimesh/box_terrain_with_two_boxes.jpg b/docs/source/_static/terrains/trimesh/box_terrain_with_two_boxes.jpg new file mode 100644 index 00000000..9d06bedb --- /dev/null +++ b/docs/source/_static/terrains/trimesh/box_terrain_with_two_boxes.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c517ab4fd58e280fbc1b97c4b85c57960e8af93cbe388c6580321a5b5ff1c362 +size 2993 diff --git a/docs/source/_static/terrains/trimesh/flat_terrain.jpg b/docs/source/_static/terrains/trimesh/flat_terrain.jpg new file mode 100644 index 00000000..26776e38 --- /dev/null +++ b/docs/source/_static/terrains/trimesh/flat_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2d0907c5a2cd6b3ee287931578b7bdba1123d968abacbd6a7f13a796a22c4e1a +size 2852 diff --git a/docs/source/_static/terrains/trimesh/floating_ring_terrain.jpg b/docs/source/_static/terrains/trimesh/floating_ring_terrain.jpg new file mode 100644 index 00000000..41d53cef --- /dev/null +++ b/docs/source/_static/terrains/trimesh/floating_ring_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:11f4f382c9ebfa3759af8471151e2840126874f26a0b3a00bad804012696a773 +size 2995 diff --git a/docs/source/_static/terrains/trimesh/gap_terrain.jpg b/docs/source/_static/terrains/trimesh/gap_terrain.jpg new file mode 100644 index 00000000..8810fcab --- /dev/null +++ b/docs/source/_static/terrains/trimesh/gap_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7aeb68558d35d186f6be3dc845d1b775a8f233fbbcc86c6a130a0ae204cac71d +size 3708 diff --git a/docs/source/_static/terrains/trimesh/inverted_pyramid_stairs_terrain.jpg b/docs/source/_static/terrains/trimesh/inverted_pyramid_stairs_terrain.jpg new file mode 100644 index 00000000..326ef95b --- /dev/null +++ b/docs/source/_static/terrains/trimesh/inverted_pyramid_stairs_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:edced4bb3f37679330ac861b710546989048bd549f5ba6aa4e23319733e89e3c +size 5796 diff --git a/docs/source/_static/terrains/trimesh/inverted_pyramid_stairs_terrain_with_holes.jpg b/docs/source/_static/terrains/trimesh/inverted_pyramid_stairs_terrain_with_holes.jpg new file mode 100644 index 00000000..96e3b0c1 --- /dev/null +++ b/docs/source/_static/terrains/trimesh/inverted_pyramid_stairs_terrain_with_holes.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:47cfdc0d011e4353e3e15c11b011444fd46b361b6445d121b12ee6696de22af8 +size 6353 diff --git a/docs/source/_static/terrains/trimesh/pit_terrain.jpg b/docs/source/_static/terrains/trimesh/pit_terrain.jpg new file mode 100644 index 00000000..51aabf2a --- /dev/null +++ b/docs/source/_static/terrains/trimesh/pit_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fc004b96d8db674643f44e93ff3aae8af4b813c9475cab7243a15ac8dfa67c4f +size 3218 diff --git a/docs/source/_static/terrains/trimesh/pit_terrain_with_two_levels.jpg b/docs/source/_static/terrains/trimesh/pit_terrain_with_two_levels.jpg new file mode 100644 index 00000000..eb5a8623 --- /dev/null +++ b/docs/source/_static/terrains/trimesh/pit_terrain_with_two_levels.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0c89bb5b34e29e6e402848e58ce79da1fb453d58eb2983be1518daf4a51966da +size 4370 diff --git a/docs/source/_static/terrains/trimesh/pyramid_stairs_terrain.jpg b/docs/source/_static/terrains/trimesh/pyramid_stairs_terrain.jpg new file mode 100644 index 00000000..be7fa692 --- /dev/null +++ b/docs/source/_static/terrains/trimesh/pyramid_stairs_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:12975e7119588cb1c0c63e77fed803552e1e5a4008af58198d01e80e89314691 +size 3595 diff --git a/docs/source/_static/terrains/trimesh/pyramid_stairs_terrain_with_holes.jpg b/docs/source/_static/terrains/trimesh/pyramid_stairs_terrain_with_holes.jpg new file mode 100644 index 00000000..836dc6e7 --- /dev/null +++ b/docs/source/_static/terrains/trimesh/pyramid_stairs_terrain_with_holes.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e63145b71405a1c57bf153f0365b82a6f35f5894fd22a9676b4c030da723b440 +size 4926 diff --git a/docs/source/_static/terrains/trimesh/rails_terrain.jpg b/docs/source/_static/terrains/trimesh/rails_terrain.jpg new file mode 100644 index 00000000..cc112e96 --- /dev/null +++ b/docs/source/_static/terrains/trimesh/rails_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:397e2a0f459c0923cc3f07e954a9835cf00a659f5ecc312465ffa0930fdee1fd +size 3326 diff --git a/docs/source/_static/terrains/trimesh/random_grid_terrain.jpg b/docs/source/_static/terrains/trimesh/random_grid_terrain.jpg new file mode 100644 index 00000000..dba91002 --- /dev/null +++ b/docs/source/_static/terrains/trimesh/random_grid_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cee6ed4e4e2e2191c1dd419c1c59d365f772bb1ff045b45771fe848dfd396ce7 +size 6378 diff --git a/docs/source/_static/terrains/trimesh/random_grid_terrain_with_holes.jpg b/docs/source/_static/terrains/trimesh/random_grid_terrain_with_holes.jpg new file mode 100644 index 00000000..3cd4f542 --- /dev/null +++ b/docs/source/_static/terrains/trimesh/random_grid_terrain_with_holes.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1be4b07068e25e7b402999ab4eb181c3115c60dc95d9f7f7a3f6296e8540167a +size 6093 diff --git a/docs/source/_static/terrains/trimesh/repeated_objects_box_terrain.jpg b/docs/source/_static/terrains/trimesh/repeated_objects_box_terrain.jpg new file mode 100644 index 00000000..4af1d1f2 --- /dev/null +++ b/docs/source/_static/terrains/trimesh/repeated_objects_box_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bf00b570b6cb5cf748c764144d4818ea3878abb48860a4202f03e112f653d46c +size 60154 diff --git a/docs/source/_static/terrains/trimesh/repeated_objects_cylinder_terrain.jpg b/docs/source/_static/terrains/trimesh/repeated_objects_cylinder_terrain.jpg new file mode 100644 index 00000000..4ebbbe0c --- /dev/null +++ b/docs/source/_static/terrains/trimesh/repeated_objects_cylinder_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ef6f85e8fe9f494d9743bdb23846b4e71c33c8998e7b91d0b5429b610cb4f6de +size 84510 diff --git a/docs/source/_static/terrains/trimesh/repeated_objects_pyramid_terrain.jpg b/docs/source/_static/terrains/trimesh/repeated_objects_pyramid_terrain.jpg new file mode 100644 index 00000000..0105dc34 --- /dev/null +++ b/docs/source/_static/terrains/trimesh/repeated_objects_pyramid_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2f000868c88452b9596ffa8f5f78a04e267ea94bb97f49ce71564efc15972d0a +size 62332 diff --git a/docs/source/_static/terrains/trimesh/star_terrain.jpg b/docs/source/_static/terrains/trimesh/star_terrain.jpg new file mode 100644 index 00000000..281916c4 --- /dev/null +++ b/docs/source/_static/terrains/trimesh/star_terrain.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:22c69f97496f7c81cb3adce0725f6f80e96c6ab78cbcd3b97929e87b4397ed24 +size 9751 diff --git a/docs/source/_static/vscode_tasks.png b/docs/source/_static/vscode_tasks.png new file mode 100644 index 00000000..b5b0d5c2 --- /dev/null +++ b/docs/source/_static/vscode_tasks.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:70b5d31518826c6e57c4182b20abbb88f0b247d596a616de4e5f7e5790f79b2c +size 983040 diff --git a/docs/source/api/index.rst b/docs/source/api/index.rst new file mode 100644 index 00000000..da2d41a1 --- /dev/null +++ b/docs/source/api/index.rst @@ -0,0 +1,17 @@ +API Reference +============= + +This page gives an overview of all the modules and classes in the Isaac Lab extensions. + +uwlab extension +------------------------ + +The following modules are available in the ``uwlab`` extension: + +.. currentmodule:: uwlab + +.. autosummary:: + :toctree: lab + + controllers + devices diff --git a/docs/source/api/lab/uwlab.controllers.rst b/docs/source/api/lab/uwlab.controllers.rst new file mode 100644 index 00000000..ee54df4f --- /dev/null +++ b/docs/source/api/lab/uwlab.controllers.rst @@ -0,0 +1,25 @@ +ï»żuwlab.controllers +==================== + +.. automodule:: uwlab.controllers + + .. rubric:: Classes + + .. autosummary:: + + MultiConstraintDifferentialIKController + MultiConstraintDifferentialIKControllerCfg + +Differential Inverse Kinematics With Multiple Bodies +----------------------------------------------------- + +.. autoclass:: MultiConstraintDifferentialIKController + :members: + :inherited-members: + :show-inheritance: + +.. autoclass:: MultiConstraintDifferentialIKControllerCfg + :members: + :inherited-members: + :show-inheritance: + :exclude-members: __init__, __new__, class_type diff --git a/docs/source/api/lab/uwlab.devices.rst b/docs/source/api/lab/uwlab.devices.rst new file mode 100644 index 00000000..20f58a4c --- /dev/null +++ b/docs/source/api/lab/uwlab.devices.rst @@ -0,0 +1,36 @@ +ï»żuwlab.devices +=============== + +.. automodule:: uwlab.devices + + .. rubric:: Classes + + .. autosummary:: + + Se3Keyboard + RealsenseT265 + RokokoGlove + + +Keyboard +-------- + +.. autoclass:: Se3Keyboard + :members: + :inherited-members: + :show-inheritance: + +Real Sense +---------- +.. autoclass:: RealsenseT265 + :members: + :inherited-members: + :show-inheritance: + +Rokoko Gloves +------------- + +.. autoclass:: RokokoGlove + :members: + :inherited-members: + :show-inheritance: diff --git a/docs/source/overview/isaac_environments.rst b/docs/source/overview/isaac_environments.rst new file mode 100644 index 00000000..33a4ffaf --- /dev/null +++ b/docs/source/overview/isaac_environments.rst @@ -0,0 +1,662 @@ +.. _environments: + +Available Isaac Environments +============================ + +The following lists comprises of all the RL tasks implementations that are available in Isaac Lab. +While we try to keep this list up-to-date, you can always get the latest list of environments by +running the following command: + +.. tab-set:: + :sync-group: os + + .. tab-item:: :icon:`fa-brands fa-linux` Linux + :sync: linux + + .. code:: bash + + ./isaaclab.sh -p scripts/environments/list_envs.py + + .. tab-item:: :icon:`fa-brands fa-windows` Windows + :sync: windows + + .. code:: batch + + isaaclab.bat -p scripts\environments\list_envs.py + +We are actively working on adding more environments to the list. If you have any environments that +you would like to add to Isaac Lab, please feel free to open a pull request! + +Single-agent +------------ + +Classic +~~~~~~~ + +Classic environments that are based on IsaacGymEnvs implementation of MuJoCo-style environments. + +.. table:: + :widths: 33 37 30 + + +------------------+-----------------------------+-------------------------------------------------------------------------+ + | World | Environment ID | Description | + +==================+=============================+=========================================================================+ + | |humanoid| | |humanoid-link| | Move towards a direction with the MuJoCo humanoid robot | + | | | | + | | |humanoid-direct-link| | | + +------------------+-----------------------------+-------------------------------------------------------------------------+ + | |ant| | |ant-link| | Move towards a direction with the MuJoCo ant robot | + | | | | + | | |ant-direct-link| | | + +------------------+-----------------------------+-------------------------------------------------------------------------+ + | |cartpole| | |cartpole-link| | Move the cart to keep the pole upwards in the classic cartpole control | + | | | | + | | |cartpole-direct-link| | | + +------------------+-----------------------------+-------------------------------------------------------------------------+ + | |cartpole| | |cartpole-rgb-link| | Move the cart to keep the pole upwards in the classic cartpole control | + | | | and perceptive inputs | + | | |cartpole-depth-link| | | + | | | | + | | |cartpole-rgb-direct-link| | | + | | | | + | | |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 +.. |cartpole| image:: ../_static/tasks/classic/cartpole.jpg + +.. |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-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 `__ +.. |cartpole-direct-link| replace:: `Isaac-Cartpole-Direct-v0 `__ +.. |cartpole-rgb-direct-link| replace:: `Isaac-Cartpole-RGB-Camera-Direct-v0 `__ +.. |cartpole-depth-direct-link| replace:: `Isaac-Cartpole-Depth-Camera-Direct-v0 `__ + +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 lift-cube environment: + +* |lift-cube-link|: Franka arm with joint position control +* |lift-cube-ik-abs-link|: Franka arm with absolute IK control +* |lift-cube-ik-rel-link|: Franka arm with relative IK control + +.. table:: + :widths: 33 37 30 + + +--------------------+-------------------------+-----------------------------------------------------------------------------+ + | World | Environment ID | Description | + +====================+=========================+=============================================================================+ + | |reach-franka| | |reach-franka-link| | Move the end-effector to a sampled target pose with the Franka robot | + +--------------------+-------------------------+-----------------------------------------------------------------------------+ + | |reach-ur10| | |reach-ur10-link| | Move the end-effector to a sampled target pose with the UR10 robot | + +--------------------+-------------------------+-----------------------------------------------------------------------------+ + | |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| | | + +--------------------+-------------------------+-----------------------------------------------------------------------------+ + | |cube-allegro| | |cube-allegro-link| | In-hand reorientation of a cube using Allegro hand | + | | | | + | | |allegro-direct-link| | | + +--------------------+-------------------------+-----------------------------------------------------------------------------+ + | |cube-shadow| | |cube-shadow-link| | In-hand reorientation of a cube using Shadow hand | + | | | | + | | |cube-shadow-ff-link| | | + | | | | + | | |cube-shadow-lstm-link| | | + +--------------------+-------------------------+-----------------------------------------------------------------------------+ + | |cube-shadow| | |cube-shadow-vis-link| | In-hand reorientation of a cube using Shadow hand using perceptive inputs | + +--------------------+-------------------------+-----------------------------------------------------------------------------+ + +.. |reach-franka| image:: ../_static/tasks/manipulation/franka_reach.jpg +.. |reach-ur10| image:: ../_static/tasks/manipulation/ur10_reach.jpg +.. |lift-cube| image:: ../_static/tasks/manipulation/franka_lift.jpg +.. |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 `__ +.. |lift-cube-link| replace:: `Isaac-Lift-Cube-Franka-v0 `__ +.. |lift-cube-ik-abs-link| replace:: `Isaac-Lift-Cube-Franka-IK-Abs-v0 `__ +.. |lift-cube-ik-rel-link| replace:: `Isaac-Lift-Cube-Franka-IK-Rel-v0 `__ +.. |cabi-franka-link| replace:: `Isaac-Open-Drawer-Franka-v0 `__ +.. |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 +~~~~~~~~~~ + +Environments based on legged locomotion tasks. + +.. table:: + :widths: 33 37 30 + + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | World | Environment ID | Description | + +==============================+==============================================+==============================================================================+ + | |velocity-flat-anymal-b| | |velocity-flat-anymal-b-link| | Track a velocity command on flat terrain with the Anymal B robot | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | |velocity-rough-anymal-b| | |velocity-rough-anymal-b-link| | Track a velocity command on rough terrain with the Anymal B robot | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | |velocity-flat-anymal-c| | |velocity-flat-anymal-c-link| | Track a velocity command on flat terrain with the Anymal C robot | + | | | | + | | |velocity-flat-anymal-c-direct-link| | | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | |velocity-rough-anymal-c| | |velocity-rough-anymal-c-link| | Track a velocity command on rough terrain with the Anymal C robot | + | | | | + | | |velocity-rough-anymal-c-direct-link| | | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | |velocity-flat-anymal-d| | |velocity-flat-anymal-d-link| | Track a velocity command on flat terrain with the Anymal D robot | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | |velocity-rough-anymal-d| | |velocity-rough-anymal-d-link| | Track a velocity command on rough terrain with the Anymal D robot | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | |velocity-flat-unitree-a1| | |velocity-flat-unitree-a1-link| | Track a velocity command on flat terrain with the Unitree A1 robot | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | |velocity-rough-unitree-a1| | |velocity-rough-unitree-a1-link| | Track a velocity command on rough terrain with the Unitree A1 robot | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | |velocity-flat-unitree-go1| | |velocity-flat-unitree-go1-link| | Track a velocity command on flat terrain with the Unitree Go1 robot | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | |velocity-rough-unitree-go1| | |velocity-rough-unitree-go1-link| | Track a velocity command on rough terrain with the Unitree Go1 robot | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | |velocity-flat-unitree-go2| | |velocity-flat-unitree-go2-link| | Track a velocity command on flat terrain with the Unitree Go2 robot | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | |velocity-rough-unitree-go2| | |velocity-rough-unitree-go2-link| | Track a velocity command on rough terrain with the Unitree Go2 robot | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | |velocity-flat-spot| | |velocity-flat-spot-link| | Track a velocity command on flat terrain with the Boston Dynamics Spot robot | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | |velocity-flat-h1| | |velocity-flat-h1-link| | Track a velocity command on flat terrain with the Unitree H1 robot | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | |velocity-rough-h1| | |velocity-rough-h1-link| | Track a velocity command on rough terrain with the Unitree H1 robot | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | |velocity-flat-g1| | |velocity-flat-g1-link| | Track a velocity command on flat terrain with the Unitree G1 robot | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | |velocity-rough-g1| | |velocity-rough-g1-link| | Track a velocity command on rough terrain with the Unitree G1 robot | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + +.. |velocity-flat-anymal-b-link| replace:: `Isaac-Velocity-Flat-Anymal-B-v0 `__ +.. |velocity-rough-anymal-b-link| replace:: `Isaac-Velocity-Rough-Anymal-B-v0 `__ + +.. |velocity-flat-anymal-c-link| replace:: `Isaac-Velocity-Flat-Anymal-C-v0 `__ +.. |velocity-rough-anymal-c-link| replace:: `Isaac-Velocity-Rough-Anymal-C-v0 `__ + +.. |velocity-flat-anymal-c-direct-link| replace:: `Isaac-Velocity-Flat-Anymal-C-Direct-v0 `__ +.. |velocity-rough-anymal-c-direct-link| replace:: `Isaac-Velocity-Rough-Anymal-C-Direct-v0 `__ + +.. |velocity-flat-anymal-d-link| replace:: `Isaac-Velocity-Flat-Anymal-D-v0 `__ +.. |velocity-rough-anymal-d-link| replace:: `Isaac-Velocity-Rough-Anymal-D-v0 `__ + +.. |velocity-flat-unitree-a1-link| replace:: `Isaac-Velocity-Flat-Unitree-A1-v0 `__ +.. |velocity-rough-unitree-a1-link| replace:: `Isaac-Velocity-Rough-Unitree-A1-v0 `__ + +.. |velocity-flat-unitree-go1-link| replace:: `Isaac-Velocity-Flat-Unitree-Go1-v0 `__ +.. |velocity-rough-unitree-go1-link| replace:: `Isaac-Velocity-Rough-Unitree-Go1-v0 `__ + +.. |velocity-flat-unitree-go2-link| replace:: `Isaac-Velocity-Flat-Unitree-Go2-v0 `__ +.. |velocity-rough-unitree-go2-link| replace:: `Isaac-Velocity-Rough-Unitree-Go2-v0 `__ + +.. |velocity-flat-spot-link| replace:: `Isaac-Velocity-Flat-Spot-v0 `__ + +.. |velocity-flat-h1-link| replace:: `Isaac-Velocity-Flat-H1-v0 `__ +.. |velocity-rough-h1-link| replace:: `Isaac-Velocity-Rough-H1-v0 `__ + +.. |velocity-flat-g1-link| replace:: `Isaac-Velocity-Flat-G1-v0 `__ +.. |velocity-rough-g1-link| replace:: `Isaac-Velocity-Rough-G1-v0 `__ + + +.. |velocity-flat-anymal-b| image:: ../_static/tasks/locomotion/anymal_b_flat.jpg +.. |velocity-rough-anymal-b| image:: ../_static/tasks/locomotion/anymal_b_rough.jpg +.. |velocity-flat-anymal-c| image:: ../_static/tasks/locomotion/anymal_c_flat.jpg +.. |velocity-rough-anymal-c| image:: ../_static/tasks/locomotion/anymal_c_rough.jpg +.. |velocity-flat-anymal-d| image:: ../_static/tasks/locomotion/anymal_d_flat.jpg +.. |velocity-rough-anymal-d| image:: ../_static/tasks/locomotion/anymal_d_rough.jpg +.. |velocity-flat-unitree-a1| image:: ../_static/tasks/locomotion/a1_flat.jpg +.. |velocity-rough-unitree-a1| image:: ../_static/tasks/locomotion/a1_rough.jpg +.. |velocity-flat-unitree-go1| image:: ../_static/tasks/locomotion/go1_flat.jpg +.. |velocity-rough-unitree-go1| image:: ../_static/tasks/locomotion/go1_rough.jpg +.. |velocity-flat-unitree-go2| image:: ../_static/tasks/locomotion/go2_flat.jpg +.. |velocity-rough-unitree-go2| image:: ../_static/tasks/locomotion/go2_rough.jpg +.. |velocity-flat-spot| image:: ../_static/tasks/locomotion/spot_flat.jpg +.. |velocity-flat-h1| image:: ../_static/tasks/locomotion/h1_flat.jpg +.. |velocity-rough-h1| image:: ../_static/tasks/locomotion/h1_rough.jpg +.. |velocity-flat-g1| image:: ../_static/tasks/locomotion/g1_flat.jpg +.. |velocity-rough-g1| image:: ../_static/tasks/locomotion/g1_rough.jpg + +Navigation +~~~~~~~~~~ + +.. table:: + :widths: 33 37 30 + + +----------------+---------------------+-----------------------------------------------------------------------------+ + | World | Environment ID | Description | + +================+=====================+=============================================================================+ + | |anymal_c_nav| | |anymal_c_nav-link| | Navigate towards a target x-y position and heading with the ANYmal C robot. | + +----------------+---------------------+-----------------------------------------------------------------------------+ + +.. |anymal_c_nav-link| replace:: `Isaac-Navigation-Flat-Anymal-C-v0 `__ + +.. |anymal_c_nav| image:: ../_static/tasks/navigation/anymal_c_nav.jpg + + +Others +~~~~~~ + +.. note:: + + Adversarial Motion Priors (AMP) training is only available with the `skrl` library, as it is the only one of the currently + integrated libraries that supports it out-of-the-box (for the other libraries, it is necessary to implement the algorithm and architectures). + See the `skrl's AMP Documentation `_ for more information. + The AMP algorithm can be activated by adding the command line input ``--algorithm AMP`` to the train/play script. + + For evaluation, the play script's command line input ``--real-time`` allows the interaction loop between the environment and the agent to run in real time, if possible. + +.. table:: + :widths: 33 37 30 + + +----------------+---------------------------+-----------------------------------------------------------------------------+ + | World | Environment ID | Description | + +================+===========================+=============================================================================+ + | |quadcopter| | |quadcopter-link| | Fly and hover the Crazyflie copter at a goal point by applying thrust. | + +----------------+---------------------------+-----------------------------------------------------------------------------+ + | |humanoid_amp| | |humanoid_amp_dance-link| | Move a humanoid robot by imitating different pre-recorded human animations | + | | | (Adversarial Motion Priors). | + | | |humanoid_amp_run-link| | | + | | | | + | | |humanoid_amp_walk-link| | | + +----------------+---------------------------+-----------------------------------------------------------------------------+ + +.. |quadcopter-link| replace:: `Isaac-Quadcopter-Direct-v0 `__ +.. |humanoid_amp_dance-link| replace:: `Isaac-Humanoid-AMP-Dance-Direct-v0 `__ +.. |humanoid_amp_run-link| replace:: `Isaac-Humanoid-AMP-Run-Direct-v0 `__ +.. |humanoid_amp_walk-link| replace:: `Isaac-Humanoid-AMP-Walk-Direct-v0 `__ + +.. |quadcopter| image:: ../_static/tasks/others/quadcopter.jpg +.. |humanoid_amp| image:: ../_static/tasks/others/humanoid_amp.jpg + + +Multi-agent +------------ + +.. note:: + + True mutli-agent training is only available with the `skrl` library, see the `Multi-Agents Documentation `_ for more information. + It supports the `IPPO` and `MAPPO` algorithms, which can be activated by adding the command line input ``--algorithm IPPO`` or ``--algorithm MAPPO`` to the train/play script. + If these environments are run with other libraries or without the `IPPO` or `MAPPO` flags, they will be converted to single-agent environments under the hood. + + +Classic +~~~~~~~ + +.. table:: + :widths: 33 37 30 + + +------------------------+------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ + | World | Environment ID | Description | + +========================+====================================+=======================================================================================================================+ + | |cart-double-pendulum| | |cart-double-pendulum-direct-link| | Move the cart and the pendulum to keep the last one upwards in the classic inverted double pendulum on a cart control | + +------------------------+------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ + +.. |cart-double-pendulum| image:: ../_static/tasks/classic/cart_double_pendulum.jpg + +.. |cart-double-pendulum-direct-link| replace:: `Isaac-Cart-Double-Pendulum-Direct-v0 `__ + +Manipulation +~~~~~~~~~~~~ + +Environments based on fixed-arm manipulation tasks. + +.. table:: + :widths: 33 37 30 + + +----------------------+--------------------------------+--------------------------------------------------------+ + | World | Environment ID | Description | + +======================+================================+========================================================+ + | |shadow-hand-over| | |shadow-hand-over-direct-link| | Passing an object from one hand over to the other hand | + +----------------------+--------------------------------+--------------------------------------------------------+ + +.. |shadow-hand-over| image:: ../_static/tasks/manipulation/shadow_hand_over.jpg + +.. |shadow-hand-over-direct-link| replace:: `Isaac-Shadow-Hand-Over-Direct-v0 `__ + +| + +Comprehensive List of Environments +================================== + +.. list-table:: + :widths: 33 25 19 25 + + * - **Task Name** + - **Inference Task Name** + - **Workflow** + - **RL Library** + * - Isaac-Ant-Direct-v0 + - + - Direct + - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Ant-v0 + - + - Manager Based + - **rsl_rl** (PPO), **rl_games** (PPO), **skrl** (PPO), **sb3** (PPO) + * - Isaac-Cart-Double-Pendulum-Direct-v0 + - + - Direct + - **rl_games** (PPO), **skrl** (IPPO, PPO, MAPPO) + * - Isaac-Cartpole-Depth-Camera-Direct-v0 + - + - Direct + - **rl_games** (PPO), **skrl** (PPO) + * - Isaac-Cartpole-Depth-v0 + - + - Manager Based + - **rl_games** (PPO) + * - Isaac-Cartpole-Direct-v0 + - + - Direct + - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO) + * - Isaac-Cartpole-RGB-Camera-Direct-v0 + - + - Direct + - **rl_games** (PPO), **skrl** (PPO) + * - Isaac-Cartpole-RGB-ResNet18-v0 + - + - Manager Based + - **rl_games** (PPO) + * - Isaac-Cartpole-RGB-TheiaTiny-v0 + - + - Manager Based + - **rl_games** (PPO) + * - Isaac-Cartpole-RGB-v0 + - + - Manager Based + - **rl_games** (PPO) + * - Isaac-Cartpole-v0 + - + - 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 + - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Humanoid-AMP-Dance-Direct-v0 + - + - Direct + - **skrl** (AMP) + * - Isaac-Humanoid-AMP-Run-Direct-v0 + - + - Direct + - **skrl** (AMP) + * - Isaac-Humanoid-AMP-Walk-Direct-v0 + - + - Direct + - **skrl** (AMP) + * - Isaac-Humanoid-Direct-v0 + - + - Direct + - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Humanoid-v0 + - + - Manager Based + - **rsl_rl** (PPO), **rl_games** (PPO), **skrl** (PPO), **sb3** (PPO) + * - Isaac-Lift-Cube-Franka-IK-Abs-v0 + - + - Manager Based + - + * - Isaac-Lift-Cube-Franka-IK-Rel-v0 + - + - Manager Based + - + * - Isaac-Lift-Cube-Franka-v0 + - Isaac-Lift-Cube-Franka-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **skrl** (PPO), **rl_games** (PPO), **sb3** (PPO) + * - Isaac-Lift-Teddy-Bear-Franka-IK-Abs-v0 + - + - Manager Based + - + * - Isaac-Navigation-Flat-Anymal-C-v0 + - Isaac-Navigation-Flat-Anymal-C-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Open-Drawer-Franka-IK-Abs-v0 + - + - Manager Based + - + * - Isaac-Open-Drawer-Franka-IK-Rel-v0 + - + - Manager Based + - + * - Isaac-Open-Drawer-Franka-v0 + - Isaac-Open-Drawer-Franka-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **rl_games** (PPO), **skrl** (PPO) + * - Isaac-Quadcopter-Direct-v0 + - + - Direct + - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Reach-Franka-IK-Abs-v0 + - + - Manager Based + - + * - Isaac-Reach-Franka-IK-Rel-v0 + - + - 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 + - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Reach-UR10-v0 + - Isaac-Reach-UR10-Play-v0 + - Manager Based + - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Repose-Cube-Allegro-Direct-v0 + - + - Direct + - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Repose-Cube-Allegro-NoVelObs-v0 + - Isaac-Repose-Cube-Allegro-NoVelObs-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **rl_games** (PPO), **skrl** (PPO) + * - Isaac-Repose-Cube-Allegro-v0 + - Isaac-Repose-Cube-Allegro-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **rl_games** (PPO), **skrl** (PPO) + * - Isaac-Repose-Cube-Shadow-Direct-v0 + - + - Direct + - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Repose-Cube-Shadow-OpenAI-FF-Direct-v0 + - + - Direct + - **rl_games** (FF), **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Repose-Cube-Shadow-OpenAI-LSTM-Direct-v0 + - + - Direct + - **rl_games** (LSTM) + * - Isaac-Repose-Cube-Shadow-Vision-Direct-v0 + - Isaac-Repose-Cube-Shadow-Vision-Direct-Play-v0 + - Direct + - **rsl_rl** (PPO), **rl_games** (VISION) + * - Isaac-Shadow-Hand-Over-Direct-v0 + - + - Direct + - **rl_games** (PPO), **skrl** (IPPO, PPO, MAPPO) + * - Isaac-Stack-Cube-Franka-IK-Rel-v0 + - + - Manager Based + - + * - Isaac-Stack-Cube-Franka-v0 + - + - Manager Based + - + * - Isaac-Stack-Cube-Instance-Randomize-Franka-IK-Rel-v0 + - + - Manager Based + - + * - Isaac-Stack-Cube-Instance-Randomize-Franka-v0 + - + - Manager Based + - + * - Isaac-Velocity-Flat-Anymal-B-v0 + - Isaac-Velocity-Flat-Anymal-B-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Velocity-Flat-Anymal-C-Direct-v0 + - + - Direct + - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Velocity-Flat-Anymal-C-v0 + - Isaac-Velocity-Flat-Anymal-C-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **rl_games** (PPO), **skrl** (PPO) + * - Isaac-Velocity-Flat-Anymal-D-v0 + - Isaac-Velocity-Flat-Anymal-D-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Velocity-Flat-Cassie-v0 + - Isaac-Velocity-Flat-Cassie-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Velocity-Flat-G1-v0 + - Isaac-Velocity-Flat-G1-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Velocity-Flat-H1-v0 + - Isaac-Velocity-Flat-H1-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Velocity-Flat-Spot-v0 + - Isaac-Velocity-Flat-Spot-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Velocity-Flat-Unitree-A1-v0 + - Isaac-Velocity-Flat-Unitree-A1-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Velocity-Flat-Unitree-Go1-v0 + - Isaac-Velocity-Flat-Unitree-Go1-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Velocity-Flat-Unitree-Go2-v0 + - Isaac-Velocity-Flat-Unitree-Go2-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Velocity-Rough-Anymal-B-v0 + - Isaac-Velocity-Rough-Anymal-B-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Velocity-Rough-Anymal-C-Direct-v0 + - + - Direct + - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Velocity-Rough-Anymal-C-v0 + - Isaac-Velocity-Rough-Anymal-C-Play-v0 + - Manager Based + - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Velocity-Rough-Anymal-D-v0 + - Isaac-Velocity-Rough-Anymal-D-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Velocity-Rough-Cassie-v0 + - Isaac-Velocity-Rough-Cassie-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Velocity-Rough-G1-v0 + - Isaac-Velocity-Rough-G1-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Velocity-Rough-H1-v0 + - Isaac-Velocity-Rough-H1-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Velocity-Rough-Unitree-A1-v0 + - Isaac-Velocity-Rough-Unitree-A1-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Velocity-Rough-Unitree-Go1-v0 + - Isaac-Velocity-Rough-Unitree-Go1-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **skrl** (PPO) + * - Isaac-Velocity-Rough-Unitree-Go2-v0 + - Isaac-Velocity-Rough-Unitree-Go2-Play-v0 + - Manager Based + - **rsl_rl** (PPO), **skrl** (PPO) diff --git a/docs/source/overview/uw_environments.rst b/docs/source/overview/uw_environments.rst new file mode 100644 index 00000000..3108f54c --- /dev/null +++ b/docs/source/overview/uw_environments.rst @@ -0,0 +1,166 @@ +.. _environments: + +Available UW Environments +=========================== + +The following lists comprises of all the RL tasks implementations that are available in UW Lab. +While we try to keep this list up-to-date, you can always get the latest list of environments by +running the following command: + +.. tab-set:: + :sync-group: os + + .. tab-item:: :icon:`fa-brands fa-linux` Linux + :sync: linux + + .. code:: bash + + python scripts/environments/list_envs.py + + .. tab-item:: :icon:`fa-brands fa-windows` Windows + :sync: windows + + .. code:: batch + + python scripts\environments\list_envs.py + + +Single-agent +------------ + +Manipulation +~~~~~~~~~~~~ + +Environments based on fixed-arm manipulation tasks. + +.. table:: + :widths: 33 37 30 + + +--------------------------------+------------------------------------------------+------------------------------------------------------------------------------+ + | World | Environment ID | Description | + +================================+================================================+==============================================================================+ + | |track-goal-ur5| | |track-goal-ur5-link| | Goal tracking with Ur5 robot with Robotiq gripper | + +--------------------------------+------------------------------------------------+------------------------------------------------------------------------------+ + | |track-goal-tycho| | |track-goal-tycho-link| | Goal tracking with Tycho robot | + +--------------------------------+------------------------------------------------+------------------------------------------------------------------------------+ + | |track-goal-xarm-leap| | |track-goal-xarm-leap-link| | Goal tracking with Xarm with Leap Hand robot | + +--------------------------------+------------------------------------------------+------------------------------------------------------------------------------+ + | |ext-nut-thread-franka| | |ext-nut-thread-franka-link| | Threading nut on to bolt on nist board | + +--------------------------------+------------------------------------------------+------------------------------------------------------------------------------+ + | |ext-gear-mesh-franka| | |ext-gear-mesh-franka-link| | Inserting gear on to gear base on nist board | + +--------------------------------+------------------------------------------------+------------------------------------------------------------------------------+ + | |ext-peg-insert-franka| | |ext-peg-insert-franka-link| | Inserting peg rod into hole on nist board | + +--------------------------------+------------------------------------------------+------------------------------------------------------------------------------+ + +.. |track-goal-ur5| image:: ../_static/tasks/manipulation/ur5_track_goal.jpg +.. |track-goal-tycho| image:: ../_static/tasks/manipulation/tycho_track_goal.jpg +.. |track-goal-xarm-leap| image:: ../_static/tasks/manipulation/xarm_leap_track_goal.jpg +.. |ext-nut-thread-franka| image:: ../_static/tasks/manipulation/factory_ext/nut_thread_ext.jpg +.. |ext-gear-mesh-franka| image:: ../_static/tasks/manipulation/factory_ext/gear_mesh_ext.jpg +.. |ext-peg-insert-franka| image:: ../_static/tasks/manipulation/factory_ext/peg_insert_ext.jpg + +.. |track-goal-ur5-link| replace:: `UW-Track-Goal-Ur5-v0 `__ +.. |track-goal-tycho-link| replace:: `UW-Track-Goal-Tycho-v0 `__ +.. |track-goal-xarm-leap-link| replace:: `UW-Track-Goal-Xarm-Leap-v0 `__ +.. |ext-nut-thread-franka-link| replace:: `UW-Nut-Thread-Franka-v0 `__ +.. |ext-gear-mesh-franka-link| replace:: `UW-Gear-Mesh-Franka-v0 `__ +.. |ext-peg-insert-franka-link| replace:: `UW-Peg-Insert-Franka-v0 `__ + +Locomotion +~~~~~~~~~~ + +Environments based on legged locomotion tasks. + +.. table:: + :widths: 33 37 30 + + +--------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | World | Environment ID | Description | + +================================+==============================================+==============================================================================+ + | |position-gap-spot| | |position-gap-spot-link| | Track a position command on gap terrain with the Spot robot | + +--------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | |position-pit-spot| | |position-pit-spot-link| | Track a position command on pit terrain with the Spot robot | + +--------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | |position-stepping-stone-spot| | |position-stepping-stone-spot-link| | Track a position command on stepping stone terrain with the Spot robot | + +--------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | |position-inv-slope-spot| | |position-inv-slope-spot-link| | Track a position command on inverse slope terrain with the Spot robot | + +--------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + | |position-obstacle-spot| | |position-obstacle-spot-link| | Track a position command on obstacle terrain with the Spot robot | + +--------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ + +.. |position-gap-spot-link| replace:: `UW-Position-Gap-Spot-v0 `__ +.. |position-pit-spot-link| replace:: `UW-Position-Pit-Spot-v0 `__ +.. |position-stepping-stone-spot-link| replace:: `UW-Position-Stepping-Stone-Spot-v0 `__ +.. |position-obstacle-spot-link| replace:: `UW-Position-Obstacle-Spot-v0 `__ +.. |position-inv-slope-spot-link| replace:: `UW-Position-Inv-Slope-Spot-v0 `__ + +.. |position-gap-spot| image:: ../_static/tasks/locomotion/spot_gap.jpg +.. |position-pit-spot| image:: ../_static/tasks/locomotion/spot_pit.jpg +.. |position-stepping-stone-spot| image:: ../_static/tasks/locomotion/spot_stepping_stone.jpg +.. |position-obstacle-spot| image:: ../_static/tasks/locomotion/spot_obstacle.jpg +.. |position-inv-slope-spot| image:: ../_static/tasks/locomotion/spot_slope.jpg + + +.. raw:: html + +
+
+
+ + + +Comprehensive List of Environments +================================== + + +.. list-table:: + :widths: 33 25 19 25 + + * - **Task Name** + - **Inference Task Name** + - **Workflow** + - **RL Library** + * - UW-Track-Goal-Ur5-v0 + - + - Manager Based + - **rsl_rl** (PPO) + * - UW-Track-Goal-Tycho-v0 + - + - Manager Based + - **rsl_rl** (PPO) + * - UW-Track-Goal-Xarm-Leap-v0 + - + - Manager Based + - **rsl_rl** (PPO) + * - UW-Nut-Thread-Franka-v0 + - + - Manager Based + - **rsl_rl** (PPO) + * - UW-Gear-Mesh-Franka-v0 + - + - Manager Based + - **rsl_rl** (PPO) + * - UW-Peg-Insert-Franka-v0 + - + - Manager Based + - **rsl_rl** (PPO) + * - UW-Position-Gap-Spot-v0 + - + - Manager Based + - **rsl_rl** (PPO) + * - UW-Position-Pit-Spot-v0 + - + - Manager Based + - **rsl_rl** (PPO) + * - UW-Position-Stepping-Stone-Spot-v0 + - + - Manager Based + - **rsl_rl** (PPO) + * - UW-Position-Obstacle-Spot-v0 + - + - Manager Based + - **rsl_rl** (PPO) + * - UW-Position-Inv-Slope-Spot-v0 + - + - Manager Based + - **rsl_rl** (PPO) diff --git a/docs/source/publications/pg1.rst b/docs/source/publications/pg1.rst new file mode 100644 index 00000000..c09dec9e --- /dev/null +++ b/docs/source/publications/pg1.rst @@ -0,0 +1,45 @@ +Parental Guidance(PG1) +====================== + +Links +----- +- **Paper on OpenReview:** `Parental Guidance: Efficient Lifelong Learning through Evolutionary Distillation `_ +- **GitHub Repository:** `UW Lab GitHub `_ + +Authors +------- +**Zhengyu Zhang** †, **Quanquan Peng** ‡, **Rosario Scalise** †, **Byron Boots** † + +† Paul G Allen School, University of Washington +‡ Shanghai Jiao Tong University + +.. image:: ../../source/_static/publications/pg1/pg1.png + :alt: Research Illustration + :align: center + +Abstract +-------- +Developing robotic agents that can perform well in diverse environments while showing a variety of behaviors is +a key challenge in AI and robotics. Traditional reinforcement learning (RL) methods often create agents that specialize +in narrow tasks, limiting their adaptability and diversity. To overcome this, we propose a preliminary, +evolution-inspired framework that includes a reproduction module, similar to natural species reproduction, +balancing diversity and specialization. By integrating RL, imitation learning (IL), and a coevolutionary agent-terrain +curriculum, our system evolves agents continuously through complex tasks. This approach promotes adaptability, +inheritance of useful traits, and continual learning. Agents not only refine inherited skills but also surpass +their predecessors. Our initial experiments show that this method improves exploration efficiency and supports +open-ended learning, offering a scalable solution where sparse reward coupled with diverse terrain environments +induces a multi-task setting. + + +BibTex +---------- +.. code:: bibtex + + @inproceedings{ + zhang2024blending, + title={Blending Reinforcement Learning and Imitation Learning for Evolutionary Continual Learning}, + author={Zhengyu Zhang and Quanquan Peng and Rosario Scalise and Byron Boots}, + booktitle={[CoRL 2024] Morphology-Aware Policy and Design Learning Workshop (MAPoDeL)}, + year={2024}, + url={https://openreview.net/forum?id=d2VTtWOCMm} + } diff --git a/docs/source/refs/license.rst b/docs/source/refs/license.rst new file mode 100644 index 00000000..0eb154a2 --- /dev/null +++ b/docs/source/refs/license.rst @@ -0,0 +1,51 @@ +.. _license: + +License +======== + +NVIDIA Isaac Sim is available freely under `individual license +`_. For more information +about its license terms, please check `here `_. +The license files for all its dependencies and included assets are available in its +`documentation `_. + + +The Isaac Lab framework is open-sourced under the +`BSD-3-Clause license `_. + + +The UW Lab framework is open-sourced under the +`BSD-3-Clause license `_. + + +.. code-block:: text + + Copyright (c) 2022-2025, The UW Lab Project Developers. + All rights reserved. + + SPDX-License-Identifier: BSD-3-Clause + + Redistribution and use in source and binary forms, with or without modification, + are permitted provided that the following conditions are met: + + 1. Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + + 2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + + 3. Neither the name of the copyright holder nor the names of its contributors + may be used to endorse or promote products derived from this software without + specific prior written permission. + + THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND + ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED + WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR + ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES + (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; + LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND + ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS + SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/docs/source/setup/installation/local_installation.rst b/docs/source/setup/installation/local_installation.rst new file mode 100644 index 00000000..74d0faac --- /dev/null +++ b/docs/source/setup/installation/local_installation.rst @@ -0,0 +1,67 @@ +Installing UW Lab +=================== + +UW Lab builds on top of IsaacLab and IsaacSim. Please follow the below instructions to install UW Lab. + + +.. note:: + + If you use Conda, we recommend using `Miniconda `_. + +Install Isaac Lab and Isaac Sim +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. note:: + Please follow one of three ways to install Isaac Lab and Isaac Sim: + For best experience with vscode development, we recommend using Binary Installation. + For easy and quick installation, we recommend using pip installation. + For advanced users, we recommend using IsaacSim and IsaacLab pip installation. + + `IsaacLab with IsaacSim pip installation `_ + + + `IsaacLab with IsaacSim binary Installation `_ + + + `IsaacSim and IsaacLab pip installation `_ + + +Please also go through **Verifying the Isaac Lab Installation:** section in above links, +If the simulator does not run or crashes while following the above +instructions, it means that something is incorrectly configured. To +debug and troubleshoot, please check Isaac Sim +`documentation `__ +and the +`forums `__. + + +Install UW Lab +~~~~~~~~~~~~~~~~ + +- Make sure that your virtual environment is activated (if applicable) + +- Install UW Lab by cloning the repository and running the installation script + + .. code-block:: bash + + git clone https://github.com/UW-Lab/UWLab.git + +- Pip Install UW Lab in edible mode + + .. code-block:: bash + + cd UWLab + ./uwlab.sh -i + + +Verify UW Lab Installation +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Try running the following command to verify that UW Lab is installed correctly: + +.. code:: bash + + python scripts/reinforcement_learning/rsl_rl/train.py --task UW-Position-Pit-Spot-v0 --num_envs 1024 + + +Congratulations! You have successfully installed UW Lab. diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 00000000..aa7e6b31 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,100 @@ +[tool.isort] + +py_version = 310 +line_length = 120 +group_by_package = true + +# Files to skip +skip_glob = ["docs/*", "logs/*", "_isaac_sim/*", ".vscode/*"] + +# Order of imports +sections = [ + "FUTURE", + "STDLIB", + "THIRDPARTY", + "ASSETS_FIRSTPARTY", + "FIRSTPARTY", + "EXTRA_FIRSTPARTY", + "TASK_FIRSTPARTY", + "LOCALFOLDER", +] + +# Extra standard libraries considered as part of python (permissive licenses +extra_standard_library = [ + "numpy", + "h5py", + "open3d", + "torch", + "tensordict", + "bpy", + "matplotlib", + "gymnasium", + "gym", + "scipy", + "hid", + "yaml", + "prettytable", + "toml", + "trimesh", + "tqdm", + "torchvision", + "transformers", + "einops" # Needed for transformers, doesn't always auto-install +] +# Imports from Isaac Sim and Omniverse +known_third_party = [ + "isaacsim.core.api", + "isaacsim.replicator.common", + "omni.replicator.core", + "pxr", + "omni.kit.*", + "warp", + "carb", + "Semantics", +] +# Imports from this repository +known_first_party = ["isaaclab", "uwlab"] +known_assets_firstparty = ["isaaclab_assets", "uwlab_assets"] +known_extra_firstparty = [ + "isaaclab_rl", + "isaaclab_mimic", + "uwlab_rl", + "uwlab_apps" +] +known_task_firstparty = ["isaaclab_tasks", "uwlab_tasks"] +# Imports from the local folder +known_local_folder = "config" + +[tool.pyright] + +include = ["source", "scripts"] +exclude = [ + "**/__pycache__", + "**/_isaac_sim", + "**/docs", + "**/logs", + ".git", + ".vscode", +] + +typeCheckingMode = "basic" +pythonVersion = "3.10" +pythonPlatform = "Linux" +enableTypeIgnoreComments = true + +# This is required as the CI pre-commit does not download the module (i.e. numpy, torch, prettytable) +# Therefore, we have to ignore missing imports +reportMissingImports = "none" +# This is required to ignore for type checks of modules with stubs missing. +reportMissingModuleSource = "none" # -> most common: prettytable in mdp managers + +reportGeneralTypeIssues = "none" # -> raises 218 errors (usage of literal MISSING in dataclasses) +reportOptionalMemberAccess = "warning" # -> raises 8 errors +reportPrivateUsage = "warning" + + +[tool.codespell] +skip = '*.usd,*.svg,*.png,_isaac_sim*,*.bib,*.css,*/_build' +quiet-level = 0 +# the world list should always have words in lower case +ignore-words-list = "haa,slq,collapsable,buss" diff --git a/scripts/environments/list_envs.py b/scripts/environments/list_envs.py new file mode 100644 index 00000000..be757337 --- /dev/null +++ b/scripts/environments/list_envs.py @@ -0,0 +1,65 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +""" +Script to print all the available environments in Isaac Lab. + +The script iterates over all registered environments and stores the details in a table. +It prints the name of the environment, the entry point and the config file. + +All the environments are registered in the `isaaclab_tasks` extension. They start +with `Isaac` in their name. +""" + +"""Launch Isaac Sim Simulator first.""" + +from isaaclab.app import AppLauncher + +# launch omniverse app +app_launcher = AppLauncher(headless=True) +simulation_app = app_launcher.app + + +"""Rest everything follows.""" + +import gymnasium as gym +from prettytable import PrettyTable + +import isaaclab_tasks # noqa: F401 +import uwlab_tasks # noqa: F401 + + +def main(): + """Print all environments registered in `isaaclab_tasks` extension.""" + # print all the available environments + table = PrettyTable(["S. No.", "Task Name", "Entry Point", "Config"]) + table.title = "Available Environments in Isaac Lab" + # set alignment of table columns + table.align["Task Name"] = "l" + table.align["Entry Point"] = "l" + table.align["Config"] = "l" + + # count of environments + index = 0 + # acquire all Isaac environments names + for task_spec in gym.registry.values(): + if "Isaac" in task_spec.id or "UW" in task_spec.id or "UW" in task_spec.id: + # add details to table + table.add_row([index + 1, task_spec.id, task_spec.entry_point, task_spec.kwargs["env_cfg_entry_point"]]) + # increment count + index += 1 + + print(table) + + +if __name__ == "__main__": + try: + # run the main function + main() + except Exception as e: + raise e + finally: + # close the app + simulation_app.close() diff --git a/scripts/environments/random_agent.py b/scripts/environments/random_agent.py new file mode 100644 index 00000000..5d0d0b3c --- /dev/null +++ b/scripts/environments/random_agent.py @@ -0,0 +1,70 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Script to an environment with random action agent.""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Random agent for Isaac Lab environments.") +parser.add_argument( + "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations." +) +parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +# 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 gymnasium as gym +import torch + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.utils import parse_env_cfg + + +def main(): + """Random actions agent with Isaac Lab environment.""" + # create environment configuration + env_cfg = parse_env_cfg( + args_cli.task, device=args_cli.device, num_envs=args_cli.num_envs, use_fabric=not args_cli.disable_fabric + ) + # create environment + env = gym.make(args_cli.task, cfg=env_cfg) + + # print info (this is vectorized environment) + print(f"[INFO]: Gym observation space: {env.observation_space}") + print(f"[INFO]: Gym action space: {env.action_space}") + # reset environment + env.reset() + # simulate environment + while simulation_app.is_running(): + # run everything in inference mode + with torch.inference_mode(): + # sample actions from -1 to 1 + actions = 2 * torch.rand(env.action_space.shape, device=env.unwrapped.device) - 1 + # apply actions + env.step(actions) + + # close the simulator + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/environments/zero_agent.py b/scripts/environments/zero_agent.py new file mode 100644 index 00000000..be7cc0f5 --- /dev/null +++ b/scripts/environments/zero_agent.py @@ -0,0 +1,70 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Script to run an environment with zero action agent.""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Zero agent for Isaac Lab environments.") +parser.add_argument( + "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations." +) +parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +# 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 gymnasium as gym +import torch + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.utils import parse_env_cfg + + +def main(): + """Zero actions agent with Isaac Lab environment.""" + # parse configuration + env_cfg = parse_env_cfg( + args_cli.task, device=args_cli.device, num_envs=args_cli.num_envs, use_fabric=not args_cli.disable_fabric + ) + # create environment + env = gym.make(args_cli.task, cfg=env_cfg) + + # print info (this is vectorized environment) + print(f"[INFO]: Gym observation space: {env.observation_space}") + print(f"[INFO]: Gym action space: {env.action_space}") + # reset environment + env.reset() + # simulate environment + while simulation_app.is_running(): + # run everything in inference mode + with torch.inference_mode(): + # compute zero actions + actions = torch.zeros(env.action_space.shape, device=env.unwrapped.device) + # apply actions + env.step(actions) + + # close the simulator + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py b/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py new file mode 100644 index 00000000..e5d2667a --- /dev/null +++ b/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py @@ -0,0 +1,303 @@ +# Copyright (c) 2024-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Script to add mimic annotations to demos to be used as source demos for mimic dataset generation. +""" + +# Launching Isaac Sim Simulator first. + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Annotate demonstrations for Isaac Lab environments.") +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +parser.add_argument( + "--input_file", type=str, default="./datasets/dataset.hdf5", help="File name of the dataset to be annotated." +) +parser.add_argument( + "--output_file", + type=str, + default="./datasets/dataset_annotated.hdf5", + help="File name of the annotated output dataset file.", +) +parser.add_argument("--auto", action="store_true", default=False, help="Automatically annotate subtasks.") +parser.add_argument( + "--signals", + type=str, + nargs="+", + default=[], + help="Sequence of subtask termination signals for all except last subtask", +) + +# 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 torch + +import isaaclab_mimic.envs # noqa: F401 + +# Only enables inputs if this script is NOT headless mode +if not args_cli.headless and not os.environ.get("HEADLESS", 0): + from isaaclab.devices import Se3Keyboard +from isaaclab.envs import ManagerBasedRLMimicEnv +from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg +from isaaclab.managers import RecorderTerm, RecorderTermCfg +from isaaclab.utils import configclass +from isaaclab.utils.datasets import HDF5DatasetFileHandler + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.utils.parse_cfg import parse_env_cfg + +is_paused = False +current_action_index = 0 +subtask_indices = [] + + +def play_cb(): + global is_paused + is_paused = False + + +def pause_cb(): + global is_paused + is_paused = True + + +def mark_subtask_cb(): + global current_action_index, subtask_indices + subtask_indices.append(current_action_index) + print(f"Marked subtask at action index: {current_action_index}") + + +class PreStepDatagenInfoRecorder(RecorderTerm): + """Recorder term that records the datagen info data in each step.""" + + def record_pre_step(self): + eef_pose_dict = {} + for eef_name in self._env.cfg.subtask_configs.keys(): + eef_pose_dict[eef_name] = self._env.get_robot_eef_pose(eef_name) + + datagen_info = { + "object_pose": self._env.get_object_poses(), + "eef_pose": eef_pose_dict, + "target_eef_pose": self._env.action_to_target_eef_pose(self._env.action_manager.action), + } + return "obs/datagen_info", datagen_info + + +@configclass +class PreStepDatagenInfoRecorderCfg(RecorderTermCfg): + """Configuration for the datagen info recorder term.""" + + class_type: type[RecorderTerm] = PreStepDatagenInfoRecorder + + +class PreStepSubtaskTermsObservationsRecorder(RecorderTerm): + """Recorder term that records the subtask completion observations in each step.""" + + def record_pre_step(self): + return "obs/datagen_info/subtask_term_signals", self._env.get_subtask_term_signals() + + +@configclass +class PreStepSubtaskTermsObservationsRecorderCfg(RecorderTermCfg): + """Configuration for the step subtask terms observation recorder term.""" + + class_type: type[RecorderTerm] = PreStepSubtaskTermsObservationsRecorder + + +@configclass +class MimicRecorderManagerCfg(ActionStateRecorderManagerCfg): + """Mimic specific recorder terms.""" + + record_pre_step_datagen_info = PreStepDatagenInfoRecorderCfg() + record_pre_step_subtask_term_signals = PreStepSubtaskTermsObservationsRecorderCfg() + + +def main(): + """Add Isaac Lab Mimic annotations to the given demo dataset file.""" + global is_paused, current_action_index, subtask_indices + + if not args_cli.auto and len(args_cli.signals) == 0: + if len(args_cli.signals) == 0: + raise ValueError("Subtask signals should be provided for manual mode.") + + # Load input dataset to be annotated + if not os.path.exists(args_cli.input_file): + raise FileNotFoundError(f"The input dataset file {args_cli.input_file} does not exist.") + dataset_file_handler = HDF5DatasetFileHandler() + dataset_file_handler.open(args_cli.input_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() + + # get output directory path and file name (without extension) from cli arguments + output_dir = os.path.dirname(args_cli.output_file) + output_file_name = os.path.splitext(os.path.basename(args_cli.output_file))[0] + # create output directory if it does not exist + if not os.path.exists(output_dir): + os.makedirs(output_dir) + + 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.") + + env_cfg = parse_env_cfg(env_name, device=args_cli.device, num_envs=1) + + env_cfg.env_name = args_cli.task + + # extract success checking function to invoke manually + success_term = None + if hasattr(env_cfg.terminations, "success"): + success_term = env_cfg.terminations.success + env_cfg.terminations.success = None + else: + raise NotImplementedError("No success termination term was found in the environment.") + + # Disable all termination terms + env_cfg.terminations = None + + # Set up recorder terms for mimic annotations + env_cfg.recorders: MimicRecorderManagerCfg = MimicRecorderManagerCfg() + if not args_cli.auto: + # disable subtask term signals recorder term if in manual mode + env_cfg.recorders.record_pre_step_subtask_term_signals = None + env_cfg.recorders.dataset_export_dir_path = output_dir + env_cfg.recorders.dataset_filename = output_file_name + + # create environment from loaded config + env = gym.make(args_cli.task, cfg=env_cfg).unwrapped + + if not isinstance(env.unwrapped, ManagerBasedRLMimicEnv): + raise ValueError("The environment should be derived from ManagerBasedRLMimicEnv") + + if args_cli.auto: + # check if the mimic API env.unwrapped.get_subtask_term_signals() is implemented + if env.unwrapped.get_subtask_term_signals.__func__ is ManagerBasedRLMimicEnv.get_subtask_term_signals: + raise NotImplementedError( + "The environment does not implement the get_subtask_term_signals method required " + "to run automatic annotations." + ) + + # reset environment + env.reset() + + # Only enables inputs if this script is NOT headless mode + if not args_cli.headless and not os.environ.get("HEADLESS", 0): + keyboard_interface = Se3Keyboard(pos_sensitivity=0.1, rot_sensitivity=0.1) + keyboard_interface.add_callback("N", play_cb) + keyboard_interface.add_callback("B", pause_cb) + if not args_cli.auto: + keyboard_interface.add_callback("S", mark_subtask_cb) + keyboard_interface.reset() + + # simulate environment -- run everything in inference mode + exported_episode_count = 0 + processed_episode_count = 0 + with contextlib.suppress(KeyboardInterrupt) and torch.inference_mode(): + while simulation_app.is_running() and not simulation_app.is_exiting(): + # Iterate over the episodes in the loaded dataset file + for episode_index, episode_name in enumerate(dataset_file_handler.get_episode_names()): + processed_episode_count += 1 + subtask_indices = [] + print(f"\nAnnotating episode #{episode_index} ({episode_name})") + episode = dataset_file_handler.load_episode(episode_name, env.unwrapped.device) + episode_data = episode.data + + # read initial state from the loaded episode + initial_state = episode_data["initial_state"] + env.unwrapped.recorder_manager.reset() + env.unwrapped.reset_to(initial_state, None, is_relative=True) + + # replay actions from this episode + actions = episode_data["actions"] + first_action = True + for action_index, action in enumerate(actions): + current_action_index = action_index + if first_action: + first_action = False + else: + while is_paused: + env.unwrapped.sim.render() + continue + action_tensor = torch.Tensor(action).reshape([1, action.shape[0]]) + env.step(torch.Tensor(action_tensor)) + + is_episode_annotated_successfully = False + if not args_cli.auto: + print(f"\tSubtasks marked at action indices: {subtask_indices}") + if len(args_cli.signals) != len(subtask_indices): + raise ValueError( + f"Number of annotated subtask signals {len(subtask_indices)} should be equal " + f" to number of subtasks {len(args_cli.signals)}" + ) + annotated_episode = env.unwrapped.recorder_manager.get_episode(0) + for subtask_index in range(len(args_cli.signals)): + # subtask termination signal is false until subtask is complete, and true afterwards + subtask_signals = torch.ones(len(actions), dtype=torch.bool) + subtask_signals[: subtask_indices[subtask_index]] = False + annotated_episode.add( + f"obs/datagen_info/subtask_term_signals/{args_cli.signals[subtask_index]}", subtask_signals + ) + is_episode_annotated_successfully = True + else: + # check if all the subtask term signals are annotated + annotated_episode = env.unwrapped.recorder_manager.get_episode(0) + subtask_term_signal_dict = annotated_episode.data["obs"]["datagen_info"]["subtask_term_signals"] + is_episode_annotated_successfully = True + for signal_name, signal_flags in subtask_term_signal_dict.items(): + if not torch.any(signal_flags): + is_episode_annotated_successfully = False + print(f'\tDid not detect completion for the subtask "{signal_name}".') + + if not bool(success_term.func(env, **success_term.params)[0]): + is_episode_annotated_successfully = False + print("\tThe final task was not completed.") + + if is_episode_annotated_successfully: + # set success to the recorded episode data and export to file + env.unwrapped.recorder_manager.set_success_to_episodes( + None, torch.tensor([[True]], dtype=torch.bool, device=env.unwrapped.device) + ) + env.unwrapped.recorder_manager.export_episodes() + exported_episode_count += 1 + print("\tExported the annotated episode.") + else: + print("\tSkipped exporting the episode due to incomplete subtask annotations.") + break + + print( + f"\nExported {exported_episode_count} (out of {processed_episode_count}) annotated" + f" episode{'s' if exported_episode_count > 1 else ''}." + ) + print("Exiting the app.") + + # Close environment after annotation is complete + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/imitation_learning/isaaclab_mimic/consolidated_demo.py b/scripts/imitation_learning/isaaclab_mimic/consolidated_demo.py new file mode 100644 index 00000000..e444d35d --- /dev/null +++ b/scripts/imitation_learning/isaaclab_mimic/consolidated_demo.py @@ -0,0 +1,473 @@ +# Copyright (c) 2024-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Script to record teleoperated demos and run mimic dataset generation in real-time. +""" + +# Launching Isaac Sim Simulator first. + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser( + description="Record demonstrations and run mimic dataset generation for Isaac Lab environments." +) +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +parser.add_argument( + "--num_demos", type=int, default=0, help="Number of demonstrations to record. Set to 0 for infinite." +) +parser.add_argument( + "--num_success_steps", + type=int, + default=10, + help="Number of continuous steps with task success for concluding a demo as successful. Default is 10.", +) +parser.add_argument( + "--num_envs", + type=int, + default=5, + help=( + "Number of environments to instantiate to test recording and generating datasets. The environment specified by" + " `teleop_env_index` will be used for teleoperation and recording while the remaining environments will be used" + " for real-time data generation. Default is 5." + ), +) +parser.add_argument( + "--teleop_env_index", + type=int, + default=0, + help="Index of the environment to be used for teleoperation. Set -1 for disabling the teleop robot. Default is 0.", +) +parser.add_argument("--teleop_device", type=str, default="keyboard", help="Device for interacting with environment.") +parser.add_argument( + "--step_hz", type=int, default=0, help="Environment stepping rate in Hz. Set to 0 for maximum speed." +) +parser.add_argument("--input_file", type=str, default=None, help="File path to the source demo dataset file.") +parser.add_argument( + "--output_file", + type=str, + default="./datasets/output_dataset.hdf5", + help="File path to export recorded episodes.", +) +parser.add_argument( + "--generated_output_file", + type=str, + default=None, + help="File path to export generated episodes by mimic.", +) +# 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 asyncio +import contextlib +import gymnasium as gym +import numpy as np +import os +import random +import time +import torch + +from isaaclab.devices import Se3Keyboard, Se3SpaceMouse +from isaaclab.envs import ManagerBasedRLMimicEnv +from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg +from isaaclab.managers import DatasetExportMode, RecorderTerm, RecorderTermCfg +from isaaclab.utils import configclass +from isaaclab.utils.datasets import HDF5DatasetFileHandler + +import isaaclab_mimic.envs # noqa: F401 +from isaaclab_mimic.datagen.data_generator import DataGenerator +from isaaclab_mimic.datagen.datagen_info_pool import DataGenInfoPool + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.utils.parse_cfg import parse_env_cfg + +# global variable to keep track of the data generation statistics +num_recorded = 0 +num_success = 0 +num_failures = 0 +num_attempts = 0 + + +class PreStepDatagenInfoRecorder(RecorderTerm): + """Recorder term that records the datagen info data in each step.""" + + def record_pre_step(self): + eef_pose_dict = {} + for eef_name in self._env.cfg.subtask_configs.keys(): + eef_pose_dict[eef_name] = self._env.get_robot_eef_pose(eef_name) + + datagen_info = { + "object_pose": self._env.get_object_poses(), + "eef_pose": eef_pose_dict, + "target_eef_pose": self._env.action_to_target_eef_pose(self._env.action_manager.action), + } + return "obs/datagen_info", datagen_info + + +@configclass +class PreStepDatagenInfoRecorderCfg(RecorderTermCfg): + """Configuration for the datagen info recorder term.""" + + class_type: type[RecorderTerm] = PreStepDatagenInfoRecorder + + +class PreStepSubtaskTermsObservationsRecorder(RecorderTerm): + """Recorder term that records the subtask completion observations in each step.""" + + def record_pre_step(self): + return "obs/datagen_info/subtask_term_signals", self._env.get_subtask_term_signals() + + +@configclass +class PreStepSubtaskTermsObservationsRecorderCfg(RecorderTermCfg): + """Configuration for the step subtask terms observation recorder term.""" + + class_type: type[RecorderTerm] = PreStepSubtaskTermsObservationsRecorder + + +@configclass +class MimicRecorderManagerCfg(ActionStateRecorderManagerCfg): + """Mimic specific recorder terms.""" + + record_pre_step_datagen_info = PreStepDatagenInfoRecorderCfg() + record_pre_step_subtask_term_signals = PreStepSubtaskTermsObservationsRecorderCfg() + + +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) + + +async def run_teleop_robot( + env, env_id, env_action_queue, shared_datagen_info_pool, success_term, exported_dataset_path, teleop_interface=None +): + """Run teleop robot.""" + global num_recorded + should_reset_teleop_instance = False + # create controller if needed + if teleop_interface is None: + 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 + def reset_teleop_instance(): + nonlocal should_reset_teleop_instance + should_reset_teleop_instance = True + + teleop_interface.add_callback("R", reset_teleop_instance) + + teleop_interface.reset() + print(teleop_interface) + + recorded_episode_dataset_file_handler = HDF5DatasetFileHandler() + recorded_episode_dataset_file_handler.create(exported_dataset_path, env_name=env.unwrapped.cfg.env_name) + + env_id_tensor = torch.tensor([env_id], dtype=torch.int64, device=env.device) + success_step_count = 0 + num_recorded = 0 + while True: + if should_reset_teleop_instance: + env.unwrapped.recorder_manager.reset(env_id_tensor) + env.unwrapped.reset(env_ids=env_id_tensor) + should_reset_teleop_instance = False + success_step_count = 0 + + # 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(1, 1) + # compute actions based on environment + teleop_action = pre_process_actions(delta_pose, gripper_command) + + await env_action_queue.put((env_id, teleop_action)) + await env_action_queue.join() + + if success_term is not None: + if bool(success_term.func(env, **success_term.params)[env_id]): + success_step_count += 1 + if success_step_count >= args_cli.num_success_steps: + env.recorder_manager.set_success_to_episodes( + env_id_tensor, torch.tensor([[True]], dtype=torch.bool, device=env.device) + ) + teleop_episode = env.unwrapped.recorder_manager.get_episode(env_id) + await shared_datagen_info_pool.add_episode(teleop_episode) + + recorded_episode_dataset_file_handler.write_episode(teleop_episode) + recorded_episode_dataset_file_handler.flush() + env.recorder_manager.reset(env_id_tensor) + num_recorded += 1 + should_reset_teleop_instance = True + else: + success_step_count = 0 + + +async def run_data_generator( + env, env_id, env_action_queue, shared_datagen_info_pool, success_term, pause_subtask=False, export_demo=True +): + """Run data generator.""" + global num_success, num_failures, num_attempts + data_generator = DataGenerator(env=env.unwrapped, src_demo_datagen_info_pool=shared_datagen_info_pool) + idle_action = torch.zeros(env.unwrapped.action_space.shape)[0] + while True: + while data_generator.src_demo_datagen_info_pool.num_datagen_infos < 1: + await env_action_queue.put((env_id, idle_action)) + await env_action_queue.join() + + results = await data_generator.generate( + env_id=env_id, + success_term=success_term, + env_action_queue=env_action_queue, + select_src_per_subtask=env.unwrapped.cfg.datagen_config.generation_select_src_per_subtask, + transform_first_robot_pose=env.unwrapped.cfg.datagen_config.generation_transform_first_robot_pose, + interpolate_from_last_target_pose=env.unwrapped.cfg.datagen_config.generation_interpolate_from_last_target_pose, + pause_subtask=pause_subtask, + export_demo=export_demo, + ) + if bool(results["success"]): + num_success += 1 + else: + num_failures += 1 + num_attempts += 1 + + +def env_loop(env, env_action_queue, shared_datagen_info_pool, asyncio_event_loop): + """Main loop for the environment.""" + global num_recorded, num_success, num_failures, num_attempts + prev_num_attempts = 0 + prev_num_recorded = 0 + + rate_limiter = None + if args_cli.step_hz > 0: + rate_limiter = RateLimiter(args_cli.step_hz) + + # simulate environment -- run everything in inference mode + is_first_print = True + with contextlib.suppress(KeyboardInterrupt) and torch.inference_mode(): + while True: + actions = torch.zeros(env.unwrapped.action_space.shape) + + # get actions from all the data generators + for i in range(env.unwrapped.num_envs): + # an async-blocking call to get an action from a data generator + env_id, action = asyncio_event_loop.run_until_complete(env_action_queue.get()) + actions[env_id] = action + + # perform action on environment + env.step(actions) + + # mark done so the data generators can continue with the step results + for i in range(env.unwrapped.num_envs): + env_action_queue.task_done() + + if prev_num_attempts != num_attempts or prev_num_recorded != num_recorded: + prev_num_attempts = num_attempts + prev_num_recorded = num_recorded + generated_sucess_rate = 100 * num_success / num_attempts if num_attempts > 0 else 0.0 + if is_first_print: + is_first_print = False + else: + print("\r", "\033[F" * 5, end="") + print("") + print("*" * 50, "\033[K") + print(f"{num_recorded} teleoperated demos recorded\033[K") + print( + f"{num_success}/{num_attempts} ({generated_sucess_rate:.1f}%) successful demos generated by" + " mimic\033[K" + ) + print("*" * 50, "\033[K") + + if args_cli.num_demos > 0 and num_recorded >= args_cli.num_demos: + print(f"All {args_cli.num_demos} demonstrations recorded. Exiting the app.") + break + + # check that simulation is stopped or not + if env.unwrapped.sim.is_stopped(): + break + + if rate_limiter: + rate_limiter.sleep(env.unwrapped) + env.close() + + +def main(): + num_envs = args_cli.num_envs + + # create output directory for recorded episodes if it does not exist + recorded_output_dir = os.path.dirname(args_cli.output_file) + if not os.path.exists(recorded_output_dir): + os.makedirs(recorded_output_dir) + + # check if the given input dataset file exists + if args_cli.input_file and not os.path.exists(args_cli.input_file): + raise FileNotFoundError(f"The dataset file {args_cli.input_file} does not exist.") + + # get the environment name + if args_cli.task is not None: + env_name = args_cli.task + elif args_cli.input_file: + # if the environment name is not specified, try to get it from the dataset file + dataset_file_handler = HDF5DatasetFileHandler() + dataset_file_handler.open(args_cli.input_file) + env_name = dataset_file_handler.get_env_name() + else: + raise ValueError("Task/env name was not specified nor found in the dataset.") + + # parse configuration + env_cfg = parse_env_cfg(env_name, device=args_cli.device, num_envs=num_envs) + env_cfg.env_name = env_name + + # extract success checking function to invoke manually + success_term = None + if hasattr(env_cfg.terminations, "success"): + success_term = env_cfg.terminations.success + env_cfg.terminations.success = None + else: + raise NotImplementedError("No success termination term was found in the environment.") + + # data generator is in charge of resetting the environment + env_cfg.terminations = None + + env_cfg.observations.policy.concatenate_terms = False + + env_cfg.recorders = MimicRecorderManagerCfg() + + env_cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_NONE + if args_cli.generated_output_file: + # create output directory for generated episodes if it does not exist + generated_output_dir = os.path.dirname(args_cli.generated_output_file) + if not os.path.exists(generated_output_dir): + os.makedirs(generated_output_dir) + generated_output_file_name = os.path.splitext(os.path.basename(args_cli.generated_output_file))[0] + env_cfg.recorders.dataset_export_dir_path = generated_output_dir + env_cfg.recorders.dataset_filename = generated_output_file_name + env_cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_SUCCEEDED_ONLY + + # create environment + env = gym.make(env_name, cfg=env_cfg) + + if not isinstance(env.unwrapped, ManagerBasedRLMimicEnv): + raise ValueError("The environment should be derived from ManagerBasedRLMimicEnv") + + # check if the mimic API env.unwrapped.get_subtask_term_signals() is implemented + if env.unwrapped.get_subtask_term_signals.__func__ is ManagerBasedRLMimicEnv.get_subtask_term_signals: + raise NotImplementedError( + "The environment does not implement the get_subtask_term_signals method required to run this script." + ) + + # set seed for generation + random.seed(env.unwrapped.cfg.datagen_config.seed) + np.random.seed(env.unwrapped.cfg.datagen_config.seed) + torch.manual_seed(env.unwrapped.cfg.datagen_config.seed) + + # reset before starting + env.reset() + + # Set up asyncio stuff + asyncio_event_loop = asyncio.get_event_loop() + env_action_queue = asyncio.Queue() + + shared_datagen_info_pool_lock = asyncio.Lock() + shared_datagen_info_pool = DataGenInfoPool( + env.unwrapped, env.unwrapped.cfg, env.unwrapped.device, asyncio_lock=shared_datagen_info_pool_lock + ) + if args_cli.input_file: + shared_datagen_info_pool.load_from_dataset_file(args_cli.input_file) + print(f"Loaded {shared_datagen_info_pool.num_datagen_infos} to datagen info pool") + + # make data generator object + data_generator_asyncio_tasks = [] + for i in range(num_envs): + if args_cli.teleop_env_index is not None and i == args_cli.teleop_env_index: + data_generator_asyncio_tasks.append( + asyncio_event_loop.create_task( + run_teleop_robot( + env, i, env_action_queue, shared_datagen_info_pool, success_term, args_cli.output_file + ) + ) + ) + continue + data_generator_asyncio_tasks.append( + asyncio_event_loop.create_task( + run_data_generator( + env, + i, + env_action_queue, + shared_datagen_info_pool, + success_term, + export_demo=bool(args_cli.generated_output_file), + ) + ) + ) + + try: + asyncio.ensure_future(asyncio.gather(*data_generator_asyncio_tasks)) + except asyncio.CancelledError: + print("Tasks were cancelled.") + + env_loop(env, env_action_queue, shared_datagen_info_pool, asyncio_event_loop) + + +if __name__ == "__main__": + try: + main() + except KeyboardInterrupt: + print("\nProgram interrupted by user. Exiting...") + # close sim app + simulation_app.close() diff --git a/scripts/imitation_learning/isaaclab_mimic/generate_dataset.py b/scripts/imitation_learning/isaaclab_mimic/generate_dataset.py new file mode 100644 index 00000000..37e99375 --- /dev/null +++ b/scripts/imitation_learning/isaaclab_mimic/generate_dataset.py @@ -0,0 +1,109 @@ +# Copyright (c) 2024-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Main data generation script. +""" + +# Launching Isaac Sim Simulator first. + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Generate demonstrations for Isaac Lab environments.") +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +parser.add_argument("--generation_num_trials", type=int, help="Number of demos to be generated.", default=None) +parser.add_argument( + "--num_envs", type=int, default=1, help="Number of environments to instantiate for generating datasets." +) +parser.add_argument("--input_file", type=str, default=None, required=True, help="File path to the source dataset file.") +parser.add_argument( + "--output_file", + type=str, + default="./datasets/output_dataset.hdf5", + help="File path to export recorded and generated episodes.", +) +parser.add_argument( + "--pause_subtask", + action="store_true", + help="pause after every subtask during generation for debugging - only useful with render flag", +) +# 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 asyncio +import gymnasium as gym +import numpy as np +import random +import torch + +import isaaclab_mimic.envs # noqa: F401 +from isaaclab_mimic.datagen.generation import env_loop, setup_async_generation, setup_env_config +from isaaclab_mimic.datagen.utils import get_env_name_from_dataset, setup_output_paths + +import isaaclab_tasks # noqa: F401 + + +def main(): + num_envs = args_cli.num_envs + + # Setup output paths and get env name + output_dir, output_file_name = setup_output_paths(args_cli.output_file) + env_name = args_cli.task or get_env_name_from_dataset(args_cli.input_file) + + # Configure environment + env_cfg, success_term = setup_env_config( + env_name=env_name, + output_dir=output_dir, + output_file_name=output_file_name, + num_envs=num_envs, + device=args_cli.device, + generation_num_trials=args_cli.generation_num_trials, + ) + + # create environment + env = gym.make(env_name, cfg=env_cfg) + + # set seed for generation + random.seed(env.unwrapped.cfg.datagen_config.seed) + np.random.seed(env.unwrapped.cfg.datagen_config.seed) + torch.manual_seed(env.unwrapped.cfg.datagen_config.seed) + + # reset before starting + env.reset() + + # Setup and run async data generation + async_components = setup_async_generation( + env=env, + num_envs=args_cli.num_envs, + input_file=args_cli.input_file, + success_term=success_term, + pause_subtask=args_cli.pause_subtask, + ) + + try: + asyncio.ensure_future(asyncio.gather(*async_components["tasks"])) + env_loop(env, async_components["action_queue"], async_components["info_pool"], async_components["event_loop"]) + except asyncio.CancelledError: + print("Tasks were cancelled.") + + +if __name__ == "__main__": + try: + main() + except KeyboardInterrupt: + print("\nProgram interrupted by user. Exiting...") + # close sim app + simulation_app.close() diff --git a/scripts/imitation_learning/robomimic/play.py b/scripts/imitation_learning/robomimic/play.py new file mode 100644 index 00000000..10377252 --- /dev/null +++ b/scripts/imitation_learning/robomimic/play.py @@ -0,0 +1,123 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Script to play and evaluate a trained policy from robomimic.""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Evaluate robomimic policy for Isaac Lab environment.") +parser.add_argument( + "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations." +) +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +parser.add_argument("--checkpoint", type=str, default=None, help="Pytorch model checkpoint to load.") +parser.add_argument("--horizon", type=int, default=800, help="Step horizon of each rollout.") +parser.add_argument("--num_rollouts", type=int, default=1, help="Number of rollouts.") +parser.add_argument("--seed", type=int, default=101, help="Random seed.") + +# 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 gymnasium as gym +import torch + +import robomimic.utils.file_utils as FileUtils +import robomimic.utils.torch_utils as TorchUtils + +from isaaclab_tasks.utils import parse_env_cfg + + +def rollout(policy, env, horizon, device): + policy.start_episode + obs_dict, _ = env.reset() + traj = dict(actions=[], obs=[], next_obs=[]) + + for i in range(horizon): + # Prepare observations + obs = obs_dict["policy"] + for ob in obs: + obs[ob] = torch.squeeze(obs[ob]) + traj["obs"].append(obs) + + # Compute actions + actions = policy(obs) + actions = torch.from_numpy(actions).to(device=device).view(1, env.action_space.shape[1]) + + # Apply actions + obs_dict, _, terminated, truncated, _ = env.step(actions) + obs = obs_dict["policy"] + + # Record trajectory + traj["actions"].append(actions.tolist()) + traj["next_obs"].append(obs) + + if terminated: + return True, traj + elif truncated: + return False, traj + + return False, traj + + +def main(): + """Run a trained policy from robomimic with Isaac Lab environment.""" + # parse configuration + env_cfg = parse_env_cfg(args_cli.task, device=args_cli.device, num_envs=1, use_fabric=not args_cli.disable_fabric) + + # Set observations to dictionary mode for Robomimic + env_cfg.observations.policy.concatenate_terms = False + + # Set termination conditions + env_cfg.terminations.time_out = None + + # Disable recorder + env_cfg.recorders = None + + # Create environment + env = gym.make(args_cli.task, cfg=env_cfg).unwrapped + + # Set seed + torch.manual_seed(args_cli.seed) + env.seed(args_cli.seed) + + # Acquire device + device = TorchUtils.get_torch_device(try_to_use_cuda=True) + + # Load policy + policy, _ = FileUtils.policy_from_checkpoint(ckpt_path=args_cli.checkpoint, device=device, verbose=True) + + # Run policy + results = [] + for trial in range(args_cli.num_rollouts): + print(f"[INFO] Starting trial {trial}") + terminated, traj = rollout(policy, env, args_cli.horizon, device) + results.append(terminated) + print(f"[INFO] Trial {trial}: {terminated}\n") + + print(f"\nSuccessful trials: {results.count(True)}, out of {len(results)} trials") + print(f"Success rate: {results.count(True) / len(results)}") + print(f"Trial Results: {results}\n") + + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/imitation_learning/robomimic/train.py b/scripts/imitation_learning/robomimic/train.py new file mode 100644 index 00000000..377a5fc4 --- /dev/null +++ b/scripts/imitation_learning/robomimic/train.py @@ -0,0 +1,360 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# MIT License +# +# Copyright (c) 2021 Stanford Vision and Learning Lab +# +# 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. + +""" +The main entry point for training policies from pre-collected data. + +Args: + algo: name of the algorithm to run. + task: name of the environment. + name: if provided, override the experiment name defined in the config + dataset: if provided, override the dataset path defined in the config + +This file has been modified from the original version in the following ways: + +""" + +"""Launch Isaac Sim Simulator first.""" + +from isaaclab.app import AppLauncher + +# launch omniverse app +app_launcher = AppLauncher(headless=True) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import argparse +import gymnasium as gym +import json +import numpy as np +import os +import sys +import time +import torch +import traceback +from collections import OrderedDict +from torch.utils.data import DataLoader + +import psutil +import robomimic.utils.env_utils as EnvUtils +import robomimic.utils.file_utils as FileUtils +import robomimic.utils.obs_utils as ObsUtils +import robomimic.utils.torch_utils as TorchUtils +import robomimic.utils.train_utils as TrainUtils +from robomimic.algo import algo_factory +from robomimic.config import config_factory +from robomimic.utils.log_utils import DataLogger, PrintLogger + +# Needed so that environment is registered +import isaaclab_tasks # noqa: F401 + + +def train(config, device): + """Train a model using the algorithm.""" + # first set seeds + np.random.seed(config.train.seed) + torch.manual_seed(config.train.seed) + + print("\n============= New Training Run with Config =============") + print(config) + print("") + log_dir, ckpt_dir, video_dir = TrainUtils.get_exp_dir(config) + + print(f">>> Saving logs into directory: {log_dir}") + print(f">>> Saving checkpoints into directory: {ckpt_dir}") + print(f">>> Saving videos into directory: {video_dir}") + + if config.experiment.logging.terminal_output_to_txt: + # log stdout and stderr to a text file + logger = PrintLogger(os.path.join(log_dir, "log.txt")) + sys.stdout = logger + sys.stderr = logger + + # read config to set up metadata for observation modalities (e.g. detecting rgb observations) + ObsUtils.initialize_obs_utils_with_config(config) + + # make sure the dataset exists + dataset_path = os.path.expanduser(config.train.data) + if not os.path.exists(dataset_path): + raise FileNotFoundError(f"Dataset at provided path {dataset_path} not found!") + + # load basic metadata from training file + print("\n============= Loaded Environment Metadata =============") + env_meta = FileUtils.get_env_metadata_from_dataset(dataset_path=config.train.data) + shape_meta = FileUtils.get_shape_metadata_from_dataset( + dataset_path=config.train.data, all_obs_keys=config.all_obs_keys, verbose=True + ) + + if config.experiment.env is not None: + env_meta["env_name"] = config.experiment.env + print("=" * 30 + "\n" + "Replacing Env to {}\n".format(env_meta["env_name"]) + "=" * 30) + + # create environment + envs = OrderedDict() + if config.experiment.rollout.enabled: + # create environments for validation runs + env_names = [env_meta["env_name"]] + + if config.experiment.additional_envs is not None: + for name in config.experiment.additional_envs: + env_names.append(name) + + for env_name in env_names: + env = EnvUtils.create_env_from_metadata( + env_meta=env_meta, + env_name=env_name, + render=False, + render_offscreen=config.experiment.render_video, + use_image_obs=shape_meta["use_images"], + ) + envs[env.name] = env + print(envs[env.name]) + + print("") + + # setup for a new training run + data_logger = DataLogger(log_dir, config=config, log_tb=config.experiment.logging.log_tb) + model = algo_factory( + algo_name=config.algo_name, + config=config, + obs_key_shapes=shape_meta["all_shapes"], + ac_dim=shape_meta["ac_dim"], + device=device, + ) + + # save the config as a json file + with open(os.path.join(log_dir, "..", "config.json"), "w") as outfile: + json.dump(config, outfile, indent=4) + + print("\n============= Model Summary =============") + print(model) # print model summary + print("") + + # load training data + trainset, validset = TrainUtils.load_data_for_training(config, obs_keys=shape_meta["all_obs_keys"]) + train_sampler = trainset.get_dataset_sampler() + print("\n============= Training Dataset =============") + print(trainset) + print("") + + # maybe retrieve statistics for normalizing observations + obs_normalization_stats = None + if config.train.hdf5_normalize_obs: + obs_normalization_stats = trainset.get_obs_normalization_stats() + + # initialize data loaders + train_loader = DataLoader( + dataset=trainset, + sampler=train_sampler, + batch_size=config.train.batch_size, + shuffle=(train_sampler is None), + num_workers=config.train.num_data_workers, + drop_last=True, + ) + + if config.experiment.validate: + # cap num workers for validation dataset at 1 + num_workers = min(config.train.num_data_workers, 1) + valid_sampler = validset.get_dataset_sampler() + valid_loader = DataLoader( + dataset=validset, + sampler=valid_sampler, + batch_size=config.train.batch_size, + shuffle=(valid_sampler is None), + num_workers=num_workers, + drop_last=True, + ) + else: + valid_loader = None + + # main training loop + best_valid_loss = None + last_ckpt_time = time.time() + + # number of learning steps per epoch (defaults to a full dataset pass) + train_num_steps = config.experiment.epoch_every_n_steps + valid_num_steps = config.experiment.validation_epoch_every_n_steps + + for epoch in range(1, config.train.num_epochs + 1): # epoch numbers start at 1 + step_log = TrainUtils.run_epoch(model=model, data_loader=train_loader, epoch=epoch, num_steps=train_num_steps) + model.on_epoch_end(epoch) + + # setup checkpoint path + epoch_ckpt_name = f"model_epoch_{epoch}" + + # check for recurring checkpoint saving conditions + should_save_ckpt = False + if config.experiment.save.enabled: + time_check = (config.experiment.save.every_n_seconds is not None) and ( + time.time() - last_ckpt_time > config.experiment.save.every_n_seconds + ) + epoch_check = ( + (config.experiment.save.every_n_epochs is not None) + and (epoch > 0) + and (epoch % config.experiment.save.every_n_epochs == 0) + ) + epoch_list_check = epoch in config.experiment.save.epochs + should_save_ckpt = time_check or epoch_check or epoch_list_check + ckpt_reason = None + if should_save_ckpt: + last_ckpt_time = time.time() + ckpt_reason = "time" + + print(f"Train Epoch {epoch}") + print(json.dumps(step_log, sort_keys=True, indent=4)) + for k, v in step_log.items(): + if k.startswith("Time_"): + data_logger.record(f"Timing_Stats/Train_{k[5:]}", v, epoch) + else: + data_logger.record(f"Train/{k}", v, epoch) + + # Evaluate the model on validation set + if config.experiment.validate: + with torch.no_grad(): + step_log = TrainUtils.run_epoch( + model=model, data_loader=valid_loader, epoch=epoch, validate=True, num_steps=valid_num_steps + ) + for k, v in step_log.items(): + if k.startswith("Time_"): + data_logger.record(f"Timing_Stats/Valid_{k[5:]}", v, epoch) + else: + data_logger.record(f"Valid/{k}", v, epoch) + + print(f"Validation Epoch {epoch}") + print(json.dumps(step_log, sort_keys=True, indent=4)) + + # save checkpoint if achieve new best validation loss + valid_check = "Loss" in step_log + if valid_check and (best_valid_loss is None or (step_log["Loss"] <= best_valid_loss)): + best_valid_loss = step_log["Loss"] + if config.experiment.save.enabled and config.experiment.save.on_best_validation: + epoch_ckpt_name += f"_best_validation_{best_valid_loss}" + should_save_ckpt = True + ckpt_reason = "valid" if ckpt_reason is None else ckpt_reason + + # Save model checkpoints based on conditions (success rate, validation loss, etc) + if should_save_ckpt: + TrainUtils.save_model( + model=model, + config=config, + env_meta=env_meta, + shape_meta=shape_meta, + ckpt_path=os.path.join(ckpt_dir, epoch_ckpt_name + ".pth"), + obs_normalization_stats=obs_normalization_stats, + ) + + # Finally, log memory usage in MB + process = psutil.Process(os.getpid()) + mem_usage = int(process.memory_info().rss / 1000000) + data_logger.record("System/RAM Usage (MB)", mem_usage, epoch) + print(f"\nEpoch {epoch} Memory Usage: {mem_usage} MB\n") + + # terminate logging + data_logger.close() + + +def main(args): + """Train a model on a task using a specified algorithm.""" + # load config + if args.task is not None: + # obtain the configuration entry point + cfg_entry_point_key = f"robomimic_{args.algo}_cfg_entry_point" + + print(f"Loading configuration for task: {args.task}") + print(gym.envs.registry.keys()) + print(" ") + cfg_entry_point_file = gym.spec(args.task).kwargs.pop(cfg_entry_point_key) + # check if entry point exists + if cfg_entry_point_file is None: + raise ValueError( + f"Could not find configuration for the environment: '{args.task}'." + f" Please check that the gym registry has the entry point: '{cfg_entry_point_key}'." + ) + + with open(cfg_entry_point_file) as f: + ext_cfg = json.load(f) + config = config_factory(ext_cfg["algo_name"]) + # update config with external json - this will throw errors if + # the external config has keys not present in the base algo config + with config.values_unlocked(): + config.update(ext_cfg) + else: + raise ValueError("Please provide a task name through CLI arguments.") + + if args.dataset is not None: + config.train.data = args.dataset + + if args.name is not None: + config.experiment.name = args.name + + # change location of experiment directory + config.train.output_dir = os.path.abspath(os.path.join("./logs", args.log_dir, args.task)) + + # get torch device + device = TorchUtils.get_torch_device(try_to_use_cuda=config.train.cuda) + + config.lock() + + # catch error during training and print it + res_str = "finished run successfully!" + try: + train(config, device=device) + except Exception as e: + res_str = f"run failed with error:\n{e}\n\n{traceback.format_exc()}" + print(res_str) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + # Experiment Name (for tensorboard, saving models, etc.) + parser.add_argument( + "--name", + type=str, + default=None, + help="(optional) if provided, override the experiment name defined in the config", + ) + + # Dataset path, to override the one in the config + parser.add_argument( + "--dataset", + type=str, + default=None, + help="(optional) if provided, override the dataset path defined in the config", + ) + + parser.add_argument("--task", type=str, default=None, help="Name of the task.") + parser.add_argument("--algo", type=str, default=None, help="Name of the algorithm.") + parser.add_argument("--log_dir", type=str, default="robomimic", help="Path to log directory") + + args = parser.parse_args() + + # run training + main(args) + # close sim app + simulation_app.close() diff --git a/scripts/reinforcement_learning/ray/cluster_configs/Dockerfile b/scripts/reinforcement_learning/ray/cluster_configs/Dockerfile new file mode 100644 index 00000000..e75b7560 --- /dev/null +++ b/scripts/reinforcement_learning/ray/cluster_configs/Dockerfile @@ -0,0 +1,23 @@ +FROM isaac-lab-base:latest + +# WGet is needed so that GCS or other cloud providers can mark the container as ready. +# Otherwise the Ray liveliness checks fail. +RUN apt-get update && apt-get install wget + +# 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/scripts/reinforcement_learning/ray/cluster_configs/google_cloud/kuberay.yaml.jinja b/scripts/reinforcement_learning/ray/cluster_configs/google_cloud/kuberay.yaml.jinja new file mode 100644 index 00000000..40ccccf7 --- /dev/null +++ b/scripts/reinforcement_learning/ray/cluster_configs/google_cloud/kuberay.yaml.jinja @@ -0,0 +1,203 @@ +# 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" + 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: {{ 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={{name}}-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/scripts/reinforcement_learning/ray/grok_cluster_with_kubectl.py b/scripts/reinforcement_learning/ray/grok_cluster_with_kubectl.py new file mode 100644 index 00000000..dcbeb680 --- /dev/null +++ b/scripts/reinforcement_learning/ray/grok_cluster_with_kubectl.py @@ -0,0 +1,274 @@ +# Copyright (c) 2024-2025, The UW 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 + + python3 scripts/reinforcement_learning/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 (skip workers) + if "head" in pod_name: + base_name = match.group(1).split("-head")[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.") + mlflow_name = f"{cluster_prefix}-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 + # This needs to be http to be resolved. HTTPS can't be resolved + # This should be fine as it is on a subnet on the cluster regardless + 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/scripts/reinforcement_learning/ray/hyperparameter_tuning/vision_cartpole_cfg.py b/scripts/reinforcement_learning/ray/hyperparameter_tuning/vision_cartpole_cfg.py new file mode 100644 index 00000000..2127e0a3 --- /dev/null +++ b/scripts/reinforcement_learning/ray/hyperparameter_tuning/vision_cartpole_cfg.py @@ -0,0 +1,50 @@ +# Copyright (c) 2024-2025, The UW 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/scripts/reinforcement_learning/ray/hyperparameter_tuning/vision_cfg.py b/scripts/reinforcement_learning/ray/hyperparameter_tuning/vision_cfg.py new file mode 100644 index 00000000..0430af72 --- /dev/null +++ b/scripts/reinforcement_learning/ray/hyperparameter_tuning/vision_cfg.py @@ -0,0 +1,154 @@ +# Copyright (c) 2024-2025, The UW 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/scripts/reinforcement_learning/ray/launch.py b/scripts/reinforcement_learning/ray/launch.py new file mode 100644 index 00000000..d04d56b8 --- /dev/null +++ b/scripts/reinforcement_learning/ray/launch.py @@ -0,0 +1,180 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import argparse +import pathlib +import subprocess +import yaml + +import util +from jinja2 import Environment, FileSystemLoader +from kubernetes import config + +"""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 + python3 scripts/reinforcement_learning/ray/launch.py -h + + # Examples + + # The following creates 8 GPUx1 nvidia l4 workers + python3 scripts/reinforcement_learning/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 + python3 scripts/reinforcement_learning/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, autoescape=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: + print(cleaned_yaml_string) + 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/scripts/reinforcement_learning/ray/mlflow_to_local_tensorboard.py b/scripts/reinforcement_learning/ray/mlflow_to_local_tensorboard.py new file mode 100644 index 00000000..11999b20 --- /dev/null +++ b/scripts/reinforcement_learning/ray/mlflow_to_local_tensorboard.py @@ -0,0 +1,149 @@ +# Copyright (c) 2024-2025, The UW 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/scripts/reinforcement_learning/ray/submit_job.py b/scripts/reinforcement_learning/ray/submit_job.py new file mode 100644 index 00000000..7fe59f02 --- /dev/null +++ b/scripts/reinforcement_learning/ray/submit_job.py @@ -0,0 +1,148 @@ +# Copyright (c) 2024-2025, The UW 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 + python3 scripts/reinforcement_learning/ray/submit_job.py \ + --aggregate_jobs /workspace/isaaclab/scripts/reinforcement_learning/ray/tuner.py \ + --cfg_file hyperparameter_tuning/vision_cartpole_cfg.py \ + --cfg_class CartpoleTheiaJobCfg --mlflow_uri + + # Example: Submitting resource wrapped job + python3 scripts/reinforcement_learning/ray/submit_job.py --aggregate_jobs wrap_resources.py --test + + # For all command line arguments + python3 scripts/reinforcement_learning/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/scripts/reinforcement_learning/ray/tuner.py b/scripts/reinforcement_learning/ray/tuner.py new file mode 100644 index 00000000..35002325 --- /dev/null +++ b/scripts/reinforcement_learning/ray/tuner.py @@ -0,0 +1,358 @@ +# Copyright (c) 2024-2025, The UW 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, one worker is created for each GPU-enabled node in the cluster for each individual job. +To use more than one worker per node (likely the case for multi-GPU machines), supply the +num_workers_per_node argument. + +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 scripts/reinforcement_learning/ray/tuner.py -h + + # Examples + # Local + ./isaaclab.sh -p scripts/reinforcement_learning/ray/tuner.py --run_mode local \ + --cfg_file scripts/reinforcement_learning/ray/hyperparameter_tuning/vision_cartpole_cfg.py \ + --cfg_class CartpoleTheiaJobCfg + # Remote (run grok cluster or create config file mentioned in :file:`submit_job.py`) + ./isaaclab.sh -p scripts/reinforcement_learning/ray/submit_job.py \ + --aggregate_jobs tuner.py \ + --cfg_file hyperparameter_tuning/vision_cartpole_cfg.py \ + --cfg_class CartpoleTheiaJobCfg --mlflow_uri + +""" + +DOCKER_PREFIX = "/workspace/isaaclab/" +BASE_DIR = os.path.expanduser("~") +PYTHON_EXEC = "./isaaclab.sh -p" +WORKFLOW = "scripts/reinforcement_learning/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: dict): + """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 + "/" + self.experiment_name + self.done = False + + if self.proc is None: + raise ValueError("Could not start trial.") + proc_status = self.proc.poll() + if proc_status is not None: # process finished, signal finish + self.data["done"] = True + print(f"[INFO]: Process finished with {proc_status}, 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: dict): + """ + Runner args include command line arguments passed to the task. + For example: + cfg["runner_args"]["headless_singleton"] = "--headless" + cfg["runner_args"]["enable_cameras_singleton"] = "--enable_cameras" + """ + assert "runner_args" in cfg, "No runner arguments specified." + """ + Task is the desired task to train on. For example: + cfg["runner_args"]["--task"] = tune.choice(["Isaac-Cartpole-RGB-TheiaTiny-v0"]) + """ + assert "--task" in cfg["runner_args"], "No task specified." + """ + Hydra args define the hyperparameters varied within the sweep. For example: + cfg["hydra_args"]["agent.params.network.cnn.activation"] = tune.choice(["relu", "elu"]) + """ + assert "hydra_args" in cfg, "No hyperparameters 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/scripts/reinforcement_learning/ray/util.py b/scripts/reinforcement_learning/ray/util.py new file mode 100644 index 00000000..b5c80a61 --- /dev/null +++ b/scripts/reinforcement_learning/ray/util.py @@ -0,0 +1,432 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import argparse +import os +import re +import subprocess +import threading +from datetime import datetime +from math import isclose + +import ray +from tensorboard.backend.event_processing.directory_watcher import DirectoryDeletedError +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. If the logs can't be + found, check the summaries sublevel. + + 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 + def get_latest_scalars(path: str) -> dict: + event_acc = EventAccumulator(path, size_guidance={"scalars": 1}) + try: + event_acc.Reload() + if event_acc.Tags()["scalars"]: + return { + tag: event_acc.Scalars(tag)[-1].value + for tag in event_acc.Tags()["scalars"] + if event_acc.Scalars(tag) + } + except (KeyError, OSError, RuntimeError, DirectoryDeletedError): + return {} + + scalars = get_latest_scalars(directory) + return scalars or get_latest_scalars(os.path.join(directory, "summaries")) + + +def get_invocation_command_from_cfg( + cfg: dict, + python_cmd: str = "/workspace/isaaclab/isaaclab.sh -p", + workflow: str = "scripts/reinforcement_learning/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 (.+)$") + + def stream_reader(stream, identifier_string, result_details): + for line in iter(stream.readline, ""): + line = line.strip() + result_details.append(f"{identifier_string}: {line}\n") + if log_all_output: + print(f"{identifier_string}: {line}") + + # Read stdout until we find experiment info + # Do some careful handling prevent overflowing the pipe reading buffer with error 141 + for line in iter(process.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: + # Start stderr reader after finding experiment info + stderr_thread = threading.Thread( + target=stream_reader, args=(process.stderr, identifier_string, result_details) + ) + stderr_thread.daemon = True + stderr_thread.start() + + # Start stdout reader to continue reading to flush buffer + stdout_thread = threading.Thread( + target=stream_reader, args=(process.stdout, identifier_string, result_details) + ) + stdout_thread.daemon = True + stdout_thread.start() + + return { + "experiment_name": experiment_name, + "logdir": logdir, + "proc": process, + "result": " ".join(result_details), + } + process.wait() + 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/scripts/reinforcement_learning/ray/wrap_resources.py b/scripts/reinforcement_learning/ray/wrap_resources.py new file mode 100644 index 00000000..96336a51 --- /dev/null +++ b/scripts/reinforcement_learning/ray/wrap_resources.py @@ -0,0 +1,150 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import argparse + +import ray +import util +from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy + +""" +This script dispatches sub-job(s) (individual jobs, use :file:`tuner.py` for tuning 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 scripts/reinforcement_learning/ray/wrap_resources.py -h + # No resource isolation; no parallelization: + ./isaaclab.sh -p scripts/reinforcement_learning/ray/wrap_resources.py + --sub_jobs ++ + # Automatic Resource Isolation; Example A: needed for parallelization + ./isaaclab.sh -p scripts/reinforcement_learning/ray/wrap_resources.py \ + --num_workers \ + --sub_jobs + + # Manual Resource Isolation; Example B: needed for parallelization + ./isaaclab.sh -p scripts/reinforcement_learning/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 scripts/reinforcement_learning/ray/wrap_resources.py --num_cpu_per_worker \ + --gpu_per_worker --ram_gb_per_worker --sub_jobs + + # to see all arguments + ./isaaclab.sh -p scripts/reinforcement_learning/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/scripts/reinforcement_learning/rl_games/play.py b/scripts/reinforcement_learning/rl_games/play.py new file mode 100644 index 00000000..14070482 --- /dev/null +++ b/scripts/reinforcement_learning/rl_games/play.py @@ -0,0 +1,206 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Script to play a checkpoint if an RL agent from RL-Games.""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Play a checkpoint of an RL agent from RL-Games.") +parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.") +parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).") +parser.add_argument( + "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations." +) +parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +parser.add_argument("--checkpoint", type=str, default=None, help="Path to model checkpoint.") +parser.add_argument( + "--use_pretrained_checkpoint", + action="store_true", + help="Use the pre-trained checkpoint from Nucleus.", +) +parser.add_argument( + "--use_last_checkpoint", + action="store_true", + help="When no checkpoint provided, use the last saved model. Otherwise use the best saved model.", +) +parser.add_argument("--real-time", action="store_true", default=False, help="Run in real-time, if possible.") +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli = parser.parse_args() +# always enable cameras to record video +if args_cli.video: + args_cli.enable_cameras = True + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + + +import gymnasium as gym +import math +import os +import time +import torch + +from rl_games.common import env_configurations, vecenv +from rl_games.common.player import BasePlayer +from rl_games.torch_runner import Runner + +from isaaclab.envs import DirectMARLEnv, multi_agent_to_single_agent +from isaaclab.utils.assets import retrieve_file_path +from isaaclab.utils.dict import print_dict +from isaaclab.utils.pretrained_checkpoint import get_published_pretrained_checkpoint + +from isaaclab_rl.rl_games import RlGamesGpuEnv, RlGamesVecEnvWrapper + +import isaaclab_tasks # noqa: F401 +import uwlab_tasks # noqa: F401 +from isaaclab_tasks.utils import get_checkpoint_path, load_cfg_from_registry, parse_env_cfg + + +def main(): + """Play with RL-Games agent.""" + # parse env configuration + env_cfg = parse_env_cfg( + args_cli.task, device=args_cli.device, num_envs=args_cli.num_envs, use_fabric=not args_cli.disable_fabric + ) + agent_cfg = load_cfg_from_registry(args_cli.task, "rl_games_cfg_entry_point") + + # specify directory for logging experiments + log_root_path = os.path.join("logs", "rl_games", agent_cfg["params"]["config"]["name"]) + log_root_path = os.path.abspath(log_root_path) + print(f"[INFO] Loading experiment from directory: {log_root_path}") + # find checkpoint + if args_cli.use_pretrained_checkpoint: + resume_path = get_published_pretrained_checkpoint("rl_games", args_cli.task) + if not resume_path: + print("[INFO] Unfortunately a pre-trained checkpoint is currently unavailable for this task.") + return + elif args_cli.checkpoint is None: + # specify directory for logging runs + run_dir = agent_cfg["params"]["config"].get("full_experiment_name", ".*") + # specify name of checkpoint + if args_cli.use_last_checkpoint: + checkpoint_file = ".*" + else: + # this loads the best checkpoint + checkpoint_file = f"{agent_cfg['params']['config']['name']}.pth" + # get path to previous checkpoint + resume_path = get_checkpoint_path(log_root_path, run_dir, checkpoint_file, other_dirs=["nn"]) + else: + resume_path = retrieve_file_path(args_cli.checkpoint) + log_dir = os.path.dirname(os.path.dirname(resume_path)) + + # wrap around environment for rl-games + rl_device = agent_cfg["params"]["config"]["device"] + clip_obs = agent_cfg["params"]["env"].get("clip_observations", math.inf) + clip_actions = agent_cfg["params"]["env"].get("clip_actions", math.inf) + + # create isaac environment + env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) + + # convert to single-agent instance if required by the RL algorithm + if isinstance(env.unwrapped, DirectMARLEnv): + env = multi_agent_to_single_agent(env) + + # wrap for video recording + if args_cli.video: + video_kwargs = { + "video_folder": os.path.join(log_root_path, log_dir, "videos", "play"), + "step_trigger": lambda step: step == 0, + "video_length": args_cli.video_length, + "disable_logger": True, + } + print("[INFO] Recording videos during training.") + print_dict(video_kwargs, nesting=4) + env = gym.wrappers.RecordVideo(env, **video_kwargs) + + # wrap around environment for rl-games + env = RlGamesVecEnvWrapper(env, rl_device, clip_obs, clip_actions) + + # register the environment to rl-games registry + # note: in agents configuration: environment name must be "rlgpu" + vecenv.register( + "IsaacRlgWrapper", lambda config_name, num_actors, **kwargs: RlGamesGpuEnv(config_name, num_actors, **kwargs) + ) + env_configurations.register("rlgpu", {"vecenv_type": "IsaacRlgWrapper", "env_creator": lambda **kwargs: env}) + + # load previously trained model + agent_cfg["params"]["load_checkpoint"] = True + agent_cfg["params"]["load_path"] = resume_path + print(f"[INFO]: Loading model checkpoint from: {agent_cfg['params']['load_path']}") + + # set number of actors into agent config + agent_cfg["params"]["config"]["num_actors"] = env.unwrapped.num_envs + # create runner from rl-games + runner = Runner() + runner.load(agent_cfg) + # obtain the agent from the runner + agent: BasePlayer = runner.create_player() + agent.restore(resume_path) + agent.reset() + + dt = env.unwrapped.physics_dt + + # reset environment + obs = env.reset() + if isinstance(obs, dict): + obs = obs["obs"] + timestep = 0 + # required: enables the flag for batched observations + _ = agent.get_batch_size(obs, 1) + # initialize RNN states if used + if agent.is_rnn: + agent.init_rnn() + # simulate environment + # note: We simplified the logic in rl-games player.py (:func:`BasePlayer.run()`) function in an + # attempt to have complete control over environment stepping. However, this removes other + # operations such as masking that is used for multi-agent learning by RL-Games. + while simulation_app.is_running(): + start_time = time.time() + # run everything in inference mode + with torch.inference_mode(): + # convert obs to agent format + obs = agent.obs_to_torch(obs) + # agent stepping + actions = agent.get_action(obs, is_deterministic=agent.is_deterministic) + # env stepping + obs, _, dones, _ = env.step(actions) + + # perform operations for terminated episodes + if len(dones) > 0: + # reset rnn state for terminated episodes + if agent.is_rnn and agent.states is not None: + for s in agent.states: + s[:, dones, :] = 0.0 + if args_cli.video: + timestep += 1 + # Exit the play loop after recording one video + if timestep == args_cli.video_length: + break + + # time delay for real-time evaluation + sleep_time = dt - (time.time() - start_time) + if args_cli.real_time and sleep_time > 0: + time.sleep(sleep_time) + + # close the simulator + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/reinforcement_learning/rl_games/train.py b/scripts/reinforcement_learning/rl_games/train.py new file mode 100644 index 00000000..bc56825b --- /dev/null +++ b/scripts/reinforcement_learning/rl_games/train.py @@ -0,0 +1,183 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Script to train RL agent with RL-Games.""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse +import sys + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Train an RL agent with RL-Games.") +parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.") +parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).") +parser.add_argument("--video_interval", type=int, default=2000, help="Interval between video recordings (in steps).") +parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment") +parser.add_argument( + "--distributed", action="store_true", default=False, help="Run training with multiple GPUs or nodes." +) +parser.add_argument("--checkpoint", type=str, default=None, help="Path to model checkpoint.") +parser.add_argument("--sigma", type=str, default=None, help="The policy's initial standard deviation.") +parser.add_argument("--max_iterations", type=int, default=None, help="RL Policy training iterations.") + +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli, hydra_args = parser.parse_known_args() +# always enable cameras to record video +if args_cli.video: + args_cli.enable_cameras = True + +# clear out sys.argv for Hydra +sys.argv = [sys.argv[0]] + hydra_args + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import gymnasium as gym +import math +import os +import random +from datetime import datetime + +from rl_games.common import env_configurations, vecenv +from rl_games.common.algo_observer import IsaacAlgoObserver +from rl_games.torch_runner import Runner + +from isaaclab.envs import ( + DirectMARLEnv, + DirectMARLEnvCfg, + DirectRLEnvCfg, + ManagerBasedRLEnvCfg, + multi_agent_to_single_agent, +) +from isaaclab.utils.assets import retrieve_file_path +from isaaclab.utils.dict import print_dict +from isaaclab.utils.io import dump_pickle, dump_yaml + +from isaaclab_rl.rl_games import RlGamesGpuEnv, RlGamesVecEnvWrapper + +import isaaclab_tasks # noqa: F401 +import uwlab_tasks # noqa: F401 +from isaaclab_tasks.utils.hydra import hydra_task_config + + +@hydra_task_config(args_cli.task, "rl_games_cfg_entry_point") +def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: dict): + """Train with RL-Games agent.""" + # override configurations with non-hydra CLI arguments + env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs + env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device + + # randomly sample a seed if seed = -1 + if args_cli.seed == -1: + args_cli.seed = random.randint(0, 10000) + + agent_cfg["params"]["seed"] = args_cli.seed if args_cli.seed is not None else agent_cfg["params"]["seed"] + agent_cfg["params"]["config"]["max_epochs"] = ( + args_cli.max_iterations if args_cli.max_iterations is not None else agent_cfg["params"]["config"]["max_epochs"] + ) + if args_cli.checkpoint is not None: + resume_path = retrieve_file_path(args_cli.checkpoint) + agent_cfg["params"]["load_checkpoint"] = True + agent_cfg["params"]["load_path"] = resume_path + print(f"[INFO]: Loading model checkpoint from: {agent_cfg['params']['load_path']}") + train_sigma = float(args_cli.sigma) if args_cli.sigma is not None else None + + # multi-gpu training config + if args_cli.distributed: + agent_cfg["params"]["seed"] += app_launcher.global_rank + agent_cfg["params"]["config"]["device"] = f"cuda:{app_launcher.local_rank}" + agent_cfg["params"]["config"]["device_name"] = f"cuda:{app_launcher.local_rank}" + agent_cfg["params"]["config"]["multi_gpu"] = True + # update env config device + env_cfg.sim.device = f"cuda:{app_launcher.local_rank}" + + # set the environment seed (after multi-gpu config for updated rank from agent seed) + # note: certain randomizations occur in the environment initialization so we set the seed here + env_cfg.seed = agent_cfg["params"]["seed"] + + # specify directory for logging experiments + log_root_path = os.path.join("logs", "rl_games", agent_cfg["params"]["config"]["name"]) + log_root_path = os.path.abspath(log_root_path) + print(f"[INFO] Logging experiment in directory: {log_root_path}") + # specify directory for logging runs + log_dir = agent_cfg["params"]["config"].get("full_experiment_name", datetime.now().strftime("%Y-%m-%d_%H-%M-%S")) + # set directory into agent config + # logging directory path: / + agent_cfg["params"]["config"]["train_dir"] = log_root_path + agent_cfg["params"]["config"]["full_experiment_name"] = log_dir + + # dump the configuration into log-directory + dump_yaml(os.path.join(log_root_path, log_dir, "params", "env.yaml"), env_cfg) + dump_yaml(os.path.join(log_root_path, log_dir, "params", "agent.yaml"), agent_cfg) + dump_pickle(os.path.join(log_root_path, log_dir, "params", "env.pkl"), env_cfg) + dump_pickle(os.path.join(log_root_path, log_dir, "params", "agent.pkl"), agent_cfg) + + # read configurations about the agent-training + rl_device = agent_cfg["params"]["config"]["device"] + clip_obs = agent_cfg["params"]["env"].get("clip_observations", math.inf) + clip_actions = agent_cfg["params"]["env"].get("clip_actions", math.inf) + + # create isaac environment + env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) + + # convert to single-agent instance if required by the RL algorithm + if isinstance(env.unwrapped, DirectMARLEnv): + env = multi_agent_to_single_agent(env) + + # wrap for video recording + if args_cli.video: + video_kwargs = { + "video_folder": os.path.join(log_root_path, log_dir, "videos", "train"), + "step_trigger": lambda step: step % args_cli.video_interval == 0, + "video_length": args_cli.video_length, + "disable_logger": True, + } + print("[INFO] Recording videos during training.") + print_dict(video_kwargs, nesting=4) + env = gym.wrappers.RecordVideo(env, **video_kwargs) + + # wrap around environment for rl-games + env = RlGamesVecEnvWrapper(env, rl_device, clip_obs, clip_actions) + + # register the environment to rl-games registry + # note: in agents configuration: environment name must be "rlgpu" + vecenv.register( + "IsaacRlgWrapper", lambda config_name, num_actors, **kwargs: RlGamesGpuEnv(config_name, num_actors, **kwargs) + ) + env_configurations.register("rlgpu", {"vecenv_type": "IsaacRlgWrapper", "env_creator": lambda **kwargs: env}) + + # set number of actors into agent config + agent_cfg["params"]["config"]["num_actors"] = env.unwrapped.num_envs + # create runner from rl-games + runner = Runner(IsaacAlgoObserver()) + runner.load(agent_cfg) + + # reset the agent and env + runner.reset() + # train the agent + if args_cli.checkpoint is not None: + runner.run({"train": True, "play": False, "sigma": train_sigma, "checkpoint": resume_path}) + else: + runner.run({"train": True, "play": False, "sigma": train_sigma}) + + # close the simulator + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/reinforcement_learning/rsl_rl/cli_args.py b/scripts/reinforcement_learning/rsl_rl/cli_args.py new file mode 100644 index 00000000..3a6e58b9 --- /dev/null +++ b/scripts/reinforcement_learning/rsl_rl/cli_args.py @@ -0,0 +1,91 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import argparse +import random +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg + + +def add_rsl_rl_args(parser: argparse.ArgumentParser): + """Add RSL-RL arguments to the parser. + + Args: + parser: The parser to add the arguments to. + """ + # create a new argument group + arg_group = parser.add_argument_group("rsl_rl", description="Arguments for RSL-RL agent.") + # -- experiment arguments + arg_group.add_argument( + "--experiment_name", type=str, default=None, help="Name of the experiment folder where logs will be stored." + ) + arg_group.add_argument("--run_name", type=str, default=None, help="Run name suffix to the log directory.") + # -- load arguments + arg_group.add_argument("--resume", type=bool, default=None, help="Whether to resume from a checkpoint.") + arg_group.add_argument("--load_run", type=str, default=None, help="Name of the run folder to resume from.") + arg_group.add_argument("--checkpoint", type=str, default=None, help="Checkpoint file to resume from.") + # -- logger arguments + arg_group.add_argument( + "--logger", type=str, default=None, choices={"wandb", "tensorboard", "neptune"}, help="Logger module to use." + ) + arg_group.add_argument( + "--log_project_name", type=str, default=None, help="Name of the logging project when using wandb or neptune." + ) + + +def parse_rsl_rl_cfg(task_name: str, args_cli: argparse.Namespace) -> RslRlOnPolicyRunnerCfg: + """Parse configuration for RSL-RL agent based on inputs. + + Args: + task_name: The name of the environment. + args_cli: The command line arguments. + + Returns: + The parsed configuration for RSL-RL agent based on inputs. + """ + from isaaclab_tasks.utils.parse_cfg import load_cfg_from_registry + + # load the default configuration + rslrl_cfg: RslRlOnPolicyRunnerCfg = load_cfg_from_registry(task_name, "rsl_rl_cfg_entry_point") + rslrl_cfg = update_rsl_rl_cfg(rslrl_cfg, args_cli) + return rslrl_cfg + + +def update_rsl_rl_cfg(agent_cfg: RslRlOnPolicyRunnerCfg, args_cli: argparse.Namespace): + """Update configuration for RSL-RL agent based on inputs. + + Args: + agent_cfg: The configuration for RSL-RL agent. + args_cli: The command line arguments. + + Returns: + The updated configuration for RSL-RL agent based on inputs. + """ + # override the default configuration with CLI arguments + if hasattr(args_cli, "seed") and args_cli.seed is not None: + # randomly sample a seed if seed = -1 + if args_cli.seed == -1: + args_cli.seed = random.randint(0, 10000) + agent_cfg.seed = args_cli.seed + if args_cli.resume is not None: + agent_cfg.resume = args_cli.resume + if args_cli.load_run is not None: + agent_cfg.load_run = args_cli.load_run + if args_cli.checkpoint is not None: + agent_cfg.load_checkpoint = args_cli.checkpoint + if args_cli.run_name is not None: + agent_cfg.run_name = args_cli.run_name + if args_cli.logger is not None: + agent_cfg.logger = args_cli.logger + # set the project name for wandb and neptune + if agent_cfg.logger in {"wandb", "neptune"} and args_cli.log_project_name: + agent_cfg.wandb_project = args_cli.log_project_name + agent_cfg.neptune_project = args_cli.log_project_name + + return agent_cfg diff --git a/scripts/reinforcement_learning/rsl_rl/play.py b/scripts/reinforcement_learning/rsl_rl/play.py new file mode 100644 index 00000000..4b371679 --- /dev/null +++ b/scripts/reinforcement_learning/rsl_rl/play.py @@ -0,0 +1,162 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Script to play a checkpoint if an RL agent from RSL-RL.""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse + +from isaaclab.app import AppLauncher + +# local imports +import cli_args # isort: skip + +# add argparse arguments +parser = argparse.ArgumentParser(description="Train an RL agent with RSL-RL.") +parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.") +parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).") +parser.add_argument( + "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations." +) +parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +parser.add_argument( + "--use_pretrained_checkpoint", + action="store_true", + help="Use the pre-trained checkpoint from Nucleus.", +) +parser.add_argument("--real-time", action="store_true", default=False, help="Run in real-time, if possible.") +# append RSL-RL cli arguments +cli_args.add_rsl_rl_args(parser) +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +args_cli = parser.parse_args() +# always enable cameras to record video +if args_cli.video: + args_cli.enable_cameras = True + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import gymnasium as gym +import os +import time +import torch + +from rsl_rl.runners import OnPolicyRunner + +from isaaclab.envs import DirectMARLEnv, multi_agent_to_single_agent +from isaaclab.utils.assets import retrieve_file_path +from isaaclab.utils.dict import print_dict +from isaaclab.utils.pretrained_checkpoint import get_published_pretrained_checkpoint + +from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper, export_policy_as_jit, export_policy_as_onnx + +import isaaclab_tasks # noqa: F401 +import uwlab_tasks # noqa: F401 +from isaaclab_tasks.utils import get_checkpoint_path, parse_env_cfg + + +def main(): + """Play with RSL-RL agent.""" + # parse configuration + env_cfg = parse_env_cfg( + args_cli.task, device=args_cli.device, num_envs=args_cli.num_envs, use_fabric=not args_cli.disable_fabric + ) + agent_cfg: RslRlOnPolicyRunnerCfg = cli_args.parse_rsl_rl_cfg(args_cli.task, args_cli) + + # specify directory for logging experiments + log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name) + log_root_path = os.path.abspath(log_root_path) + print(f"[INFO] Loading experiment from directory: {log_root_path}") + if args_cli.use_pretrained_checkpoint: + resume_path = get_published_pretrained_checkpoint("rsl_rl", args_cli.task) + if not resume_path: + print("[INFO] Unfortunately a pre-trained checkpoint is currently unavailable for this task.") + return + elif args_cli.checkpoint: + resume_path = retrieve_file_path(args_cli.checkpoint) + else: + resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint) + + log_dir = os.path.dirname(resume_path) + + # create isaac environment + env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) + + # convert to single-agent instance if required by the RL algorithm + if isinstance(env.unwrapped, DirectMARLEnv): + env = multi_agent_to_single_agent(env) + + # wrap for video recording + if args_cli.video: + video_kwargs = { + "video_folder": os.path.join(log_dir, "videos", "play"), + "step_trigger": lambda step: step == 0, + "video_length": args_cli.video_length, + "disable_logger": True, + } + print("[INFO] Recording videos during training.") + print_dict(video_kwargs, nesting=4) + env = gym.wrappers.RecordVideo(env, **video_kwargs) + + # wrap around environment for rsl-rl + env = RslRlVecEnvWrapper(env) + + print(f"[INFO]: Loading model checkpoint from: {resume_path}") + # load previously trained model + ppo_runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device) + ppo_runner.load(resume_path) + + # obtain the trained policy for inference + policy = ppo_runner.get_inference_policy(device=env.unwrapped.device) + + # export policy to onnx/jit + export_model_dir = os.path.join(os.path.dirname(resume_path), "exported") + export_policy_as_jit( + ppo_runner.alg.actor_critic, ppo_runner.obs_normalizer, path=export_model_dir, filename="policy.pt" + ) + export_policy_as_onnx( + ppo_runner.alg.actor_critic, normalizer=ppo_runner.obs_normalizer, path=export_model_dir, filename="policy.onnx" + ) + + dt = env.unwrapped.physics_dt + + # reset environment + obs, _ = env.get_observations() + timestep = 0 + # simulate environment + while simulation_app.is_running(): + start_time = time.time() + # run everything in inference mode + with torch.inference_mode(): + # agent stepping + actions = policy(obs) + # env stepping + obs, _, _, _ = env.step(actions) + if args_cli.video: + timestep += 1 + # Exit the play loop after recording one video + if timestep == args_cli.video_length: + break + + # time delay for real-time evaluation + sleep_time = dt - (time.time() - start_time) + if args_cli.real_time and sleep_time > 0: + time.sleep(sleep_time) + + # close the simulator + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/reinforcement_learning/rsl_rl/train.py b/scripts/reinforcement_learning/rsl_rl/train.py new file mode 100644 index 00000000..9daaf6a4 --- /dev/null +++ b/scripts/reinforcement_learning/rsl_rl/train.py @@ -0,0 +1,200 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Script to train RL agent with RSL-RL.""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse +import sys + +from isaaclab.app import AppLauncher + +# local imports +import cli_args # isort: skip + + +# add argparse arguments +parser = argparse.ArgumentParser(description="Train an RL agent with RSL-RL.") +parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.") +parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).") +parser.add_argument("--video_interval", type=int, default=2000, help="Interval between video recordings (in steps).") +parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment") +parser.add_argument("--max_iterations", type=int, default=None, help="RL Policy training iterations.") +# append RSL-RL cli arguments +cli_args.add_rsl_rl_args(parser) +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +args_cli, hydra_args = parser.parse_known_args() + +# always enable cameras to record video +if args_cli.video: + args_cli.enable_cameras = True + +# clear out sys.argv for Hydra +sys.argv = [sys.argv[0]] + hydra_args + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import gymnasium as gym +import os +import torch +from datetime import datetime + +from rsl_rl.runners import OnPolicyRunner + +from isaaclab.envs import ( + DirectMARLEnv, + DirectMARLEnvCfg, + DirectRLEnvCfg, + ManagerBasedRLEnvCfg, + multi_agent_to_single_agent, +) +from isaaclab.utils.dict import print_dict +from isaaclab.utils.io import dump_pickle, dump_yaml + +from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper + +import isaaclab_tasks # noqa: F401 +import uwlab_tasks # noqa: F401 +from isaaclab_tasks.utils import get_checkpoint_path +from isaaclab_tasks.utils.hydra import hydra_task_config + +torch.backends.cuda.matmul.allow_tf32 = True +torch.backends.cudnn.allow_tf32 = True +torch.backends.cudnn.deterministic = False +torch.backends.cudnn.benchmark = False + + +def process_agent_cfg(env_cfg, agent_cfg): + if hasattr(agent_cfg.algorithm, "symmetry_cfg") and agent_cfg.algorithm.symmetry_cfg is None: + del agent_cfg.algorithm.symmetry_cfg + + if hasattr(agent_cfg.algorithm, "behavior_cloning_cfg"): + if agent_cfg.algorithm.behavior_cloning_cfg is None: + del agent_cfg.algorithm.behavior_cloning_cfg + else: + bc_cfg = agent_cfg.algorithm.behavior_cloning_cfg + if bc_cfg.experts_observation_group_cfg is not None: + import importlib + + # resolve path to the module location + mod_name, attr_name = bc_cfg.experts_observation_group_cfg.split(":") + mod = importlib.import_module(mod_name) + cfg_cls = mod + for attr in attr_name.split("."): + cfg_cls = getattr(cfg_cls, attr) + cfg = cfg_cls() + setattr(env_cfg.observations, "expert_obs", cfg) + + if hasattr(agent_cfg.algorithm, "offline_algorithm_cfg"): + if agent_cfg.algorithm.offline_algorithm_cfg is None: + del agent_cfg.algorithm.offline_algorithm_cfg + else: + if agent_cfg.algorithm.offline_algorithm_cfg.behavior_cloning_cfg is None: + del agent_cfg.algorithm.offline_algorithm_cfg.behavior_cloning_cfg + else: + bc_cfg = agent_cfg.algorithm.offline_algorithm_cfg.behavior_cloning_cfg + if bc_cfg.experts_observation_group_cfg is not None: + import importlib + + # resolve path to the module location + mod_name, attr_name = bc_cfg.experts_observation_group_cfg.split(":") + mod = importlib.import_module(mod_name) + cfg_cls = mod + for attr in attr_name.split("."): + cfg_cls = getattr(cfg_cls, attr) + cfg = cfg_cls() + setattr(env_cfg.observations, "expert_obs", cfg) + return agent_cfg + + +@hydra_task_config(args_cli.task, "rsl_rl_cfg_entry_point") +def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlOnPolicyRunnerCfg): + """Train with RSL-RL agent.""" + # override configurations with non-hydra CLI arguments + agent_cfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli) + env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs + agent_cfg.max_iterations = ( + args_cli.max_iterations if args_cli.max_iterations is not None else agent_cfg.max_iterations + ) + + # set the environment seed + # note: certain randomizations occur in the environment initialization so we set the seed here + env_cfg.seed = agent_cfg.seed + env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device + agent_cfg = process_agent_cfg(env_cfg, agent_cfg) + # specify directory for logging experiments + log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name) + log_root_path = os.path.abspath(log_root_path) + print(f"[INFO] Logging experiment in directory: {log_root_path}") + # specify directory for logging runs: {time-stamp}_{run_name} + log_dir = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + # This way, the Ray Tune workflow can extract experiment name. + print(f"Exact experiment name requested from command line: {log_dir}") + if agent_cfg.run_name: + log_dir += f"_{agent_cfg.run_name}" + log_dir = os.path.join(log_root_path, log_dir) + + # create isaac environment + env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) + + # convert to single-agent instance if required by the RL algorithm + if isinstance(env.unwrapped, DirectMARLEnv): + env = multi_agent_to_single_agent(env) + + # save resume path before creating a new log_dir + if agent_cfg.resume: + resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint) + + # wrap for video recording + if args_cli.video: + video_kwargs = { + "video_folder": os.path.join(log_dir, "videos", "train"), + "step_trigger": lambda step: step % args_cli.video_interval == 0, + "video_length": args_cli.video_length, + "disable_logger": True, + } + print("[INFO] Recording videos during training.") + print_dict(video_kwargs, nesting=4) + env = gym.wrappers.RecordVideo(env, **video_kwargs) + + # wrap around environment for rsl-rl + env = RslRlVecEnvWrapper(env) + + # create runner from rsl-rl + runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device) + # write git state to logs + runner.add_git_repo_to_log(__file__) + # load the checkpoint + if agent_cfg.resume: + print(f"[INFO]: Loading model checkpoint from: {resume_path}") + # load previously trained model + runner.load(resume_path) + + # dump the configuration into log-directory + dump_yaml(os.path.join(log_dir, "params", "env.yaml"), env_cfg) + dump_yaml(os.path.join(log_dir, "params", "agent.yaml"), agent_cfg) + dump_pickle(os.path.join(log_dir, "params", "env.pkl"), env_cfg) + dump_pickle(os.path.join(log_dir, "params", "agent.pkl"), agent_cfg) + + # run training + runner.learn(num_learning_iterations=agent_cfg.max_iterations, init_at_random_ep_len=True) + + # close the simulator + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/reinforcement_learning/sb3/play.py b/scripts/reinforcement_learning/sb3/play.py new file mode 100644 index 00000000..ac63d218 --- /dev/null +++ b/scripts/reinforcement_learning/sb3/play.py @@ -0,0 +1,169 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Script to play a checkpoint if an RL agent from Stable-Baselines3.""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Play a checkpoint of an RL agent from Stable-Baselines3.") +parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.") +parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).") +parser.add_argument( + "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations." +) +parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +parser.add_argument("--checkpoint", type=str, default=None, help="Path to model checkpoint.") +parser.add_argument( + "--use_pretrained_checkpoint", + action="store_true", + help="Use the pre-trained checkpoint from Nucleus.", +) +parser.add_argument( + "--use_last_checkpoint", + action="store_true", + help="When no checkpoint provided, use the last saved model. Otherwise use the best saved model.", +) +parser.add_argument("--real-time", action="store_true", default=False, help="Run in real-time, if possible.") +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli = parser.parse_args() +# always enable cameras to record video +if args_cli.video: + args_cli.enable_cameras = True + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import gymnasium as gym +import numpy as np +import os +import time +import torch + +from stable_baselines3 import PPO +from stable_baselines3.common.vec_env import VecNormalize + +from isaaclab.envs import DirectMARLEnv, multi_agent_to_single_agent +from isaaclab.utils.dict import print_dict +from isaaclab.utils.pretrained_checkpoint import get_published_pretrained_checkpoint + +from isaaclab_rl.sb3 import Sb3VecEnvWrapper, process_sb3_cfg + +import isaaclab_tasks # noqa: F401 +import uwlab_tasks # noqa: F401 +from isaaclab_tasks.utils.parse_cfg import get_checkpoint_path, load_cfg_from_registry, parse_env_cfg + + +def main(): + """Play with stable-baselines agent.""" + # parse configuration + env_cfg = parse_env_cfg( + args_cli.task, device=args_cli.device, num_envs=args_cli.num_envs, use_fabric=not args_cli.disable_fabric + ) + agent_cfg = load_cfg_from_registry(args_cli.task, "sb3_cfg_entry_point") + + # directory for logging into + log_root_path = os.path.join("logs", "sb3", args_cli.task) + log_root_path = os.path.abspath(log_root_path) + # checkpoint and log_dir stuff + if args_cli.use_pretrained_checkpoint: + checkpoint_path = get_published_pretrained_checkpoint("sb3", args_cli.task) + if not checkpoint_path: + print("[INFO] Unfortunately a pre-trained checkpoint is currently unavailable for this task.") + return + elif args_cli.checkpoint is None: + if args_cli.use_last_checkpoint: + checkpoint = "model_.*.zip" + else: + checkpoint = "model.zip" + checkpoint_path = get_checkpoint_path(log_root_path, ".*", checkpoint) + else: + checkpoint_path = args_cli.checkpoint + log_dir = os.path.dirname(checkpoint_path) + + # post-process agent configuration + agent_cfg = process_sb3_cfg(agent_cfg) + + # create isaac environment + env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) + + # convert to single-agent instance if required by the RL algorithm + if isinstance(env.unwrapped, DirectMARLEnv): + env = multi_agent_to_single_agent(env) + + # wrap for video recording + if args_cli.video: + video_kwargs = { + "video_folder": os.path.join(log_dir, "videos", "play"), + "step_trigger": lambda step: step == 0, + "video_length": args_cli.video_length, + "disable_logger": True, + } + print("[INFO] Recording videos during training.") + print_dict(video_kwargs, nesting=4) + env = gym.wrappers.RecordVideo(env, **video_kwargs) + # wrap around environment for stable baselines + env = Sb3VecEnvWrapper(env) + + # normalize environment (if needed) + if "normalize_input" in agent_cfg: + env = VecNormalize( + env, + training=True, + norm_obs="normalize_input" in agent_cfg and agent_cfg.pop("normalize_input"), + norm_reward="normalize_value" in agent_cfg and agent_cfg.pop("normalize_value"), + clip_obs="clip_obs" in agent_cfg and agent_cfg.pop("clip_obs"), + gamma=agent_cfg["gamma"], + clip_reward=np.inf, + ) + + # create agent from stable baselines + print(f"Loading checkpoint from: {checkpoint_path}") + agent = PPO.load(checkpoint_path, env, print_system_info=True) + + dt = env.unwrapped.physics_dt + + # reset environment + obs = env.reset() + timestep = 0 + # simulate environment + while simulation_app.is_running(): + start_time = time.time() + # run everything in inference mode + with torch.inference_mode(): + # agent stepping + actions, _ = agent.predict(obs, deterministic=True) + # env stepping + obs, _, _, _ = env.step(actions) + if args_cli.video: + timestep += 1 + # Exit the play loop after recording one video + if timestep == args_cli.video_length: + break + + # time delay for real-time evaluation + sleep_time = dt - (time.time() - start_time) + if args_cli.real_time and sleep_time > 0: + time.sleep(sleep_time) + + # close the simulator + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/reinforcement_learning/sb3/train.py b/scripts/reinforcement_learning/sb3/train.py new file mode 100644 index 00000000..7a9f1bf1 --- /dev/null +++ b/scripts/reinforcement_learning/sb3/train.py @@ -0,0 +1,165 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Script to train RL agent with Stable Baselines3. + +Since Stable-Baselines3 does not support buffers living on GPU directly, +we recommend using smaller number of environments. Otherwise, +there will be significant overhead in GPU->CPU transfer. +""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse +import sys + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Train an RL agent with Stable-Baselines3.") +parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.") +parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).") +parser.add_argument("--video_interval", type=int, default=2000, help="Interval between video recordings (in steps).") +parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment") +parser.add_argument("--max_iterations", type=int, default=None, help="RL Policy training iterations.") +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli, hydra_args = parser.parse_known_args() +# always enable cameras to record video +if args_cli.video: + args_cli.enable_cameras = True + +# clear out sys.argv for Hydra +sys.argv = [sys.argv[0]] + hydra_args + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import gymnasium as gym +import numpy as np +import os +import random +from datetime import datetime + +from stable_baselines3 import PPO +from stable_baselines3.common.callbacks import CheckpointCallback +from stable_baselines3.common.logger import configure +from stable_baselines3.common.vec_env import VecNormalize + +from isaaclab.envs import ( + DirectMARLEnv, + DirectMARLEnvCfg, + DirectRLEnvCfg, + ManagerBasedRLEnvCfg, + multi_agent_to_single_agent, +) +from isaaclab.utils.dict import print_dict +from isaaclab.utils.io import dump_pickle, dump_yaml + +from isaaclab_rl.sb3 import Sb3VecEnvWrapper, process_sb3_cfg + +import isaaclab_tasks # noqa: F401 +import uwlab_tasks # noqa: F401 +from isaaclab_tasks.utils.hydra import hydra_task_config + + +@hydra_task_config(args_cli.task, "sb3_cfg_entry_point") +def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: dict): + """Train with stable-baselines agent.""" + # randomly sample a seed if seed = -1 + if args_cli.seed == -1: + args_cli.seed = random.randint(0, 10000) + + # override configurations with non-hydra CLI arguments + env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs + agent_cfg["seed"] = args_cli.seed if args_cli.seed is not None else agent_cfg["seed"] + # max iterations for training + if args_cli.max_iterations is not None: + agent_cfg["n_timesteps"] = args_cli.max_iterations * agent_cfg["n_steps"] * env_cfg.scene.num_envs + + # set the environment seed + # note: certain randomizations occur in the environment initialization so we set the seed here + env_cfg.seed = agent_cfg["seed"] + env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device + + # directory for logging into + run_info = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + log_root_path = os.path.abspath(os.path.join("logs", "sb3", args_cli.task)) + print(f"[INFO] Logging experiment in directory: {log_root_path}") + print(f"Exact experiment name requested from command line: {run_info}") + log_dir = os.path.join(log_root_path, run_info) + # dump the configuration into log-directory + dump_yaml(os.path.join(log_dir, "params", "env.yaml"), env_cfg) + dump_yaml(os.path.join(log_dir, "params", "agent.yaml"), agent_cfg) + dump_pickle(os.path.join(log_dir, "params", "env.pkl"), env_cfg) + dump_pickle(os.path.join(log_dir, "params", "agent.pkl"), agent_cfg) + + # post-process agent configuration + agent_cfg = process_sb3_cfg(agent_cfg) + # read configurations about the agent-training + policy_arch = agent_cfg.pop("policy") + n_timesteps = agent_cfg.pop("n_timesteps") + + # create isaac environment + env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) + + # convert to single-agent instance if required by the RL algorithm + if isinstance(env.unwrapped, DirectMARLEnv): + env = multi_agent_to_single_agent(env) + + # wrap for video recording + if args_cli.video: + video_kwargs = { + "video_folder": os.path.join(log_dir, "videos", "train"), + "step_trigger": lambda step: step % args_cli.video_interval == 0, + "video_length": args_cli.video_length, + "disable_logger": True, + } + print("[INFO] Recording videos during training.") + print_dict(video_kwargs, nesting=4) + env = gym.wrappers.RecordVideo(env, **video_kwargs) + + # wrap around environment for stable baselines + env = Sb3VecEnvWrapper(env) + + if "normalize_input" in agent_cfg: + env = VecNormalize( + env, + training=True, + norm_obs="normalize_input" in agent_cfg and agent_cfg.pop("normalize_input"), + norm_reward="normalize_value" in agent_cfg and agent_cfg.pop("normalize_value"), + clip_obs="clip_obs" in agent_cfg and agent_cfg.pop("clip_obs"), + gamma=agent_cfg["gamma"], + clip_reward=np.inf, + ) + + # create agent from stable baselines + agent = PPO(policy_arch, env, verbose=1, **agent_cfg) + # configure the logger + new_logger = configure(log_dir, ["stdout", "tensorboard"]) + agent.set_logger(new_logger) + + # callbacks for agent + checkpoint_callback = CheckpointCallback(save_freq=1000, save_path=log_dir, name_prefix="model", verbose=2) + # train the agent + agent.learn(total_timesteps=n_timesteps, callback=checkpoint_callback) + # save the final model + agent.save(os.path.join(log_dir, "model")) + + # close the simulator + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/reinforcement_learning/skrl/play.py b/scripts/reinforcement_learning/skrl/play.py new file mode 100644 index 00000000..43b59c52 --- /dev/null +++ b/scripts/reinforcement_learning/skrl/play.py @@ -0,0 +1,210 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +""" +Script to play a checkpoint of an RL agent from skrl. + +Visit the skrl documentation (https://skrl.readthedocs.io) to see the examples structured in +a more user-friendly way. +""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Play a checkpoint of an RL agent from skrl.") +parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.") +parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).") +parser.add_argument( + "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations." +) +parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +parser.add_argument("--checkpoint", type=str, default=None, help="Path to model checkpoint.") +parser.add_argument( + "--use_pretrained_checkpoint", + action="store_true", + help="Use the pre-trained checkpoint from Nucleus.", +) +parser.add_argument( + "--ml_framework", + type=str, + default="torch", + choices=["torch", "jax", "jax-numpy"], + help="The ML framework used for training the skrl agent.", +) +parser.add_argument( + "--algorithm", + type=str, + default="PPO", + choices=["AMP", "PPO", "IPPO", "MAPPO"], + help="The RL algorithm used for training the skrl agent.", +) +parser.add_argument("--real-time", action="store_true", default=False, help="Run in real-time, if possible.") + +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +args_cli = parser.parse_args() +# always enable cameras to record video +if args_cli.video: + args_cli.enable_cameras = True + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import gymnasium as gym +import os +import time +import torch + +import skrl +from packaging import version + +# check for minimum supported skrl version +SKRL_VERSION = "1.4.1" +if version.parse(skrl.__version__) < version.parse(SKRL_VERSION): + skrl.logger.error( + f"Unsupported skrl version: {skrl.__version__}. " + f"Install supported version using 'pip install skrl>={SKRL_VERSION}'" + ) + exit() + +if args_cli.ml_framework.startswith("torch"): + from skrl.utils.runner.torch import Runner +elif args_cli.ml_framework.startswith("jax"): + from skrl.utils.runner.jax import Runner + +from isaaclab.envs import DirectMARLEnv, multi_agent_to_single_agent +from isaaclab.utils.dict import print_dict +from isaaclab.utils.pretrained_checkpoint import get_published_pretrained_checkpoint + +from isaaclab_rl.skrl import SkrlVecEnvWrapper + +import isaaclab_tasks # noqa: F401 +import uwlab_tasks # noqa: F401 +from isaaclab_tasks.utils import get_checkpoint_path, load_cfg_from_registry, parse_env_cfg + +# config shortcuts +algorithm = args_cli.algorithm.lower() + + +def main(): + """Play with skrl agent.""" + # configure the ML framework into the global skrl variable + if args_cli.ml_framework.startswith("jax"): + skrl.config.jax.backend = "jax" if args_cli.ml_framework == "jax" else "numpy" + + # parse configuration + env_cfg = parse_env_cfg( + args_cli.task, device=args_cli.device, num_envs=args_cli.num_envs, use_fabric=not args_cli.disable_fabric + ) + try: + experiment_cfg = load_cfg_from_registry(args_cli.task, f"skrl_{algorithm}_cfg_entry_point") + except ValueError: + experiment_cfg = load_cfg_from_registry(args_cli.task, "skrl_cfg_entry_point") + + # specify directory for logging experiments (load checkpoint) + log_root_path = os.path.join("logs", "skrl", experiment_cfg["agent"]["experiment"]["directory"]) + log_root_path = os.path.abspath(log_root_path) + print(f"[INFO] Loading experiment from directory: {log_root_path}") + # get checkpoint path + if args_cli.use_pretrained_checkpoint: + resume_path = get_published_pretrained_checkpoint("skrl", args_cli.task) + if not resume_path: + print("[INFO] Unfortunately a pre-trained checkpoint is currently unavailable for this task.") + return + elif args_cli.checkpoint: + resume_path = os.path.abspath(args_cli.checkpoint) + else: + resume_path = get_checkpoint_path( + log_root_path, run_dir=f".*_{algorithm}_{args_cli.ml_framework}", other_dirs=["checkpoints"] + ) + log_dir = os.path.dirname(os.path.dirname(resume_path)) + + # create isaac environment + env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) + + # convert to single-agent instance if required by the RL algorithm + if isinstance(env.unwrapped, DirectMARLEnv) and algorithm in ["ppo"]: + env = multi_agent_to_single_agent(env) + + # get environment (physics) dt for real-time evaluation + try: + dt = env.physics_dt + except AttributeError: + dt = env.unwrapped.physics_dt + + # wrap for video recording + if args_cli.video: + video_kwargs = { + "video_folder": os.path.join(log_dir, "videos", "play"), + "step_trigger": lambda step: step == 0, + "video_length": args_cli.video_length, + "disable_logger": True, + } + print("[INFO] Recording videos during training.") + print_dict(video_kwargs, nesting=4) + env = gym.wrappers.RecordVideo(env, **video_kwargs) + + # wrap around environment for skrl + env = SkrlVecEnvWrapper(env, ml_framework=args_cli.ml_framework) # same as: `wrap_env(env, wrapper="auto")` + + # configure and instantiate the skrl runner + # https://skrl.readthedocs.io/en/latest/api/utils/runner.html + experiment_cfg["trainer"]["close_environment_at_exit"] = False + experiment_cfg["agent"]["experiment"]["write_interval"] = 0 # don't log to TensorBoard + experiment_cfg["agent"]["experiment"]["checkpoint_interval"] = 0 # don't generate checkpoints + runner = Runner(env, experiment_cfg) + + print(f"[INFO] Loading model checkpoint from: {resume_path}") + runner.agent.load(resume_path) + # set agent to evaluation mode + runner.agent.set_running_mode("eval") + + # reset environment + obs, _ = env.reset() + timestep = 0 + # simulate environment + while simulation_app.is_running(): + start_time = time.time() + + # run everything in inference mode + with torch.inference_mode(): + # agent stepping + outputs = runner.agent.act(obs, timestep=0, timesteps=0) + # - multi-agent (deterministic) actions + if hasattr(env, "possible_agents"): + actions = {a: outputs[-1][a].get("mean_actions", outputs[0][a]) for a in env.possible_agents} + # - single-agent (deterministic) actions + else: + actions = outputs[-1].get("mean_actions", outputs[0]) + # env stepping + obs, _, _, _, _ = env.step(actions) + if args_cli.video: + timestep += 1 + # exit the play loop after recording one video + if timestep == args_cli.video_length: + break + + # time delay for real-time evaluation + sleep_time = dt - (time.time() - start_time) + if args_cli.real_time and sleep_time > 0: + time.sleep(sleep_time) + + # close the simulator + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/reinforcement_learning/skrl/train.py b/scripts/reinforcement_learning/skrl/train.py new file mode 100644 index 00000000..55757f1c --- /dev/null +++ b/scripts/reinforcement_learning/skrl/train.py @@ -0,0 +1,202 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +""" +Script to train RL agent with skrl. + +Visit the skrl documentation (https://skrl.readthedocs.io) to see the examples structured in +a more user-friendly way. +""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse +import sys + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Train an RL agent with skrl.") +parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.") +parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).") +parser.add_argument("--video_interval", type=int, default=2000, help="Interval between video recordings (in steps).") +parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment") +parser.add_argument( + "--distributed", action="store_true", default=False, help="Run training with multiple GPUs or nodes." +) +parser.add_argument("--checkpoint", type=str, default=None, help="Path to model checkpoint to resume training.") +parser.add_argument("--max_iterations", type=int, default=None, help="RL Policy training iterations.") +parser.add_argument( + "--ml_framework", + type=str, + default="torch", + choices=["torch", "jax", "jax-numpy"], + help="The ML framework used for training the skrl agent.", +) +parser.add_argument( + "--algorithm", + type=str, + default="PPO", + choices=["AMP", "PPO", "IPPO", "MAPPO"], + help="The RL algorithm used for training the skrl agent.", +) + +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli, hydra_args = parser.parse_known_args() +# always enable cameras to record video +if args_cli.video: + args_cli.enable_cameras = True + +# clear out sys.argv for Hydra +sys.argv = [sys.argv[0]] + hydra_args + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import gymnasium as gym +import os +import random +from datetime import datetime + +import skrl +from packaging import version + +# check for minimum supported skrl version +SKRL_VERSION = "1.4.1" +if version.parse(skrl.__version__) < version.parse(SKRL_VERSION): + skrl.logger.error( + f"Unsupported skrl version: {skrl.__version__}. " + f"Install supported version using 'pip install skrl>={SKRL_VERSION}'" + ) + exit() + +if args_cli.ml_framework.startswith("torch"): + from skrl.utils.runner.torch import Runner +elif args_cli.ml_framework.startswith("jax"): + from skrl.utils.runner.jax import Runner + +from isaaclab.envs import ( + DirectMARLEnv, + DirectMARLEnvCfg, + DirectRLEnvCfg, + ManagerBasedRLEnvCfg, + multi_agent_to_single_agent, +) +from isaaclab.utils.assets import retrieve_file_path +from isaaclab.utils.dict import print_dict +from isaaclab.utils.io import dump_pickle, dump_yaml + +from isaaclab_rl.skrl import SkrlVecEnvWrapper + +import isaaclab_tasks # noqa: F401 +import uwlab_tasks # noqa: F401 +from isaaclab_tasks.utils.hydra import hydra_task_config + +# config shortcuts +algorithm = args_cli.algorithm.lower() +agent_cfg_entry_point = "skrl_cfg_entry_point" if algorithm in ["ppo"] else f"skrl_{algorithm}_cfg_entry_point" + + +@hydra_task_config(args_cli.task, agent_cfg_entry_point) +def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: dict): + """Train with skrl agent.""" + # override configurations with non-hydra CLI arguments + env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs + env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device + + # multi-gpu training config + if args_cli.distributed: + env_cfg.sim.device = f"cuda:{app_launcher.local_rank}" + # max iterations for training + if args_cli.max_iterations: + agent_cfg["trainer"]["timesteps"] = args_cli.max_iterations * agent_cfg["agent"]["rollouts"] + agent_cfg["trainer"]["close_environment_at_exit"] = False + # configure the ML framework into the global skrl variable + if args_cli.ml_framework.startswith("jax"): + skrl.config.jax.backend = "jax" if args_cli.ml_framework == "jax" else "numpy" + + # randomly sample a seed if seed = -1 + if args_cli.seed == -1: + args_cli.seed = random.randint(0, 10000) + + # set the agent and environment seed from command line + # note: certain randomization occur in the environment initialization so we set the seed here + agent_cfg["seed"] = args_cli.seed if args_cli.seed is not None else agent_cfg["seed"] + env_cfg.seed = agent_cfg["seed"] + + # specify directory for logging experiments + log_root_path = os.path.join("logs", "skrl", agent_cfg["agent"]["experiment"]["directory"]) + log_root_path = os.path.abspath(log_root_path) + print(f"[INFO] Logging experiment in directory: {log_root_path}") + # specify directory for logging runs: {time-stamp}_{run_name} + log_dir = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + f"_{algorithm}_{args_cli.ml_framework}" + print(f"Exact experiment name requested from command line {log_dir}") + if agent_cfg["agent"]["experiment"]["experiment_name"]: + log_dir += f'_{agent_cfg["agent"]["experiment"]["experiment_name"]}' + # set directory into agent config + agent_cfg["agent"]["experiment"]["directory"] = log_root_path + agent_cfg["agent"]["experiment"]["experiment_name"] = log_dir + # update log_dir + log_dir = os.path.join(log_root_path, log_dir) + + # dump the configuration into log-directory + dump_yaml(os.path.join(log_dir, "params", "env.yaml"), env_cfg) + dump_yaml(os.path.join(log_dir, "params", "agent.yaml"), agent_cfg) + dump_pickle(os.path.join(log_dir, "params", "env.pkl"), env_cfg) + dump_pickle(os.path.join(log_dir, "params", "agent.pkl"), agent_cfg) + + # get checkpoint path (to resume training) + resume_path = retrieve_file_path(args_cli.checkpoint) if args_cli.checkpoint else None + + # create isaac environment + env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) + + # convert to single-agent instance if required by the RL algorithm + if isinstance(env.unwrapped, DirectMARLEnv) and algorithm in ["ppo"]: + env = multi_agent_to_single_agent(env) + + # wrap for video recording + if args_cli.video: + video_kwargs = { + "video_folder": os.path.join(log_dir, "videos", "train"), + "step_trigger": lambda step: step % args_cli.video_interval == 0, + "video_length": args_cli.video_length, + "disable_logger": True, + } + print("[INFO] Recording videos during training.") + print_dict(video_kwargs, nesting=4) + env = gym.wrappers.RecordVideo(env, **video_kwargs) + + # wrap around environment for skrl + env = SkrlVecEnvWrapper(env, ml_framework=args_cli.ml_framework) # same as: `wrap_env(env, wrapper="auto")` + + # configure and instantiate the skrl runner + # https://skrl.readthedocs.io/en/latest/api/utils/runner.html + runner = Runner(env, agent_cfg) + + # load checkpoint (if specified) + if resume_path: + print(f"[INFO] Loading model checkpoint from: {resume_path}") + runner.agent.load(resume_path) + + # run training + runner.run() + + # close the simulator + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/tools/blender_obj.py b/scripts/tools/blender_obj.py new file mode 100644 index 00000000..aaef0a87 --- /dev/null +++ b/scripts/tools/blender_obj.py @@ -0,0 +1,99 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +""" +Convert a mesh file to `.obj` using blender. + +This file processes a given dae mesh file and saves the resulting mesh file in obj format. + +It needs to be called using the python packaged with blender, i.e.: + + blender --background --python blender_obj.py -- -in_file FILE -out_file FILE + +For more information: https://docs.blender.org/api/current/index.html + +The script was tested on Blender 3.2 on Ubuntu 20.04LTS. +""" + +import bpy +import os +import sys + + +def parse_cli_args(): + """Parse the input command line arguments.""" + import argparse + + # get the args passed to blender after "--", all of which are ignored by + # blender so scripts may receive their own arguments + argv = sys.argv + + if "--" not in argv: + argv = [] # as if no args are passed + else: + argv = argv[argv.index("--") + 1 :] # get all args after "--" + + # When --help or no args are given, print this help + usage_text = ( + f"Run blender in background mode with this script:\n\tblender --background --python {__file__} -- [options]" + ) + parser = argparse.ArgumentParser(description=usage_text) + # Add arguments + parser.add_argument("-i", "--in_file", metavar="FILE", type=str, required=True, help="Path to input OBJ file.") + parser.add_argument("-o", "--out_file", metavar="FILE", type=str, required=True, help="Path to output OBJ file.") + args = parser.parse_args(argv) + # Check if any arguments provided + if not argv or not args.in_file or not args.out_file: + parser.print_help() + return None + # return arguments + return args + + +def convert_to_obj(in_file: str, out_file: str, save_usd: bool = False): + """Convert a mesh file to `.obj` using blender. + + Args: + in_file: Input mesh file to process. + out_file: Path to store output obj file. + """ + # check valid input file + if not os.path.exists(in_file): + raise FileNotFoundError(in_file) + # add ending of file format + if not out_file.endswith(".obj"): + out_file += ".obj" + # create directory if it doesn't exist for destination file + if not os.path.exists(os.path.dirname(out_file)): + os.makedirs(os.path.dirname(out_file), exist_ok=True) + # reset scene to empty + bpy.ops.wm.read_factory_settings(use_empty=True) + # load object into scene + if in_file.endswith(".dae"): + bpy.ops.wm.collada_import(filepath=in_file) + elif in_file.endswith(".stl") or in_file.endswith(".STL"): + bpy.ops.import_mesh.stl(filepath=in_file) + else: + raise ValueError(f"Input file not in dae/stl format: {in_file}") + # convert to obj format and store with z up + # TODO: Read the convention from dae file instead of manually fixing it. + # Reference: https://docs.blender.org/api/2.79/bpy.ops.export_scene.html + bpy.ops.export_scene.obj( + filepath=out_file, check_existing=False, axis_forward="Y", axis_up="Z", global_scale=1, path_mode="RELATIVE" + ) + # save it as usd as well + if save_usd: + out_file = out_file.replace("obj", "usd") + bpy.ops.wm.usd_export(filepath=out_file, check_existing=False) + + +if __name__ == "__main__": + # read arguments + cli_args = parse_cli_args() + # check CLI args + if cli_args is None: + sys.exit() + # process via blender + convert_to_obj(cli_args.in_file, cli_args.out_file) diff --git a/scripts/tools/check_instanceable.py b/scripts/tools/check_instanceable.py new file mode 100644 index 00000000..5e5cddda --- /dev/null +++ b/scripts/tools/check_instanceable.py @@ -0,0 +1,133 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +""" +This script uses the cloner API to check if asset has been instanced properly. + +Usage with different inputs (replace `` and `` with the path to the +original asset and the instanced asset respectively): + +```bash +./isaaclab.sh -p source/tools/check_instanceable.py -n 4096 --headless --physics +./isaaclab.sh -p source/tools/check_instanceable.py -n 4096 --headless --physics +./isaaclab.sh -p source/tools/check_instanceable.py -n 4096 --headless +./isaaclab.sh -p source/tools/check_instanceable.py -n 4096 --headless +``` + +Output from the above commands: + +```bash +>>> Cloning time (cloner.clone): 0.648198 seconds +>>> Setup time (sim.reset): : 5.843589 seconds +[#clones: 4096, physics: True] Asset: : 6.491870 seconds + +>>> Cloning time (cloner.clone): 0.693133 seconds +>>> Setup time (sim.reset): 50.860526 seconds +[#clones: 4096, physics: True] Asset: : 51.553743 seconds + +>>> Cloning time (cloner.clone) : 0.687201 seconds +>>> Setup time (sim.reset) : 6.302215 seconds +[#clones: 4096, physics: False] Asset: : 6.989500 seconds + +>>> Cloning time (cloner.clone) : 0.678150 seconds +>>> Setup time (sim.reset) : 52.854054 seconds +[#clones: 4096, physics: False] Asset: : 53.532287 seconds +``` + +""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse +import contextlib +import os + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser("Utility to empirically check if asset in instanced properly.") +parser.add_argument("input", type=str, help="The path to the USD file.") +parser.add_argument("-n", "--num_clones", type=int, default=128, help="Number of clones to spawn.") +parser.add_argument("-s", "--spacing", type=float, default=1.5, help="Spacing between instances in a grid.") +parser.add_argument("-p", "--physics", action="store_true", default=False, help="Clone assets using physics cloner.") +# 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 isaacsim.core.utils.prims as prim_utils +from isaacsim.core.api.simulation_context import SimulationContext +from isaacsim.core.cloner import GridCloner +from isaacsim.core.utils.carb import set_carb_setting + +from isaaclab.utils import Timer +from isaaclab.utils.assets import check_file_path + + +def main(): + """Spawns the USD asset robot and clones it using Isaac Gym Cloner API.""" + # check valid file path + if not check_file_path(args_cli.input): + raise ValueError(f"Invalid file path: {args_cli.input}") + # Load kit helper + sim = SimulationContext( + stage_units_in_meters=1.0, physics_dt=0.01, rendering_dt=0.01, backend="torch", device="cuda:0" + ) + # enable fabric which avoids passing data over to USD structure + # this speeds up the read-write operation of GPU buffers + if sim.get_physics_context().use_gpu_pipeline: + sim.get_physics_context().enable_fabric(True) + # increase GPU buffer dimensions + sim.get_physics_context().set_gpu_found_lost_aggregate_pairs_capacity(2**25) + sim.get_physics_context().set_gpu_total_aggregate_pairs_capacity(2**21) + # enable hydra scene-graph instancing + # this is needed to visualize the scene when fabric is enabled + set_carb_setting(sim._settings, "/persistent/omnihydra/useSceneGraphInstancing", True) + + # Create interface to clone the scene + cloner = GridCloner(spacing=args_cli.spacing) + cloner.define_base_env("/World/envs") + prim_utils.define_prim("/World/envs/env_0") + # Spawn things into stage + prim_utils.create_prim("/World/Light", "DistantLight") + + # Everything under the namespace "/World/envs/env_0" will be cloned + prim_utils.create_prim("/World/envs/env_0/Asset", "Xform", usd_path=os.path.abspath(args_cli.input)) + # Clone the scene + num_clones = args_cli.num_clones + + # Create a timer to measure the cloning time + with Timer(f"[#clones: {num_clones}, physics: {args_cli.physics}] Asset: {args_cli.input}"): + # Clone the scene + with Timer(">>> Cloning time (cloner.clone)"): + cloner.define_base_env("/World/envs") + envs_prim_paths = cloner.generate_paths("/World/envs/env", num_paths=num_clones) + _ = cloner.clone( + source_prim_path="/World/envs/env_0", prim_paths=envs_prim_paths, replicate_physics=args_cli.physics + ) + # Play the simulator + with Timer(">>> Setup time (sim.reset)"): + sim.reset() + + # Simulate scene (if not headless) + if not args_cli.headless: + with contextlib.suppress(KeyboardInterrupt): + while sim.is_playing(): + # perform step + sim.step() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/tools/convert_mesh.py b/scripts/tools/convert_mesh.py new file mode 100644 index 00000000..03c0f509 --- /dev/null +++ b/scripts/tools/convert_mesh.py @@ -0,0 +1,175 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +""" +Utility to convert a OBJ/STL/FBX into USD format. + +The OBJ file format is a simple data-format that represents 3D geometry alone — namely, the position +of each vertex, the UV position of each texture coordinate vertex, vertex normals, and the faces that +make each polygon defined as a list of vertices, and texture vertices. + +An STL file describes a raw, unstructured triangulated surface by the unit normal and vertices (ordered +by the right-hand rule) of the triangles using a three-dimensional Cartesian coordinate system. + +FBX files are a type of 3D model file created using the Autodesk FBX software. They can be designed and +modified in various modeling applications, such as Maya, 3ds Max, and Blender. Moreover, FBX files typically +contain mesh, material, texture, and skeletal animation data. +Link: https://www.autodesk.com/products/fbx/overview + + +This script uses the asset converter extension from Isaac Sim (``omni.kit.asset_converter``) to convert a +OBJ/STL/FBX asset into USD format. It is designed as a convenience script for command-line use. + + +positional arguments: + input The path to the input mesh (.OBJ/.STL/.FBX) file. + output The path to store the USD file. + +optional arguments: + -h, --help Show this help message and exit + --make-instanceable, Make the asset instanceable for efficient cloning. (default: False) + --collision-approximation The method used for approximating collision mesh. Defaults to convexDecomposition. + Set to \"none\" to not add a collision mesh to the converted mesh. (default: convexDecomposition) + --mass The mass (in kg) to assign to the converted asset. (default: None) + +""" + +"""Launch Isaac Sim Simulator first.""" + + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Utility to convert a mesh file into USD format.") +parser.add_argument("input", type=str, help="The path to the input mesh file.") +parser.add_argument("output", type=str, help="The path to store the USD file.") +parser.add_argument( + "--make-instanceable", + action="store_true", + default=False, + help="Make the asset instanceable for efficient cloning.", +) +parser.add_argument( + "--collision-approximation", + type=str, + default="convexDecomposition", + choices=["convexDecomposition", "convexHull", "none"], + help=( + 'The method used for approximating collision mesh. Set to "none" ' + "to not add a collision mesh to the converted mesh." + ), +) +parser.add_argument( + "--mass", + type=float, + default=None, + help="The mass (in kg) to assign to the converted asset. If not provided, then no mass is added.", +) +# 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 contextlib +import os + +import carb +import isaacsim.core.utils.stage as stage_utils +import omni.kit.app + +from isaaclab.sim.converters import MeshConverter, MeshConverterCfg +from isaaclab.sim.schemas import schemas_cfg +from isaaclab.utils.assets import check_file_path +from isaaclab.utils.dict import print_dict + + +def main(): + # check valid file path + mesh_path = args_cli.input + if not os.path.isabs(mesh_path): + mesh_path = os.path.abspath(mesh_path) + if not check_file_path(mesh_path): + raise ValueError(f"Invalid mesh file path: {mesh_path}") + + # create destination path + dest_path = args_cli.output + if not os.path.isabs(dest_path): + dest_path = os.path.abspath(dest_path) + + # Mass properties + if args_cli.mass is not None: + mass_props = schemas_cfg.MassPropertiesCfg(mass=args_cli.mass) + rigid_props = schemas_cfg.RigidBodyPropertiesCfg() + else: + mass_props = None + rigid_props = None + + # Collision properties + collision_props = schemas_cfg.CollisionPropertiesCfg(collision_enabled=args_cli.collision_approximation != "none") + + # Create Mesh converter config + mesh_converter_cfg = MeshConverterCfg( + mass_props=mass_props, + rigid_props=rigid_props, + collision_props=collision_props, + asset_path=mesh_path, + force_usd_conversion=True, + usd_dir=os.path.dirname(dest_path), + usd_file_name=os.path.basename(dest_path), + make_instanceable=args_cli.make_instanceable, + collision_approximation=args_cli.collision_approximation, + ) + + # Print info + print("-" * 80) + print("-" * 80) + print(f"Input Mesh file: {mesh_path}") + print("Mesh importer config:") + print_dict(mesh_converter_cfg.to_dict(), nesting=0) + print("-" * 80) + print("-" * 80) + + # Create Mesh converter and import the file + mesh_converter = MeshConverter(mesh_converter_cfg) + # print output + print("Mesh importer output:") + print(f"Generated USD file: {mesh_converter.usd_path}") + print("-" * 80) + print("-" * 80) + + # Determine if there is a GUI to update: + # acquire settings interface + carb_settings_iface = carb.settings.get_settings() + # read flag for whether a local GUI is enabled + local_gui = carb_settings_iface.get("/app/window/enabled") + # read flag for whether livestreaming GUI is enabled + livestream_gui = carb_settings_iface.get("/app/livestream/enabled") + + # Simulate scene (if not headless) + if local_gui or livestream_gui: + # Open the stage with USD + stage_utils.open_stage(mesh_converter.usd_path) + # Reinitialize the simulation + app = omni.kit.app.get_app_interface() + # Run simulation + with contextlib.suppress(KeyboardInterrupt): + while app.is_running(): + # perform step + app.update() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/tools/convert_mjcf.py b/scripts/tools/convert_mjcf.py new file mode 100644 index 00000000..079bd8db --- /dev/null +++ b/scripts/tools/convert_mjcf.py @@ -0,0 +1,139 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +""" +Utility to convert a MJCF into USD format. + +MuJoCo XML Format (MJCF) is an XML file format used in MuJoCo to describe all elements of a robot. +For more information, see: http://www.mujoco.org/book/XMLreference.html + +This script uses the MJCF importer extension from Isaac Sim (``isaacsim.asset.importer.mjcf``) to convert +a MJCF asset into USD format. It is designed as a convenience script for command-line use. For more information +on the MJCF importer, see the documentation for the extension: +https://docs.isaacsim.omniverse.nvidia.com/latest/robot_setup/ext_isaacsim_asset_importer_mjcf.html + + +positional arguments: + input The path to the input URDF file. + output The path to store the USD file. + +optional arguments: + -h, --help Show this help message and exit + --fix-base Fix the base to where it is imported. (default: False) + --import-sites Import sites by parse tag. (default: True) + --make-instanceable Make the asset instanceable for efficient cloning. (default: False) + +""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Utility to convert a MJCF into USD format.") +parser.add_argument("input", type=str, help="The path to the input MJCF file.") +parser.add_argument("output", type=str, help="The path to store the USD file.") +parser.add_argument("--fix-base", action="store_true", default=False, help="Fix the base to where it is imported.") +parser.add_argument( + "--import-sites", action="store_true", default=False, help="Import sites by parsing the tag." +) +parser.add_argument( + "--make-instanceable", + action="store_true", + default=False, + help="Make the asset instanceable for efficient cloning.", +) + +# 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 contextlib +import os + +import carb +import isaacsim.core.utils.stage as stage_utils +import omni.kit.app + +from isaaclab.sim.converters import MjcfConverter, MjcfConverterCfg +from isaaclab.utils.assets import check_file_path +from isaaclab.utils.dict import print_dict + + +def main(): + # check valid file path + mjcf_path = args_cli.input + if not os.path.isabs(mjcf_path): + mjcf_path = os.path.abspath(mjcf_path) + if not check_file_path(mjcf_path): + raise ValueError(f"Invalid file path: {mjcf_path}") + # create destination path + dest_path = args_cli.output + if not os.path.isabs(dest_path): + dest_path = os.path.abspath(dest_path) + + # create the converter configuration + mjcf_converter_cfg = MjcfConverterCfg( + asset_path=mjcf_path, + usd_dir=os.path.dirname(dest_path), + usd_file_name=os.path.basename(dest_path), + fix_base=args_cli.fix_base, + import_sites=args_cli.import_sites, + force_usd_conversion=True, + make_instanceable=args_cli.make_instanceable, + ) + + # Print info + print("-" * 80) + print("-" * 80) + print(f"Input MJCF file: {mjcf_path}") + print("MJCF importer config:") + print_dict(mjcf_converter_cfg.to_dict(), nesting=0) + print("-" * 80) + print("-" * 80) + + # Create mjcf converter and import the file + mjcf_converter = MjcfConverter(mjcf_converter_cfg) + # print output + print("MJCF importer output:") + print(f"Generated USD file: {mjcf_converter.usd_path}") + print("-" * 80) + print("-" * 80) + + # Determine if there is a GUI to update: + # acquire settings interface + carb_settings_iface = carb.settings.get_settings() + # read flag for whether a local GUI is enabled + local_gui = carb_settings_iface.get("/app/window/enabled") + # read flag for whether livestreaming GUI is enabled + livestream_gui = carb_settings_iface.get("/app/livestream/enabled") + + # Simulate scene (if not headless) + if local_gui or livestream_gui: + # Open the stage with USD + stage_utils.open_stage(mjcf_converter.usd_path) + # Reinitialize the simulation + app = omni.kit.app.get_app_interface() + # Run simulation + with contextlib.suppress(KeyboardInterrupt): + while app.is_running(): + # perform step + app.update() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/tools/convert_urdf.py b/scripts/tools/convert_urdf.py new file mode 100644 index 00000000..9721fbab --- /dev/null +++ b/scripts/tools/convert_urdf.py @@ -0,0 +1,163 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +""" +Utility to convert a URDF into USD format. + +Unified Robot Description Format (URDF) is an XML file format used in ROS to describe all elements of +a robot. For more information, see: http://wiki.ros.org/urdf + +This script uses the URDF importer extension from Isaac Sim (``isaacsim.asset.importer.urdf``) to convert a +URDF asset into USD format. It is designed as a convenience script for command-line use. For more +information on the URDF importer, see the documentation for the extension: +https://docs.isaacsim.omniverse.nvidia.com/latest/robot_setup/ext_isaacsim_asset_importer_urdf.html + + +positional arguments: + input The path to the input URDF file. + output The path to store the USD file. + +optional arguments: + -h, --help Show this help message and exit + --merge-joints Consolidate links that are connected by fixed joints. (default: False) + --fix-base Fix the base to where it is imported. (default: False) + --joint-stiffness The stiffness of the joint drive. (default: 100.0) + --joint-damping The damping of the joint drive. (default: 1.0) + --joint-target-type The type of control to use for the joint drive. (default: "position") + +""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Utility to convert a URDF into USD format.") +parser.add_argument("input", type=str, help="The path to the input URDF file.") +parser.add_argument("output", type=str, help="The path to store the USD file.") +parser.add_argument( + "--merge-joints", + action="store_true", + default=False, + help="Consolidate links that are connected by fixed joints.", +) +parser.add_argument("--fix-base", action="store_true", default=False, help="Fix the base to where it is imported.") +parser.add_argument( + "--joint-stiffness", + type=float, + default=100.0, + help="The stiffness of the joint drive.", +) +parser.add_argument( + "--joint-damping", + type=float, + default=1.0, + help="The damping of the joint drive.", +) +parser.add_argument( + "--joint-target-type", + type=str, + default="position", + choices=["position", "velocity", "none"], + help="The type of control to use for the joint drive.", +) + +# 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 contextlib +import os + +import carb +import isaacsim.core.utils.stage as stage_utils +import omni.kit.app + +from isaaclab.sim.converters import UrdfConverter, UrdfConverterCfg +from isaaclab.utils.assets import check_file_path +from isaaclab.utils.dict import print_dict + + +def main(): + # check valid file path + urdf_path = args_cli.input + if not os.path.isabs(urdf_path): + urdf_path = os.path.abspath(urdf_path) + if not check_file_path(urdf_path): + raise ValueError(f"Invalid file path: {urdf_path}") + # create destination path + dest_path = args_cli.output + if not os.path.isabs(dest_path): + dest_path = os.path.abspath(dest_path) + + # Create Urdf converter config + urdf_converter_cfg = UrdfConverterCfg( + asset_path=urdf_path, + usd_dir=os.path.dirname(dest_path), + usd_file_name=os.path.basename(dest_path), + fix_base=args_cli.fix_base, + merge_fixed_joints=args_cli.merge_joints, + force_usd_conversion=True, + joint_drive=UrdfConverterCfg.JointDriveCfg( + gains=UrdfConverterCfg.JointDriveCfg.PDGainsCfg( + stiffness=args_cli.joint_stiffness, + damping=args_cli.joint_damping, + ), + target_type=args_cli.joint_target_type, + ), + ) + + # Print info + print("-" * 80) + print("-" * 80) + print(f"Input URDF file: {urdf_path}") + print("URDF importer config:") + print_dict(urdf_converter_cfg.to_dict(), nesting=0) + print("-" * 80) + print("-" * 80) + + # Create Urdf converter and import the file + urdf_converter = UrdfConverter(urdf_converter_cfg) + # print output + print("URDF importer output:") + print(f"Generated USD file: {urdf_converter.usd_path}") + print("-" * 80) + print("-" * 80) + + # Determine if there is a GUI to update: + # acquire settings interface + carb_settings_iface = carb.settings.get_settings() + # read flag for whether a local GUI is enabled + local_gui = carb_settings_iface.get("/app/window/enabled") + # read flag for whether livestreaming GUI is enabled + livestream_gui = carb_settings_iface.get("/app/livestream/enabled") + + # Simulate scene (if not headless) + if local_gui or livestream_gui: + # Open the stage with USD + stage_utils.open_stage(urdf_converter.usd_path) + # Reinitialize the simulation + app = omni.kit.app.get_app_interface() + # Run simulation + with contextlib.suppress(KeyboardInterrupt): + while app.is_running(): + # perform step + app.update() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/tools/merge_hdf5_datasets.py b/scripts/tools/merge_hdf5_datasets.py new file mode 100644 index 00000000..f7b8cfda --- /dev/null +++ b/scripts/tools/merge_hdf5_datasets.py @@ -0,0 +1,47 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import argparse +import h5py +import os + +parser = argparse.ArgumentParser(description="Merge a set of HDF5 datasets.") +parser.add_argument( + "--input_files", + type=str, + nargs="+", + default=[], + help="A list of paths to HDF5 files to merge.", +) +parser.add_argument("--output_file", type=str, default="merged_dataset.hdf5", help="File path to merged output.") + +args_cli = parser.parse_args() + + +def merge_datasets(): + for filepath in args_cli.input_files: + if not os.path.exists(filepath): + raise FileNotFoundError(f"The dataset file {filepath} does not exist.") + + with h5py.File(args_cli.output_file, "w") as output: + episode_idx = 0 + copy_attributes = True + + for filepath in args_cli.input_files: + + with h5py.File(filepath, "r") as input: + for episode, data in input["data"].items(): + input.copy(f"data/{episode}", output, f"data/demo_{episode_idx}") + episode_idx += 1 + + if copy_attributes: + output["data"].attrs["env_args"] = input["data"].attrs["env_args"] + copy_attributes = False + + print(f"Merged dataset saved to {args_cli.output_file}") + + +if __name__ == "__main__": + merge_datasets() diff --git a/scripts/tools/pretrained_checkpoint.py b/scripts/tools/pretrained_checkpoint.py new file mode 100644 index 00000000..9476cf98 --- /dev/null +++ b/scripts/tools/pretrained_checkpoint.py @@ -0,0 +1,376 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +""" +Script to manage pretrained checkpoints for our environments. +""" + +import argparse + +from isaaclab.app import AppLauncher + +# Initialize the parser +parser = argparse.ArgumentParser( + description=""" +Script used for the training and publishing of pre-trained checkpoints for Isaac Lab. + +Examples : + # Train an agent using the rl_games workflow on the Isaac-Cartpole-v0 environment. + pretrained_checkpoint.py --train rl_games:Isaac-Cartpole-v0 + # Train and publish the checkpoints for all workflows on only the direct Cartpole environments. + pretrained_checkpoint.py -tp "*:Isaac-Cartpole-*Direct-v0" \\ + --/persistent/isaaclab/asset_root/pretrained_checkpoints="/some/path" + # Review all repose cube jobs, excluding the Play tasks and skrl + pretrained_checkpoint.py -r "*:*Repose-Cube*" --exclude "*:*Play*" --exclude skrl:* + # Publish all results (that have been reviewed and approved). + pretrained_checkpoint.py --publish --all \\ + --/persistent/isaaclab/asset_root/pretrained_checkpoints="/some/path" +""", + formatter_class=argparse.RawTextHelpFormatter, +) + +# Add positional arguments that can accept zero or more values +parser.add_argument( + "jobs", + nargs="*", + help=""" +A job consists of a workflow and a task name separated by a colon (wildcards optional), for example : + rl_games:Isaac-Humanoid-*v0 + rsl_rl:Isaac-Ant-*-v0 + *:Isaac-Velocity-Flat-Spot-v0 +""", +) +parser.add_argument("-t", "--train", action="store_true", help="Train checkpoints for later publishing.") +parser.add_argument("-p", "--publish_checkpoint", action="store_true", help="Publish pre-trained checkpoints.") +parser.add_argument("-r", "--review", action="store_true", help="Review checkpoints.") +parser.add_argument("-l", "--list", action="store_true", help="List all available environments and workflows.") +parser.add_argument("-f", "--force", action="store_true", help="Force training when results already exist.") +parser.add_argument("-a", "--all", action="store_true", help="Run all valid workflow task pairs.") +parser.add_argument( + "-E", + "--exclude", + action="append", + type=str, + default=[], + help="Excludes jobs matching the argument, with wildcard support.", +) +parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") +parser.add_argument("--force_review", action="store_true", help="Forces review when one already exists.") +parser.add_argument("--force_publish", action="store_true", help="Publish checkpoints without review.") +parser.add_argument("--headless", action="store_true", help="Run training without the UI.") + +args, _ = parser.parse_known_args() + +# Need something to do +if len(args.jobs) == 0 and not args.all: + parser.error("Jobs must be provided, or --all.") + +# Must train, publish, review or list +if not (args.train or args.publish_checkpoint or args.review or args.list): + parser.error("A train, publish, review or list flag must be given.") + +# List excludes train and publish +if args.list and (args.train or args.publish_checkpoint): + parser.error("Can't train or publish when listing.") + +# launch omniverse app +app_launcher = AppLauncher(headless=True) +simulation_app = app_launcher.app + + +import csv + +# Now everything else +import fnmatch +import gymnasium as gym +import json +import numpy as np +import os +import subprocess +import sys + +import omni.client +from omni.client._omniclient import CopyBehavior + +from isaaclab.utils.pretrained_checkpoint import ( + WORKFLOW_EXPERIMENT_NAME_VARIABLE, + WORKFLOW_PLAYER, + WORKFLOW_TRAINER, + WORKFLOWS, + get_log_root_path, + get_pretrained_checkpoint_path, + get_pretrained_checkpoint_publish_path, + get_pretrained_checkpoint_review, + get_pretrained_checkpoint_review_path, + has_pretrained_checkpoint_job_finished, + has_pretrained_checkpoint_job_run, + has_pretrained_checkpoints_asset_root_dir, +) + +# Need somewhere to publish +if args.publish_checkpoint and not has_pretrained_checkpoints_asset_root_dir(): + raise Exception("A /persistent/isaaclab/asset_root/pretrained_checkpoints setting is required to publish.") + + +def train_job(workflow, task_name, headless=False, force=False, num_envs=None): + """ + This trains a task using the workflow's train.py script, overriding the experiment name to ensure unique + log directories. By default it will return if an experiment has already been run. + + Args: + workflow: The workflow. + task_name: The task name. + headless: Should the training run without the UI. + force: Run training even if previous experiments have been run. + num_envs: How many simultaneous environments to simulate, overriding the config. + """ + + log_root_path = get_log_root_path(workflow, task_name) + + # We already ran this + if not force and os.path.exists(log_root_path) and len(os.listdir(log_root_path)) > 0: + print(f"Skipping training of {workflow}:{task_name}, already has been run") + return + + print(f"Training {workflow}:{task_name}") + + # Construct our command + cmd = [ + sys.executable, + WORKFLOW_TRAINER[workflow], + "--task", + task_name, + "--enable_cameras", + ] + + # Changes the directory name for logging + if WORKFLOW_EXPERIMENT_NAME_VARIABLE[workflow]: + cmd.append(f"{WORKFLOW_EXPERIMENT_NAME_VARIABLE[workflow]}={task_name}") + + if headless: + cmd.append("--headless") + if num_envs: + cmd.extend(["--num_envs", str(num_envs)]) + + print("Running : " + " ".join(cmd)) + + subprocess.run(cmd) + + +def review_pretrained_checkpoint(workflow, task_name, force_review=False, num_envs=None): + """ + This initiates a review of the pretrained checkpoint. The play.py script for the workflow is run, and the user + inspects the results. When done they close the simulator and will be prompted for their review. + + Args: + workflow: The workflow. + task_name: The task name. + force_review: Performs the review even if a review already exists. + num_envs: How many simultaneous environments to simulate, overriding the config. + """ + + # This workflow task pair hasn't been trained + if not has_pretrained_checkpoint_job_run(workflow, task_name): + print(f"Skipping review of {workflow}:{task_name}, hasn't been trained yet") + return + + # Couldn't find the checkpoint + if not has_pretrained_checkpoint_job_finished(workflow, task_name): + print(f"Training not complete for {workflow}:{task_name}") + return + + review = get_pretrained_checkpoint_review(workflow, task_name) + + if not force_review and review and review["reviewed"]: + print(f"Review already complete for {workflow}:{task_name}") + return + + print(f"Reviewing {workflow}:{task_name}") + + # Construct our command + cmd = [ + sys.executable, + WORKFLOW_PLAYER[workflow], + "--task", + task_name, + "--checkpoint", + get_pretrained_checkpoint_path(workflow, task_name), + "--enable_cameras", + ] + + if num_envs: + cmd.extend(["--num_envs", str(num_envs)]) + + print("Running : " + " ".join(cmd)) + + subprocess.run(cmd) + + # Give user a chance to leave the old review + if force_review and review and review["reviewed"]: + result = review["result"] + notes = review.get("notes") + print(f"A review already exists for {workflow}:{task_name}, it was marked as '{result}'.") + print(f" Notes: {notes}") + answer = input("Would you like to replace it? Please answer yes or no (y/n) [n]: ").strip().lower() + if answer != "y": + return + + # Get the verdict from the user + print(f"Do you accept this checkpoint for {workflow}:{task_name}?") + + answer = input("Please answer yes, no or undetermined (y/n/u) [u]: ").strip().lower() + if answer not in {"y", "n", "u"}: + answer = "u" + answer_map = { + "y": "accepted", + "n": "rejected", + "u": "undetermined", + } + + # Create the review dict + review = { + "reviewed": True, + "result": answer_map[answer], + } + + # Maybe add some notes + notes = input("Please add notes or hit enter: ").strip().lower() + if notes: + review["notes"] = notes + + # Save the review JSON file + path = get_pretrained_checkpoint_review_path(workflow, task_name) + if not path: + raise Exception("This shouldn't be possible, something went very wrong.") + + with open(path, "w") as f: + json.dump(review, f, indent=4) + + +def publish_pretrained_checkpoint(workflow, task_name, force_publish=False): + """ + This publishes the pretrained checkpoint to Nucleus using the asset path in the + /persistent/isaaclab/asset_root/pretrained_checkpoints Carb variable. + + Args: + workflow: The workflow. + task_name: The task name. + force_publish: Publish without review. + """ + + # This workflow task pair hasn't been trained + if not has_pretrained_checkpoint_job_run(workflow, task_name): + print(f"Skipping publishing of {workflow}:{task_name}, hasn't been trained yet") + return + + # Couldn't find the checkpoint + if not has_pretrained_checkpoint_job_finished(workflow, task_name): + print(f"Training not complete for {workflow}:{task_name}") + return + + # Get local pretrained checkpoint path + local_path = get_pretrained_checkpoint_path(workflow, task_name) + if not local_path: + raise Exception("This shouldn't be possible, something went very wrong.") + + # Not forcing, need to check review results + if not force_publish: + + # Grab the review if it exists + review = get_pretrained_checkpoint_review(workflow, task_name) + + if not review or not review["reviewed"]: + print(f"Skipping publishing of {workflow}:{task_name}, hasn't been reviewed yet") + return + + result = review["result"] + if result != "accepted": + print(f'Skipping publishing of {workflow}:{task_name}, review result was "{result}"') + return + + print(f"Publishing {workflow}:{task_name}") + + # Copy the file + publish_path = get_pretrained_checkpoint_publish_path(workflow, task_name) + omni.client.copy_file(local_path, publish_path, CopyBehavior.OVERWRITE) + + +def get_job_summary_row(workflow, task_name): + """Returns a single row summary of the job""" + + has_run = has_pretrained_checkpoint_job_run(workflow, task_name) + has_finished = has_pretrained_checkpoint_job_finished(workflow, task_name) + review = get_pretrained_checkpoint_review(workflow, task_name) + + if review: + result = review.get("result", "undetermined") + notes = review.get("notes", "") + else: + result = "" + notes = "" + + return [workflow, task_name, has_run, has_finished, result, notes] + + +def main(): + + # Figure out what workflows and tasks we'll be using + if args.all: + jobs = ["*:*"] + else: + jobs = args.jobs + + if args.list: + print() + print("# Workflow, Task, Ran, Finished, Review, Notes") + + summary_rows = [] + + # Could be implemented more efficiently, but the performance gain would be inconsequential + for workflow in WORKFLOWS: + for task_spec in sorted(gym.registry.values(), key=lambda t: t.id): + job_id = f"{workflow}:{task_spec.id}" + + # We've excluded this job + if any(fnmatch.fnmatch(job_id, e) for e in args.exclude): + continue + + # None of our jobs match this pair + if not np.any(np.array([fnmatch.fnmatch(job_id, job) for job in jobs])): + continue + + # No config for this workflow + if workflow + "_cfg_entry_point" not in task_spec.kwargs: + continue + + if args.list: + summary_rows.append(get_job_summary_row(workflow, task_spec.id)) + continue + + # Training reviewing and publishing + if args.train: + train_job(workflow, task_spec.id, args.headless, args.force, args.num_envs) + + if args.review: + review_pretrained_checkpoint(workflow, task_spec.id, args.force_review, args.num_envs) + + if args.publish_checkpoint: + publish_pretrained_checkpoint(workflow, task_spec.id, args.force_publish) + + if args.list: + writer = csv.writer(sys.stdout, quotechar='"', quoting=csv.QUOTE_MINIMAL) + writer.writerows(summary_rows) + + +if __name__ == "__main__": + + try: + # Run the main function + main() + except Exception as e: + raise e + finally: + # Close the app + simulation_app.close() diff --git a/scripts/tools/process_meshes_to_obj.py b/scripts/tools/process_meshes_to_obj.py new file mode 100644 index 00000000..17b8a765 --- /dev/null +++ b/scripts/tools/process_meshes_to_obj.py @@ -0,0 +1,90 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Convert all mesh files to `.obj` in given folders.""" + +import argparse +import os +import shutil +import subprocess + +# Constants +# Path to blender +BLENDER_EXE_PATH = shutil.which("blender") + + +def parse_cli_args(): + """Parse the input command line arguments.""" + # add argparse arguments + parser = argparse.ArgumentParser("Utility to convert all mesh files to `.obj` in given folders.") + parser.add_argument("input_dir", type=str, help="The input directory from which to load meshes.") + parser.add_argument( + "-o", + "--output_dir", + type=str, + default=None, + help="The output directory to save converted meshes into. Default is same as input directory.", + ) + args_cli = parser.parse_args() + # resolve output directory + if args_cli.output_dir is None: + args_cli.output_dir = args_cli.input_dir + # return arguments + return args_cli + + +def run_blender_convert2obj(in_file: str, out_file: str): + """Calls the python script using `subprocess` to perform processing of mesh file. + + Args: + in_file: Input mesh file. + out_file: Output obj file. + """ + # resolve for python file + tools_dirname = os.path.dirname(os.path.abspath(__file__)) + script_file = os.path.join(tools_dirname, "blender_obj.py") + # complete command + command_exe = f"{BLENDER_EXE_PATH} --background --python {script_file} -- -i {in_file} -o {out_file}" + # break command into list + command_exe_list = command_exe.split(" ") + # run command + subprocess.run(command_exe_list) + + +def convert_meshes(source_folders: list[str], destination_folders: list[str]): + """Processes all mesh files of supported format into OBJ file using blender. + + Args: + source_folders: List of directories to search for meshes. + destination_folders: List of directories to dump converted files. + """ + # create folder for corresponding destination + for folder in destination_folders: + os.makedirs(folder, exist_ok=True) + # iterate over each folder + for in_folder, out_folder in zip(source_folders, destination_folders): + # extract all dae files in the directory + mesh_filenames = [f for f in os.listdir(in_folder) if f.endswith("dae")] + mesh_filenames += [f for f in os.listdir(in_folder) if f.endswith("stl")] + mesh_filenames += [f for f in os.listdir(in_folder) if f.endswith("STL")] + # print status + print(f"Found {len(mesh_filenames)} files to process in directory: {in_folder}") + # iterate over each OBJ file + for mesh_file in mesh_filenames: + # extract meshname + mesh_name = os.path.splitext(mesh_file)[0] + # complete path of input and output files + in_file_path = os.path.join(in_folder, mesh_file) + out_file_path = os.path.join(out_folder, mesh_name + ".obj") + # perform blender processing + print("Processing: ", in_file_path) + run_blender_convert2obj(in_file_path, out_file_path) + + +if __name__ == "__main__": + # Parse command line arguments + args = parse_cli_args() + # Run conversion + convert_meshes([args.input_dir], [args.output_dir]) diff --git a/scripts/tools/record_demos.py b/scripts/tools/record_demos.py new file mode 100644 index 00000000..63cd65d7 --- /dev/null +++ b/scripts/tools/record_demos.py @@ -0,0 +1,256 @@ +# Copyright (c) 2024-2025, 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) + --num_demos Number of demonstrations to record. (default: 0) + --num_success_steps Number of continuous steps with task success for concluding a demo as successful. (default: 10) +""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse +import os + +from isaaclab.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.") +parser.add_argument( + "--num_demos", type=int, default=0, help="Number of demonstrations to record. Set to 0 for infinite." +) +parser.add_argument( + "--num_success_steps", + type=int, + default=10, + help="Number of continuous steps with task success for concluding a demo as successful. Default is 10.", +) +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli = parser.parse_args() + +if args_cli.teleop_device.lower() == "handtracking": + vars(args_cli)["experience"] = f'{os.environ["ISAACLAB_PATH"]}/apps/isaaclab.python.xr.openxr.kit' + +# launch the simulator +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import contextlib +import gymnasium as gym +import time +import torch + +import omni.log + +from isaaclab.devices import Se3HandTracking, Se3Keyboard, Se3SpaceMouse +from isaaclab.envs import ViewerCfg +from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg +from isaaclab.envs.ui import ViewportCameraController + +import isaaclab_tasks # noqa: F401 +import uwlab_tasks # noqa: F401 +from isaaclab_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.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.""" + + # if handtracking is selected, rate limiting is achieved via OpenXR + if args_cli.teleop_device.lower() == "handtracking": + rate_limiter = None + else: + 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 + + # extract success checking function to invoke in the main loop + success_term = None + if hasattr(env_cfg.terminations, "success"): + success_term = env_cfg.terminations.success + env_cfg.terminations.success = None + else: + omni.log.warn( + "No success termination term was found in the environment." + " Will not be able to mark recorded demos as successful." + ) + + # modify configuration such that the environment runs indefinitely until + # the goal is reached or other termination conditions are met + 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).unwrapped + + # 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 + + # 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) + elif args_cli.teleop_device.lower() == "handtracking": + from isaacsim.xr.openxr import OpenXRSpec + + teleop_interface = Se3HandTracking(OpenXRSpec.XrHandEXT.XR_HAND_RIGHT_EXT, False, True) + teleop_interface.add_callback("RESET", reset_recording_instance) + viewer = ViewerCfg(eye=(-0.25, -0.3, 0.5), lookat=(0.6, 0, 0), asset_name="viewer") + ViewportCameraController(env, viewer) + else: + raise ValueError( + f"Invalid device interface '{args_cli.teleop_device}'. Supported: 'keyboard', 'spacemouse', 'handtracking'." + ) + + 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 + current_recorded_demo_count = 0 + success_step_count = 0 + 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 success_term is not None: + if bool(success_term.func(env, **success_term.params)[0]): + success_step_count += 1 + if success_step_count >= args_cli.num_success_steps: + env.recorder_manager.record_pre_reset([0], force_export_or_skip=False) + env.recorder_manager.set_success_to_episodes( + [0], torch.tensor([[True]], dtype=torch.bool, device=env.device) + ) + env.recorder_manager.export_episodes([0]) + should_reset_recording_instance = True + else: + success_step_count = 0 + + if should_reset_recording_instance: + env.recorder_manager.reset() + env.reset() + should_reset_recording_instance = False + success_step_count = 0 + + # print out the current demo count if it has changed + if env.recorder_manager.exported_successful_episode_count > current_recorded_demo_count: + current_recorded_demo_count = env.recorder_manager.exported_successful_episode_count + print(f"Recorded {current_recorded_demo_count} successful demonstrations.") + + if args_cli.num_demos > 0 and env.recorder_manager.exported_successful_episode_count >= args_cli.num_demos: + print(f"All {args_cli.num_demos} demonstrations recorded. Exiting the app.") + break + + # check that simulation is stopped or not + if env.sim.is_stopped(): + break + + if rate_limiter: + rate_limiter.sleep(env) + + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/tools/replay_demos.py b/scripts/tools/replay_demos.py new file mode 100644 index 00000000..74eeaf67 --- /dev/null +++ b/scripts/tools/replay_demos.py @@ -0,0 +1,227 @@ +# Copyright (c) 2024-2025, 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 isaaclab.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 isaaclab.devices import Se3Keyboard +from isaaclab.utils.datasets import EpisodeData, HDF5DatasetFileHandler + +import isaaclab_tasks # noqa: F401 +from isaaclab_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).unwrapped + + 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('Press "B" to pause and "N" to resume the replayed actions.') + + # 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.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.device + ) + env_episode_data_map[env_id] = episode_data + # Set initial state for the new episode + initial_state = episode_data.get_initial_state() + env.reset_to(initial_state, torch.tensor([env_id], device=env.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.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.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/uwlab/config/extension.toml b/source/uwlab/config/extension.toml new file mode 100644 index 00000000..ee2aacd7 --- /dev/null +++ b/source/uwlab/config/extension.toml @@ -0,0 +1,44 @@ +[package] + +# Semantic Versioning is used: https://semver.org/ +version = "0.8.0" + +# Description +title = "UW Lab framework for Robot Learning" +description="Extension providing main framework interfaces and abstractions for robot learning." +readme = "docs/README.md" +repository = "https://github.com/UW-Lab/UWLab" +category = "robotics, metaworld" +keywords = ["kit", "robotics", "learning", "ai"] + +[dependencies] +"omni.isaac.core" = {} +"omni.isaac.ml_archive" = {} +"omni.replicator.core" = {} +"isaaclab" = {} +"isaaclab_assets" = {} +"isaaclab_tasks" = {} + +[python.pipapi] +requirements = [ + "numpy", + "prettytable==3.3.0", + "toml", + "hidapi", + "gymnasium==0.29.0", + "trimesh" +] + +modules = [ + "numpy", + "prettytable", + "toml", + "hid", + "gymnasium", + "trimesh" +] + +use_online_index=true + +[[python.module]] +name = "uwlab" diff --git a/source/uwlab/docs/CHANGELOG.rst b/source/uwlab/docs/CHANGELOG.rst new file mode 100644 index 00000000..cbbdfbf0 --- /dev/null +++ b/source/uwlab/docs/CHANGELOG.rst @@ -0,0 +1,408 @@ +Changelog +--------- + +0.8.0 (2024-11-14) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* renamed and standardized genome, gene in :file:`uwlab.genes.genome` + + +0.7.2 (2024-11-10) +~~~~~~~~~~~~~~~~~~ + +Removed +^^^^^^^ + +* Deprecating :folder:`uwlab.envs.assets.deformable` related deformable modules + + +0.7.1 (2024-10-24) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* fixed :class:`uwlab.device.RokokoGloveT265` to support up to date teleoperation pipeline +* fixed :class:`uwlab.device.RokokoGloveKeyboard` to support up to date teleoperation pipeline + +0.7.0 (2024-10-20) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* added :file:`uwlab.scene.large_scene_cfg` to support importing prebuilt scene in usd that has +* configured articulation, rigidbody, deformable, and so on. + + +0.6.1 (2024-10-20) +~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Dropping DeformableInteractiveScene from :file:`uwlab.scene.deformable_interactive_scene_cfg` as +official deformable has been added to Isaac Lab + + +0.6.0 (2024-10-20) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* support encodable terrain in :file:`uwlab.lab.terrains.descriptive_terrain` +* with current encoding to be behavior and terrain type with an example located at +* :file:`uwlab.lab.terrains.config.descriptive_terrain` + +0.5.5 (2024-09-09) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* adding property is_closed to :class:`uwlab.envs.UWManagerBasedRl` + +0.5.4 (2024-09-06) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Transferred Experimental Evolution code into lab extension as :dir:`uwlab.evolution_system` + +0.5.3 (2024-09-02) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Adding event that reset from demonstration :func:`uwlab.envs.mdp.events.reset_from_demonstration` +* Adding event that record state of simulation:func:`uwlab.envs.mdp.events.record_state_configuration` + +0.5.2 (2024-09-01) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Adding event that make viewport camera follows robot at :func:`uwlab.envs.mdp.events.viewport_follow_robot` + + +0.5.1 (2024-08-23) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Bug fix for :func:`uwlab.envs.UWManagerBasedRl.step` where data manager existence is not queried correctly + + +0.5.0 (2024-08-06) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added features that support obj typed sub-terrain, and custom supply of the spawning locations + please check :folder:`uwlab.lab.terrains` + + +0.4.3 (2024-08-06) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Removed :file:`uwlab.terrains.enhanced_terrain_importer.py` as it is ended up not being a solution + + +0.4.2 (2024-08-06) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Removed :func:`uwlab.envs.mdp.events.reset_root_state_uniform` instead, reset_root_state_uniform is imported + from isaac lab + + +0.4.1 (2024-08-06) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Bug fix for :func:`uwlab.envs.UWManagerBasedRl.close` is self.extensions not self.extension + + +0.4.0 (2024-07-29) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* updated dependency and meta information to isaac sim 4.1.0 + + +0.3.0 (2024-07-28) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^^^ +Added experiment feature categorical command type for commanding anything that can be represented +by integer at :folder:`uwlab.envs.mdp.commands` + + +0.2.7 (2024-07-28) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ +* MultiConstraintDifferentialIKController now supports multi environments parallel computes +* added reward :func:`uwlab.envs.mdp.position_command_error` +* added reward :func:`uwlab.envs.mdp.position_command_error_tanh` +* added reward :func:`uwlab.envs.mdp.orientation_command_error_tanh` +* removed :func:`uwlab.envs.mdp.track_interpolated_lin_vel_xy_exp` as this is fetching task specific +* removed :func:`uwlab.envs.mdp.track_interpolated_ang_vel_z_exp` as this is fetching task specific + + +0.2.6 (2024-07-27) +~~~~~~~~~~~~~~~~~~ +Added +^^^^^ +* Added reward term :func:`uwlab.envs.mdp.reward_body1_body2_within_distance` for reward proximity + two objects proximity + +Changed +^^^^^^^ +* Updating default rough terrain tiling configuration at :class:`uwlab.terrains.config` + + +0.2.5 (2024-07-27) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ +* Removed dependency on ``import os`` to support custom extension in :class:`uwlab.actuators.EffortMotor` + + +0.2.4 (2024-07-26) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ +* Changed :class:`uwlab.actuators.EffortMotor` inherites and uses super classes stiffness, + damping, effort limit instead of redefining a redundant field as of :class:`uwlab.actuators.HebiEffortMotor` + +* Changed : :class:`uwlab.actuators.EffortMotorCfg` added to support above change + +* Changed : :class:`uwlab.actuators.__init__` added to support above change + + +0.2.3 (2024-07-20) +~~~~~~~~~~~~~~~~~~ + + +Added +^^^^^ +* Added debug :func:`uwlab.devices.RokokoGloveKeyboard.debug_advance_all_joint_data.` + for glove data visualization + +Changed +^^^^^^^ +* Changed :class:`uwlab.devices.RokokoGloveKeyboard.` class requires + input initial command pose to correctly set robot reset command target + +* Edited Thumb scaling input in :class:`uwlab.devices.RokokoGlove` that corrects + thumb length mismatch in teleoperation + + +0.2.2 (2024-07-15) +~~~~~~~~~~~~~~~~~~ + + +Changed +^^^^^^^ +* Changed :func:`uwlab.sim.spawners.from_files.from_files_cfg.MultiAssetCfg` to support + multi objects scaling. +* Changed :func:`uwlab.sim.spawners.from_files.from_files.spawn_multi_object_randomly_sdf` + to support multi objects scaling. + + +0.2.1 (2024-07-14) +~~~~~~~~~~~~~~~~~~ + + +Added +^^^^^ +* UW lab now support multi assets spawning +* Added :func:`uwlab.sim.spawners.from_files.from_files.spawn_multi_object_randomly_sdf` + and :func:`uwlab.sim.spawners.from_files.from_files.spawn_multi_object_randomly` +* Added :func:`uwlab.sim.spawners.from_files.from_files_cfg.MultiAssetCfg` + + +0.2.0 (2024-07-10) +~~~~~~~~~~~~~~~~~~ + + +Changed +^^^^^^^ + +* Added Reward Term :func:`uwlab.envs.mdp.rewards.reward_body1_frame2_distance` +* Let Keyboard device accepts initial transform pose input :class:`uwlab.devices.Se3Keyboard` + + +0.1.9 (2024-07-10) +~~~~~~~~~~~~~~~~~~ + + +Changed +^^^^^^^ + +* Documented :class:`uwlab.controllers.MultiConstraintDifferentialIKController`, + :class:`uwlab.controllers.MultiConstraintDifferentialIKControllerCfg` + + +0.1.8 (2024-07-09) +~~~~~~~~~~~~~~~~~~ + + +Changed +^^^^^^^ + +* Documented :class:`uwlab.devices.RokokoGlove`, + :class:`uwlab.devices.RokokoGloveKeyboard`, :class:`uwlab.devices.Se3Keyboard` + + + +0.1.7 (2024-07-08) +~~~~~~~~~~~~~~~~~~ + + +Changed +^^^^^^^ + +* Added proximal distance scaling in :class:`uwlab.devices.rokoko_glove.RokokoGlove` +* Fixed the order checking for the :class:`uwlab.controllers.differential_ik.MultiConstraintDifferentialIKController` + + +Added +^^^^^ +* Added combined control that separates pose and finger joints in + :class:`uwlab.devices.rokoko_glove_keyboard.RokokoGloveKeyboard` + + +0.1.6 (2024-07-06) +~~~~~~~~~~~~~~~~~~ + + +Changed +^^^^^^^ + +* :class:`uwlab.actuators.actuator_cfg.HebiStrategy3ActuatorCfg` added the field that scales position_p and effort_p +* :class:`uwlab.actuators.actuator_cfg.HebiStrategy4ActuatorCfg` added the field that scales position_p and effort_p +* :class:`uwlab.actuators.actuator_pd.py.HebiStrategy3Actuator` reflected the field that scales position_p and effort_p +* :class:`uwlab.actuators.actuator_pd.py.HebiStrategy4Actuator` reflected the field that scales position_p and effort_p +* Improved Reuseability :class:`uwlab.devices.rokoko_udp_receiver.Rokoko_Glove` such that the returned joint position respects the +order user inputs. Added debug visualization. Plan to add scale by knuckle width to match the leap hand knuckle width + +0.1.5 (2024-07-04) +~~~~~~~~~~~~~~~~~~ + + +Changed +^^^^^^^ +* :meth:`uwlab.envs.data_manager_based_rl.step` the actual environment update rate now becomes +decimation square, as square allows a nice property that tuning decimation creates minimal effect on the learning +behavior. + + +0.1.4 (2024-06-29) +~~~~~~~~~~~~~~~~~~ + + +Changed +^^^^^^^ +* allow user input specific tracking name :meth:`uwlab.device.rokoko_udp_receiver.Rokoko_Glove.__init__` to address + inefficiency when left or right has tracking is unnecessary, and future need in increasing, decreasing number of track + parts with ease. In addition, the order which parts are outputted is now ordered by user's list input, removing the need + of manually reorder the output when the output is fixed + +0.1.3 (2024-06-28) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added :class:`uwlab.envs.mdp.actions.MultiConstraintsDifferentialInverseKinematicsActionCfg` + + +Changed +^^^^^^^ +* cleaned, memory preallocated :class:`uwlab.device.rokoko_udp_receiver.Rokoko_Glove` so it is much more readable and efficient + + +0.1.2 (2024-06-27) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added :class:`uwlab.envs.mdp.actions.MultiConstraintsDifferentialInverseKinematicsActionCfg` + + +Changed +^^^^^^^ +* Removed duplicate functions in :class:`uwlab.envs.mdp.actions.actions_cfg` already defined in Isaac lab +* Removed :file:`uwlab.envs.mdp.actions.binary_joint_actions.py` as it completely duplicates Isaac lab implementation +* Removed :file:`uwlab.envs.mdp.actions.joint_actions.py` as it completely duplicates Isaac lab implementation +* Removed :file:`uwlab.envs.mdp.actions.non_holonomic_actions.py` as it completely duplicates Isaac lab implementation +* Cleaned :class:`uwlab.controllers.differential_ik.DifferentialIKController` + +0.1.1 (2024-06-26) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Rokoko smart glove device reading +* separation of :class:`uwlab.envs.mdp.actions.MultiConstraintDifferentialInverseKinematicsAction` + from :class:`isaaclab.envs.mdp.actions.DifferentialInverseKinematicsAction` + +* separation of :class:`uwlab.envs.mdp.actions.MultiConstraintDifferentialIKController` + from :class:`isaaclab.envs.mdp.actions.DifferentialIKController` + +* separation of :class:`uwlab.envs.mdp.actions.MultiConstraintDifferentialIKControllerCfg` + from :class:`isaaclab.envs.mdp.actions.DifferentialIKControllerCfg` + + +Changed +^^^^^^^ +* Changed :func:`uwlab.envs.mdp.events.reset_tycho_to_default` to :func:`uwlab.envs.mdp.events.reset_robot_to_default` +* Changed :func:`uwlab.envs.mdp.events.update_joint_positions` to :func:`uwlab.envs.mdp.events.update_joint_target_positions_to_current` +* Removed unnecessary import in :class:`uwlab.envs.mdp.events` +* Removed unnecessary import in :class:`uwlab.envs.mdp.rewards` +* Removed unnecessary import in :class:`uwlab.envs.mdp.terminations` + + +Updated +^^^^^^^ + +* Updated :meth:`uwlab.envs.DeformableBasedEnv.__init__` up to date with :meth:`isaaclab.envs.ManagerBasedEnv.__init__` +* Updated :class:`uwlab.envs.HebiRlEnvCfg` to :class:`uwlab.envs.UWManagerBasedRlCfg` +* Updated :class:`uwlab.envs.HebiRlEnv` to :class:`uwlab.envs.UWManagerBasedRl` + + +0.1.0 (2024-06-11) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Performed uwlab refactorization. Tested to work alone, and also with tycho +* Updated README Instruction +* Plan to do: check out not duplicate logic, clean up this repository. diff --git a/source/uwlab/pyproject.toml b/source/uwlab/pyproject.toml new file mode 100644 index 00000000..d90ac353 --- /dev/null +++ b/source/uwlab/pyproject.toml @@ -0,0 +1,3 @@ +[build-system] +requires = ["setuptools", "wheel", "toml"] +build-backend = "setuptools.build_meta" diff --git a/source/uwlab/setup.py b/source/uwlab/setup.py new file mode 100644 index 00000000..d7cea5f4 --- /dev/null +++ b/source/uwlab/setup.py @@ -0,0 +1,41 @@ +"""Installation script for the 'uwlab' python package.""" + +import os +import toml + +from setuptools import setup + +# Obtain the extension data from the extension.toml file +EXTENSION_PATH = os.path.dirname(os.path.realpath(__file__)) +# Read the extension.toml file +EXTENSION_TOML_DATA = toml.load(os.path.join(EXTENSION_PATH, "config", "extension.toml")) + +# Minimum dependencies required prior to installation +INSTALL_REQUIRES = [ + "pybullet", +] + +PYTORCH_INDEX_URL = ["https://download.pytorch.org/whl/cu118"] + +# Installation operation +setup( + name="uwlab", + author="UW and Isaac Lab Project Developers", + maintainer="UW and Isaac Lab Project Developers", + url=EXTENSION_TOML_DATA["package"]["repository"], + version=EXTENSION_TOML_DATA["package"]["version"], + description=EXTENSION_TOML_DATA["package"]["description"], + keywords=EXTENSION_TOML_DATA["package"]["keywords"], + license="BSD-3-Clause", + include_package_data=True, + python_requires=">=3.10", + install_requires=INSTALL_REQUIRES, + dependency_links=PYTORCH_INDEX_URL, + packages=["uwlab"], + classifiers=[ + "Natural Language :: English", + "Programming Language :: Python :: 3.10", + "Isaac Sim :: 4.5.0", + ], + zip_safe=False, +) diff --git a/source/uwlab/uwlab/actuators/__init__.py b/source/uwlab/uwlab/actuators/__init__.py new file mode 100644 index 00000000..bc9b0785 --- /dev/null +++ b/source/uwlab/uwlab/actuators/__init__.py @@ -0,0 +1,32 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-package for different actuator models. + +Actuator models are used to model the behavior of the actuators in an articulation. These +are usually meant to be used in simulation to model different actuator dynamics and delays. + +There are two main categories of actuator models that are supported: + +- **Implicit**: Motor model with ideal PD from the physics engine. This is similar to having a continuous time + PD controller. The motor model is implicit in the sense that the motor model is not explicitly defined by the user. +- **Explicit**: Motor models based on physical drive models. + + - **Physics-based**: Derives the motor models based on first-principles. + - **Neural Network-based**: Learned motor models from actuator data. + +Every actuator model inherits from the :class:`isaaclab.actuators.ActuatorBase` class, +which defines the common interface for all actuator models. The actuator models are handled +and called by the :class:`isaaclab.assets.Articulation` class. +""" + +from .actuator_cfg import ( + EffortMotorCfg, + HebiDCMotorCfg, + HebiEffortMotorCfg, + HebiStrategy3ActuatorCfg, + HebiStrategy4ActuatorCfg, +) +from .actuator_pd import EffortMotor, HebiDCMotor, HebiEffortMotor, HebiStrategy3Actuator, HebiStrategy4Actuator diff --git a/source/uwlab/uwlab/actuators/actuator_cfg.py b/source/uwlab/uwlab/actuators/actuator_cfg.py new file mode 100644 index 00000000..498a4d39 --- /dev/null +++ b/source/uwlab/uwlab/actuators/actuator_cfg.py @@ -0,0 +1,16 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from isaaclab.actuators.actuator_cfg import ActuatorBaseCfg +from isaaclab.utils import configclass + +from . import actuator_pd + + +@configclass +class EffortMotorCfg(ActuatorBaseCfg): + class_type: type = actuator_pd.EffortMotor diff --git a/source/uwlab/uwlab/actuators/actuator_pd.py b/source/uwlab/uwlab/actuators/actuator_pd.py new file mode 100644 index 00000000..50fdebd7 --- /dev/null +++ b/source/uwlab/uwlab/actuators/actuator_pd.py @@ -0,0 +1,34 @@ +# Copyright (c) 2024-2025, The UW 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 + +from isaaclab.actuators.actuator_base import ActuatorBase +from isaaclab.utils.types import ArticulationActions + +if TYPE_CHECKING: + from .actuator_cfg import EffortMotorCfg + + +class EffortMotor(ActuatorBase): + cfg: EffortMotorCfg + + def __init__(self, cfg: EffortMotorCfg, *args, **kwargs): + super().__init__(cfg, *args, **kwargs) + + def compute(self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor): + self.computed_effort = torch.clip(control_action.joint_efforts, -self.effort_limit, self.effort_limit) + self.applied_effort = self.computed_effort + control_action.joint_efforts = self.applied_effort + control_action.joint_positions = None + control_action.joint_velocities = None + return control_action + + def reset(self, env_ids: Sequence[int]): + pass diff --git a/source/uwlab/uwlab/assets/__init__.py b/source/uwlab/uwlab/assets/__init__.py new file mode 100644 index 00000000..bed6d2dc --- /dev/null +++ b/source/uwlab/uwlab/assets/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .articulation import ArticulationCfg, ArticulationData, UniversalArticulation +from .asset_base import AssetBase +from .asset_base_cfg import AssetBaseCfg diff --git a/source/uwlab/uwlab/assets/articulation/__init__.py b/source/uwlab/uwlab/assets/articulation/__init__.py new file mode 100644 index 00000000..01281185 --- /dev/null +++ b/source/uwlab/uwlab/assets/articulation/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .articulation import UniversalArticulation +from .articulation_cfg import ArticulationCfg +from .articulation_data import ArticulationData +from .articulation_view import * diff --git a/source/uwlab/uwlab/assets/articulation/articulation.py b/source/uwlab/uwlab/assets/articulation/articulation.py new file mode 100644 index 00000000..d338fc61 --- /dev/null +++ b/source/uwlab/uwlab/assets/articulation/articulation.py @@ -0,0 +1,1535 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Flag for pyright to ignore type errors in this file. +# pyright: reportPrivateUsage=false + +from __future__ import annotations + +import torch +from collections.abc import Sequence +from prettytable import PrettyTable +from typing import TYPE_CHECKING + +import isaacsim.core.utils.stage as stage_utils +import omni.log +from pxr import PhysxSchema, UsdPhysics + +import isaaclab.utils.math as math_utils +import isaaclab.utils.string as string_utils +from isaaclab.actuators import ActuatorBase, ActuatorBaseCfg, ImplicitActuator +from isaaclab.utils.types import ArticulationActions + +from ..asset_base import AssetBase +from .articulation_data import ArticulationData +from .articulation_view import ArticulationView + +if TYPE_CHECKING: + from .articulation_cfg import ArticulationCfg + + +class UniversalArticulation(AssetBase): + """An articulation asset class. + + An articulation is a collection of rigid bodies connected by joints. The joints can be either + fixed or actuated. The joints can be of different types, such as revolute, prismatic, D-6, etc. + However, the articulation class has currently been tested with revolute and prismatic joints. + The class supports both floating-base and fixed-base articulations. The type of articulation + is determined based on the root joint of the articulation. If the root joint is fixed, then + the articulation is considered a fixed-base system. Otherwise, it is considered a floating-base + system. This can be checked using the :attr:`Articulation.is_fixed_base` attribute. + + For an asset to be considered an articulation, the root prim of the asset must have the + `USD ArticulationRootAPI`_. This API is used to define the sub-tree of the articulation using + the reduced coordinate formulation. On playing the simulation, the physics engine parses the + articulation root prim and creates the corresponding articulation in the physics engine. The + articulation root prim can be specified using the :attr:`AssetBaseCfg.prim_path` attribute. + + The articulation class also provides the functionality to augment the simulation of an articulated + system with custom actuator models. These models can either be explicit or implicit, as detailed in + the :mod:`isaaclab.actuators` module. The actuator models are specified using the + :attr:`ArticulationCfg.actuators` attribute. These are then parsed and used to initialize the + corresponding actuator models, when the simulation is played. + + During the simulation step, the articulation class first applies the actuator models to compute + the joint commands based on the user-specified targets. These joint commands are then applied + into the simulation. The joint commands can be either position, velocity, or effort commands. + As an example, the following snippet shows how this can be used for position commands: + + .. code-block:: python + + # an example instance of the articulation class + my_articulation = Articulation(cfg) + + # set joint position targets + my_articulation.set_joint_position_target(position) + # propagate the actuator models and apply the computed commands into the simulation + my_articulation.write_data_to_sim() + + # step the simulation using the simulation context + sim_context.step() + + # update the articulation state, where dt is the simulation time step + my_articulation.update(dt) + + .. _`USD ArticulationRootAPI`: https://openusd.org/dev/api/class_usd_physics_articulation_root_a_p_i.html + + """ + + cfg: ArticulationCfg + """Configuration instance for the articulations.""" + + actuators: dict[str, ActuatorBase] + """Dictionary of actuator instances for the articulation. + + The keys are the actuator names and the values are the actuator instances. The actuator instances + are initialized based on the actuator configurations specified in the :attr:`ArticulationCfg.actuators` + attribute. They are used to compute the joint commands during the :meth:`write_data_to_view` function. + """ + + def __init__(self, cfg: ArticulationCfg): + """Initialize the articulation. + + 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 + + self._root_state_dep_warn = False + self._root_pose_dep_warn = False + self._root_vel_dep_warn = False + + """ + Properties + """ + + @property + def data(self) -> ArticulationData: + return self._data + + @property + def num_instances(self) -> int: + return self.view.count + + @property + def is_fixed_base(self) -> bool: + """Whether the articulation is a fixed-base or floating-base system.""" + return self.view.fixed_base + + @property + def num_joints(self) -> int: + """Number of joints in articulation.""" + return self.view.dof_count + + @property + def num_fixed_tendons(self) -> int: + """Number of fixed tendons in articulation.""" + return self.view.max_fixed_tendons + + @property + def num_bodies(self) -> int: + """Number of bodies in articulation.""" + return self.view.dof_count + + @property + def joint_names(self) -> list[str]: + """Ordered names of joints in articulation.""" + return self.view.joint_names + + @property + def fixed_tendon_names(self) -> list[str]: + """Ordered names of fixed tendons in articulation.""" + return self._fixed_tendon_names + + @property + def body_names(self) -> list[str]: + """Ordered names of bodies in articulation.""" + return self.view.body_names + + @property + def view(self) -> ArticulationView: + """Articulation view for the asset (PhysX). + + Note: + Use this view with caution. It requires handling of tensors in a specific way. + """ + return self._view + + """ + Operations. + """ + + def reset(self, env_ids: Sequence[int] | None = None): + # use ellipses object to skip initial indices. + if env_ids is None: + env_ids = slice(None) + # reset actuators + for actuator in self.actuators.values(): + actuator.reset(env_ids) + # reset external wrench + self._external_force_b[env_ids] = 0.0 + self._external_torque_b[env_ids] = 0.0 + + def write_data_to_sim(self): + """Write external wrenches and joint commands to the simulation. + + If any explicit actuators are present, then the actuator models are used to compute the + joint commands. Otherwise, the joint commands are directly set into 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.view.apply_forces_and_torques_at_position( + force_data=self._external_force_b.view(-1, 3), + torque_data=self._external_torque_b.view(-1, 3), + position_data=None, + indices=self._ALL_INDICES, + is_global=False, + ) + + # apply actuator models + self._apply_actuator_model() + # write actions into simulation + self.view.set_dof_actuation_forces(self._joint_effort_target_sim, self._ALL_INDICES) + # position and velocity targets only for implicit actuators + if self._has_implicit_actuators: + self.view.set_dof_position_targets(self._joint_pos_target_sim, self._ALL_INDICES) + self.view.set_dof_velocity_targets(self._joint_vel_target_sim, self._ALL_INDICES) + + def update(self, dt: float): + self._data.update(dt) + # self._view.update(dt) + + """ + Operations - Finders. + """ + + def find_bodies(self, name_keys: str | Sequence[str], preserve_order: bool = False) -> tuple[list[int], list[str]]: + """Find bodies in the articulation based on the name keys. + + Please check the :meth:`isaaclab.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 body names. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the body indices and names. + """ + return string_utils.resolve_matching_names(name_keys, self.body_names, preserve_order) + + def find_joints( + self, name_keys: str | Sequence[str], joint_subset: list[str] | None = None, preserve_order: bool = False + ) -> tuple[list[int], list[str]]: + """Find joints in the articulation based on the name keys. + + Please see the :func:`isaaclab.utils.string.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 joint names. + joint_subset: A subset of joints to search for. Defaults to None, which means all joints + in the articulation are searched. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the joint indices and names. + """ + if joint_subset is None: + joint_subset = self.joint_names + # find joints + return string_utils.resolve_matching_names(name_keys, joint_subset, preserve_order) + + def find_fixed_tendons( + self, name_keys: str | Sequence[str], tendon_subsets: list[str] | None = None, preserve_order: bool = False + ) -> tuple[list[int], list[str]]: + """Find fixed tendons in the articulation based on the name keys. + + Please see the :func:`isaaclab.utils.string.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 joint + names with fixed tendons. + tendon_subsets: A subset of joints with fixed tendons to search for. Defaults to None, which means + all joints in the articulation are searched. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the tendon indices and names. + """ + if tendon_subsets is None: + # tendons follow the joint names they are attached to + tendon_subsets = self.fixed_tendon_names + # find tendons + return string_utils.resolve_matching_names(name_keys, tendon_subsets, preserve_order) + + """ + Operations - Writers. + """ + + def write_root_state_to_sim(self, root_state: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root 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. + """ + + # deprecation warning + if not self._root_state_dep_warn: + omni.log.warn( + "DeprecationWarning: Articluation.write_root_state_to_view will be removed in a future release. Please" + " use write_root_link_state_to_view or write_root_com_state_to_view 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_view will be removed in a future release. Please" + " use write_root_link_pose_to_view or write_root_com_pose_to_view 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. + """ + # resolve all indices + physx_env_ids = env_ids + if env_ids is None: + env_ids = slice(None) + 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_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_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.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 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_view will be removed in a future release." + " Please use write_root_link_velocity_to_view or write_root_com_velocity_to_view 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: + env_ids = slice(None) + 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_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.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, + position: torch.Tensor, + velocity: torch.Tensor, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | slice | None = None, + ): + """Write joint positions and velocities to the simulation. + + Args: + position: Joint positions. Shape is (len(env_ids), len(joint_ids)). + velocity: Joint velocities. Shape is (len(env_ids), len(joint_ids)). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets 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] + # set into internal buffers + self._data.joint_pos[env_ids, joint_ids] = position + self._data.joint_vel[env_ids, joint_ids] = velocity + self._data._previous_joint_vel[env_ids, joint_ids] = velocity + 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.view.set_dof_positions(self._data.joint_pos, indices=physx_env_ids) + self.view.set_dof_velocities(self._data.joint_vel, indices=physx_env_ids) + + def write_joint_stiffness_to_sim( + self, + stiffness: torch.Tensor | float, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ): + """Write joint stiffness into the simulation. + + Args: + stiffness: Joint stiffness. Shape is (len(env_ids), len(joint_ids)). + joint_ids: The joint indices to set the stiffness for. Defaults to None (all joints). + env_ids: The environment indices to set the stiffness for. Defaults to None (all environments). + """ + # note: This function isn't setting the values for actuator models. (#128) + # 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] + # set into internal buffers + self._data.joint_stiffness[env_ids, joint_ids] = stiffness + # set into simulation + self.view.set_dof_stiffnesses(self._data.joint_stiffness.cpu(), indices=physx_env_ids.cpu()) + + def write_joint_damping_to_sim( + self, + damping: torch.Tensor | float, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ): + """Write joint damping into the simulation. + + Args: + damping: Joint damping. Shape is (len(env_ids), len(joint_ids)). + joint_ids: The joint indices to set the damping for. + Defaults to None (all joints). + env_ids: The environment indices to set the damping for. + Defaults to None (all environments). + """ + # note: This function isn't setting the values for actuator models. (#128) + # 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] + # set into internal buffers + self._data.joint_damping[env_ids, joint_ids] = damping + # set into simulation + self.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_vel_limits = self.view.get_dof_max_velocities().to(self.device) + self._data.joint_vel_limits[env_ids, joint_ids] = limits + # set into simulation + self.view.set_dof_max_velocities(self._data.joint_vel_limits.cpu(), indices=physx_env_ids.cpu()) + + def write_joint_effort_limit_to_sim( + self, + limits: torch.Tensor | float, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ): + """Write joint effort limits into the simulation. + + Args: + limits: Joint torque limits. Shape is (len(env_ids), len(joint_ids)). + joint_ids: The joint indices to set the joint torque limits for. Defaults to None (all joints). + env_ids: The environment indices to set the joint torque limits for. Defaults to None (all environments). + """ + # note: This function isn't setting the values for actuator models. (#128) + # 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.cpu() + # set into internal buffers + torque_limit_all = self.view.get_dof_max_forces() + torque_limit_all[env_ids, joint_ids] = limits + # set into simulation + self.view.set_dof_max_forces(torque_limit_all.cpu(), indices=physx_env_ids.cpu()) + + def write_joint_armature_to_sim( + self, + armature: torch.Tensor | float, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ): + """Write joint armature into the simulation. + + Args: + armature: Joint armature. Shape is (len(env_ids), len(joint_ids)). + joint_ids: The joint indices to set the joint torque limits for. Defaults to None (all joints). + env_ids: The environment indices to set the joint torque limits 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] + # set into internal buffers + self._data.joint_armature[env_ids, joint_ids] = armature + # set into simulation + self.view.set_dof_armatures(self._data.joint_armature.cpu(), indices=physx_env_ids.cpu()) + + def write_joint_friction_to_sim( + self, + joint_friction: torch.Tensor | float, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ): + """Write joint friction into the simulation. + + Args: + joint_friction: Joint friction. Shape is (len(env_ids), len(joint_ids)). + joint_ids: The joint indices to set the joint torque limits for. Defaults to None (all joints). + env_ids: The environment indices to set the joint torque limits 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] + # set into internal buffers + self._data.joint_friction[env_ids, joint_ids] = joint_friction + # set into simulation + self.view.set_dof_friction_coefficients(self._data.joint_friction.cpu(), indices=physx_env_ids.cpu()) + + def write_joint_limits_to_sim( + self, + limits: torch.Tensor | float, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ): + """Write joint limits into the simulation. + + Args: + limits: Joint limits. Shape is (len(env_ids), len(joint_ids), 2). + joint_ids: The joint indices to set the limits for. Defaults to None (all joints). + env_ids: The environment indices to set the limits for. Defaults to None (all environments). + """ + # note: This function isn't setting the values for actuator models. (#128) + # 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] + # 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[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." + ) + # set into simulation + self.view.set_dof_limits(self._data.joint_limits.cpu(), indices=physx_env_ids.cpu()) + + """ + Operations - Setters. + """ + + def set_external_force_and_torque( + self, + forces: torch.Tensor, + torques: torch.Tensor, + body_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ): + """Set external force and torque to apply on the asset's 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, 3), torques=torch.zeros(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_view` function + right before the simulation step. + + Args: + forces: External forces in bodies' local frame. Shape is (len(env_ids), len(body_ids), 3). + torques: External torques in bodies' local frame. Shape is (len(env_ids), len(body_ids), 3). + body_ids: Body indices to apply external wrench to. Defaults to None (all bodies). + 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_INDICES + elif not isinstance(env_ids, torch.Tensor): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + # -- body_ids + if body_ids is None: + body_ids = torch.arange(self.num_bodies, dtype=torch.long, device=self.device) + elif isinstance(body_ids, slice): + body_ids = torch.arange(self.num_bodies, dtype=torch.long, device=self.device)[body_ids] + elif not isinstance(body_ids, torch.Tensor): + body_ids = torch.tensor(body_ids, dtype=torch.long, device=self.device) + + # note: we need to do this complicated indexing since torch doesn't support multi-indexing + # create global body indices from env_ids and env_body_ids + # (env_id * total_bodies_per_env) + body_id + indices = body_ids.repeat(len(env_ids), 1) + env_ids.unsqueeze(1) * self.num_bodies + indices = indices.view(-1) + # set into internal buffers + # note: these are applied in the write_to_view function + self._external_force_b.flatten(0, 1)[indices] = forces.flatten(0, 1) + self._external_torque_b.flatten(0, 1)[indices] = torques.flatten(0, 1) + else: + self.has_external_wrench = False + + def set_joint_position_target( + self, target: torch.Tensor, joint_ids: Sequence[int] | slice | None = None, env_ids: Sequence[int] | None = None + ): + """Set joint position targets into internal buffers. + + .. note:: + This function does not apply the joint targets to the simulation. It only fills the buffers with + the desired values. To apply the joint targets, call the :meth:`write_data_to_view` function. + + Args: + target: Joint position targets. Shape is (len(env_ids), len(joint_ids)). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + """ + # resolve indices + if env_ids is None: + env_ids = slice(None) + 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] + # set targets + self._data.joint_pos_target[env_ids, joint_ids] = target + + def set_joint_velocity_target( + self, target: torch.Tensor, joint_ids: Sequence[int] | slice | None = None, env_ids: Sequence[int] | None = None + ): + """Set joint velocity targets into internal buffers. + + .. note:: + This function does not apply the joint targets to the simulation. It only fills the buffers with + the desired values. To apply the joint targets, call the :meth:`write_data_to_view` function. + + Args: + target: Joint velocity targets. Shape is (len(env_ids), len(joint_ids)). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + """ + # resolve indices + if env_ids is None: + env_ids = slice(None) + 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] + # set targets + self._data.joint_vel_target[env_ids, joint_ids] = target + + def set_joint_effort_target( + self, target: torch.Tensor, joint_ids: Sequence[int] | slice | None = None, env_ids: Sequence[int] | None = None + ): + """Set joint efforts into internal buffers. + + .. note:: + This function does not apply the joint targets to the simulation. It only fills the buffers with + the desired values. To apply the joint targets, call the :meth:`write_data_to_view` function. + + Args: + target: Joint effort targets. Shape is (len(env_ids), len(joint_ids)). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + """ + # resolve indices + if env_ids is None: + env_ids = slice(None) + 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] + # set targets + self._data.joint_effort_target[env_ids, joint_ids] = target + + """ + Operations - Tendons. + """ + + def set_fixed_tendon_stiffness( + self, + stiffness: torch.Tensor, + fixed_tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ): + """Set fixed tendon stiffness into internal buffers. + + .. note:: + This function does not apply the tendon stiffness to the simulation. It only fills the buffers with + the desired values. To apply the tendon stiffness, call the :meth:`write_fixed_tendon_properties_to_view` function. + + Args: + stiffness: Fixed tendon stiffness. Shape is (len(env_ids), len(fixed_tendon_ids)). + fixed_tendon_ids: The tendon indices to set the stiffness for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the stiffness for. Defaults to None (all environments). + """ + # resolve indices + if env_ids is None: + env_ids = slice(None) + if fixed_tendon_ids is None: + fixed_tendon_ids = slice(None) + if env_ids != slice(None) and fixed_tendon_ids != slice(None): + env_ids = env_ids[:, None] + # set stiffness + self._data.fixed_tendon_stiffness[env_ids, fixed_tendon_ids] = stiffness + + def set_fixed_tendon_damping( + self, + damping: torch.Tensor, + fixed_tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ): + """Set fixed tendon damping into internal buffers. + + .. note:: + This function does not apply the tendon damping to the simulation. It only fills the buffers with + the desired values. To apply the tendon damping, call the :meth:`write_fixed_tendon_properties_to_view` function. + + Args: + damping: Fixed tendon damping. Shape is (len(env_ids), len(fixed_tendon_ids)). + fixed_tendon_ids: The tendon indices to set the damping for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the damping for. Defaults to None (all environments). + """ + # resolve indices + if env_ids is None: + env_ids = slice(None) + if fixed_tendon_ids is None: + fixed_tendon_ids = slice(None) + if env_ids != slice(None) and fixed_tendon_ids != slice(None): + env_ids = env_ids[:, None] + # set damping + self._data.fixed_tendon_damping[env_ids, fixed_tendon_ids] = damping + + def set_fixed_tendon_limit_stiffness( + self, + limit_stiffness: torch.Tensor, + fixed_tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ): + """Set fixed tendon limit stiffness efforts into internal buffers. + + .. note:: + This function does not apply the tendon limit stiffness to the simulation. It only fills the buffers with + the desired values. To apply the tendon limit stiffness, call the :meth:`write_fixed_tendon_properties_to_view` function. + + Args: + limit_stiffness: Fixed tendon limit stiffness. Shape is (len(env_ids), len(fixed_tendon_ids)). + fixed_tendon_ids: The tendon indices to set the limit stiffness for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the limit stiffness for. Defaults to None (all environments). + """ + # resolve indices + if env_ids is None: + env_ids = slice(None) + if fixed_tendon_ids is None: + fixed_tendon_ids = slice(None) + if env_ids != slice(None) and fixed_tendon_ids != slice(None): + env_ids = env_ids[:, None] + # set limit_stiffness + self._data.fixed_tendon_limit_stiffness[env_ids, fixed_tendon_ids] = limit_stiffness + + def set_fixed_tendon_limit( + self, + limit: torch.Tensor, + fixed_tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ): + """Set fixed tendon limit efforts into internal buffers. + + .. note:: + This function does not apply the tendon limit to the simulation. It only fills the buffers with + the desired values. To apply the tendon limit, call the :meth:`write_fixed_tendon_properties_to_view` function. + + Args: + limit: Fixed tendon limit. Shape is (len(env_ids), len(fixed_tendon_ids)). + fixed_tendon_ids: The tendon indices to set the limit for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the limit for. Defaults to None (all environments). + """ + # resolve indices + if env_ids is None: + env_ids = slice(None) + if fixed_tendon_ids is None: + fixed_tendon_ids = slice(None) + if env_ids != slice(None) and fixed_tendon_ids != slice(None): + env_ids = env_ids[:, None] + # set limit + self._data.fixed_tendon_limit[env_ids, fixed_tendon_ids] = limit + + def set_fixed_tendon_rest_length( + self, + rest_length: torch.Tensor, + fixed_tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ): + """Set fixed tendon rest length efforts into internal buffers. + + .. note:: + This function does not apply the tendon rest length to the simulation. It only fills the buffers with + the desired values. To apply the tendon rest length, call the :meth:`write_fixed_tendon_properties_to_view` function. + + Args: + rest_length: Fixed tendon rest length. Shape is (len(env_ids), len(fixed_tendon_ids)). + fixed_tendon_ids: The tendon indices to set the rest length for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the rest length for. Defaults to None (all environments). + """ + # resolve indices + if env_ids is None: + env_ids = slice(None) + if fixed_tendon_ids is None: + fixed_tendon_ids = slice(None) + if env_ids != slice(None) and fixed_tendon_ids != slice(None): + env_ids = env_ids[:, None] + # set rest_length + self._data.fixed_tendon_rest_length[env_ids, fixed_tendon_ids] = rest_length + + def set_fixed_tendon_offset( + self, + offset: torch.Tensor, + fixed_tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ): + """Set fixed tendon offset efforts into internal buffers. + + .. note:: + This function does not apply the tendon offset to the simulation. It only fills the buffers with + the desired values. To apply the tendon offset, call the :meth:`write_fixed_tendon_properties_to_view` function. + + Args: + offset: Fixed tendon offset. Shape is (len(env_ids), len(fixed_tendon_ids)). + fixed_tendon_ids: The tendon indices to set the offset for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the offset for. Defaults to None (all environments). + """ + # resolve indices + if env_ids is None: + env_ids = slice(None) + if fixed_tendon_ids is None: + fixed_tendon_ids = slice(None) + if env_ids != slice(None) and fixed_tendon_ids != slice(None): + env_ids = env_ids[:, None] + # set offset + self._data.fixed_tendon_offset[env_ids, fixed_tendon_ids] = offset + + def write_fixed_tendon_properties_to_sim( + self, + fixed_tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ): + """Write fixed tendon properties into the simulation. + + Args: + fixed_tendon_ids: The fixed tendon indices to set the limits for. Defaults to None (all fixed tendons). + env_ids: The environment indices to set the limits for. Defaults to None (all environments). + """ + # resolve indices + physx_env_ids = env_ids + if env_ids is None: + physx_env_ids = self._ALL_INDICES + if fixed_tendon_ids is None: + fixed_tendon_ids = slice(None) + + # set into simulation + self.view.set_fixed_tendon_properties( + self._data.fixed_tendon_stiffness, + self._data.fixed_tendon_damping, + self._data.fixed_tendon_limit_stiffness, + self._data.fixed_tendon_limit, + self._data.fixed_tendon_rest_length, + self._data.fixed_tendon_offset, + indices=physx_env_ids, + ) + + """ + Internal helper. + """ + + def _initialize_impl(self, device: str = None): + # create simulation view + self._device = device if device else "cpu" + self._view: ArticulationView = self.cfg.articulation_view_cfg.class_type( + cfg=self.cfg.articulation_view_cfg, device=self.device + ) + + # log information about the articulation + print(f"Is fixed root: {self.is_fixed_base}") + print(f"Number of bodies: {self.num_bodies}") + print(f"Body names: {self.body_names}") + print(f"Number of joints: {self.num_joints}") + print(f"Joint names: {self.joint_names}") + print(f"Number of fixed tendons: {self.num_fixed_tendons}") + + # container for data access + self._data = ArticulationData(self.view, self.device) + + # create buffers + self._create_buffers() + # process configuration + self._process_cfg() + self._process_actuators_cfg() + self._process_fixed_tendons() + # validate configuration + self._validate_cfg() + # update the robot data + self.update(0.0) + # log joint information + self._log_articulation_joint_info() + self.initialize_default_data() + + self._view.play() + + def initialize_default_data(self): + self._data.joint_pos_target = self._data.default_joint_pos.clone() + self._data.joint_vel_target = self._data.default_joint_vel.clone() + self._view.set_dof_position_targets(self._data.joint_pos_target, indices=self._ALL_INDICES) + self._view.set_dof_velocity_targets(self._data.joint_vel_target, indices=self._ALL_INDICES) + + def _create_buffers(self): + # constants + self._ALL_INDICES = torch.arange(self.num_instances, 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_bodies, 3), device=self.device) + self._external_torque_b = torch.zeros_like(self._external_force_b) + + # asset data + # -- properties + self._data.joint_names = self.joint_names + self._data.body_names = self.body_names + + # -- bodies + self._data.default_mass = self.view.get_masses().clone() + self._data.default_inertia = self.view.get_inertias().clone() + + # -- default joint state + self._data.default_joint_pos = torch.zeros(self.num_instances, self.num_joints, device=self.device) + self._data.default_joint_vel = torch.zeros_like(self._data.default_joint_pos) + + # -- joint commands + self._data.joint_pos_target = torch.zeros_like(self._data.default_joint_pos) + self._data.joint_vel_target = torch.zeros_like(self._data.default_joint_pos) + self._data.joint_effort_target = torch.zeros_like(self._data.default_joint_pos) + self._data.joint_stiffness = torch.zeros_like(self._data.default_joint_pos) + self._data.joint_damping = torch.zeros_like(self._data.default_joint_pos) + self._data.joint_armature = torch.zeros_like(self._data.default_joint_pos) + self._data.joint_friction = torch.zeros_like(self._data.default_joint_pos) + self._data.joint_limits = torch.zeros(self.num_instances, self.num_joints, 2, device=self.device) + + # -- joint commands (explicit) + self._data.computed_torque = torch.zeros_like(self._data.default_joint_pos) + self._data.applied_torque = torch.zeros_like(self._data.default_joint_pos) + + # -- tendons + if self.num_fixed_tendons > 0: + self._data.fixed_tendon_stiffness = torch.zeros( + self.num_instances, self.num_fixed_tendons, device=self.device + ) + self._data.fixed_tendon_damping = torch.zeros( + self.num_instances, self.num_fixed_tendons, device=self.device + ) + self._data.fixed_tendon_limit_stiffness = torch.zeros( + self.num_instances, self.num_fixed_tendons, device=self.device + ) + self._data.fixed_tendon_limit = torch.zeros( + self.num_instances, self.num_fixed_tendons, 2, device=self.device + ) + self._data.fixed_tendon_rest_length = torch.zeros( + self.num_instances, self.num_fixed_tendons, device=self.device + ) + self._data.fixed_tendon_offset = torch.zeros(self.num_instances, self.num_fixed_tendons, device=self.device) + + # -- other data + self._data.soft_joint_pos_limits = torch.zeros(self.num_instances, self.num_joints, 2, device=self.device) + self._data.soft_joint_vel_limits = torch.zeros(self.num_instances, self.num_joints, device=self.device) + self._data.gear_ratio = torch.ones(self.num_instances, self.num_joints, device=self.device) + + # -- initialize default buffers related to joint properties + self._data.default_joint_stiffness = torch.zeros(self.num_instances, self.num_joints, device=self.device) + self._data.default_joint_damping = torch.zeros(self.num_instances, self.num_joints, device=self.device) + self._data.default_joint_armature = torch.zeros(self.num_instances, self.num_joints, device=self.device) + self._data.default_joint_friction = torch.zeros(self.num_instances, self.num_joints, device=self.device) + self._data.default_joint_limits = torch.zeros(self.num_instances, self.num_joints, 2, device=self.device) + + # -- initialize default buffers related to fixed tendon properties + if self.num_fixed_tendons > 0: + self._data.default_fixed_tendon_stiffness = torch.zeros( + self.num_instances, self.num_fixed_tendons, device=self.device + ) + self._data.default_fixed_tendon_damping = torch.zeros( + self.num_instances, self.num_fixed_tendons, device=self.device + ) + self._data.default_fixed_tendon_limit_stiffness = torch.zeros( + self.num_instances, self.num_fixed_tendons, device=self.device + ) + self._data.default_fixed_tendon_limit = torch.zeros( + self.num_instances, self.num_fixed_tendons, 2, device=self.device + ) + self._data.default_fixed_tendon_rest_length = torch.zeros( + self.num_instances, self.num_fixed_tendons, device=self.device + ) + self._data.default_fixed_tendon_offset = torch.zeros( + self.num_instances, self.num_fixed_tendons, device=self.device + ) + + # soft joint position limits (recommended not to be too close to limits). + joint_pos_limits = self.view.get_dof_limits() + joint_pos_mean = (joint_pos_limits[..., 0] + joint_pos_limits[..., 1]) / 2 + joint_pos_range = joint_pos_limits[..., 1] - joint_pos_limits[..., 0] + soft_limit_factor = self.cfg.soft_joint_pos_limit_factor + # add to data + self._data.soft_joint_pos_limits[..., 0] = joint_pos_mean - 0.5 * joint_pos_range * soft_limit_factor + self._data.soft_joint_pos_limits[..., 1] = joint_pos_mean + 0.5 * joint_pos_range * soft_limit_factor + + # create buffers to store processed actions from actuator models + self._joint_pos_target_sim = torch.zeros_like(self._data.joint_pos_target) + self._joint_vel_target_sim = torch.zeros_like(self._data.joint_pos_target) + self._joint_effort_target_sim = torch.zeros_like(self._data.joint_pos_target) + + def _process_cfg(self): + """Post processing of configuration parameters.""" + # default state + # -- root state + # note: we cast to tuple to avoid torch/numpy type mismatch. + default_root_state = ( + tuple(self.cfg.init_state.pos) + + tuple(self.cfg.init_state.rot) + + tuple(self.cfg.init_state.lin_vel) + + tuple(self.cfg.init_state.ang_vel) + ) + default_root_state = torch.tensor(default_root_state, dtype=torch.float, device=self.device) + self._data.default_root_state = default_root_state.repeat(self.num_instances, 1) + # -- joint state + # joint pos + indices_list, _, values_list = string_utils.resolve_matching_names_values( + self.cfg.init_state.joint_pos, self.joint_names + ) + self._data.default_joint_pos[:, indices_list] = torch.tensor(values_list, device=self.device) + # joint vel + indices_list, _, values_list = string_utils.resolve_matching_names_values( + self.cfg.init_state.joint_vel, self.joint_names + ) + self._data.default_joint_vel[:, indices_list] = torch.tensor(values_list, device=self.device) + + # -- joint limits + self._data.default_joint_limits = self.view.get_dof_limits().to(device=self.device).clone() + self._data.joint_limits = self._data.default_joint_limits.clone() + + """ + 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._view = None + + """ + Internal helpers -- Actuators. + """ + + def _process_actuators_cfg(self): + """Process and apply articulation joint properties.""" + # create actuators + self.actuators = dict() + # flag for implicit actuators + # if this is false, we by-pass certain checks when doing actuator-related operations + self._has_implicit_actuators = False + + # cache the values coming from the usd + self._data.default_joint_stiffness = self.view.get_dof_stiffnesses().to(self.device).clone() + self._data.default_joint_damping = self.view.get_dof_dampings().to(self.device).clone() + self._data.default_joint_armature = self.view.get_dof_armatures().to(self.device).clone() + self._data.default_joint_friction = self.view.get_dof_friction_coefficients().to(self.device).clone() + + # iterate over all actuator configurations + for actuator_name, actuator_cfg in self.cfg.actuators.items(): + # type annotation for type checkers + actuator_cfg: ActuatorBaseCfg + # create actuator group + joint_ids, joint_names = self.find_joints(actuator_cfg.joint_names_expr) + # check if any joints are found + if len(joint_names) == 0: + raise ValueError( + f"No joints found for actuator group: {actuator_name} with joint name expression:" + f" {actuator_cfg.joint_names_expr}." + ) + # create actuator collection + # note: for efficiency avoid indexing when over all indices + actuator: ActuatorBase = actuator_cfg.class_type( + cfg=actuator_cfg, + joint_names=joint_names, + joint_ids=( + slice(None) if len(joint_names) == self.num_joints else torch.tensor(joint_ids, device=self.device) + ), + num_envs=self.num_instances, + device=self.device, + 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.view.get_dof_max_forces().to(self.device).clone()[:, joint_ids], + velocity_limit=self.view.get_dof_max_velocities().to(self.device).clone()[:, joint_ids], + ) + # log information on actuator groups + omni.log.info( + f"Actuator collection: {actuator_name} with model '{actuator_cfg.class_type.__name__}' and" + f" joint names: {joint_names} [{joint_ids}]." + ) + # store actuator group + self.actuators[actuator_name] = actuator + # set the passed gains and limits into the simulation + if isinstance(actuator, ImplicitActuator): + self._has_implicit_actuators = True + # the gains and limits are set into the simulation since actuator model is implicit + 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: + # the gains and limits are processed by the actuator model + # we set gains to zero, and torque limit to a high value in simulation to avoid any interference + 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) + # 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()) + if total_act_joints != (self.num_joints - self.num_fixed_tendons): + omni.log.warn( + "Not all actuators are configured! Total number of actuated joints not equal to number of" + f" joints available: {total_act_joints} != {self.num_joints - self.num_fixed_tendons}." + ) + + def _process_fixed_tendons(self): + """Process fixed tendons.""" + # create a list to store the fixed tendon names + self._fixed_tendon_names = list() + + # parse fixed tendons properties if they exist + if self.num_fixed_tendons > 0: + stage = stage_utils.get_current_stage() + joint_paths = self.view.dof_paths[0] + + # iterate over all joints to find tendons attached to them + for j in range(self.num_joints): + 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): + joint_name = usd_joint_path.split("/")[-1] + self._fixed_tendon_names.append(joint_name) + + self._data.fixed_tendon_names = self._fixed_tendon_names + self._data.default_fixed_tendon_stiffness = self.view.get_fixed_tendon_stiffnesses().clone() + self._data.default_fixed_tendon_damping = self.view.get_fixed_tendon_dampings().clone() + self._data.default_fixed_tendon_limit_stiffness = self.view.get_fixed_tendon_limit_stiffnesses().clone() + self._data.default_fixed_tendon_limit = self.view.get_fixed_tendon_limits().clone() + self._data.default_fixed_tendon_rest_length = self.view.get_fixed_tendon_rest_lengths().clone() + self._data.default_fixed_tendon_offset = self.view.get_fixed_tendon_offsets().clone() + + def _apply_actuator_model(self): + """Processes joint commands for the articulation by forwarding them to the actuators. + + The actions are first processed using actuator models. Depending on the robot configuration, + the actuator models compute the joint level simulation commands and sets them into the PhysX buffers. + """ + # process actions per group + for actuator in self.actuators.values(): + # prepare input for actuator model based on cached data + # TODO : A tensor dict would be nice to do the indexing of all tensors together + control_action = ArticulationActions( + joint_positions=self._data.joint_pos_target[:, actuator.joint_indices], + joint_velocities=self._data.joint_vel_target[:, actuator.joint_indices], + joint_efforts=self._data.joint_effort_target[:, actuator.joint_indices], + joint_indices=actuator.joint_indices, + ) + # compute joint command from the actuator model + control_action = actuator.compute( + control_action, + joint_pos=self._data.joint_pos[:, actuator.joint_indices], + joint_vel=self._data.joint_vel[:, actuator.joint_indices], + ) + # update targets (these are set into the simulation) + if control_action.joint_positions is not None: + self._joint_pos_target_sim[:, actuator.joint_indices] = control_action.joint_positions + if control_action.joint_velocities is not None: + self._joint_vel_target_sim[:, actuator.joint_indices] = control_action.joint_velocities + if control_action.joint_efforts is not None: + self._joint_effort_target_sim[:, actuator.joint_indices] = control_action.joint_efforts + # update state of the actuator model + # -- torques + self._data.computed_torque[:, actuator.joint_indices] = actuator.computed_effort + self._data.applied_torque[:, actuator.joint_indices] = actuator.applied_effort + # -- actuator data + self._data.soft_joint_vel_limits[:, actuator.joint_indices] = actuator.velocity_limit + # TODO: find a cleaner way to handle gear ratio. Only needed for variable gear ratio actuators. + if hasattr(actuator, "gear_ratio"): + self._data.gear_ratio[:, actuator.joint_indices] = actuator.gear_ratio + + """ + Internal helpers -- Debugging. + """ + + def _validate_cfg(self): + """Validate the configuration after processing. + + Note: + This function should be called only after the configuration has been processed and the buffers have been + created. Otherwise, some settings that are altered during processing may not be validated. + For instance, the actuator models may change the joint max velocity limits. + """ + # check that the default values are within the limits + joint_pos_limits = self.view.get_dof_limits()[0].to(self.device) + out_of_range = self._data.default_joint_pos[0] < joint_pos_limits[:, 0] + out_of_range |= self._data.default_joint_pos[0] > joint_pos_limits[:, 1] + violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) + # throw error if any of the default joint positions are out of the limits + if len(violated_indices) > 0: + # prepare message for violated joints + msg = "The following joints have default positions out of the limits: \n" + for idx in violated_indices: + joint_name = self.data.joint_names[idx] + joint_limits = joint_pos_limits[idx] + joint_pos = self.data.default_joint_pos[0, idx] + # add to message + msg += f"\t- '{joint_name}': {joint_pos:.3f} not in [{joint_limits[0]:.3f}, {joint_limits[1]:.3f}]\n" + raise ValueError(msg) + + # check that the default joint velocities are within the limits + joint_max_vel = self.view.get_dof_max_velocities()[0].to(self.device) + out_of_range = torch.abs(self._data.default_joint_vel[0]) > joint_max_vel + violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) + if len(violated_indices) > 0: + # prepare message for violated joints + msg = "The following joints have default velocities out of the limits: \n" + for idx in violated_indices: + joint_name = self.data.joint_names[idx] + joint_limits = [-joint_max_vel[idx], joint_max_vel[idx]] + joint_vel = self.data.default_joint_vel[0, idx] + # add to message + msg += f"\t- '{joint_name}': {joint_vel:.3f} not in [{joint_limits[0]:.3f}, {joint_limits[1]:.3f}]\n" + raise ValueError(msg) + + def _log_articulation_joint_info(self): + """Log information about the articulation's simulated joints.""" + # read out all joint parameters from simulation + # -- gains + stiffnesses = self.view.get_dof_stiffnesses()[0].tolist() + dampings = self.view.get_dof_dampings()[0].tolist() + # -- properties + armatures = self.view.get_dof_armatures()[0].tolist() + frictions = self.view.get_dof_friction_coefficients()[0].tolist() + # -- limits + position_limits = self.view.get_dof_limits()[0].tolist() + velocity_limits = self.view.get_dof_max_velocities()[0].tolist() + effort_limits = self.view.get_dof_max_forces()[0].tolist() + # create table for term information + table = PrettyTable(float_format=".3f") + table.field_names = [ + "Index", + "Name", + "Stiffness", + "Damping", + "Armature", + "Friction", + "Position Limits", + "Velocity Limits", + "Effort Limits", + ] + # set alignment of table columns + table.align["Name"] = "l" + # add info on each term + for index, name in enumerate(self.joint_names): + table.add_row( + [ + index, + name, + stiffnesses[index], + dampings[index], + armatures[index], + frictions[index], + position_limits[index], + velocity_limits[index], + effort_limits[index], + ] + ) + + # read out all tendon parameters from simulation + if self.num_fixed_tendons > 0: + # -- gains + ft_stiffnesses = self.view.get_fixed_tendon_stiffnesses()[0].tolist() + ft_dampings = self.view.get_fixed_tendon_dampings()[0].tolist() + # -- limits + ft_limit_stiffnesses = self.view.get_fixed_tendon_limit_stiffnesses()[0].tolist() + ft_limits = self.view.get_fixed_tendon_limits()[0].tolist() + ft_rest_lengths = self.view.get_fixed_tendon_rest_lengths()[0].tolist() + ft_offsets = self.view.get_fixed_tendon_offsets()[0].tolist() + # create table for term information + tendon_table = PrettyTable(float_format=".3f") + tendon_table.title = f"Simulation Tendon Information (Prim path: {self.cfg.prim_path})" + tendon_table.field_names = [ + "Index", + "Stiffness", + "Damping", + "Limit Stiffness", + "Limit", + "Rest Length", + "Offset", + ] + # add info on each term + for index in range(self.num_fixed_tendons): + tendon_table.add_row( + [ + index, + ft_stiffnesses[index], + ft_dampings[index], + ft_limit_stiffnesses[index], + ft_limits[index], + ft_rest_lengths[index], + ft_offsets[index], + ] + ) diff --git a/source/uwlab/uwlab/assets/articulation/articulation_cfg.py b/source/uwlab/uwlab/assets/articulation/articulation_cfg.py new file mode 100644 index 00000000..6b74986f --- /dev/null +++ b/source/uwlab/uwlab/assets/articulation/articulation_cfg.py @@ -0,0 +1,55 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import MISSING + +from isaaclab.actuators import ActuatorBaseCfg +from isaaclab.utils import configclass + +from ..asset_base_cfg import AssetBaseCfg +from .articulation import UniversalArticulation +from .articulation_view import ArticulationViewCfg + + +@configclass +class ArticulationCfg(AssetBaseCfg): + """Configuration parameters for an articulation.""" + + @configclass + class InitialStateCfg(AssetBaseCfg.InitialStateCfg): + """Initial state of the articulation.""" + + # root velocity + lin_vel: tuple[float, float, float] = (0.0, 0.0, 0.0) + """Linear velocity of the root in simulation world frame. Defaults to (0.0, 0.0, 0.0).""" + ang_vel: tuple[float, float, float] = (0.0, 0.0, 0.0) + """Angular velocity of the root in simulation world frame. Defaults to (0.0, 0.0, 0.0).""" + + # joint state + joint_pos: dict[str, float] = {".*": 0.0} + """Joint positions of the joints. Defaults to 0.0 for all joints.""" + joint_vel: dict[str, float] = {".*": 0.0} + """Joint velocities of the joints. Defaults to 0.0 for all joints.""" + + ## + # Initialize configurations. + ## + + class_type: type[UniversalArticulation] = UniversalArticulation + + articulation_view_cfg: ArticulationViewCfg = MISSING + + init_state: InitialStateCfg = InitialStateCfg() + """Initial state of the articulated object. Defaults to identity pose with zero velocity and zero joint state.""" + + soft_joint_pos_limit_factor: float = 1.0 + """Fraction specifying the range of DOF position limits (parsed from the asset) to use. Defaults to 1.0. + + The joint position limits are scaled by this factor to allow for a limited range of motion. + This is accessible in the articulation data through :attr:`ArticulationData.soft_joint_pos_limits` attribute. + """ + + actuators: dict[str, ActuatorBaseCfg] = MISSING + """Actuators for the robot with corresponding joint names.""" diff --git a/source/uwlab/uwlab/assets/articulation/articulation_data.py b/source/uwlab/uwlab/assets/articulation/articulation_data.py new file mode 100644 index 00000000..ebfb2b07 --- /dev/null +++ b/source/uwlab/uwlab/assets/articulation/articulation_data.py @@ -0,0 +1,858 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch +import weakref + +import omni.log + +import isaaclab.utils.math as math_utils +from isaaclab.utils.buffers import TimestampedBuffer + +from .articulation_view import ArticulationView + + +class ArticulationData: + """Data container for an articulation. + + This class contains the data for an articulation in the simulation. The data includes the state of + the root rigid body, the state of all the bodies in the articulation, and the joint state. The data is + stored in the simulation world frame unless otherwise specified. + + An articulation is comprised of multiple rigid bodies or links. 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, the two frames may not coincide with each other. In the robotics sense, the actor frame + can be interpreted as the link frame. + """ + + def __init__(self, view: ArticulationView, device: str): + """Initializes the articulation data. + + Args: + root_physx_view: The root articulation view. + device: The device used for processing. + """ + # Set the parameters + self.device = device + # Set the root articulation 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.view: ArticulationView = weakref.proxy(view) + + # Set initial time stamp + self._timestamp = 0.0 + + # Convert to direction vector + gravity_dir = torch.tensor((0.0, 0.0, -9.8), device=self.device) + gravity_dir = math_utils.normalize(gravity_dir.unsqueeze(0)).squeeze(0) + + # Initialize constants + self.GRAVITY_VEC_W = gravity_dir.repeat(self.view.count, 1) + self.FORWARD_VEC_B = torch.tensor((1.0, 0.0, 0.0), device=self.device).repeat(self.view.count, 1) + + # Initialize history for finite differencing + self._previous_joint_vel = self.view.get_dof_velocities().clone() + + # 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._timestamp += dt + # Trigger an update of the joint acceleration buffer at a higher frequency + # since we do finite differencing. + self.joint_acc + + def update_articulations_kinematic(self): + """Update the kinematic state of all articulations.""" + pass + + ## + # Names. + ## + + body_names: list[str] = None + """Body names in the order parsed by the simulation view.""" + + joint_names: list[str] = None + """Joint names in the order parsed by the simulation view.""" + + fixed_tendon_names: list[str] = None + """Fixed tendon names in the order parsed by the simulation view.""" + + ## + # Defaults. + ## + + default_root_state: torch.Tensor = None + """Default root state ``[pos, quat, lin_vel, ang_vel]`` in local environment 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 its center of mass frame. + """ + + default_mass: torch.Tensor = None + """Default mass read from the simulation. Shape is (num_instances, num_bodies).""" + + default_inertia: torch.Tensor = None + """Default inertia read from the simulation. Shape is (num_instances, num_bodies, 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}]`. + """ + + default_joint_pos: torch.Tensor = None + """Default joint positions of all joints. Shape is (num_instances, num_joints).""" + + default_joint_vel: torch.Tensor = None + """Default joint velocities of all joints. Shape is (num_instances, num_joints).""" + + default_joint_stiffness: torch.Tensor = None + """Default joint stiffness of all joints. Shape is (num_instances, num_joints).""" + + default_joint_damping: torch.Tensor = None + """Default joint damping of all joints. Shape is (num_instances, num_joints).""" + + default_joint_armature: torch.Tensor = None + """Default joint armature of all joints. Shape is (num_instances, num_joints).""" + + default_joint_friction: torch.Tensor = None + """Default joint friction of all joints. Shape is (num_instances, num_joints).""" + + default_joint_limits: torch.Tensor = None + """Default joint limits of all joints. Shape is (num_instances, num_joints, 2).""" + + default_fixed_tendon_stiffness: torch.Tensor = None + """Default tendon stiffness of all tendons. Shape is (num_instances, num_fixed_tendons).""" + + default_fixed_tendon_damping: torch.Tensor = None + """Default tendon damping of all tendons. Shape is (num_instances, num_fixed_tendons).""" + + default_fixed_tendon_limit_stiffness: torch.Tensor = None + """Default tendon limit stiffness of all tendons. Shape is (num_instances, num_fixed_tendons).""" + + default_fixed_tendon_rest_length: torch.Tensor = None + """Default tendon rest length of all tendons. Shape is (num_instances, num_fixed_tendons).""" + + default_fixed_tendon_offset: torch.Tensor = None + """Default tendon offset of all tendons. Shape is (num_instances, num_fixed_tendons).""" + + default_fixed_tendon_limit: torch.Tensor = None + """Default tendon limits of all tendons. Shape is (num_instances, num_fixed_tendons, 2).""" + + ## + # Joint commands -- Set into simulation. + ## + + joint_pos_target: torch.Tensor = None + """Joint position targets commanded by the user. Shape is (num_instances, num_joints). + + For an implicit actuator model, the targets are directly set into the simulation. + For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), + which are then set into the simulation. + """ + + joint_vel_target: torch.Tensor = None + """Joint velocity targets commanded by the user. Shape is (num_instances, num_joints). + + For an implicit actuator model, the targets are directly set into the simulation. + For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), + which are then set into the simulation. + """ + + joint_effort_target: torch.Tensor = None + """Joint effort targets commanded by the user. Shape is (num_instances, num_joints). + + For an implicit actuator model, the targets are directly set into the simulation. + For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), + which are then set into the simulation. + """ + + ## + # Joint commands -- Explicit actuators. + ## + + computed_torque: torch.Tensor = None + """Joint torques computed from the actuator model (before clipping). Shape is (num_instances, num_joints). + + This quantity is the raw torque output from the actuator mode, before any clipping is applied. + It is exposed for users who want to inspect the computations inside the actuator model. + For instance, to penalize the learning agent for a difference between the computed and applied torques. + + Note: The torques are zero for implicit actuator models. + """ + + applied_torque: torch.Tensor = None + """Joint torques applied from the actuator model (after clipping). Shape is (num_instances, num_joints). + + These torques are set into the simulation, after clipping the :attr:`computed_torque` based on the + actuator model. + + Note: The torques are zero for implicit actuator models. + """ + + ## + # Joint properties. + ## + + joint_stiffness: torch.Tensor = None + """Joint stiffness provided to simulation. Shape is (num_instances, num_joints).""" + + joint_damping: torch.Tensor = None + """Joint damping 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_vel_limits: torch.Tensor = None + """Joint maximum velocity provided to simulation. Shape is (num_instances, num_joints).""" + + ## + # Fixed tendon properties. + ## + + fixed_tendon_stiffness: torch.Tensor = None + """Fixed tendon stiffness provided to simulation. Shape is (num_instances, num_fixed_tendons).""" + + fixed_tendon_damping: torch.Tensor = None + """Fixed tendon damping provided to simulation. Shape is (num_instances, num_fixed_tendons).""" + + fixed_tendon_limit_stiffness: torch.Tensor = None + """Fixed tendon limit stiffness provided to simulation. Shape is (num_instances, num_fixed_tendons).""" + + fixed_tendon_rest_length: torch.Tensor = None + """Fixed tendon rest length provided to simulation. Shape is (num_instances, num_fixed_tendons).""" + + fixed_tendon_offset: torch.Tensor = None + """Fixed tendon offset provided to simulation. Shape is (num_instances, num_fixed_tendons).""" + + fixed_tendon_limit: torch.Tensor = None + """Fixed tendon limits provided to simulation. Shape is (num_instances, num_fixed_tendons, 2).""" + + ## + # Other Data. + ## + + soft_joint_pos_limits: torch.Tensor = None + """Joint positions limits for all joints. Shape is (num_instances, num_joints, 2).""" + + soft_joint_vel_limits: torch.Tensor = None + """Joint velocity limits for all joints. Shape is (num_instances, num_joints).""" + + gear_ratio: torch.Tensor = None + """Gear ratio for relating motor torques to applied Joint torques. Shape is (num_instances, num_joints).""" + + ## + # Properties. + ## + + @property + 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 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._timestamp: + # read data from simulation + pose = self.view.get_root_transforms().clone() + pose[:, 3:7] = math_utils.convert_quat(pose[:, 3:7], to="wxyz") + velocity = self.view.get_root_velocities() + # set the buffer data and timestamp + self._root_state_w.data = torch.cat((pose, velocity), dim=-1) + self._root_state_w.timestamp = self._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._timestamp: + # read data from simulation + pose = self.view.get_root_transforms().clone() + pose[:, 3:7] = math_utils.convert_quat(pose[:, 3:7], to="wxyz") + velocity = self.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._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._timestamp: + # read data from simulation (pose is of link) + pose = self.view.get_root_transforms().clone() + pose[:, 3:7] = math_utils.convert_quat(pose[:, 3:7], to="wxyz") + velocity = self.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._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, num_bodies, 13). + + 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._timestamp: + self.update_articulations_kinematic() + # read data from simulation + poses = self.view.get_link_transforms().clone() + poses[..., 3:7] = math_utils.convert_quat(poses[..., 3:7], to="wxyz") + velocities = self.view.get_link_velocities() + # set the buffer data and timestamp + self._body_state_w.data = torch.cat((poses, velocities), dim=-1) + self._body_state_w.timestamp = self._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._timestamp: + self.update_articulations_kinematic() + # read data from simulation + pose = self.view.get_link_transforms().clone() + pose[..., 3:7] = math_utils.convert_quat(pose[..., 3:7], to="wxyz") + velocity = self.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._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._timestamp: + self.update_articulations_kinematic() + # read data from simulation (pose is of link) + pose = self.view.get_link_transforms().clone() + pose[..., 3:7] = math_utils.convert_quat(pose[..., 3:7], to="wxyz") + velocity = self.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._timestamp + return self._body_com_state_w.data + + @property + def body_acc_w(self): + """Acceleration of all bodies (center of mass). Shape is (num_instances, num_bodies, 6). + + All values are relative to the world. + """ + if self._body_acc_w.timestamp < self._timestamp: + # read data from simulation and set the buffer data and timestamp + self._body_acc_w.data = self.view.get_link_accelerations() + + self._body_acc_w.timestamp = self._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_link_quat_w, self.GRAVITY_VEC_W) + + @property + def heading_w(self): + """Yaw heading of the base frame (in radians). Shape is (num_instances,). + + 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.root_link_quat_w, self.FORWARD_VEC_B) + return torch.atan2(forward_w[:, 1], forward_w[:, 0]) + + @property + def joint_pos(self): + """Joint positions of all joints. Shape is (num_instances, num_joints).""" + if self._joint_pos.timestamp < self._timestamp: + # read data from simulation and set the buffer data and timestamp + self._joint_pos.data = self.view.get_dof_positions() + self._joint_pos.timestamp = self._timestamp + return self._joint_pos.data + + @property + def joint_vel(self): + """Joint velocities of all joints. Shape is (num_instances, num_joints).""" + if self._joint_vel.timestamp < self._timestamp: + # read data from simulation and set the buffer data and timestamp + self._joint_vel.data = self.view.get_dof_velocities() + self._joint_vel.timestamp = self._timestamp + return self._joint_vel.data + + @property + def joint_acc(self): + """Joint acceleration of all joints. Shape is (num_instances, num_joints).""" + if self._joint_acc.timestamp < self._timestamp: + # note: we use finite differencing to compute acceleration + time_elapsed = self._timestamp - self._joint_acc.timestamp + self._joint_acc.data = (self.joint_vel - self._previous_joint_vel) / time_elapsed + self._joint_acc.timestamp = self._timestamp + # update the previous joint velocity + self._previous_joint_vel[:] = self.joint_vel + return self._joint_acc.data + + ## + # Derived properties. + ## + + @property + 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 relative to the world. + """ + return self.root_state_w[:, :3] + + @property + 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 relative to the world. + """ + return self.root_state_w[:, 3:7] + + @property + 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 + relative to the world. + """ + return self.root_state_w[:, 7:13] + + @property + 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 relative to the world. + """ + return self.root_state_w[:, 7:10] + + @property + 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 relative to the world. + """ + return self.root_state_w[:, 10:13] + + @property + 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 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) + + @property + 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 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._timestamp: + # read data from simulation (pose is of link) + pose = self.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._timestamp: + # read data from simulation (pose is of link) + pose = self.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._timestamp: + # read data from simulation (pose is of link) + velocity = self.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._timestamp: + # read data from simulation (pose is of link) + velocity = self.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._timestamp: + self.update_articulations_kinematic() + # read data from simulation (pose is of link) + velocity = self.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_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 relative + to the world. + """ + return self.body_state_w[..., 7:13] + + @property + 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 relative to the world. + """ + return self.body_state_w[..., 7:10] + + @property + 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 relative to the world. + """ + return self.body_state_w[..., 10:13] + + @property + 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 relative to the world. + """ + return self.body_acc_w[..., 0:3] + + @property + 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 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._timestamp: + self.update_articulations_kinematic() + # read data from simulation + pose = self.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._timestamp: + self.update_articulations_kinematic() + # read data from simulation + pose = self.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._timestamp: + self.update_articulations_kinematic() + # read data from simulation (velocity is of com) + velocity = self.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._timestamp: + self.update_articulations_kinematic() + # read data from simulation (velocity is of com) + velocity = self.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._timestamp: + self.update_articulations_kinematic() + # read data from simulation (velocity is of com) + velocity = self.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.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.view.get_coms().to(self.device)[..., 3:7] + return math_utils.convert_quat(quat, to="wxyz") diff --git a/source/uwlab/uwlab/assets/articulation/articulation_drive/__init__.py b/source/uwlab/uwlab/assets/articulation/articulation_drive/__init__.py new file mode 100644 index 00000000..44e81d7e --- /dev/null +++ b/source/uwlab/uwlab/assets/articulation/articulation_drive/__init__.py @@ -0,0 +1,4 @@ +from .articulation_drive import ArticulationDrive +from .articulation_drive_cfg import ArticulationDriveCfg +from .articulation_drive_data import ArticulationDriveData +from .articulation_drive_process import ArticulationDriveDedicatedProcess diff --git a/source/uwlab/uwlab/assets/articulation/articulation_drive/articulation_drive.py b/source/uwlab/uwlab/assets/articulation/articulation_drive/articulation_drive.py new file mode 100644 index 00000000..ec5f3f45 --- /dev/null +++ b/source/uwlab/uwlab/assets/articulation/articulation_drive/articulation_drive.py @@ -0,0 +1,60 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from abc import abstractmethod +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from . import ArticulationDriveCfg + + # TODO: in next release remove this and should be completely + + +class ArticulationDrive: + def __init__(self, cfg: ArticulationDriveCfg): + raise NotImplementedError + + """ + Below method should be implemented by the child class + """ + + @property + def ordered_joint_names(self) -> list[str]: + raise NotImplementedError + + @abstractmethod + def close(self) -> None: + raise NotImplementedError + + @abstractmethod + def read_dof_states(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + raise NotImplementedError + + @abstractmethod + def write_dof_targets(self, pos_target: torch.Tensor, vel_target: torch.Tensor, eff_target: torch.Tensor): + raise NotImplementedError + + @abstractmethod + def set_dof_stiffnesses(self, stiffnesses: torch.Tensor) -> None: + raise NotImplementedError + + @abstractmethod + def set_dof_armatures(self, armatures: torch.Tensor) -> None: + raise NotImplementedError + + @abstractmethod + def set_dof_frictions(self, frictions: torch.Tensor) -> None: + raise NotImplementedError + + @abstractmethod + def set_dof_dampings(self, dampings: torch.Tensor) -> None: + raise NotImplementedError + + @abstractmethod + def set_dof_limits(self, limits: torch.Tensor) -> None: + raise NotImplementedError diff --git a/source/uwlab/uwlab/assets/articulation/articulation_drive/articulation_drive_cfg.py b/source/uwlab/uwlab/assets/articulation/articulation_drive/articulation_drive_cfg.py new file mode 100644 index 00000000..a9523ae0 --- /dev/null +++ b/source/uwlab/uwlab/assets/articulation/articulation_drive/articulation_drive_cfg.py @@ -0,0 +1,24 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import MISSING +from typing import Callable + +from isaaclab.utils import configclass + +from .articulation_drive import ArticulationDrive + + +@configclass +class ArticulationDriveCfg: + """Configuration parameters for an articulation view.""" + + class_type: Callable[..., ArticulationDrive] = MISSING # type: ignore + + use_multiprocessing: bool = False + + dt = 0.01 + + device: str = "cpu" diff --git a/source/uwlab/uwlab/assets/articulation/articulation_drive/articulation_drive_data.py b/source/uwlab/uwlab/assets/articulation/articulation_drive/articulation_drive_data.py new file mode 100644 index 00000000..58a1aa8c --- /dev/null +++ b/source/uwlab/uwlab/assets/articulation/articulation_drive/articulation_drive_data.py @@ -0,0 +1,35 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch +from typing import List, TypedDict + + +class ArticulationDriveData(TypedDict): + is_running: bool + close: bool + link_names: List[str] + dof_names: List[str] + dof_types: List[str] + pos: torch.Tensor + vel: torch.Tensor + torque: torch.Tensor + pos_target: torch.Tensor + vel_target: torch.Tensor + eff_target: torch.Tensor + link_transforms: torch.Tensor + link_velocities: torch.Tensor + link_mass: torch.Tensor + link_inertia: torch.Tensor + link_coms: torch.Tensor + mass_matrix: torch.Tensor + dof_stiffness: torch.Tensor + dof_armatures: torch.Tensor + dof_frictions: torch.Tensor + dof_damping: torch.Tensor + dof_limits: torch.Tensor + dof_max_forces: torch.Tensor + dof_max_velocity: torch.Tensor + jacobians: torch.Tensor diff --git a/source/uwlab/uwlab/assets/articulation/articulation_drive/articulation_drive_process.py b/source/uwlab/uwlab/assets/articulation/articulation_drive/articulation_drive_process.py new file mode 100644 index 00000000..c9ea25f9 --- /dev/null +++ b/source/uwlab/uwlab/assets/articulation/articulation_drive/articulation_drive_process.py @@ -0,0 +1,183 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import queue +import threading +import time +import torch +import torch.multiprocessing as mp +from abc import abstractmethod +from typing import TYPE_CHECKING + +from .articulation_drive import ArticulationDrive + +if TYPE_CHECKING: + from . import ArticulationDriveCfg, ArticulationDriveData + + +# This class provides a sketch for dedicated process for handling hardware communication in an ArticulationDrive, +# rather than forcing every drive in a “view” to share a single process. + +# Why? +# If you have multiple drives (e.g. an arm with two hands), each drive might otherwise block or slow down others when talking to hardware. +# By creating a separate process per drive, each drive handles its own hardware I/O without impacting the rest of the system. + +# How it Works: +# On creation, every drive class spawns either a multiprocessing Process or a thread (depending on configuration). +# This separate worker loop regularly reads from and writes to the hardware (position, velocity, torque, etc.). +# Commands from the main thread (such as setting stiffness or limits) go through a queue and get handled in the worker process. + +# Benefits: +# No single bottleneck from multiple drives all sharing one process. +# More responsive hardware communication for each drive. +# Cleaner separation of logic—each drive manages its own thread/process. + +# Before this Design +# - ArticulationView1(Process1) - ArticulationView2(Process2) +# - ArticulationDrive1 - ArticulationDrive1 +# - ArticulationDrive2 +# - ArticulationDrive3 + +# as you may see, ArticulationView2 maybe fine, but ArticulationView1 may have a bottleneck due +# multiple hardware calls and some maybe blocking. + +# After this Design: +# - ArticulationView1(Process1) - ArticulationView2(Process4 thread1) +# - ArticulationDrive1(Process2) - ArticulationDrive(Process4 thread2) +# - ArticulationDrive2(Process3 thread1) +# - ArticulationDrive3(Process3 thread2) + + +# Use this design when you suspect a single thread/process for all drives within a articulation View +# might become slow or blocking due to heavy hardware calls or large simulations. + + +class ArticulationDriveDedicatedProcess(ArticulationDrive): + def __init__(self, cfg: ArticulationDriveCfg): + self._dt = cfg.dt + if cfg.use_multiprocessing: + self.manager = mp.Manager() + self.init_event = mp.Event() + self.cmd_queue = mp.Queue() + self.ack_queue = mp.Queue() + # Create a shared dictionary to communicate with the child process + self.shared_data: ArticulationDriveData = self.manager.dict() # type: ignore + self._proc = mp.Process(target=self._run, args=(self.init_event,)) + else: + # Threading shares memory within the same process, so no need for mp.Manager + self.manager = None + self.init_event = threading.Event() + self.cmd_queue = queue.Queue() + self.ack_queue = queue.Queue() + self.shared_data: ArticulationDriveData = {} # type: ignore + self._proc = threading.Thread(target=self._run, daemon=True, args=(self.init_event,)) + + """ + Below method should be implemented by the child class + """ + + @property + def ordered_joint_names(self) -> list[str]: + raise NotImplementedError + + @abstractmethod + def close(self) -> None: + raise NotImplementedError + + @abstractmethod + def read_dof_states(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + raise NotImplementedError + + @abstractmethod + def write_dof_targets(self, pos_target: torch.Tensor, vel_target: torch.Tensor, eff_target: torch.Tensor): + raise NotImplementedError + + @abstractmethod + def set_dof_stiffnesses(self, stiffnesses: torch.Tensor) -> None: + raise NotImplementedError + + @abstractmethod + def set_dof_armatures(self, armatures: torch.Tensor) -> None: + raise NotImplementedError + + @abstractmethod + def set_dof_frictions(self, frictions: torch.Tensor) -> None: + raise NotImplementedError + + @abstractmethod + def set_dof_dampings(self, dampings: torch.Tensor) -> None: + raise NotImplementedError + + @abstractmethod + def set_dof_limits(self, limits: torch.Tensor) -> None: + raise NotImplementedError + + """Below methods should not be called by any other class + """ + + def _run(self, init_event: threading.Event): + init_event.set() + + next_poll_time = time.time() + self._dt + while True: + if self.shared_data["close"]: + break + + self._process_blocking_commands() + + if self.shared_data["is_running"]: + now = time.time() + if now >= next_poll_time: + next_poll_time = now + self._dt + # 1) read pos/vel from hardware + pos, vel, eff = self.read_dof_states() + self.shared_data["pos"][:] = pos + self.shared_data["vel"][:] = vel + self.shared_data["torque"][:] = eff + + # 2) write target to hardware and kinematic + self.write_dof_targets( + pos_target=self.shared_data["pos_target"], + vel_target=self.shared_data["vel_target"], + eff_target=self.shared_data["eff_target"], + ) + # Sleep till next poll time + time.sleep(max(next_poll_time - time.time(), 0)) + else: + # now is less than next poll time, sleep till next poll time + time.sleep(next_poll_time - now) + else: + # If not running, keep sleeping + time.sleep(self._dt) + + # Done, disconnect hardware + self.close() + print("DynamixelWorker: Child process stopped") + + def _process_blocking_commands(self): + if not self.cmd_queue.empty(): + cmd = self.cmd_queue.get() + match cmd: + case "set_dof_stiffnesses": + self.set_dof_stiffnesses(self.shared_data["dof_stiffness"]) + self.ack_queue.put({"status": "OK", "command": "set_dof_stiffnesses"}) + case "set_dof_armatures": + self.set_dof_armatures(self.shared_data["dof_armatures"]) + self.ack_queue.put({"status": "OK", "command": "set_dof_armatures"}) + case "set_dof_frictions": + self.set_dof_frictions(self.shared_data["dof_frictions"]) + self.ack_queue.put({"status": "OK", "command": "set_dof_frictions"}) + case "set_dof_dampings": + self.set_dof_dampings(self.shared_data["dof_damping"]) + self.ack_queue.put({"status": "OK", "command": "set_dof_dampings"}) + case "set_dof_limits": + self.set_dof_limits(self.shared_data["dof_limits"]) + self.ack_queue.put({"status": "OK", "command": "set_dof_limits"}) + + """ + Below methods are for users to call + """ diff --git a/source/uwlab/uwlab/assets/articulation/articulation_view/__init__.py b/source/uwlab/uwlab/assets/articulation/articulation_view/__init__.py new file mode 100644 index 00000000..69aa3a8c --- /dev/null +++ b/source/uwlab/uwlab/assets/articulation/articulation_view/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from ..articulation_drive import ArticulationDrive, ArticulationDriveCfg +from .articulation_view import ArticulationView, SharedDataSchema +from .articulation_view_cfg import ArticulationViewCfg, BulletArticulationViewCfg diff --git a/source/uwlab/uwlab/assets/articulation/articulation_view/articulation_view.py b/source/uwlab/uwlab/assets/articulation/articulation_view/articulation_view.py new file mode 100644 index 00000000..02057b4a --- /dev/null +++ b/source/uwlab/uwlab/assets/articulation/articulation_view/articulation_view.py @@ -0,0 +1,657 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch +from abc import abstractmethod +from typing import List, TypedDict + + +class SharedDataSchema(TypedDict): + is_running: bool + close: bool + link_names: List[str] + dof_names: List[str] + dof_types: List[str] + pos: torch.Tensor + vel: torch.Tensor + torque: torch.Tensor + pos_target: torch.Tensor + vel_target: torch.Tensor + eff_target: torch.Tensor + link_transforms: torch.Tensor + link_velocities: torch.Tensor + link_mass: torch.Tensor + link_inertia: torch.Tensor + link_coms: torch.Tensor + mass_matrix: torch.Tensor + dof_stiffness: torch.Tensor + dof_armatures: torch.Tensor + dof_frictions: torch.Tensor + dof_damping: torch.Tensor + dof_limits: torch.Tensor + dof_max_forces: torch.Tensor + dof_max_velocity: torch.Tensor + jacobians: torch.Tensor + + +class ArticulationView: + """ + An abstract interface for querying and setting the state of an articulated mechanism. + This interface aims to provide a unified set of methods for interacting with: + - Physics engines (e.g., PhysX, MuJoCo, Bullet, etc.) + - Real robot hardware + - Simulated or software-only kinematics/dynamics backends + + The subclass implementing this interface must handle how these queries and commands + are reflected in the underlying system. + """ + + def __init__(self): + """ + Initializes the articulation view. + + Subclasses must implement: + - Initialization and any required backend setup. + - Necessary internal references for subsequent get/set methods. + + Raises: + NotImplementedError: If not overridden by subclasses. + """ + raise NotImplementedError + + @abstractmethod + def play(self): + """ + Play the articulation(s) in the view. + """ + raise NotImplementedError + + @abstractmethod + def pause(self): + """ + Pause the articulation(s) in the view. + """ + raise NotImplementedError + + @abstractmethod + def close(self): + """ + Pause the articulation(s) in the view. + """ + raise NotImplementedError + + @property + def count(self) -> int: + """ + Number of articulation instances being managed by this view. + + E.g., if your environment has N separate copies of the same robot, + this property would return N. + + Returns: + int: The total number of articulation instances. + """ + raise NotImplementedError + + @property + def fixed_base(self) -> bool: + """ + Indicates whether the articulation(s) in this view has a fixed base. + + A fixed-base articulation does not move freely in space + (e.g., a robotic arm bolted to a table), + whereas a floating-base articulation can move freely. + + Returns: + bool: True if the articulation is fixed-base, False otherwise. + """ + raise NotImplementedError + + @property + def dof_count(self) -> int: + """ + Number of degrees of freedom (DOFs) for the articulation(s) in this view. + + Returns: + int: Count of all DOFs for the managed articulation(s). + """ + raise NotImplementedError + + @property + def max_fixed_tendons(self) -> int: + """ + Maximum number of 'fixed tendons' (sometimes known as cables or passively constrained joints). + + Returns: + int: The number of fixed tendon connections in the articulation(s). + """ + raise NotImplementedError + + @property + def num_bodies(self) -> int: + """ + Number of rigid bodies (links) in the articulation(s). + + Returns: + int: Total number of rigid bodies in the articulation(s). + """ + raise NotImplementedError + + @property + def joint_names(self) -> list[str]: + """ + Ordered list of joint names in the articulation(s). + + Returns: + list[str]: Names of the joints in order of their DOF indices. + """ + raise NotImplementedError + + @property + def fixed_tendon_names(self) -> list[str]: + """ + Ordered list of names for the fixed tendons (cables) in the articulation(s). + + Returns: + list[str]: Names of the fixed tendons, if any. + """ + raise NotImplementedError + + @property + def body_names(self) -> list[str]: + """ + Ordered list of body (link) names in the articulation(s). + + Returns: + list[str]: Names of the rigid bodies in order of their indices. + """ + raise NotImplementedError + + @abstractmethod + def get_root_transforms(self) -> torch.Tensor: + """ + Get the root poses (position + orientation) of each articulation instance. + + This typically refers to the base link or root transform + of a floating/fixed-base articulation in world coordinates. + + Returns: + torch.Tensor: A tensor of shape (count, 7), + each row containing [px, py, pz, qw, qx, qy, qz]. + """ + raise NotImplementedError + + @abstractmethod + def get_root_velocities(self) -> torch.Tensor: + """ + Get the linear and angular velocities of the root link(s) of each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, 6), + each row containing [vx, vy, vz, wx, wy, wz]. + """ + raise NotImplementedError + + @abstractmethod + def get_link_accelerations(self) -> torch.Tensor: + """ + Get the link accelerations for each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, num_bodies, 6), + containing [ax, ay, az, alpha_x, alpha_y, alpha_z] for each link. + """ + raise NotImplementedError + + @abstractmethod + def get_link_transforms(self) -> torch.Tensor: + """ + Get the transforms (position + orientation) of each link for each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, num_bodies, 7), + for [px, py, pz, qx, qy, qz, qw] per link. + """ + raise NotImplementedError + + @abstractmethod + def get_link_velocities(self) -> torch.Tensor: + """ + Get the linear and angular velocities of each link in the articulation(s). + + Returns: + torch.Tensor: A tensor of shape (count, num_bodies, 6), + containing [vx, vy, vz, wx, wy, wz] per link. + """ + raise NotImplementedError + + @abstractmethod + def get_coms(self) -> torch.Tensor: + """ + Get the center-of-mass (COM) positions of each link or the entire articulation. + + Depending on the implementation, this can mean: + - Per-link COM (shape = (count, num_bodies, 3)) + - Single COM for the entire system (shape = (count, 3)) + + Returns: + torch.Tensor: COM positions in the world or local coordinate frame + depending on the backend's convention. + """ + raise NotImplementedError + + @abstractmethod + def get_masses(self) -> torch.Tensor: + """ + Get the masses for each link in the articulation(s). + + Returns: + torch.Tensor: A tensor of shape (count, num_bodies), + containing the mass of each link. + """ + raise NotImplementedError + + @abstractmethod + def get_inertias(self) -> torch.Tensor: + """ + Get the inertial tensors (often expressed in link-local frames) for each link. + + Returns: + torch.Tensor: A tensor of shape (count, num_bodies, 3, 3), + containing the inertia matrix for each link. + """ + raise NotImplementedError + + @abstractmethod + def get_dof_positions(self) -> torch.Tensor: + """ + Get the joint positions for each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, dof_count) + with joint angles or positions. + """ + raise NotImplementedError + + @abstractmethod + def get_dof_velocities(self) -> torch.Tensor: + """ + Get the joint velocities for each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, dof_count), + with joint velocity values. + """ + raise NotImplementedError + + @abstractmethod + def get_dof_max_velocities(self) -> torch.Tensor: + """ + Get the maximum velocity limits for each joint in each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, dof_count), + containing velocity limits per joint. + """ + raise NotImplementedError + + @abstractmethod + def get_dof_max_forces(self) -> torch.Tensor: + """ + Get the maximum force (torque) limits for each joint in each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, dof_count), + with torque/force limits per joint. + """ + raise NotImplementedError + + @abstractmethod + def get_dof_stiffnesses(self) -> torch.Tensor: + """ + Get the joint stiffness values for each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, dof_count), + containing the stiffness for each DOF. + """ + raise NotImplementedError + + @abstractmethod + def get_dof_dampings(self) -> torch.Tensor: + """ + Get the joint damping values for each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, dof_count), + containing the damping for each DOF. + """ + raise NotImplementedError + + @abstractmethod + def get_dof_armatures(self) -> torch.Tensor: + """ + Get the armature values for each joint in each articulation instance. + + The 'armature' is sometimes used to represent an inertia-like term + used in certain simulation backends or real hardware compensations. + + Returns: + torch.Tensor: A tensor of shape (count, dof_count). + """ + raise NotImplementedError + + @abstractmethod + def get_dof_friction_coefficients(self) -> torch.Tensor: + """ + Get the friction coefficients for each joint in each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, dof_count). + """ + raise NotImplementedError + + @abstractmethod + def get_dof_limits(self) -> torch.Tensor: + """ + Get the joint position limits for each joint in each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, dof_count, 2), + where the last dimension stores [lower_limit, upper_limit]. + """ + raise NotImplementedError + + @abstractmethod + def get_fixed_tendon_stiffnesses(self) -> torch.Tensor: + """ + Get the stiffness values for each fixed tendon across the articulation(s). + + Returns: + torch.Tensor: A tensor of shape (count, max_fixed_tendons), + with the stiffness value for each tendon. + """ + raise NotImplementedError + + @abstractmethod + def get_fixed_tendon_dampings(self) -> torch.Tensor: + """ + Get the damping values for each fixed tendon across the articulation(s). + + Returns: + torch.Tensor: A tensor of shape (count, max_fixed_tendons), + with the damping value for each tendon. + """ + raise NotImplementedError + + @abstractmethod + def get_fixed_tendon_limit_stiffnesses(self) -> torch.Tensor: + """ + Get the limit stiffness values for each fixed tendon. + + Returns: + torch.Tensor: A tensor of shape (count, max_fixed_tendons), + with the limit stiffness for each tendon. + """ + raise NotImplementedError + + @abstractmethod + def get_fixed_tendon_limits(self) -> torch.Tensor: + """ + Get the limit range for each fixed tendon, which might represent + min/max constraints on the tendon length or tension. + + Returns: + torch.Tensor: A tensor of shape (count, max_fixed_tendons, 2), + containing [lower_limit, upper_limit] for each tendon. + """ + raise NotImplementedError + + @abstractmethod + def get_fixed_tendon_rest_lengths(self) -> torch.Tensor: + """ + Get the rest lengths for each fixed tendon across the articulation(s). + + Returns: + torch.Tensor: A tensor of shape (count, max_fixed_tendons). + """ + raise NotImplementedError + + @abstractmethod + def get_fixed_tendon_offsets(self) -> torch.Tensor: + """ + Get the offset values for each fixed tendon across the articulation(s). + + Returns: + torch.Tensor: A tensor of shape (count, max_fixed_tendons). + """ + raise NotImplementedError + + @abstractmethod + def set_dof_actuation_forces(self, forces: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the actuation forces (torques) for the specified articulation instances. + + Args: + forces (torch.Tensor): A tensor of shape (len(indices), dof_count) + specifying the commanded forces/torques. + indices (torch.Tensor): A tensor of indices specifying which + articulation instances to apply these forces. + """ + raise NotImplementedError + + @abstractmethod + def set_dof_position_targets(self, positions: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the position targets for the specified articulation instances, + if the underlying controller or simulation uses position-based control. + + Args: + positions (torch.Tensor): A tensor of shape (len(indices), dof_count) + specifying desired joint positions. + indices (torch.Tensor): Indices of articulation instances to apply these targets. + """ + raise NotImplementedError + + @abstractmethod + def set_dof_positions(self, positions: torch.Tensor, indices: torch.Tensor) -> None: + """ + Hard-set the joint positions for the specified articulation instances. + Usually used for resetting or overriding joint states directly. + + Args: + positions (torch.Tensor): A tensor of shape (len(indices), dof_count) + specifying the new positions. + indices (torch.Tensor): Indices of articulation instances to set positions for. + """ + raise NotImplementedError + + @abstractmethod + def set_dof_velocity_targets(self, velocities: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the velocity targets for the specified articulation instances, + if the underlying controller or simulation uses velocity-based control. + + Args: + velocities (torch.Tensor): A tensor of shape (len(indices), dof_count) + specifying desired joint velocities. + indices (torch.Tensor): Indices of articulation instances to apply these targets. + """ + raise NotImplementedError + + @abstractmethod + def set_dof_velocities(self, velocities: torch.Tensor, indices: torch.Tensor) -> None: + """ + Hard-set the joint velocities for the specified articulation instances. + Usually used for resetting or overriding joint states directly. + + Args: + velocities (torch.Tensor): A tensor of shape (len(indices), dof_count) + specifying new joint velocities. + indices (torch.Tensor): Indices of articulation instances to set velocities for. + """ + raise NotImplementedError + + @abstractmethod + def set_root_transforms(self, root_poses_xyzw: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the root transforms (position + orientation) for each articulation instance. + Orientation is expected in (x, y, z, w) format. + + Args: + root_poses_xyzw (torch.Tensor): A tensor of shape (len(indices), 7), + containing [px, py, pz, qx, qy, qz, qw]. + indices (torch.Tensor): Indices of articulation instances to set transforms for. + """ + raise NotImplementedError + + @abstractmethod + def set_dof_stiffnesses(self, stiffness: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the joint stiffness values for the specified articulation instances. + + Args: + stiffness (torch.Tensor): A tensor of shape (len(indices), dof_count) + with new stiffness values. + indices (torch.Tensor): Indices specifying which articulation instances to affect. + """ + raise NotImplementedError + + @abstractmethod + def set_dof_dampings(self, damping: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the joint damping values for the specified articulation instances. + + Args: + damping (torch.Tensor): A tensor of shape (len(indices), dof_count) + with new damping values. + indices (torch.Tensor): Indices specifying which articulation instances to affect. + """ + raise NotImplementedError + + @abstractmethod + def set_dof_armatures(self, armatures: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the joint armature values for the specified articulation instances. + + Args: + armatures (torch.Tensor): A tensor of shape (len(indices), dof_count), + specifying new armature values. + indices (torch.Tensor): Indices specifying which articulation instances to affect. + """ + raise NotImplementedError + + @abstractmethod + def set_dof_friction_coefficients(self, friction_coefficients: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the friction coefficients for each joint of the specified articulation instances. + + Args: + friction_coefficients (torch.Tensor): A tensor of shape (len(indices), dof_count), + specifying new friction values. + indices (torch.Tensor): Indices specifying which articulation instances to affect. + """ + raise NotImplementedError + + @abstractmethod + def set_dof_max_velocities(self, max_velocities: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the maximum allowed velocities for each joint in the specified articulation instances. + + Args: + max_velocities (torch.Tensor): A tensor of shape (len(indices), dof_count), + specifying new velocity limits. + indices (torch.Tensor): Indices specifying which articulation instances to affect. + """ + raise NotImplementedError + + @abstractmethod + def set_dof_max_forces(self, max_forces: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the maximum allowed forces (torques) for each joint in the specified articulation instances. + + Args: + max_forces (torch.Tensor): A tensor of shape (len(indices), dof_count), + specifying new force/torque limits. + indices (torch.Tensor): Indices specifying which articulation instances to affect. + """ + raise NotImplementedError + + @abstractmethod + def set_dof_limits(self, limits: torch.Tensor, indices: torch.Tensor): + """ + Set new position limits (lower/upper) for each joint in the specified articulation instances. + + Args: + limits (torch.Tensor): A tensor of shape (len(indices), dof_count, 2), + specifying [lower_limit, upper_limit] for each joint. + indices (torch.Tensor): Indices specifying which articulation instances to affect. + """ + raise NotImplementedError + + @abstractmethod + def set_fixed_tendon_properties( + self, + fixed_tendon_stiffness: torch.Tensor, + fixed_tendon_damping: torch.Tensor, + fixed_tendon_limit_stiffness: torch.Tensor, + fixed_tendon_limit: torch.Tensor, + fixed_tendon_rest_length: torch.Tensor, + fixed_tendon_offset: torch.Tensor, + indices: torch.Tensor, + ): + """ + Set the properties of fixed tendons (cables) for the specified articulation instances. + + Args: + fixed_tendon_stiffness (torch.Tensor): A tensor of shape (len(indices), max_fixed_tendons), + specifying the stiffness for each tendon. + fixed_tendon_damping (torch.Tensor): A tensor of shape (len(indices), max_fixed_tendons), + specifying the damping for each tendon. + fixed_tendon_limit_stiffness (torch.Tensor): A tensor of shape (len(indices), max_fixed_tendons), + specifying the limit stiffness for each tendon. + fixed_tendon_limit (torch.Tensor): A tensor of shape (len(indices), max_fixed_tendons, 2), + specifying [lower_limit, upper_limit] for each tendon. + fixed_tendon_rest_length (torch.Tensor): A tensor of shape (len(indices), max_fixed_tendons), + specifying the rest length for each tendon. + fixed_tendon_offset (torch.Tensor): A tensor of shape (len(indices), max_fixed_tendons), + specifying the offset for each tendon. + indices (torch.Tensor): Indices specifying which articulation instances to affect. + """ + raise NotImplementedError + + @abstractmethod + def apply_forces_and_torques_at_position( + self, + force_data: torch.Tensor, + torque_data: torch.Tensor, + position_data: torch.Tensor, + indices: torch.Tensor, + is_global: bool, + ) -> None: + """ + Apply forces and torques to bodies at specified positions (in local or global coordinates). + + Args: + force_data (torch.Tensor): A tensor of shape (len(indices), nbodies, 3), + specifying the force to be applied. + torque_data (torch.Tensor): A tensor of shape (len(indices), nbodies, 3), + specifying the torque to be applied. + position_data (torch.Tensor): A tensor of shape (len(indices), nbodies, 3), + specifying the point of application + (in local or world coordinates). + indices (torch.Tensor): Indices specifying which articulation instances to affect. + is_global (bool): + If True, forces/torques are expressed in world/global coordinates. + Otherwise, they are in local link coordinates. + """ + raise NotImplementedError + + @abstractmethod + def _initialize_default_data(self, data): + """ + Initialize the default data for the articulation view. + """ + raise NotImplementedError diff --git a/source/uwlab/uwlab/assets/articulation/articulation_view/articulation_view_cfg.py b/source/uwlab/uwlab/assets/articulation/articulation_view/articulation_view_cfg.py new file mode 100644 index 00000000..20addbf4 --- /dev/null +++ b/source/uwlab/uwlab/assets/articulation/articulation_view/articulation_view_cfg.py @@ -0,0 +1,47 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from dataclasses import MISSING +from typing import TYPE_CHECKING, Callable + +from isaaclab.utils import configclass + +from .articulation_view import ArticulationView +from .bullet_articulation_view import BulletArticulationView + +if TYPE_CHECKING: + from ..articulation_drive import ArticulationDriveCfg + + +@configclass +class ArticulationViewCfg: + """Configuration parameters for an articulation view.""" + + class_type: Callable[[ArticulationViewCfg, str], ArticulationView] = MISSING # type: ignore + + device: str = "cpu" + + +@configclass +class BulletArticulationViewCfg(ArticulationViewCfg): + class_type: Callable[..., BulletArticulationView] = BulletArticulationView + + drive_cfg: ArticulationDriveCfg = MISSING # type: ignore + + urdf: str = MISSING # type: ignore + + debug_visualize: bool = False + + use_multiprocessing: bool = False + + isaac_joint_names: list[str] = MISSING # type: ignore + """Joint names in the Isaac Sim order, this is important when real is executing the action target from isaac sim during + sync mode. If the real environment are run independently, this field is not necessary.""" + + dummy_mode: bool = False + + dt: float = 0.02 diff --git a/source/uwlab/uwlab/assets/articulation/articulation_view/bullet_articulation_view.py b/source/uwlab/uwlab/assets/articulation/articulation_view/bullet_articulation_view.py new file mode 100644 index 00000000..1c6f0d23 --- /dev/null +++ b/source/uwlab/uwlab/assets/articulation/articulation_view/bullet_articulation_view.py @@ -0,0 +1,883 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import queue +import threading +import time +import torch +import torch.multiprocessing as mp +from typing import TYPE_CHECKING + +from . import ArticulationView +from .utils.articulation_kinematics import BulletArticulationKinematics + +if TYPE_CHECKING: + from ..articulation_data import ArticulationData + from ..articulation_drive import ArticulationDrive + from . import BulletArticulationViewCfg, SharedDataSchema + + +def init_shared_data(count: int, ndof: int, nbodies: int) -> SharedDataSchema: + data: SharedDataSchema = { + "is_running": False, + "close": False, + "link_names": [], + "dof_names": [], + "dof_types": [], + "pos": torch.zeros(count, ndof, device="cpu").share_memory_(), + "vel": torch.zeros(count, ndof, device="cpu").share_memory_(), + "torque": torch.zeros(count, ndof, device="cpu").share_memory_(), + "pos_target": torch.zeros(count, ndof, device="cpu").share_memory_(), + "vel_target": torch.zeros(count, ndof, device="cpu").share_memory_(), + "eff_target": torch.zeros(count, ndof, device="cpu").share_memory_(), + "link_transforms": torch.zeros(count, nbodies, 7, device="cpu").share_memory_(), + "link_velocities": torch.zeros(count, nbodies, 6, device="cpu").share_memory_(), + "link_mass": torch.zeros(count, nbodies, device="cpu").share_memory_(), + "link_inertia": torch.zeros(count, nbodies, 9, device="cpu").share_memory_(), + "link_coms": torch.zeros(count, nbodies, 7, device="cpu").share_memory_(), + "mass_matrix": torch.zeros(count, ndof, ndof, device="cpu").share_memory_(), + "dof_stiffness": torch.zeros(count, ndof, device="cpu").share_memory_(), + "dof_armatures": torch.zeros(count, ndof, device="cpu").share_memory_(), + "dof_frictions": torch.zeros(count, ndof, device="cpu").share_memory_(), + "dof_damping": torch.zeros(count, ndof, device="cpu").share_memory_(), + "dof_limits": torch.zeros(count, ndof, 2, device="cpu").share_memory_(), + "dof_max_forces": torch.zeros(count, ndof, device="cpu").share_memory_(), + "dof_max_velocity": torch.zeros(count, ndof, device="cpu").share_memory_(), + "jacobians": torch.zeros(count, nbodies, 6, ndof, device="cpu").share_memory_(), + } + return data + + +class BulletArticulationView(ArticulationView): + """ + LeapHand Implementation of Articulation View + """ + + def __init__( + self, + device: str, + cfg: BulletArticulationViewCfg, + ): + """ + Initializes the articulation view. + """ + # setup hand + self.device = device + self._dt = cfg.dt + self._urdf = cfg.urdf + self._debug_visualize = cfg.debug_visualize + self._isaac_joint_names = cfg.isaac_joint_names + + self._dummy_mode = cfg.dummy_mode + + if cfg.use_multiprocessing: + self.manager = mp.Manager() + self.init_event = mp.Event() + self.cmd_queue = mp.Queue() + self.ack_queue = mp.Queue() + # Create a shared dictionary to communicate with the child process + self.shared_data: SharedDataSchema = self.manager.dict() # type: ignore + self._proc = mp.Process(target=self._run, args=(self.init_event,)) + else: + # Threading shares memory within the same process, so no need for mp.Manager + self.manager = None + self.init_event = threading.Event() + self.cmd_queue = queue.Queue() + self.ack_queue = queue.Queue() + self.shared_data: SharedDataSchema = {} # type: ignore + self._proc = threading.Thread(target=self._run, daemon=True, args=(self.init_event,)) + + # store data that will never be retrieved from kinematic articulation + self.drive_cfg = cfg.drive_cfg + + # Spawn the child process + self._proc.start() + self.init_event.wait() + + def _run(self, init_event: threading.Event): + # Loop until closed + _kinematic = BulletArticulationKinematics(self._urdf, True, self._debug_visualize, dt=self._dt, device="cpu") + self.shared_data.update(init_shared_data(self.count, _kinematic.num_dof, _kinematic.num_links)) + self._populate_shared_data(_kinematic) + _drive: ArticulationDrive = None # type: ignore + if not self._dummy_mode: + _drive = self.drive_cfg.class_type(cfg=self.drive_cfg) + _drive_joint_names = _drive.ordered_joint_names + self._isaac_to_real_idx = [ + self._isaac_joint_names.index(name) for name in _drive_joint_names if name in self._isaac_joint_names + ] + self._real_to_isaac_idx = [ + _drive_joint_names.index(name) for name in self._isaac_joint_names if name in _drive_joint_names + ] + + self._bullet_to_real_idx = [ + self.shared_data["dof_names"].index(name) + for name in _drive_joint_names + if name in self.shared_data["dof_names"] + ] + self._real_to_bullet_idx = [ + _drive_joint_names.index(name) for name in self.shared_data["dof_names"] if name in _drive_joint_names + ] + + self._isaac_to_bullet_idx = [ + self._isaac_joint_names.index(name) + for name in self.shared_data["dof_names"] + if name in self._isaac_joint_names + ] + self._bullet_to_isaac_idx = [ + self.shared_data["dof_names"].index(name) + for name in self._isaac_joint_names + if name in self.shared_data["dof_names"] + ] + + init_event.set() + + next_poll_time = time.time() + self._dt + while True: + if self.shared_data["close"]: + break + + self._process_blocking_commands(_kinematic, _drive) + + if self.shared_data["is_running"]: + now = time.time() + if now >= next_poll_time: + next_poll_time = now + self._dt + if self._dummy_mode: + self._read_write_dummy_states() + _kinematic.set_dof_states( + self.shared_data["pos"][:, self._isaac_to_bullet_idx], + self.shared_data["vel"][:, self._isaac_to_bullet_idx], + self.shared_data["torque"][:, self._isaac_to_bullet_idx], + ) + + self.shared_data["link_coms"][:] = _kinematic.get_link_coms() + self.shared_data["link_transforms"][:] = _kinematic.get_link_transforms() + self.shared_data["link_velocities"][:] = _kinematic.get_link_velocities() + else: + # 1) read pos/vel from hardware + pos, vel, eff = _drive.read_dof_states() + self.shared_data["pos"][:, self._real_to_isaac_idx] = pos[:, self._real_to_isaac_idx] + self.shared_data["vel"][:, self._real_to_isaac_idx] = vel[:, self._real_to_isaac_idx] + self.shared_data["torque"][:, self._real_to_isaac_idx] = eff[:, self._real_to_isaac_idx] + + _kinematic.set_dof_states( + self.shared_data["pos"][:, self._isaac_to_bullet_idx], + self.shared_data["vel"][:, self._isaac_to_bullet_idx], + self.shared_data["torque"][:, self._isaac_to_bullet_idx], + ) + + self.shared_data["link_coms"][:] = _kinematic.get_link_coms() + self.shared_data["link_transforms"][:] = _kinematic.get_link_transforms() + self.shared_data["link_velocities"][:] = _kinematic.get_link_velocities() + + # 2) write target to hardware and kinematic + _kinematic.set_dof_targets( + self.shared_data["pos_target"][:, self._isaac_to_bullet_idx], + self.shared_data["vel_target"][:, self._isaac_to_bullet_idx], + self.shared_data["eff_target"][:, self._isaac_to_bullet_idx], + ) + _drive.write_dof_targets( + pos_target=self.shared_data["pos_target"][:, self._isaac_to_real_idx], + vel_target=self.shared_data["vel_target"][:, self._isaac_to_real_idx], + eff_target=self.shared_data["eff_target"][:, self._isaac_to_real_idx], + ) + if self._debug_visualize: + _kinematic.render() + # Sleep till next poll time + time.sleep(max(next_poll_time - time.time(), 0)) + else: + # now is less than next poll time, sleep till next poll time + time.sleep(next_poll_time - now) + else: + # If not running, keep sleeping + time.sleep(self._dt) + + # Done, disconnect hardware + _drive.close() + _kinematic.close() + print("DynamixelWorker: Child process stopped") + + def _process_blocking_commands(self, _kinematic: BulletArticulationKinematics, _drive: ArticulationDrive): + if not self.cmd_queue.empty(): + cmd = self.cmd_queue.get() + match cmd: + case "set_dof_stiffnesses": + _kinematic.set_dof_stiffnesses(self.shared_data["dof_stiffness"][:, self._isaac_to_bullet_idx]) + if not self._dummy_mode: + _drive.set_dof_stiffnesses(self.shared_data["dof_stiffness"][:, self._bullet_to_real_idx]) + self.ack_queue.put({"status": "OK", "command": "set_dof_stiffnesses"}) + case "set_dof_armatures": + _kinematic.set_dof_armatures(self.shared_data["dof_armatures"][:, self._isaac_to_bullet_idx]) + if not self._dummy_mode: + _drive.set_dof_armatures(self.shared_data["dof_armatures"][:, self._bullet_to_real_idx]) + self.ack_queue.put({"status": "OK", "command": "set_dof_armatures"}) + case "set_dof_frictions": + _kinematic.set_dof_frictions(self.shared_data["dof_frictions"][:, self._isaac_to_bullet_idx]) + if not self._dummy_mode: + _drive.set_dof_frictions(self.shared_data["dof_frictions"][:, self._bullet_to_real_idx]) + self.ack_queue.put({"status": "OK", "command": "set_dof_frictions"}) + case "set_dof_dampings": + _kinematic.set_dof_dampings(self.shared_data["dof_damping"][:, self._isaac_to_bullet_idx]) + if not self._dummy_mode: + _drive.set_dof_dampings(self.shared_data["dof_damping"][:, self._bullet_to_real_idx]) + self.ack_queue.put({"status": "OK", "command": "set_dof_dampings"}) + case "set_dof_limits": + _kinematic.set_dof_limits(self.shared_data["dof_limits"][:, self._isaac_to_bullet_idx]) + if not self._dummy_mode: + _drive.set_dof_limits(self.shared_data["dof_limits"][:, self._bullet_to_real_idx]) + self.ack_queue.put({"status": "OK", "command": "set_dof_limits"}) + + def _populate_shared_data(self, _kinematic: BulletArticulationKinematics): + self.shared_data["link_names"] = _kinematic.link_names + self.shared_data["dof_names"] = _kinematic.joint_names + + self.shared_data["pos"][:] = _kinematic.get_dof_positions(clone=False) + self.shared_data["vel"][:] = _kinematic.get_dof_velocities(clone=False) + self.shared_data["torque"][:] = _kinematic.get_dof_torques(clone=False) + self.shared_data["pos_target"][:] = _kinematic.get_dof_position_targets(clone=False) + self.shared_data["vel_target"][:] = _kinematic.get_dof_velocity_targets(clone=False) + # self.shared_data["eff_target"][:] + + self.shared_data["link_transforms"][:] = _kinematic.get_link_transforms(clone=False) + self.shared_data["link_velocities"][:] = _kinematic.get_link_velocities(clone=False) + self.shared_data["link_mass"][:] = _kinematic.get_link_masses(clone=False) + self.shared_data["link_inertia"][:] = _kinematic.get_link_inertias(clone=False) + self.shared_data["link_coms"][:] = _kinematic.get_link_coms(clone=False) + self.shared_data["mass_matrix"][:] = _kinematic.get_mass_matrix() + self.shared_data["dof_stiffness"][:] = _kinematic.get_dof_stiffnesses(clone=False) + self.shared_data["dof_armatures"][:] = _kinematic.get_dof_armatures(clone=False) + self.shared_data["dof_frictions"][:] = _kinematic.get_dof_frictions(clone=False) + self.shared_data["dof_damping"][:] = _kinematic.get_dof_dampings(clone=False) + self.shared_data["dof_limits"][:] = _kinematic.get_dof_limits(clone=False) + self.shared_data["dof_max_forces"][:] = _kinematic.get_dof_max_forces(clone=False) + self.shared_data["dof_max_velocity"][:] = _kinematic.get_dof_max_velocities(clone=False) + self.shared_data["jacobians"][:] = _kinematic.get_jacobian() + + def _read_write_dummy_states(self): + self.shared_data["pos"] = self.shared_data["pos_target"] + self.shared_data["vel"] = self.shared_data["vel_target"] + self.shared_data["torque"] = self.shared_data["eff_target"] + + # Public API + # ------------------------------------------------------------------------- + # Rules: + # Share the Data, not Object! + # Below method should not be called in the child process, namely self._run, + # Implementation of below methods should not directly access fields used for child process, e.g: + # _kinematic + # _drive + + def play(self): + self.shared_data["is_running"] = True + + def pause(self): + self.shared_data["is_running"] = False + + def close(self): + self.shared_data["is_running"] = False + self.shared_data["close"] = True + if self._proc is not None: + self._proc.join() + self._proc = None + + if self.manager is not None: + self.manager.shutdown() + self.manager = None + + @property + def count(self) -> int: + """ + Number of articulation instances being managed by this view. + + E.g., if your environment has N separate copies of the same robot, + this property would return N. + + Returns: + int: The total number of articulation instances. + """ + return 1 + + @property + def fixed_base(self) -> bool: + """ + Indicates whether the articulation(s) in this view has a fixed base. + + A fixed-base articulation does not move freely in space + (e.g., a robotic arm bolted to a table), + whereas a floating-base articulation can move freely. + + Returns: + bool: True if the articulation is fixed-base, False otherwise. + """ + return True + + @property + def dof_count(self) -> int: + """ + Number of degrees of freedom (DOFs) for the articulation(s) in this view. + + Returns: + int: Count of all DOFs for the managed articulation(s). + """ + return self.shared_data["dof_names"].__len__() + + @property + def max_fixed_tendons(self) -> int: + """ + Maximum number of 'fixed tendons' (sometimes known as cables or passively constrained joints). + + Returns: + int: The number of fixed tendon connections in the articulation(s). + """ + return 0 + + @property + def num_bodies(self) -> int: + """ + Number of rigid bodies (links) in the articulation(s). + + Returns: + int: Total number of rigid bodies in the articulation(s). + """ + return self.shared_data["link_names"].__len__() + + @property + def joint_names(self) -> list[str]: + """ + Ordered list of joint names in the articulation(s). + + Returns: + list[str]: Names of the joints in order of their DOF indices. + """ + return self.shared_data["dof_names"].copy() + + @property + def fixed_tendon_names(self) -> list[str]: + """ + Ordered list of names for the fixed tendons (cables) in the articulation(s). + + Returns: + list[str]: Names of the fixed tendons, if any. + """ + return [] + + @property + def body_names(self) -> list[str]: + """ + Ordered list of body (link) names in the articulation(s). + + Returns: + list[str]: Names of the rigid bodies in order of their indices. + """ + return self.shared_data["link_names"].copy() + + def get_root_transforms(self) -> torch.Tensor: + """ + Get the root poses (position + orientation) of each articulation instance. + + This typically refers to the base link or root transform + of a floating/fixed-base articulation in world coordinates. + + Returns: + torch.Tensor: A tensor of shape (count, 7), + each row containing [px, py, pz, qw, qx, qy, qz]. + """ + return self.shared_data["link_transforms"][:, 0, :].clone().to(self.device) + + def get_root_velocities(self) -> torch.Tensor: + """ + Get the linear and angular velocities of the root link(s) of each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, 6), + each row containing [vx, vy, vz, wx, wy, wz]. + """ + return self.shared_data["link_velocities"][:, 0, :].clone().to(self.device) + + def get_link_accelerations(self) -> torch.Tensor: + """ + Get the link accelerations for each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, num_bodies, 6), + containing [ax, ay, az, alpha_x, alpha_y, alpha_z] for each link. + """ + print("getter function: get_link_accelerations, result are place holder values, do not rely on them") + return torch.zeros((self.count, self.num_bodies, 6), dtype=torch.float32, device=self.device) + + def get_link_transforms(self) -> torch.Tensor: + """ + Get the transforms (position + orientation) of each link for each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, num_bodies, 7), + for [px, py, pz, qw, qx, qy, qz] per link. + """ + return self.shared_data["link_transforms"].clone().to(self.device) + + def get_link_velocities(self) -> torch.Tensor: + """ + Get the linear and angular velocities of each link in the articulation(s). + + Returns: + torch.Tensor: A tensor of shape (count, num_bodies, 6), + containing [vx, vy, vz, wx, wy, wz] per link. + """ + return self.shared_data["link_velocities"].clone().to(self.device) + + def get_coms(self) -> torch.Tensor: + """ + Get the center-of-mass (COM) positions of each link or the entire articulation. + + Depending on the implementation, this can mean: + - Per-link COM (shape = (count, num_bodies, 7)) + - Single COM for the entire system (shape = (count, 7)) + + Returns: + torch.Tensor: COM positions in the world or local coordinate frame + depending on the backend's convention. + """ + return self.shared_data["link_coms"].clone().to(self.device) + + def get_masses(self) -> torch.Tensor: + """ + Get the masses for each link in the articulation(s). + + Returns: + torch.Tensor: A tensor of shape (count, num_bodies), + containing the mass of each link. + """ + link_mass = self.shared_data["link_mass"].clone().to(self.device) + return link_mass + + def get_inertias(self) -> torch.Tensor: + """ + Get the inertial tensors (often expressed in link-local frames) for each link. + + Returns: + torch.Tensor: A tensor of shape (count, num_bodies, 9), + containing the inertia matrix for each link. + """ + return self.shared_data["link_inertia"].to(self.device) + + def get_dof_positions(self) -> torch.Tensor: + """ + Get the joint positions for each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, dof_count) + with joint angles or positions. + """ + return self.shared_data["pos"].clone().to(self.device) + + def get_dof_velocities(self) -> torch.Tensor: + """ + Get the joint velocities for each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, dof_count), + with joint velocity values. + """ + return self.shared_data["vel"].clone().to(self.device) + + def get_dof_torques(self) -> torch.Tensor: + """ + Get the joint torques for each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, dof_count), + with joint velocity values. + """ + return self.shared_data["torque"].clone().to(self.device) + + def get_dof_max_velocities(self) -> torch.Tensor: + """ + Get the maximum velocity limits for each joint in each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, dof_count), + containing velocity limits per joint. + """ + return self.shared_data["dof_max_velocity"].clone().to(self.device) + + def get_dof_max_forces(self) -> torch.Tensor: + """ + Get the maximum force (torque) limits for each joint in each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, dof_count), + with torque/force limits per joint. + """ + # max force expects cpu tensor + return self.shared_data["dof_max_forces"].clone().to("cpu") + + def get_dof_stiffnesses(self) -> torch.Tensor: + """ + Get the joint stiffness values for each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, dof_count), + containing the stiffness for each DOF. + """ + return self.shared_data["dof_stiffness"].clone().to(self.device) + + def get_dof_dampings(self) -> torch.Tensor: + """ + Get the joint damping values for each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, dof_count), + containing the damping for each DOF. + """ + return self.shared_data["dof_damping"].clone().to(self.device) + + def get_dof_armatures(self) -> torch.Tensor: + """ + Get the armature values for each joint in each articulation instance. + + The 'armature' is sometimes used to represent an inertia-like term + used in certain simulation backends or real hardware compensations. + + Returns: + torch.Tensor: A tensor of shape (count, dof_count). + """ + return self.shared_data["dof_armatures"].clone().to(self.device) + + def get_dof_friction_coefficients(self) -> torch.Tensor: + """ + Get the friction coefficients for each joint in each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, dof_count). + """ + return self.shared_data["dof_frictions"].clone().to(self.device) + + def get_dof_limits(self) -> torch.Tensor: + """ + Get the joint position limits for each joint in each articulation instance. + + Returns: + torch.Tensor: A tensor of shape (count, dof_count, 2), + where the last dimension stores [lower_limit, upper_limit]. + """ + return self.shared_data["dof_limits"].clone().to(self.device) + + def get_fixed_tendon_stiffnesses(self) -> torch.Tensor: + """ + Get the stiffness values for each fixed tendon across the articulation(s). + + Returns: + torch.Tensor: A tensor of shape (count, max_fixed_tendons), + with the stiffness value for each tendon. + """ + print("getter function: get_fixed_tendon_stiffnesses, result are place holder values, do not rely on them") + return torch.zeros((self.count, self.dof_count), dtype=torch.float32, device=self.device) + + def get_fixed_tendon_dampings(self) -> torch.Tensor: + """ + Get the damping values for each fixed tendon across the articulation(s). + + Returns: + torch.Tensor: A tensor of shape (count, max_fixed_tendons), + with the damping value for each tendon. + """ + print("getter function: get_fixed_tendon_dampings, result are place holder values, do not rely on them") + return torch.zeros((self.count, self.dof_count), dtype=torch.float32, device=self.device) + + def get_fixed_tendon_limit_stiffnesses(self) -> torch.Tensor: + """ + Get the limit stiffness values for each fixed tendon. + + Returns: + torch.Tensor: A tensor of shape (count, max_fixed_tendons), + with the limit stiffness for each tendon. + """ + print( + "getter function: get_fixed_tendon_limit_stiffnesses, result are place holder values, do not rely on them" + ) + return torch.zeros((self.count, self.dof_count), dtype=torch.float32, device=self.device) + + def get_fixed_tendon_limits(self) -> torch.Tensor: + """ + Get the limit range for each fixed tendon, which might represent + min/max constraints on the tendon length or tension. + + Returns: + torch.Tensor: A tensor of shape (count, max_fixed_tendons, 2), + containing [lower_limit, upper_limit] for each tendon. + """ + print("getter function: get_fixed_tendon_limits, result are place holder values, do not rely on them") + return torch.zeros((self.count, self.dof_count), dtype=torch.float32, device=self.device) + + def get_fixed_tendon_rest_lengths(self) -> torch.Tensor: + """ + Get the rest lengths for each fixed tendon across the articulation(s). + + Returns: + torch.Tensor: A tensor of shape (count, max_fixed_tendons). + """ + print("getter function: get_fixed_tendon_rest_lengths, result are place holder values, do not rely on them") + return torch.zeros((self.count, self.dof_count), dtype=torch.float32, device=self.device) + + def get_fixed_tendon_offsets(self) -> torch.Tensor: + """ + Get the offset values for each fixed tendon across the articulation(s). + + Returns: + torch.Tensor: A tensor of shape (count, max_fixed_tendons). + """ + print("getter function: get_fixed_tendon_offsets, result are place holder values, do not rely on them") + return torch.zeros((self.count, self.dof_count), dtype=torch.float32, device=self.device) + + def set_dof_actuation_forces(self, forces: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the actuation forces (torques) for the specified articulation instances. + + Args: + forces (torch.Tensor): A tensor of shape (len(indices), dof_count) + specifying the commanded forces/torques. + indices (torch.Tensor): A tensor of indices specifying which + articulation instances to apply these forces. + """ + if torch.any(forces != 0): + print( + "calling placeholder function: set_dof_actuation_forces with non-zero values, but function is not" + " implemented" + ) + + def set_dof_position_targets(self, positions: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the position targets for the specified articulation instances, + if the underlying controller or simulation uses position-based control. + + Args: + positions (torch.Tensor): A tensor of shape (len(indices), dof_count) + specifying desired joint positions. + indices (torch.Tensor): Indices of articulation instances to apply these targets. + """ + self.shared_data["pos_target"][:] = positions.cpu() + + def set_dof_positions(self, positions: torch.Tensor, indices: torch.Tensor, threshold: float = 1e-2) -> None: + """ + Hard-set the joint positions for the specified articulation instances. + Usually used for resetting or overriding joint states directly. + + Args: + positions (torch.Tensor): A tensor of shape (len(indices), dof_count) + specifying the new positions. + indices (torch.Tensor): Indices of articulation instances to set positions for. + """ + count = 0 + while torch.sum(self.get_dof_positions() - positions).abs() > threshold: + self.set_dof_position_targets(positions, indices) + count += 1 + if count > 50: + break + + def set_dof_velocity_targets(self, velocities: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the velocity targets for the specified articulation instances, + if the underlying controller or simulation uses velocity-based control. + + Args: + velocities (torch.Tensor): A tensor of shape (len(indices), dof_count) + specifying desired joint velocities. + indices (torch.Tensor): Indices of articulation instances to apply these targets. + """ + if torch.any(velocities != 0): + print( + "calling placeholder function: set_dof_velocity_targets with non-zero values, but function is not" + " implemented" + ) + # the interface is correct but driver doesn't support velocity control yet + self.shared_data["vel_target"][:] = velocities.cpu() + + def set_dof_velocities(self, velocities: torch.Tensor, indices: torch.Tensor) -> None: + """ + Hard-set the joint velocities for the specified articulation instances. + Usually used for resetting or overriding joint states directly. + + Args: + velocities (torch.Tensor): A tensor of shape (len(indices), dof_count) + specifying new joint velocities. + indices (torch.Tensor): Indices of articulation instances to set velocities for. + """ + if torch.any(velocities != 0): + print( + "calling placeholder function: set_dof_velocities with non-zero values, but function is not implemented" + ) + pass + + def set_root_transforms(self, root_poses_xyzw: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the root transforms (position + orientation) for each articulation instance. + Orientation is expected in (x, y, z, w) format. + + Args: + root_poses_xyzw (torch.Tensor): A tensor of shape (len(indices), 7), + containing [px, py, pz, qx, qy, qz, qw]. + indices (torch.Tensor): Indices of articulation instances to set transforms for. + """ + if torch.any(root_poses_xyzw != 0): + print( + "calling placeholder function: set_root_transforms with non-zero values, but function is not" + " implemented" + ) + pass + + def set_root_velocities(self, root_velocities: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the root velocities (linvel + angvel) for each articulation instance. + + Args: + root_velocities (torch.Tensor): A tensor of shape (len(indices), 6), + containing [x, y, z, rx, ry, rz]. + indices (torch.Tensor): Indices of articulation instances to set transforms for. + """ + if torch.any(root_velocities != 0): + print( + "calling placeholder function: set_root_transforms with non-zero values, but function is not" + " implemented" + ) + pass + + def set_dof_stiffnesses(self, stiffness: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the joint stiffness values for the specified articulation instances. + + Args: + stiffness (torch.Tensor): A tensor of shape (len(indices), dof_count) + with new stiffness values. + indices (torch.Tensor): Indices specifying which articulation instances to affect. + """ + self.shared_data["dof_stiffness"][:] = stiffness.cpu() + self.cmd_queue.put("set_dof_stiffnesses") + ack = self.ack_queue.get() + if ack != "OK": + print(f"Warning: Child returned ack={ack}") + + def set_dof_dampings(self, damping: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the joint damping values for the specified articulation instances. + + Args: + damping (torch.Tensor): A tensor of shape (len(indices), dof_count) + with new damping values. + indices (torch.Tensor): Indices specifying which articulation instances to affect. + """ + self.shared_data["dof_damping"][:] = damping.cpu() + self.cmd_queue.put("set_dof_dampings") + ack = self.ack_queue.get() + if ack != "OK": + print(f"Warning: Child returned ack={ack}") + + def set_dof_armatures(self, armatures: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the joint armature values for the specified articulation instances. + + Args: + armatures (torch.Tensor): A tensor of shape (len(indices), dof_count), + specifying new armature values. + indices (torch.Tensor): Indices specifying which articulation instances to affect. + """ + self.shared_data["dof_armatures"][:] = armatures.cpu() + self.cmd_queue.put("set_dof_armatures") + ack = self.ack_queue.get() + if ack != "OK": + print(f"Warning: Child returned ack={ack}") + + def set_dof_friction_coefficients(self, friction_coefficients: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the friction coefficients for each joint of the specified articulation instances. + + Args: + friction_coefficients (torch.Tensor): A tensor of shape (len(indices), dof_count), + specifying new friction values. + indices (torch.Tensor): Indices specifying which articulation instances to affect. + """ + self.shared_data["dof_frictions"][:] = friction_coefficients.cpu() + self.cmd_queue.put("set_dof_frictions") + ack = self.ack_queue.get() + if ack != "OK": + print(f"Warning: Child returned ack={ack}") + + def set_dof_max_velocities(self, max_velocities: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the maximum allowed velocities for each joint in the specified articulation instances. + + Args: + max_velocities (torch.Tensor): A tensor of shape (len(indices), dof_count), + specifying new velocity limits. + indices (torch.Tensor): Indices specifying which articulation instances to affect. + """ + if torch.any(max_velocities != 0): + print( + "calling placeholder function: set_dof_max_velocities with non-zero values, but function is not" + " implemented" + ) + pass + + def set_dof_max_forces(self, max_forces: torch.Tensor, indices: torch.Tensor) -> None: + """ + Set the maximum allowed forces (torques) for each joint in the specified articulation instances. + + Args: + max_forces (torch.Tensor): A tensor of shape (len(indices), dof_count), + specifying new force/torque limits. + indices (torch.Tensor): Indices specifying which articulation instances to affect. + """ + if torch.any(max_forces != 0): + print( + "calling placeholder function: set_dof_max_forces with non-zero values, but function is not implemented" + ) + pass + + def set_dof_limits(self, limits: torch.Tensor, indices: torch.Tensor): + """ + Set new position limits (lower/upper) for each joint in the specified articulation instances. + + Args: + limits (torch.Tensor): A tensor of shape (len(indices), dof_count, 2), + specifying [lower_limit, upper_limit] for each joint. + indices (torch.Tensor): Indices specifying which articulation instances to affect. + """ + self.shared_data["dof_limits"][:] = limits.cpu() + self.cmd_queue.put("set_dof_limits") + ack = self.ack_queue.get() + if ack != "OK": + print(f"Warning: Child returned ack={ack}") + + def apply_forces_and_torques_at_position( + self, + force_data: torch.Tensor, + torque_data: torch.Tensor, + position_data: torch.Tensor, + indices: torch.Tensor, + is_global: bool, + ) -> None: + """ + Apply forces and torques to bodies at specified positions (in local or global coordinates). + + Args: + force_data (torch.Tensor): A tensor of shape (len(indices), nbodies, 3), + specifying the force to be applied. + torque_data (torch.Tensor): A tensor of shape (len(indices), nbodies, 3), + specifying the torque to be applied. + position_data (torch.Tensor): A tensor of shape (len(indices), nbodies, 3), + specifying the point of application + (in local or world coordinates). + indices (torch.Tensor): Indices specifying which articulation instances to affect. + is_global (bool): + If True, forces/torques are expressed in world/global coordinates. + Otherwise, they are in local link coordinates. + """ + raise NotImplementedError("Real does not support applying virtual forces.") + + def _initialize_default_data(self, data: ArticulationData): + self.set_dof_position_targets(data.default_joint_pos, indices=torch.arange(self.count)) + self.set_dof_velocity_targets(data.default_joint_vel, indices=torch.arange(self.count)) diff --git a/source/uwlab/uwlab/assets/articulation/articulation_view/utils/articulation_kinematics.py b/source/uwlab/uwlab/assets/articulation/articulation_view/utils/articulation_kinematics.py new file mode 100644 index 00000000..a7bbfa7c --- /dev/null +++ b/source/uwlab/uwlab/assets/articulation/articulation_view/utils/articulation_kinematics.py @@ -0,0 +1,637 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch +from typing import List, Union + +import pybullet as p + +from isaaclab.utils import math as math_utils + + +class BulletArticulationKinematicsData: + def __init__(self, dof_dim: int, link_dim: int, device: str): + self.dof_dim = dof_dim + self.link_dim = link_dim + self.device = device + self.reset() + + def reset(self): + self.link_names: list[str] = [] + self.dof_names: list[str] = [] + self.dof_types: list[str] = [] + self.dof_indices = torch.zeros((1, self.dof_dim), device=self.device) + + self.link_transforms = torch.zeros((1, self.link_dim, 7), device=self.device) + self.link_velocities = torch.zeros((1, self.link_dim, 6), device=self.device) + self.link_mass = torch.zeros((1, self.link_dim), device=self.device) + self.link_inertia = torch.zeros((1, self.link_dim, 9), device=self.device) + self.link_coms = torch.zeros((1, self.link_dim, 7), device=self.device) + + self.mass_matrix = torch.zeros((1, self.dof_dim, self.dof_dim), device=self.device) + + self.dof_positions = torch.zeros((1, self.dof_dim), device=self.device) + self.dof_velocities = torch.zeros((1, self.dof_dim), device=self.device) + self.dof_accelerations = torch.zeros((1, self.dof_dim), device=self.device) + self.dof_torques = torch.zeros((1, self.dof_dim), device=self.device) + + self.dof_position_target = torch.zeros((1, self.dof_dim), device=self.device) + self.dof_velocity_target = torch.zeros((1, self.dof_dim), device=self.device) + self.dof_torque_target = torch.zeros((1, self.dof_dim), device=self.device) + + self.dof_stiffness = torch.zeros((1, self.dof_dim), device=self.device) + self.dof_armatures = torch.zeros((1, self.dof_dim), device=self.device) + self.dof_frictions = torch.zeros((1, self.dof_dim), device=self.device) + self.dof_damping = torch.zeros((1, self.dof_dim), device=self.device) + self.dof_limits = torch.zeros((1, self.dof_dim, 2), device=self.device) + self.dof_max_forces = torch.zeros((1, self.dof_dim), device=self.device) + self.dof_max_velocity = torch.zeros((1, self.dof_dim), device=self.device) + + self.jacobians = torch.zeros((1, self.link_dim, 6, self.dof_dim), device=self.device) + + +dof_types_dict: dict[int, str] = { + p.JOINT_REVOLUTE: "revolute", + p.JOINT_PRISMATIC: "prismatic", + p.JOINT_SPHERICAL: "spherical", + p.JOINT_PLANAR: "planar", + p.JOINT_FIXED: "fixed", +} + + +class BulletArticulationKinematics: + def __init__(self, urdf_path, is_fixed_base, debug_visualize, dt, device): + """ + Initialize PyBullet in DIRECT mode and load the URDF at urdf_path. + """ + # Connect in DIRECT mode (no GUI) + self.debug_visualize = debug_visualize + if debug_visualize: + self.client_id = p.connect(p.GUI) + else: + self.client_id = p.connect(p.DIRECT) + + self._is_fixed_base = is_fixed_base + # Load URDF, useFixedBase=True to keep it from falling due to gravity + self.articulation = p.loadURDF(urdf_path, useFixedBase=is_fixed_base, physicsClientId=self.client_id) + # Query how many joints (which also defines how many child links). + self._num_joints = p.getNumJoints(self.articulation, physicsClientId=self.client_id) + self._num_links = self._num_joints + 1 + + # Identify which joints are actually movable (revolute or prismatic) + self._dof_indicies = [] + for j in range(self._num_joints): + info = p.getJointInfo(self.articulation, j, physicsClientId=self.client_id) + joint_type = info[2] # 0=REVOLUTE, 1=PRISMATIC, 4=FIXED, ... + if joint_type not in [p.JOINT_FIXED]: + self._dof_indicies.append(j) + + # Number of degrees of freedom (movable joints) + self._num_dofs = len(self._dof_indicies) + print("Total Links:", self._num_links, "Total joints:", self._num_joints, "Movable DoF:", self._num_dofs) + + # Initialize the state storage + self.articulation_view_data = BulletArticulationKinematicsData(self._num_dofs, self._num_links, device) + self.device = device + + self.populate_joint_state() + self.populate_link_transforms() + self.populate_link_velocities() + self.populate_link_mass() + self.populate_inertia() + + def close(self): + """ + Signal the rendering thread to stop and wait for it to join. + """ + self.pause = True + + @property + def fixed_base(self): + return self._is_fixed_base + + @property + def num_links(self): + return self._num_links + + @property + def num_dof(self): + return self._num_dofs + + @property + def joint_names(self): + return self.articulation_view_data.dof_names + + @property + def link_names(self): + return self.articulation_view_data.link_names + + def render(self): + p.stepSimulation(physicsClientId=self.client_id) + + def get_root_link_transform(self, clone: bool = True) -> torch.Tensor: + data = self.articulation_view_data.link_transforms[:, 0] + return data.clone() if clone else data + + def get_root_link_velocity(self, clone: bool = True) -> torch.Tensor: + data = self.articulation_view_data.link_velocities[:, 0] + return data.clone() if clone else data + + def get_link_transforms( + self, + body_indices: Union[List[int], torch.Tensor, slice] = slice(None), + clone: bool = True, + ) -> torch.Tensor: + data = self.articulation_view_data.link_transforms[:, body_indices] + return data.clone() if clone else data + + def get_link_velocities( + self, + body_indices: Union[List[int], torch.Tensor, slice] = slice(None), + clone: bool = True, + ) -> torch.Tensor: + data = self.articulation_view_data.link_velocities[:, body_indices] + return data.clone() if clone else data + + def get_link_coms( + self, + body_indices: Union[List[int], torch.Tensor, slice] = slice(None), + clone: bool = True, + ) -> torch.Tensor: + """in body local frames""" + data = self.articulation_view_data.link_coms[:, body_indices] + return data.clone() if clone else data + + def get_link_masses( + self, + body_indices: Union[List[int], torch.Tensor, slice] = slice(None), + clone: bool = True, + ) -> torch.Tensor: + data = self.articulation_view_data.link_mass[:, body_indices] + return data.clone() if clone else data + + def get_link_inertias( + self, + body_indices: Union[List[int], torch.Tensor, slice] = slice(None), + clone: bool = True, + ) -> torch.Tensor: + data = self.articulation_view_data.link_inertia[:, body_indices] + return data.clone() if clone else data + + def get_dof_limits( + self, + body_indices: Union[List[int], torch.Tensor, slice] = slice(None), + clone: bool = True, + ) -> torch.Tensor: + data = self.articulation_view_data.dof_limits[:, :, body_indices] + return data.clone() if clone else data + + def get_dof_positions( + self, + body_indices: Union[List[int], torch.Tensor, slice] = slice(None), + clone: bool = True, + ) -> torch.Tensor: + data = self.articulation_view_data.dof_positions[:, body_indices] + return data.clone() if clone else data + + def get_dof_position_targets( + self, + body_indices: Union[List[int], torch.Tensor, slice] = slice(None), + clone: bool = True, + ) -> torch.Tensor: + data = self.articulation_view_data.dof_position_target[:, body_indices] + return data.clone() if clone else data + + def get_dof_velocities( + self, + body_indices: Union[List[int], torch.Tensor, slice] = slice(None), + clone: bool = True, + ) -> torch.Tensor: + data = self.articulation_view_data.dof_velocities[:, body_indices] + return data.clone() if clone else data + + def get_dof_velocity_targets( + self, + body_indices: Union[List[int], torch.Tensor, slice] = slice(None), + clone: bool = True, + ) -> torch.Tensor: + data = self.articulation_view_data.dof_velocity_target[:, body_indices] + return data.clone() if clone else data + + def get_dof_max_velocities( + self, + joint_indices: Union[List[int], torch.Tensor, slice] = slice(None), + clone: bool = True, + ) -> torch.Tensor: + data = self.articulation_view_data.dof_max_velocity[:, joint_indices] + return data.clone() if clone else data + + def get_dof_torques( + self, + joint_indices: Union[List[int], torch.Tensor, slice] = slice(None), + clone: bool = True, + ) -> torch.Tensor: + data = self.articulation_view_data.dof_torques[:, joint_indices] + return data.clone() if clone else data + + def get_dof_max_forces( + self, + joint_indices: Union[List[int], torch.Tensor, slice] = slice(None), + clone: bool = True, + ) -> torch.Tensor: + data = self.articulation_view_data.dof_max_forces[:, joint_indices] + return data.clone() if clone else data + + def get_dof_stiffnesses( + self, + joint_indices: Union[List[int], torch.Tensor, slice] = slice(None), + clone: bool = True, + ): + data = self.articulation_view_data.dof_stiffness[:, joint_indices] + return data.clone() if clone else data + + def get_dof_dampings( + self, + joint_indices: Union[List[int], torch.Tensor, slice] = slice(None), + clone: bool = True, + ) -> torch.Tensor: + data = self.articulation_view_data.dof_damping[:, joint_indices] + return data.clone() if clone else data + + def get_dof_frictions( + self, + joint_indices: Union[List[int], torch.Tensor, slice] = slice(None), + clone: bool = True, + ): + data = self.articulation_view_data.dof_frictions[:, joint_indices] + return data.clone() if clone else data + + def get_dof_armatures( + self, + joint_indices: Union[List[int], torch.Tensor, slice] = slice(None), + clone: bool = True, + ): + data = self.articulation_view_data.dof_armatures[:, joint_indices] + return data.clone() if clone else data + + def get_mass_matrix(self) -> torch.Tensor: + """ + Return the mass matrix for the current state. + """ + # 1) Gather the current joint positions from PyBullet + joint_positions = self.get_dof_positions() + + # 2) Compute the mass matrix + mass_matrix = p.calculateMassMatrix( + bodyUniqueId=self.articulation, objPositions=joint_positions[0].tolist(), physicsClientId=self.client_id + ) + + # 3) Convert the mass matrix to a torch tensor + mass_matrix = torch.tensor(mass_matrix, device=self.device) + + return mass_matrix + + def get_jacobian(self) -> torch.Tensor: + """ + Return the Jacobian for each link in shape (num_links, 6, num_joints). + + - num_links + - 6 = [dPos/dq (3 rows), dRot/dq (3 rows)] + - num_joints = total joints + """ + # 1) Gather the current joint positions from PyBullet + joint_positions = self.get_dof_positions() + + # 2) For each link, calculate the Jacobian + jacobians = torch.zeros((1, self._num_links, 6, self._num_dofs), device=self.device) + for link_idx in range(0, self._num_links): + linJ, angJ = p.calculateJacobian( + bodyUniqueId=self.articulation, + linkIndex=link_idx - 1, + localPosition=[0, 0, 0], + objPositions=joint_positions[0].tolist(), + objVelocities=torch.zeros_like(joint_positions)[0].tolist(), + objAccelerations=torch.zeros_like(joint_positions)[0].tolist(), + physicsClientId=self.client_id, + ) + # linJ, angJ are (3, num_joints) + linJ = torch.tensor(linJ, device=self.device) + angJ = torch.tensor(angJ, device=self.device) + jacobians[0, link_idx, :3, :] = linJ + jacobians[0, link_idx, 3:, :] = angJ + + return jacobians + + def set_masses( + self, + masses: torch.Tensor, + indices: Union[List[int], torch.Tensor, slice] = slice(None), + link_indices: Union[List[int], torch.Tensor, slice] = slice(None), + ): + """ + Set the mass of each link in PyBullet at runtime using p.changeDynamics. + NOTE: If you pass an invalid mass for the base (index -1), PyBullet might ignore it + or throw an error. + """ + # Transfer the data from the user-supplied tensor to our internal state + # (in case you want to keep an internal copy). + self.articulation_view_data.link_mass[0, link_indices] = masses.to(self.device) + + # Actually call p.changeDynamics to update the mass in Bullet + if isinstance(link_indices, slice): + effective_indices = range(self._num_links)[link_indices] + elif isinstance(link_indices, torch.Tensor): + effective_indices = link_indices.tolist() + else: + effective_indices = link_indices + + for i, link_idx in enumerate(effective_indices): + mass_value = masses[i].item() if len(masses.shape) > 0 else masses.item() + p.changeDynamics( + bodyUniqueId=self.articulation, linkIndex=link_idx, mass=mass_value, physicsClientId=self.client_id + ) + + def set_root_velocities( + self, + root_velocities: torch.Tensor, + indices: Union[List[int], torch.Tensor, slice] = slice(None), + ): + pass + + def set_dof_limits( + self, + limits: torch.Tensor, + indices: Union[List[int], torch.Tensor, slice] = slice(None), + joint_indices: Union[List[int], torch.Tensor, slice] = slice(None), + ): + self.articulation_view_data.dof_limits[indices, joint_indices] = limits.to(self.device) + + def set_dof_positions( + self, + positions: torch.Tensor, + indices: Union[List[int], torch.Tensor, slice] = slice(None), + joint_indices: Union[List[int], torch.Tensor, slice] = slice(None), + ): + with torch.inference_mode(mode=True): + # Also store internally + self.articulation_view_data.dof_positions[indices, joint_indices] = positions.to(self.device) + + if isinstance(joint_indices, slice): + effective_indices = self._dof_indicies[joint_indices] + elif isinstance(joint_indices, torch.Tensor): + effective_indices = [self._dof_indicies[i] for i in joint_indices.tolist()] + else: + effective_indices = [self._dof_indicies[i] for i in joint_indices] + # For each movable dof + for idx, j_id in enumerate(effective_indices): + p.resetJointState( + bodyUniqueId=self.articulation, + jointIndex=j_id, + targetValue=positions[0][idx].item(), + targetVelocity=0.0, + physicsClientId=self.client_id, + ) + self.populate_link_transforms() + + def set_dof_position_targets( + self, + positions: torch.Tensor, + indices: Union[List[int], torch.Tensor, slice] = slice(None), + joint_indices: Union[List[int], torch.Tensor, slice] = slice(None), + ): + self.articulation_view_data.dof_position_target[indices, joint_indices] = positions.to(self.device) + + def set_dof_velocities( + self, + velocities: torch.Tensor, + indices: Union[List[int], torch.Tensor, slice] = slice(None), + joint_indices: Union[List[int], torch.Tensor, slice] = slice(None), + ): + self.articulation_view_data.dof_velocities[indices, joint_indices] = velocities.to(self.device) + + if isinstance(joint_indices, slice): + effective_indices = self._dof_indicies[joint_indices] + if isinstance(joint_indices, torch.Tensor): + effective_indices = [self._dof_indicies[i] for i in joint_indices.tolist()] + + for idx, j_id in enumerate(effective_indices): + # get current joint position so we don't overwrite it + cur_state = p.getJointState(self.articulation, j_id, physicsClientId=self.client_id) + cur_pos = cur_state[0] + p.resetJointState( + bodyUniqueId=self.articulation, + jointIndex=j_id, + targetValue=cur_pos, + targetVelocity=velocities[0][idx].item(), + physicsClientId=self.client_id, + ) + self.populate_link_velocities() + + def set_dof_velocity_targets( + self, + velocities: torch.Tensor, + indices: Union[List[int], torch.Tensor, slice] = slice(None), + joint_indices: Union[List[int], torch.Tensor, slice] = slice(None), + ): + self.articulation_view_data.dof_velocity_target[indices, joint_indices] = velocities.to(self.device) + + def set_dof_torques( + self, + torques: torch.Tensor, + indices: Union[List[int], torch.Tensor, slice] = slice(None), + joint_indices: Union[List[int], torch.Tensor, slice] = slice(None), + ): + self.articulation_view_data.dof_torques[indices, joint_indices] = torques.to(self.device) + + def set_dof_states( + self, + positions: torch.Tensor, + velocities: torch.Tensor, + efforts: torch.Tensor, + indices: Union[List[int], torch.Tensor, slice] = slice(None), + joint_indices: Union[List[int], torch.Tensor, slice] = slice(None), + ): + self.set_dof_positions(positions, indices, joint_indices) + self.set_dof_velocities(velocities, indices, joint_indices) + self.set_dof_torques(efforts, indices, joint_indices) + + def set_dof_stiffnesses( + self, + stiffness: torch.Tensor, + indices: Union[List[int], torch.Tensor, slice] = slice(None), + joint_indices: Union[List[int], torch.Tensor, slice] = slice(None), + ): + self.articulation_view_data.dof_stiffness[indices, joint_indices] = stiffness.to(self.device) + + def set_dof_dampings( + self, + damping: torch.Tensor, + indices: Union[List[int], torch.Tensor, slice] = slice(None), + joint_indices: Union[List[int], torch.Tensor, slice] = slice(None), + ): + self.articulation_view_data.dof_damping[indices, joint_indices] = damping.to(self.device) + + def set_dof_armatures( + self, + armatures: torch.Tensor, + indices: Union[List[int], torch.Tensor, slice] = slice(None), + joint_indices: Union[List[int], torch.Tensor, slice] = slice(None), + ): + self.articulation_view_data.dof_armatures[indices, joint_indices] = armatures.to(self.device) + + def set_dof_frictions( + self, + frictions: torch.Tensor, + indices: Union[List[int], torch.Tensor, slice] = slice(None), + joint_indices: Union[List[int], torch.Tensor, slice] = slice(None), + ): + self.articulation_view_data.dof_frictions[indices, joint_indices] = frictions.to(self.device) + + def forward_kinematics(self, positions: torch.Tensor) -> torch.Tensor: + self.current_joint_positions = self.articulation_view_data.dof_positions.clone() + # 1) Set the joint positions in PyBullet + self.set_dof_positions(positions) + + # 2) Get the link transforms + link_transforms = self.get_link_transforms() + + # 3) Reset the joint positions to the original state + self.set_dof_positions(self.current_joint_positions) + + return link_transforms + + def set_dof_targets(self, positions, velocities, torques): + self.set_dof_position_targets(positions) + self.set_dof_velocity_targets(velocities) + self.set_dof_torques(torques) + + def populate_joint_state(self): + dof_names = [] + dof_types = [] + link_names = [] + for j in range(self._num_dofs): + info = p.getJointInfo(self.articulation, self._dof_indicies[j], physicsClientId=self.client_id) + dof_index: int = info[0] + dof_name: str = info[1].decode("utf-8") + dof_type: int = info[2] + damping: float = info[6] + friction: float = info[7] + lower_limit: float = info[8] + upper_limit: float = info[9] + max_forces: float = info[10] + max_velocity: float = info[11] + link_name = info[12].decode("utf-8") + + if dof_type is not p.JOINT_FIXED: + self.articulation_view_data.dof_indices[:, j] = dof_index + dof_names.append(dof_name) + dof_types.append(dof_types_dict[dof_type]) + self.articulation_view_data.dof_damping[:, j] = damping + self.articulation_view_data.dof_frictions[:, j] = friction + self.articulation_view_data.dof_limits[:, j, 0] = lower_limit + self.articulation_view_data.dof_limits[:, j, 1] = upper_limit + self.articulation_view_data.dof_max_forces[:, j] = max_forces + self.articulation_view_data.dof_max_velocity[:, j] = max_velocity + + link_names.append(link_name) + + self.articulation_view_data.dof_names = dof_names + self.articulation_view_data.dof_types = dof_types + + base, _ = p.getBodyInfo(self.articulation, physicsClientId=self.client_id) + self.articulation_view_data.link_names = [base.decode("utf-8")] + link_names + + def populate_link_transforms(self): + link_world_pose = torch.zeros_like(self.articulation_view_data.link_transforms) + # populate the base link + base_pose = p.getBasePositionAndOrientation(self.articulation, physicsClientId=self.client_id) + pos = base_pose[0] + quat = base_pose[1] # (x,y,z,w) + self.articulation_view_data.link_transforms[:, 0, :3] = torch.tensor(pos, device=self.device) + self.articulation_view_data.link_transforms[:, 0, 3:] = torch.tensor(quat, device=self.device) + + for link_idx in range(1, self._num_links): + link_state = p.getLinkState( + self.articulation, + link_idx - 1, + computeLinkVelocity=True, + computeForwardKinematics=True, + physicsClientId=self.client_id, + ) + link_world_position: list[float] = link_state[0] + link_world_orientation: list[float] = link_state[1] # (x,y,z,w) + + link_world_pose[:, link_idx, :3] = torch.tensor(link_world_position) + link_world_pose[:, link_idx, 3:] = torch.tensor(link_world_orientation) + + self.articulation_view_data.link_transforms = link_world_pose.to(self.device) + + def populate_link_velocities(self): + link_world_velocity = torch.zeros_like(self.articulation_view_data.link_velocities) + + lin_vel, ang_vel = p.getBaseVelocity(self.articulation, physicsClientId=self.client_id) + link_world_velocity[:, 0, :3] = torch.tensor(lin_vel) + link_world_velocity[:, 0, 3:] = torch.tensor(ang_vel) + for link_idx in range(1, self._num_links): + link_state = p.getLinkState( + self.articulation, + link_idx - 1, + computeLinkVelocity=True, + computeForwardKinematics=True, + physicsClientId=self.client_id, + ) + world_link_linear_velocity: list[float] = link_state[6] + world_link_angular_velocity: list[float] = link_state[7] + link_world_velocity[:, link_idx, :3] = torch.tensor(world_link_linear_velocity) + link_world_velocity[:, link_idx, 3:] = torch.tensor(world_link_angular_velocity) + + self.articulation_view_data.link_velocities = link_world_velocity.to(self.device) + + def populate_link_mass(self): + dyn_info = p.getDynamicsInfo(self.articulation, -1, physicsClientId=self.client_id) + self.articulation_view_data.link_mass[:, 0] = dyn_info[0] + + for link_idx in range(1, self._num_links): + # 3) Get the mass of this link + dyn_info = p.getDynamicsInfo(self.articulation, link_idx - 1, physicsClientId=self.client_id) + self.articulation_view_data.link_mass[:, link_idx] = dyn_info[0] + + def populate_inertia(self): + # Initialize output: (batch=1, num_links, 9) + inertias = torch.zeros_like(self.articulation_view_data.link_inertia) + coms = torch.zeros_like(self.articulation_view_data.link_coms) + for link_idx in range(0, self._num_links): + # Retrieve the dynamics info from PyBullet + dyn_info = p.getDynamicsInfo(self.articulation, link_idx - 1, physicsClientId=self.client_id) + + local_inertia_diag = torch.tensor(dyn_info[2]) # shape [3] + local_inertial_pos = torch.tensor(dyn_info[3]) # (x, y, z) + local_inertia_quat = torch.tensor(dyn_info[4]) # shape [4], (x,y,z,w) + + # 1) Diagonal inertia in a 3×3 matrix + Idiag = torch.diag(local_inertia_diag) + + # 2) Convert PyBullet quaternion (x,y,z,w) into a rotation matrix + # Adjust for any function that expects (w, x, y, z) vs. (x, y, z, w). + # For example: + # local_inertia_quat -> (x, y, z, w) + # convert_quat(...) -> reorder to (w, x, y, z) if needed. + R = math_utils.matrix_from_quat( + math_utils.convert_quat(local_inertia_quat, to="wxyz") # type: ignore + ) # shape (3,3) + + # 3) Compute the full inertia in link frame: I = R * Idiag * R^T + Ifull = R @ Idiag @ R.transpose(0, 1) # shape (3,3) + + # 4) Flatten row-major into a length-9 vector + inertias[0, link_idx, :] = Ifull.reshape(-1) + coms[:, link_idx, :3] = local_inertial_pos + coms[:, link_idx, 3:] = local_inertia_quat + + self.articulation_view_data.link_inertia[:] = inertias.to(self.device) + self.articulation_view_data.link_coms[:] = coms.to(self.device) + + def __del__(self): + p.disconnect() diff --git a/source/uwlab/uwlab/assets/asset_base.py b/source/uwlab/uwlab/assets/asset_base.py new file mode 100644 index 00000000..74ce9219 --- /dev/null +++ b/source/uwlab/uwlab/assets/asset_base.py @@ -0,0 +1,234 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import inspect +import weakref +from abc import ABC, abstractmethod +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any + +import omni.kit.app +import omni.timeline + +import isaaclab.sim as sim_utils + +if TYPE_CHECKING: + from .asset_base_cfg import AssetBaseCfg + + +class AssetBase(ABC): + """The base interface class for assets. + + An asset corresponds to any physics-enabled object that can be spawned in the simulation. These include + rigid objects, articulated objects, deformable objects etc. The core functionality of an asset is to + provide a set of buffers that can be used to interact with the simulator. The buffers are updated + by the asset class and can be written into the simulator using the their respective ``write`` methods. + This allows a convenient way to perform post-processing operations on the buffers before writing them + into the simulator and obtaining the corresponding simulation results. + + The class handles both the spawning of the asset into the USD stage as well as initialization of necessary + physics handles to interact with the asset. Upon construction of the asset instance, the prim corresponding + to the asset is spawned into the USD stage if the spawn configuration is not None. The spawn configuration + is defined in the :attr:`AssetBaseCfg.spawn` attribute. In case the configured :attr:`AssetBaseCfg.prim_path` + is an expression, then the prim is spawned at all the matching paths. Otherwise, a single prim is spawned + at the configured path. For more information on the spawn configuration, see the + :mod:`isaaclab.sim.spawners` module. + + Unlike Isaac Sim interface, where one usually needs to call the + :meth:`isaacsim.coreprims.XFormPrimView.initialize` method to initialize the PhysX handles, the asset + class automatically initializes and invalidates the PhysX handles when the stage is played/stopped. This + is done by registering callbacks for the stage play/stop events. + + Additionally, the class registers a callback for debug visualization of the asset if a debug visualization + is implemented in the asset class. This can be enabled by setting the :attr:`AssetBaseCfg.debug_vis` attribute + to True. The debug visualization is implemented through the :meth:`_set_debug_vis_impl` and + :meth:`_debug_vis_callback` methods. + """ + + def __init__(self, cfg: AssetBaseCfg): + """Initialize the asset base. + + Args: + cfg: The configuration class for the asset. + + 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 + self._is_initialized = False + # add handle for debug visualization (this is set to a valid handle inside set_debug_vis) + self._debug_vis_handle = None + + def __del__(self): + """Unsubscribe from the callbacks.""" + # clear physics events handles + if self._initialize_handle: + self._initialize_handle.unsubscribe() + self._initialize_handle = None + if self._invalidate_initialize_handle: + self._invalidate_initialize_handle.unsubscribe() + self._invalidate_initialize_handle = None + # clear debug visualization + if self._debug_vis_handle: + self._debug_vis_handle.unsubscribe() + self._debug_vis_handle = None + + """ + Properties + """ + + @property + def is_initialized(self) -> bool: + """Whether the asset is initialized. + + Returns True if the asset is initialized, False otherwise. + """ + return self._is_initialized + + @property + @abstractmethod + def num_instances(self) -> int: + """Number of instances of the asset. + + This is equal to the number of asset instances per environment multiplied by the number of environments. + """ + return NotImplementedError + + @property + def device(self) -> str: + """Memory device for computation.""" + return self._device + + @property + @abstractmethod + def data(self) -> Any: + """Data related to the asset.""" + return NotImplementedError + + @property + def has_debug_vis_implementation(self) -> bool: + """Whether the asset 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 + + """ + Operations. + """ + + def set_debug_vis(self, debug_vis: bool) -> bool: + """Sets whether to visualize the asset data. + + Args: + debug_vis: Whether to visualize the asset data. + + Returns: + Whether the debug visualization was successfully set. False if the asset + does not support debug visualization. + """ + # 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 + if debug_vis: + # create a subscriber for the post update event if it doesn't exist + if 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: + # remove the subscriber if it exists + if self._debug_vis_handle is not None: + self._debug_vis_handle.unsubscribe() + self._debug_vis_handle = None + # return success + return True + + @abstractmethod + def reset(self, env_ids: Sequence[int] | None = None): + """Resets all internal buffers of selected environments. + + Args: + env_ids: The indices of the object to reset. Defaults to None (all instances). + """ + raise NotImplementedError + + @abstractmethod + def write_data_to_sim(self): + """Writes data to the simulator.""" + raise NotImplementedError + + @abstractmethod + def update(self, dt: float): + """Update the internal buffers. + + The time step ``dt`` is used to compute numerical derivatives of quantities such as joint + accelerations which are not provided by the simulator. + + Args: + dt: The amount of time passed from last ``update`` call. + """ + raise NotImplementedError + + """ + Implementation specific. + """ + + @abstractmethod + def _initialize_impl(self): + """Initializes the PhysX handles and internal buffers.""" + raise NotImplementedError + + def _set_debug_vis_impl(self, debug_vis: bool): + """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__}.") + + def _debug_vis_callback(self, event): + """Callback for debug visualization. + + This function calls the visualization objects and sets the data to visualize into them. + """ + raise NotImplementedError(f"Debug visualization is not implemented for {self.__class__.__name__}.") + + """ + Internal simulation callbacks. + """ + + def _initialize_callback(self, event): + """Initializes the scene elements. + + Note: + PhysX handles are only enabled once the simulator starts playing. Hence, this function needs to be + called whenever the simulator "plays" from a "stop" state. + """ + if not self._is_initialized: + # obtain simulation related information + sim = sim_utils.SimulationContext.instance() + if sim is None: + raise RuntimeError("SimulationContext is not initialized! Please initialize SimulationContext first.") + self._backend = sim.backend + self._device = sim.device + # initialize the asset + self._initialize_impl() + # set flag + self._is_initialized = True + + def _invalidate_initialize_callback(self, event): + """Invalidates the scene elements.""" + self._is_initialized = False diff --git a/source/uwlab/uwlab/assets/asset_base_cfg.py b/source/uwlab/uwlab/assets/asset_base_cfg.py new file mode 100644 index 00000000..5d39e7bd --- /dev/null +++ b/source/uwlab/uwlab/assets/asset_base_cfg.py @@ -0,0 +1,46 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.utils import configclass + +from .asset_base import AssetBase + + +@configclass +class AssetBaseCfg: + """The base configuration class for an asset's parameters. + + Please see the :class:`AssetBase` class for more information on the asset class. + """ + + @configclass + class InitialStateCfg: + """Initial state of the asset. + + This defines the default initial state of the asset when it is spawned into the simulation, as + well as the default state when the simulation is reset. + + After parsing the initial state, the asset class stores this information in the :attr:`data` + attribute of the asset class. This can then be accessed by the user to modify the state of the asset + during the simulation, for example, at resets. + """ + + # root position + pos: tuple[float, float, float] = (0.0, 0.0, 0.0) + """Position of the root in simulation world 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) of the root in simulation world frame. + Defaults to (1.0, 0.0, 0.0, 0.0). + """ + + 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:`isaaclab.assets.asset_base.AssetBase`. + """ + + init_state: InitialStateCfg = InitialStateCfg() + """Initial state of the rigid object. Defaults to identity pose.""" diff --git a/source/uwlab/uwlab/controllers/__init__.py b/source/uwlab/uwlab/controllers/__init__.py new file mode 100644 index 00000000..ff4bb75e --- /dev/null +++ b/source/uwlab/uwlab/controllers/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-package for different controllers and motion-generators. + +Controllers or motion generators are responsible for closed-loop tracking of a given command. The +controller can be a simple PID controller or a more complex controller such as impedance control +or inverse kinematics control. The controller is responsible for generating the desired joint-level +commands to be sent to the robot. +""" + +from .differential_ik import MultiConstraintDifferentialIKController +from .differential_ik_cfg import MultiConstraintDifferentialIKControllerCfg diff --git a/source/uwlab/uwlab/controllers/differential_ik.py b/source/uwlab/uwlab/controllers/differential_ik.py new file mode 100644 index 00000000..0eb4b027 --- /dev/null +++ b/source/uwlab/uwlab/controllers/differential_ik.py @@ -0,0 +1,250 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.utils.math import apply_delta_pose, compute_pose_error + +if TYPE_CHECKING: + from .differential_ik_cfg import MultiConstraintDifferentialIKControllerCfg + + +class MultiConstraintDifferentialIKController: + r"""Differential inverse kinematics (IK) controller. + + .. note:: + This controller extends Isaac Lab differential inverse kinematic and provides multi constraints solution. + + This controller is based on the concept of differential inverse kinematics [1, 2] which is a method for computing + the change in joint positions that yields the desired change in pose. + + .. math:: + + \Delta \mathbf{q} = \mathbf{J}^{\dagger} \Delta \mathbf{x} + \mathbf{q}_{\text{desired}} = \mathbf{q}_{\text{current}} + \Delta \mathbf{q} + + where :math:`\mathbf{J}^{\dagger}` is the pseudo-inverse of the Jacobian matrix :math:`\mathbf{J}`, + :math:`\Delta \mathbf{x}` is the desired change in pose, and :math:`\mathbf{q}_{\text{current}}` + is the current joint positions. + + Currently only below methods are supported : + - "dls": Damped version of Moore-Penrose pseudo-inverse (also called Levenberg-Marquardt) + + + .. caution:: + Isaac Lab non-multiconstraints version implementation contains methods such as pinv, svd, trans, however + their solution is not yet implemented even though these inputs are available. + + The controller does not assume anything about the frames of the current and desired end-effector pose, + or the joint-space velocities. It is up to the user to ensure that these quantities are given + in the correct format. + + Reference: + [1] https://ethz.ch/content/dam/ethz/special-interest/mavt/robotics-n-intelligent-systems/rsl-dam/documents/RobotDynamics2017/RD_HS2017script.pdf + [2] https://www.cs.cmu.edu/~15464-s13/lectures/lecture6/iksurvey.pdf + [3] https://isaac-sim.github.io/IsaacLab/source/api/lab/isaaclab.controllers.html + """ + + def __init__(self, cfg: MultiConstraintDifferentialIKControllerCfg, num_bodies: int, num_envs: int, device: str): + """Initialize the controller. + + Args: + cfg: The configuration for the controller. + num_envs: The number of environments. + device: The device to use for computations. + """ + # store inputs + self.cfg = cfg + self.num_envs = num_envs + self.num_bodies = num_bodies + self._device = device + # create buffers + self.ee_pos_des = torch.zeros(self.num_envs, self.num_bodies, 3, device=self._device) + self.ee_quat_des = torch.zeros(self.num_envs, self.num_bodies, 4, device=self._device) + # -- input command + self._command = torch.zeros(self.num_envs, self.action_dim, device=self._device).view( + self.num_envs, self.num_bodies, -1 + ) + + """ + Properties. + """ + + @property + def action_dim(self) -> int: + """Dimension of the controller's input command.""" + if self.cfg.command_type == "position": + return 3 * self.num_bodies # (x, y, z) + elif self.cfg.command_type == "pose" and self.cfg.use_relative_mode: + return 6 * self.num_bodies # (dx, dy, dz, droll, dpitch, dyaw) + else: + return 7 * self.num_bodies # (x, y, z, qw, qx, qy, qz) + + """ + Operations. + """ + + def reset(self, env_ids: torch.Tensor = None): # type: ignore + """Reset the internals. + + Args: + env_ids: The environment indices to reset. If None, then all environments are reset. + """ + pass + + def set_command( + self, command: torch.Tensor, ee_pos: torch.Tensor | None = None, ee_quat: torch.Tensor | None = None + ): + """Set target end-effector pose command. + + Based on the configured command type and relative mode, the method computes the desired end-effector pose. + It is up to the user to ensure that the command is given in the correct frame. The method only + applies the relative mode if the command type is ``position_rel`` or ``pose_rel``. + + Args: + command: The input command in shape (N, 3 * num_bodies) or (N, 6 * num_bodies) or (N, 7 * num_bodies). + ee_pos: The current end-effector position in shape (N, 3). + This is only needed if the command type is ``position_rel`` or ``pose_rel``. + ee_quat: The current end-effector orientation (w, x, y, z) in shape (N, 4). + This is only needed if the command type is ``position_*`` or ``pose_rel``. + + Raises: + ValueError: If the command type is ``position_*`` and :attr:`ee_quat` is None. + ValueError: If the command type is ``position_rel`` and :attr:`ee_pos` is None. + ValueError: If the command type is ``pose_rel`` and either :attr:`ee_pos` or :attr:`ee_quat` is None. + """ + # store command + self._command[:] = command.view(self._command.shape) + # compute the desired end-effector pose + if self.cfg.command_type == "position": + # we need end-effector orientation even though we are in position mode + # this is only needed for display purposes + if ee_quat is None: + raise ValueError("End-effector orientation can not be None for `position_*` command type!") + # compute targets + if self.cfg.use_relative_mode: + if ee_pos is None: + raise ValueError("End-effector position can not be None for `position_rel` command type!") + self.ee_pos_des[:] = ee_pos + self._command + self.ee_quat_des[:] = ee_quat + else: + self.ee_pos_des[:] = self._command.view(self.num_envs, -1, 3) + self.ee_quat_des[:] = ee_quat + else: + # compute targets + if self.cfg.use_relative_mode: + if ee_pos is None or ee_quat is None: + raise ValueError( + "Neither end-effector position nor orientation can be None for `pose_rel` command type!" + ) + self.ee_pos_des, self.ee_quat_des = apply_delta_pose( + ee_pos.view(-1, 3), ee_quat.view(-1, 4), self._command.view(-1, 6) + ) + else: + self.ee_pos_des = self._command[:, :, 0:3] + self.ee_quat_des = self._command[:, :, 3:7] + + def compute( + self, ee_pos: torch.Tensor, ee_quat: torch.Tensor, jacobian: torch.Tensor, joint_pos: torch.Tensor + ) -> torch.Tensor: + """Computes the target joint positions that will yield the desired end effector pose. + + Args: + ee_pos: The current end-effector position in shape (N, 3). + ee_quat: The current end-effector orientation in shape (N, 4). + jacobian: The geometric jacobian matrix in shape (N, 6, num_joints). + joint_pos: The current joint positions in shape (N, num_joints). + + Returns: + The target joint positions commands in shape (N, num_joints). + """ + # compute the delta in joint-space + if "position" in self.cfg.command_type: + position_error = self.ee_pos_des - ee_pos + jacobian_pos = jacobian[:, :, 0:3, :] + delta_joint_pos = self._compute_delta_joint_pos(delta_pose=position_error, jacobian=jacobian_pos) + else: + position_error, axis_angle_error = compute_pose_error( + ee_pos.view(-1, 3), + ee_quat.view(-1, 4), + self.ee_pos_des.view(-1, 3), + self.ee_quat_des.view(-1, 4), + rot_error_type="axis_angle", + ) + pose_error = torch.cat((position_error, axis_angle_error), dim=1) + delta_joint_pos = self._compute_delta_joint_pos(delta_pose=pose_error, jacobian=jacobian) + # return the desired joint positions + return joint_pos + delta_joint_pos + + """ + Helper functions. + """ + + def _compute_delta_joint_pos(self, delta_pose: torch.Tensor, jacobian: torch.Tensor) -> torch.Tensor: + """Computes the change in joint position that yields the desired change in pose. + + The method uses the Jacobian mapping from joint-space velocities to end-effector velocities + to compute the delta-change in the joint-space that moves the robot closer to a desired + end-effector position. + + Args: + delta_pose: The desired delta pose in shape (N, 3) or (N, 6). + jacobian: The geometric jacobian matrix in shape (N, 3, num_joints) or (N, 6, num_joints). + + Returns: + The desired delta in joint space. Shape is (N, num-jointsß). + """ + if self.cfg.ik_params is None: + raise RuntimeError(f"Inverse-kinematics parameters for method '{self.cfg.ik_method}' is not defined!") + # compute the delta in joint-space + if self.cfg.ik_method == "pinv": # Jacobian pseudo-inverse + # parameters + k_val = self.cfg.ik_params["k_val"] + # computation + jacobian_pinv = torch.linalg.pinv(jacobian) + delta_joint_pos = k_val * jacobian_pinv @ delta_pose.unsqueeze(-1) + delta_joint_pos = delta_joint_pos.squeeze(-1) + elif self.cfg.ik_method == "svd": # adaptive SVD + # parameters + k_val = self.cfg.ik_params["k_val"] + min_singular_value = self.cfg.ik_params["min_singular_value"] + # computation + # U: 6xd, S: dxd, V: d x num-joint + U, S, Vh = torch.linalg.svd(jacobian) + S_inv = 1.0 / S + 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_joint_pos = k_val * jacobian_pinv @ delta_pose.unsqueeze(-1) + delta_joint_pos = delta_joint_pos.squeeze(-1) + elif self.cfg.ik_method == "trans": # Jacobian transpose + # parameters + k_val = self.cfg.ik_params["k_val"] + # computation + jacobian_T = torch.transpose(jacobian, dim0=1, dim1=2) + delta_joint_pos = k_val * jacobian_T @ delta_pose.unsqueeze(-1) + delta_joint_pos = delta_joint_pos.squeeze(-1) + elif self.cfg.ik_method == "dls": # damped least squares + # parameters + delta_pose = delta_pose.reshape(self.num_envs, -1) + jacobian = jacobian.reshape(self.num_envs, -1, jacobian.shape[-1]) + lambda_val = self.cfg.ik_params["lambda_val"] + # computation + jacobian_T = torch.transpose(jacobian, dim0=1, dim1=2) + lambda_matrix = (lambda_val**2) * torch.eye(n=jacobian.shape[1], device=self._device) + delta_joint_pos = ( + jacobian_T @ torch.inverse(jacobian @ jacobian_T + lambda_matrix) @ delta_pose.unsqueeze(-1) + ) + delta_joint_pos = delta_joint_pos.view(self.num_envs, -1) + else: + raise ValueError(f"Unsupported inverse-kinematics method: {self.cfg.ik_method}") + + return delta_joint_pos diff --git a/source/uwlab/uwlab/controllers/differential_ik_cfg.py b/source/uwlab/uwlab/controllers/differential_ik_cfg.py new file mode 100644 index 00000000..7181f849 --- /dev/null +++ b/source/uwlab/uwlab/controllers/differential_ik_cfg.py @@ -0,0 +1,17 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.controllers import DifferentialIKControllerCfg +from isaaclab.utils import configclass + +from .differential_ik import MultiConstraintDifferentialIKController + + +@configclass +class MultiConstraintDifferentialIKControllerCfg(DifferentialIKControllerCfg): + """Configuration for multi-constraint differential inverse kinematics controller.""" + + class_type: type = MultiConstraintDifferentialIKController + """The associated controller class.""" diff --git a/source/uwlab/uwlab/devices/__init__.py b/source/uwlab/uwlab/devices/__init__.py new file mode 100644 index 00000000..a477fb3d --- /dev/null +++ b/source/uwlab/uwlab/devices/__init__.py @@ -0,0 +1,25 @@ +# Copyright (c) 2022-2025, The UW Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-package providing interfaces to different teleoperation devices. + +Currently, the following categories of devices are supported: + +* **Se3Keyboard**: Standard keyboard with WASD and arrow keys. +* **RealsenseT265**: Realsense cameras 6 degrees of freedom. +* **RokokoGlove**: Rokoko gloves that track position and quaternion of the hand. + +All device interfaces inherit from the :class:`isaaclab.devices.DeviceBase` class, which provides a +common interface for all devices. The device interface reads the input data when +the :meth:`DeviceBase.advance` method is called. It also provides the function :meth:`DeviceBase.add_callback` +to add user-defined callback functions to be called when a particular input is pressed from +the peripheral device. +""" + +from .device_cfg import * +from .realsense_t265 import RealsenseT265 +from .rokoko_glove import RokokoGlove +from .se3_keyboard import Se3Keyboard +from .teleop_cfg import TeleopCfg diff --git a/source/uwlab/uwlab/devices/device_cfg.py b/source/uwlab/uwlab/devices/device_cfg.py new file mode 100644 index 00000000..7080734b --- /dev/null +++ b/source/uwlab/uwlab/devices/device_cfg.py @@ -0,0 +1,59 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from typing import Callable, Literal + +from isaaclab.devices import DeviceBase +from isaaclab.utils import configclass + +from .realsense_t265 import RealsenseT265 +from .rokoko_glove import RokokoGlove +from .se3_keyboard import Se3Keyboard + + +@configclass +class DeviceBaseTeleopCfg: + class_type: Callable[..., DeviceBase] = DeviceBase + + +@configclass +class KeyboardCfg(DeviceBaseTeleopCfg): + class_type: Callable[..., Se3Keyboard] = Se3Keyboard + + pos_sensitivity: float = 0.01 + + rot_sensitivity: float = 0.01 + + enable_gripper_command: bool = False + + +@configclass +class RokokoGlovesCfg(DeviceBaseTeleopCfg): + class_type: Callable[..., RokokoGlove] = RokokoGlove + + UDP_IP: str = "0.0.0.0" # Listen on all available network interfaces + + UDP_PORT: int = 14043 # Make sure this matches the port used in Rokoko Studio Live + + left_hand_track: list[str] = [] + + right_hand_track: list[str] = [] + + scale: float = 1 + + proximal_offset: float = 0.3 + + thumb_scale: float = 1.1 + + command_type: Literal["pose", "pos"] = "pos" + + +@configclass +class RealsenseT265Cfg(DeviceBaseTeleopCfg): + class_type: Callable[..., RealsenseT265] = RealsenseT265 + + cam_device_id: str = "905312110639" + + device: str = "cuda:0" diff --git a/source/uwlab/uwlab/devices/realsense_t265.py b/source/uwlab/uwlab/devices/realsense_t265.py new file mode 100644 index 00000000..895e5db6 --- /dev/null +++ b/source/uwlab/uwlab/devices/realsense_t265.py @@ -0,0 +1,95 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from enum import IntEnum +from typing import TYPE_CHECKING + +from isaaclab.devices import DeviceBase +from isaaclab.utils.math import compute_pose_error +from uwlab.utils.math import create_axis_remap_function + +if TYPE_CHECKING: + from .device_cfg import RealsenseT265Cfg + + +class Preset(IntEnum): + Custom = 0 + Default = 1 + Hand = 2 + HighAccuracy = 3 + HighDensity = 4 + MediumDensity = 5 + + +class RealsenseT265(DeviceBase): + def __init__( + self, + cfg: RealsenseT265Cfg, + device="cuda:0", + ): + import pyrealsense2 as rs + + self.cam_device_id = cfg.cam_device_id + self.device = device + self._additional_callbacks = dict() + self.pipeline = rs.pipeline() # type: ignore + self.config = rs.config() # type: ignore + + self.ctx = rs.context() # type: ignore + self.t265_pipeline = rs.pipeline(self.ctx) # type: ignore + self.t265_config = rs.config() # type: ignore + self.t265_config.enable_device(self.cam_device_id) + self.t265_config.enable_stream(rs.stream.pose) # type: ignore + self.init_pose = torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0]], device=device) + self.current_pose = self.init_pose.clone() + self.t265_pipeline.start(self.t265_config) + self.axis_remap_fn = create_axis_remap_function(forward="z", left="-x", up="y", device=self.device) + + self.quat_error = None + + def reset(self): + self.current_pose = self.init_pose.clone() + t265_frames = self.t265_pipeline.wait_for_frames() + pose_frame = t265_frames.get_pose_frame() + pose_data = pose_frame.get_pose_data() + self.initial_quat = torch.tensor( + [[pose_data.rotation.w, pose_data.rotation.x, pose_data.rotation.y, pose_data.rotation.z]], + device=self.device, + ) + self.initial_pos = torch.tensor( + [[pose_data.translation.x, pose_data.translation.y, pose_data.translation.z]], device=self.device + ) + + def add_callback(self, key, func): + # camera does not have any callbacks + pass + + def advance(self): + """ + Advance the device state, read the device translational and rotational data, + then transform them xyz rpy that makes intuitive sense in Isaac Sim world coordinates. + """ + t265_frames = self.t265_pipeline.wait_for_frames() + pose_frame = t265_frames.get_pose_frame() + pose_data = pose_frame.get_pose_data() + + curr_quat = torch.tensor( + [[pose_data.rotation.w, pose_data.rotation.x, pose_data.rotation.y, pose_data.rotation.z]], + device=self.device, + ) + curr_pos = torch.tensor( + [[pose_data.translation.x, pose_data.translation.y, pose_data.translation.z]], device=self.device + ) + del_xyz, del_rpy = compute_pose_error( + curr_pos, curr_quat, self.initial_pos, self.initial_quat, rot_error_type="axis_angle" + ) + del_xyz, del_rpy = self.axis_remap_fn(del_xyz, del_rpy) + + self.current_pose[:, :3] = del_xyz + self.current_pose[:, 3:] = del_rpy + return self.current_pose diff --git a/source/uwlab/uwlab/devices/rokoko_glove.py b/source/uwlab/uwlab/devices/rokoko_glove.py new file mode 100644 index 00000000..2b806124 --- /dev/null +++ b/source/uwlab/uwlab/devices/rokoko_glove.py @@ -0,0 +1,329 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import json +import socket +import torch +from collections.abc import Callable +from typing import TYPE_CHECKING + +from isaaclab.devices import DeviceBase +from isaaclab.utils.math import quat_apply, quat_from_euler_xyz + +if TYPE_CHECKING: + from uwlab.devices import RokokoGlovesCfg + + +class RokokoGlove(DeviceBase): + """A Rokoko_Glove controller for sending SE(3) commands as absolute poses of hands individual part + This class is designed to track hands and fingers's pose from rokoko gloves. + It uses the udp network protocol to listen to Rokoko Live Studio data gathered from Rokoko smart gloves, + and process the data in form of torch Tensor. + Addressing the efficiency and ease to understand, the tracking will only be performed with user's parts + input, and all Literal of available parts is exhaustively listed in the comment under method __init__. + + available tracking literals: + LEFT_HAND: + leftHand, leftThumbProximal, leftThumbMedial, leftThumbDistal, leftThumbTip, + leftIndexProximal, leftIndexMedial, leftIndexDistal, leftIndexTip, + leftMiddleProximal, leftMiddleMedial, leftMiddleDistal, leftMiddleTip, + leftRingProximal, leftRingMedial, leftRingDistal, leftRingTip, + leftLittleProximal, leftLittleMedial, leftLittleDistal, leftLittleTip + + RIGHT_HAND: + rightHand, rightThumbProximal, rightThumbMedial, rightThumbDistal, rightThumbTip + rightIndexProximal, rightIndexMedial, rightIndexDistal, rightIndexTip, + rightMiddleProximal, rightMiddleMedial, rightMiddleDistal, rightMiddleTip + rightRingProximal, rightRingMedial, rightRingDistal, rightRingTip, + rightLittleProximal, rightLittleMedial, rightLittleDistal, rightLittleTip + """ + + def __init__( + self, + cfg: RokokoGlovesCfg, # Make sure this matches the port used in Rokoko Studio Live + device="cuda:0", + ): + """Initialize the Rokoko_Glove Controller. + Be aware that current implementation outputs pose of each hand part in the same order as input list, + but parts come from left hand always come before parts from right hand. + + Args: + UDP_IP: The IP Address of network to listen to, 0.0.0.0 refers to all available networks + UDP_PORT: The port Rokoko Studio Live sends to + left_hand_track: the tracking point of left hand this class will be tracking. + right_hand_track: the tracking point of right hand this class will be tracking. + scale: the overall scale for the hand. + proximal_offset: the inter proximal offset that shorten or widen the spread of hand. + """ + import lz4.frame + + self.lz4frame = lz4.frame + self.device = device + self._additional_callbacks = dict() + # Define the IP address and port to listen on + self.UDP_IP = cfg.UDP_IP + self.UDP_PORT = cfg.UDP_PORT + self.scale = cfg.scale + self.proximal_offset = cfg.proximal_offset + self.thumb_scale = cfg.thumb_scale + self.left_fingertip_names = cfg.left_hand_track + self.right_fingertip_names = cfg.right_hand_track + self.command_type = cfg.command_type + + # Create a UDP socket + self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 8192) + self.sock.bind((self.UDP_IP, self.UDP_PORT)) + + self.normalize_position = True + self.normalize_rotation = True + + print(f"Listening for UDP packets on {self.UDP_IP}:{self.UDP_PORT}") + + # fmt: off + self.left_hand_joint_names = [ + 'leftHand', + 'leftThumbProximal', 'leftThumbMedial', 'leftThumbDistal', 'leftThumbTip', + 'leftIndexProximal', 'leftIndexMedial', 'leftIndexDistal', 'leftIndexTip', + 'leftMiddleProximal', 'leftMiddleMedial', 'leftMiddleDistal', 'leftMiddleTip', + 'leftRingProximal', 'leftRingMedial', 'leftRingDistal', 'leftRingTip', + 'leftLittleProximal', 'leftLittleMedial', 'leftLittleDistal', 'leftLittleTip'] + + self.right_hand_joint_names = [ + 'rightHand', + 'rightThumbProximal', 'rightThumbMedial', 'rightThumbDistal', 'rightThumbTip', + 'rightIndexProximal', 'rightIndexMedial', 'rightIndexDistal', 'rightIndexTip', + 'rightMiddleProximal', 'rightMiddleMedial', 'rightMiddleDistal', 'rightMiddleTip', + 'rightRingProximal', 'rightRingMedial', 'rightRingDistal', 'rightRingTip', + 'rightLittleProximal', 'rightLittleMedial', 'rightLittleDistal', 'rightLittleTip'] + # fmt: on + + self.left_joint_dict = {self.left_hand_joint_names[i]: i for i in range(len(self.left_hand_joint_names))} + self.right_joint_dict = {self.right_hand_joint_names[i]: i for i in range(len(self.right_hand_joint_names))} + + self.left_finger_dict = {i: self.left_joint_dict[i] for i in self.left_fingertip_names} + self.right_finger_dict = {i: self.right_joint_dict[i] for i in self.right_fingertip_names} + + self.left_fingertip_poses = torch.zeros((len(self.left_hand_joint_names), 7), device=self.device) + self.right_fingertip_poses = torch.zeros((len(self.right_hand_joint_names), 7), device=self.device) + self.fingertip_poses = torch.zeros( + (len(self.left_hand_joint_names) + len(self.right_hand_joint_names), 7), device=self.device + ) + output_indices_list = [ + *[self.right_joint_dict[i] for i in self.left_fingertip_names], + *[self.right_joint_dict[i] + len(self.left_hand_joint_names) for i in self.right_fingertip_names], + ] + self.output_indices = torch.tensor(output_indices_list, device=self.device) + + # necessary joints for compute normalization + self.necessary_joints = [] + if len(self.right_fingertip_names): + self.necessary_joints.extend(["rightIndexProximal", "rightMiddleProximal", "rightHand"]) + if len(self.left_fingertip_names): + self.necessary_joints.extend(["leftIndexProximal", "leftMiddleProximal", "leftHand"]) + + def reset(self): + "Reset Internal Buffer" + self.left_fingertip_poses = torch.zeros((len(self.left_hand_joint_names), 7), device=self.device) + self.right_fingertip_poses = torch.zeros((len(self.right_hand_joint_names), 7), device=self.device) + + def advance(self): + """Provides the properly scaled, ordered, selected tracking results received from Rokoko Studio. + + Returns: + A tuple containing the 2D (n,7) pose array ordered by user inputted joint track list, and a dummy truth value. + """ + self.left_fingertip_poses = torch.zeros((len(self.left_hand_joint_names), 7), device=self.device) + self.right_fingertip_poses = torch.zeros((len(self.right_hand_joint_names), 7), device=self.device) + body_data = self._get_gloves_data() + + for joint_name in self.left_fingertip_names: + joint_data = body_data[joint_name] + joint_position = torch.tensor(list(joint_data["position"].values())) + joint_rotation = torch.tensor(list(joint_data["rotation"].values())) + self.left_fingertip_poses[self.right_joint_dict[joint_name]][:3] = joint_position + self.left_fingertip_poses[self.left_joint_dict[joint_name]][3:] = joint_rotation + + for joint_name in self.right_fingertip_names: + joint_data = body_data[joint_name] + joint_position = torch.tensor(list(joint_data["position"].values())) + joint_rotation = torch.tensor(list(joint_data["rotation"].values())) + self.right_fingertip_poses[self.right_joint_dict[joint_name]][:3] = joint_position + self.right_fingertip_poses[self.right_joint_dict[joint_name]][3:] = joint_rotation + # for normalization purpose + + for joint_name in self.necessary_joints: + joint_data = body_data[joint_name] + joint_position = torch.tensor(list(joint_data["position"].values())) + joint_rotation = torch.tensor(list(joint_data["rotation"].values())) + self.right_fingertip_poses[self.right_joint_dict[joint_name]][:3] = joint_position + self.right_fingertip_poses[self.right_joint_dict[joint_name]][3:] = joint_rotation + + left_wrist_position = self.left_fingertip_poses[0][:3] + if len(self.left_fingertip_names) > 0: + # scale + self.left_fingertip_poses[:, :3] = ( + self.left_fingertip_poses[:, :3] - left_wrist_position + ) * self.scale + left_wrist_position + # reposition + leftIndexProximalIdx = self.left_joint_dict["leftIndexProximal"] + leftMiddleProximalIdx = self.left_joint_dict["leftMiddleProximal"] + leftRingProximalIdx = self.left_joint_dict["leftRingProximal"] + leftLittleProximalIdx = self.left_joint_dict["leftLittleProximal"] + + reposition_vector = ( + self.left_fingertip_poses[leftMiddleProximalIdx][:3] + - self.left_fingertip_poses[leftIndexProximalIdx][:3] + ) + self.left_fingertip_poses[leftIndexProximalIdx:, :3] += self.proximal_offset * reposition_vector + self.left_fingertip_poses[leftMiddleProximalIdx:, :3] += self.proximal_offset * reposition_vector + self.left_fingertip_poses[leftRingProximalIdx:, :3] += self.proximal_offset * reposition_vector + self.left_fingertip_poses[leftLittleProximalIdx:, :3] += self.proximal_offset * reposition_vector + self.fingertip_poses[: len(self.left_fingertip_poses)] = self.left_fingertip_poses + + right_wrist_position = self.right_fingertip_poses[0][:3] + if len(self.right_fingertip_names) > 0: + rightThumbProximalIdx = self.right_joint_dict["rightThumbProximal"] + rightIndexProximalIdx = self.right_joint_dict["rightIndexProximal"] + rightMiddleProximalIdx = self.right_joint_dict["rightMiddleProximal"] + rightRingProximalIdx = self.right_joint_dict["rightRingProximal"] + rightLittleProximalIdx = self.right_joint_dict["rightLittleProximal"] + # normalize + rot_matrix = self.normalize_hand_positions(self.right_joint_dict) + position_normalized_poses = self.right_fingertip_poses[:, :3] - right_wrist_position + if self.normalize_rotation: + self.right_fingertip_poses[:, :3] = position_normalized_poses @ rot_matrix + # scale + self.right_fingertip_poses[:, :3] = self.right_fingertip_poses[:, :3] * self.scale + t_idx = rightThumbProximalIdx + self.right_fingertip_poses[t_idx : t_idx + 4, :3] = ( + self.right_fingertip_poses[t_idx : t_idx + 4, :3] * self.thumb_scale + ) + if not self.normalize_position: + self.right_fingertip_poses[:, :3] += right_wrist_position + # reposition + reposition_vector = ( + self.right_fingertip_poses[rightMiddleProximalIdx][:3] + - self.right_fingertip_poses[rightIndexProximalIdx][:3] + ) + self.right_fingertip_poses[rightIndexProximalIdx:, :3] += self.proximal_offset * reposition_vector + self.right_fingertip_poses[rightMiddleProximalIdx:, :3] += self.proximal_offset * reposition_vector + self.right_fingertip_poses[rightRingProximalIdx:, :3] += self.proximal_offset * reposition_vector + self.right_fingertip_poses[rightLittleProximalIdx:, :3] += self.proximal_offset * reposition_vector + self.fingertip_poses[len(self.left_fingertip_poses) :] = self.right_fingertip_poses + self.fingertip_poses[self.output_indices] = self.to_world_convention(self.fingertip_poses[self.output_indices]) + if self.command_type == "pos": + return self.fingertip_poses[self.output_indices][:, :3] + return self.fingertip_poses[self.output_indices] + + def to_world_convention(self, poses: torch.Tensor): + # # to world convention + poses[:, 1] = -poses[:, 1] + + x = torch.tensor([0], device=self.device) + y = torch.tensor([0], device=self.device) + z = torch.tensor([0.4], device=self.device) + quat = quat_from_euler_xyz(x, y, z) + poses[:, :3] = quat_apply(quat, poses[:, :3]) + return poses + + def normalize_hand_positions(self, hand_keypoints): + wrist = self.right_fingertip_poses[hand_keypoints["rightHand"]][:3] + index_proximal = self.right_fingertip_poses[hand_keypoints["rightIndexProximal"]][:3] + middle_proximal = self.right_fingertip_poses[hand_keypoints["rightMiddleProximal"]][:3] + + # Define the plane with two vectors + vec1 = index_proximal - wrist + vec2 = middle_proximal - index_proximal + + # Compute orthonormal basis for the plane + vec1_normalized = vec1 / vec1.norm() + vec2_proj = vec2 - torch.dot(vec2, vec1_normalized) * vec1_normalized + vec2_normalized = vec2_proj / vec2_proj.norm() + vec3_normalized = torch.linalg.cross(vec1_normalized, vec2_normalized) + + # Construct rotation matrix + rotation_matrix = torch.stack([vec1_normalized, vec2_normalized, vec3_normalized], dim=-1) + + return rotation_matrix + + def _get_gloves_data(self): + while True: + try: + data, addr = self.sock.recvfrom(8192) # Buffer size is 1024 bytes + break + except OSError as e: + print(f"Error: {e}") + continue + decompressed_data = self.lz4frame.decompress(data) + received_json = json.loads(decompressed_data) + # received_json = json.loads(data) + body_data = received_json["scene"]["actors"][0]["body"] + return body_data + + def add_callback(self, key: str, func: Callable): + # check keys supported by callback + if key not in ["L", "R"]: + raise ValueError(f"Only left (L) and right (R) buttons supported. Provided: {key}.") + # TODO: Improve this to allow multiple buttons on same key. + self._additional_callbacks[key] = func + + def debug_advance_all_joint_data(self): + """Provides the properly scaled, all tracking results received from Rokoko Studio. + It is intended to use a debug and visualization function inspecting all data from Rokoko Glove. + + Returns: + A tuple containing the 2D (42,7) pose array(left:0-21, right:21-42), and a dummy truth value. + """ + body_data = self._get_gloves_data() + + # for joint_name in self.left_hand_joint_names: + # joint_data = body_data[joint_name] + # joint_position = torch.tensor(list(joint_data["position"].values())) + # joint_rotation = torch.tensor(list(joint_data["rotation"].values())) + # self.left_fingertip_poses[self.left_joint_dict[joint_name]][:3] = joint_position + # self.left_fingertip_poses[self.left_joint_dict[joint_name]][3:] = joint_rotation + + for joint_name in self.right_hand_joint_names: + joint_data = body_data[joint_name] + joint_position = torch.tensor(list(joint_data["position"].values())) + joint_rotation = torch.tensor(list(joint_data["rotation"].values())) + self.right_fingertip_poses[self.right_joint_dict[joint_name]][:3] = joint_position + self.right_fingertip_poses[self.right_joint_dict[joint_name]][3:] = joint_rotation + + # left_wrist_position = self.left_fingertip_poses[0][:3] + + # self.left_fingertip_poses[:, :3] = (self.left_fingertip_poses[:, :3] - left_wrist_position) * self.scale + left_wrist_position + # self.fingertip_poses[:len(self.left_fingertip_poses)] = self.left_fingertip_poses + + right_wrist_position = self.right_fingertip_poses[0][:3] + # scale + rightThumbProximalIdx = self.right_joint_dict["rightThumbProximal"] + rightIndexProximalIdx = self.right_joint_dict["rightIndexProximal"] + rightMiddleProximalIdx = self.right_joint_dict["rightMiddleProximal"] + rightRingProximalIdx = self.right_joint_dict["rightRingProximal"] + rightLittleProximalIdx = self.right_joint_dict["rightLittleProximal"] + # scale + self.right_fingertip_poses[:, :3] = ( + self.right_fingertip_poses[:, :3] - right_wrist_position + ) * self.scale + right_wrist_position + t_idx = rightThumbProximalIdx + self.right_fingertip_poses[t_idx : t_idx + 4, :3] = ( + self.right_fingertip_poses[t_idx : t_idx + 4, :3] - right_wrist_position + ) * self.thumb_scale + right_wrist_position + # reposition + reposition_vector = ( + self.right_fingertip_poses[rightMiddleProximalIdx][:3] + - self.right_fingertip_poses[rightIndexProximalIdx][:3] + ) + self.right_fingertip_poses[rightIndexProximalIdx:, :3] += self.proximal_offset * reposition_vector + self.right_fingertip_poses[rightMiddleProximalIdx:, :3] += self.proximal_offset * reposition_vector + self.right_fingertip_poses[rightRingProximalIdx:, :3] += self.proximal_offset * reposition_vector + self.right_fingertip_poses[rightLittleProximalIdx:, :3] += self.proximal_offset * reposition_vector + self.fingertip_poses[len(self.left_fingertip_poses) :] = self.right_fingertip_poses + + return self.fingertip_poses, True diff --git a/source/uwlab/uwlab/devices/se3_keyboard.py b/source/uwlab/uwlab/devices/se3_keyboard.py new file mode 100644 index 00000000..5fdedf8e --- /dev/null +++ b/source/uwlab/uwlab/devices/se3_keyboard.py @@ -0,0 +1,109 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.devices import Se3Keyboard +from isaaclab.utils.math import matrix_from_quat + +if TYPE_CHECKING: + from uwlab.devices import KeyboardCfg + + +class Se3Keyboard(Se3Keyboard): + """A keyboard controller for sending SE(3) commands as absolute poses and binary command (open/close). + + This class is designed to provide a keyboard controller for a robotic arm with a gripper. + It uses the Omniverse keyboard interface to listen to keyboard events and map them to robot's + task-space commands. Different from Isaac Lab Se3Keyboard that adds delta command to current robot pose, + this implementation controls a target pose which robot actively tracks. this error corrective property is + advantageous to use when robot is prone to drift or under actuated + + The command comprises of two parts: + + * absolute pose: a 6D vector of (x, y, z, roll, pitch, yaw) in meters and radians. + * gripper: a binary command to open or close the gripper. + + Key bindings: + ============================== ================= ================= + Description Key (+ve axis) Key (-ve axis) + ============================== ================= ================= + Toggle gripper (open/close) K + Move along x-axis W S + Move along y-axis A D + Move along z-axis Q E + Rotate along x-axis Z X + Rotate along y-axis T G + Rotate along z-axis C V + ============================== ================= ================= + + .. seealso:: + + The official documentation for the keyboard interface: `Carb Keyboard Interface `__. + + """ + + def __init__(self, cfg: KeyboardCfg, device="cuda:0"): + """Initialize the keyboard layer. + + Args: + pos_sensitivity: Magnitude of input position command scaling. Defaults to 0.05. + rot_sensitivity: Magnitude of scale input rotation commands scaling. Defaults to 0.5. + """ + super().__init__(cfg.pos_sensitivity, cfg.rot_sensitivity) + self.device = device + self.enable_gripper_command = cfg.enable_gripper_command + self.init_pose = torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0]], device=self.device) + self.abs_pose = torch.zeros((1, 6), device=self.device) + + """ + Operations + """ + + def reset(self): + super().reset() + self.abs_pose = self.init_pose.clone() + + def advance(self): + """Provides the result from internal target command modified by keyboard event. + + Returns: + A tuple containing the absolute pose command and gripper commands. + """ + delta_pose, gripper_command = super().advance() + delta_pose = delta_pose.astype("float32") + # convert to torch + delta_pose = torch.tensor(delta_pose, device=self.device).view(1, -1) + self.abs_pose[:, :3] += delta_pose[:, :3] + self.abs_pose[:, 3:] += delta_pose[:, 3:] + + if self.enable_gripper_command: + if gripper_command: + gripper_command = torch.tensor([[1.0]], device=self.device) + else: + gripper_command = torch.tensor([[-1.0]], device=self.device) + return self.abs_pose.clone(), gripper_command + else: + return self.abs_pose.clone() + + +def apply_local_translation(current_position, local_translation, orientation_quaternion): + # Assuming matrix_from_quat correctly handles batch inputs and outputs a batch of rotation matrices + rotation_matrix = matrix_from_quat(orientation_quaternion) # Expected shape (n, 3, 3) + + # Ensure local_translation is correctly shaped for batch matrix multiplication + local_translation = local_translation.unsqueeze(-1) # Shape becomes (n, 3, 1) for matmul + + local_translation[:, [1, 2]] = -local_translation[:, [2, 1]] + # Rotate the local translation vector to align with the global frame + global_translation = torch.matmul(rotation_matrix, local_translation).squeeze(-1) # Back to shape (n, 3) + + # Apply the translated vector to the object's current position + new_position = current_position + global_translation + + return new_position diff --git a/source/uwlab/uwlab/devices/teleop.py b/source/uwlab/uwlab/devices/teleop.py new file mode 100644 index 00000000..a7bb49ff --- /dev/null +++ b/source/uwlab/uwlab/devices/teleop.py @@ -0,0 +1,225 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from dataclasses import dataclass +from typing import TYPE_CHECKING, Callable, Literal + +import isaaclab.utils.math as math_utils +from isaaclab.devices import DeviceBase +from isaaclab.managers import SceneEntityCfg +from isaaclab.markers import VisualizationMarkers +from isaaclab.markers.config import CUBOID_MARKER_CFG, FRAME_MARKER_CFG +from uwlab.utils.math import create_axis_remap_function + +if TYPE_CHECKING: + from isaaclab.assets import Articulation + + from .teleop_cfg import TeleopCfg + + +@dataclass +class TeleopState: + device_interface: DeviceBase + axis_remap: tuple[str, str, str] + command_type: Literal["position", "pose"] + attach_body: SceneEntityCfg + attach_scope: Literal["self", "descendants"] + pose_reference_body: SceneEntityCfg + debug_vis: bool = False + num_command_body: int = None + reorient_func: Callable = None + ref_pos_b: torch.Tensor = None + ref_quat_b: torch.Tensor = None + attach_pos_b: torch.Tensor = None + attach_quat_b: torch.Tensor = None + + def update_ref(self, robot: Articulation): + ref_body_id = self.pose_reference_body.body_ids + ref_pos_b, ref_quat_b = math_utils.subtract_frame_transforms( + robot.data.root_pos_w, + robot.data.root_quat_w, + robot.data.body_link_pos_w[:, ref_body_id, :].view(-1, 3), + robot.data.body_link_quat_w[:, ref_body_id, :].view(-1, 4), + ) + self.ref_pos_b = ref_pos_b.repeat_interleave(self.num_command_body, dim=0) + self.ref_quat_b = ref_quat_b.repeat_interleave(self.num_command_body, dim=0) + + def update_attach(self, robot: Articulation): + attach_body_id = self.attach_body.body_ids + self.attach_pos_b, self.attach_quat_b = math_utils.subtract_frame_transforms( + robot.data.root_pos_w, + robot.data.root_quat_w, + robot.data.body_link_pos_w[:, attach_body_id, :].view(-1, 3), + robot.data.body_link_quat_w[:, attach_body_id, :].view(-1, 4), + ) + self.attach_pos_b = self.attach_pos_b.repeat_interleave(self.num_command_body, dim=0) + self.attach_quat_b = self.attach_quat_b.repeat_interleave(self.num_command_body, dim=0) + + def combine_frame_on_root( + self, robot: Articulation, command_pos_b: torch.Tensor, command_quat_b: torch.Tensor | None + ): + command_pos_w, command_quat_w = math_utils.combine_frame_transforms( + robot.data.root_pos_w.repeat_interleave(self.num_command_body, dim=0), + robot.data.root_quat_w.repeat_interleave(self.num_command_body, dim=0), + command_pos_b, + command_quat_b, + ) + return command_pos_w, command_quat_w + + def combine_frame_on_ref_b(self, command_pos: torch.Tensor, command_rot: torch.Tensor | None): + if command_rot is not None and command_rot.shape[1] == 4: + command_rot = math_utils.axis_angle_from_quat(command_rot[:, 3:]) + + command_pos, command_rot = self.reorient_func(command_pos, command_rot) + if command_rot is not None: + command_rot = math_utils.quat_from_euler_xyz(command_rot[:, 0], command_rot[:, 1], command_rot[:, 2]) + command_pos_b, command_quat_b = math_utils.combine_frame_transforms( + t01=self.ref_pos_b, + q01=self.ref_quat_b, + t12=command_pos, + q12=command_rot, + ) + return command_pos_b, command_quat_b + + def combine_frame_on_attach_b(self, command_pos: torch.Tensor, command_quat: torch.Tensor | None): + if self.attach_scope == "descendants": + command_pos_b, command_quat_b = math_utils.combine_frame_transforms( + t01=self.attach_pos_b, + q01=self.attach_quat_b, + t12=command_pos, + q12=command_quat, + ) + else: + command_pos_b = self.attach_pos_b + command_pos + if command_quat is None: + return command_pos_b, None + command_quat_b = math_utils.quat_mul(command_quat, self.attach_quat_b) + return command_pos_b, command_quat_b + + +class Teleop: + def __init__(self, cfg: TeleopCfg, env): + self._env = env + self.num_envs = env.unwrapped.num_envs + self.cfg = cfg + devices = cfg.teleop_devices + self.teleops: list[TeleopState] = [] + for device in devices.values(): + teleop_device = TeleopState( + device_interface=device.teleop_interface_cfg.class_type(device.teleop_interface_cfg, env.device), + axis_remap=device.reference_axis_remap, + command_type=device.command_type, + attach_body=device.attach_body, + attach_scope=device.attach_scope, + pose_reference_body=device.pose_reference_body, + debug_vis=device.debug_vis, + ) + self.teleops.append(teleop_device) + teleop_device.device_interface.reset() + teleop_command = teleop_device.device_interface.advance() + teleop_command = (teleop_command,) if not isinstance(teleop_command, tuple) else teleop_command + teleop_device.num_command_body = teleop_command[0].shape[0] + + total_command_bodys = self.num_envs * teleop_device.num_command_body + teleop_device.ref_pos_b = torch.zeros((total_command_bodys, 3), device=env.device) + teleop_device.ref_quat_b = torch.zeros((total_command_bodys, 4), device=env.device) + teleop_device.attach_pos_b = torch.zeros((total_command_bodys, 3), device=env.device) + teleop_device.attach_quat_b = torch.zeros((total_command_bodys, 4), device=env.device) + + device.attach_body.resolve(env.unwrapped.scene) + device.pose_reference_body.resolve(env.unwrapped.scene) + + self.device = env.device + self.command_w = None + + frame_marker_cfg = FRAME_MARKER_CFG.copy() # type: ignore + frame_marker_cfg.markers["frame"].scale = (0.025, 0.025, 0.025) + cuboid_marker_cfg = CUBOID_MARKER_CFG.copy() # type: ignore + cuboid_marker_cfg.markers["cuboid"].size = (0.01, 0.01, 0.01) + self.pose_marker = VisualizationMarkers(frame_marker_cfg.replace(prim_path="/Visuals/device_pose")) + self.position_marker = VisualizationMarkers(cuboid_marker_cfg.replace(prim_path="/Visuals/device_position")) + self.robot = self._env.unwrapped.scene["robot"] + + @property + def command(self): + return self.command_w + + def add_callback(self, key: str, func: Callable): + # check keys supported by callback + for teleop_device in self.teleops: + teleop_device.device_interface.add_callback(key, func) + + def reset(self): + robot = self._env.unwrapped.scene["robot"] + for teleop_device in self.teleops: + teleop_device.device_interface.reset() + + teleop_device.update_ref(robot) + teleop_device.update_attach(robot) + teleop_device.reorient_func = create_axis_remap_function(*teleop_device.axis_remap, device=self.device) + + def advance(self): + self.command_b = [] + self.visualization_pose = [] + self.visualization_position = [] + for teleop_device in self.teleops: + teleop_command = teleop_device.device_interface.advance() + teleop_command = (teleop_command,) if not isinstance(teleop_command, tuple) else teleop_command + teleop_command = self._world_command_to_robot_command(teleop_device, *teleop_command) + self.command_b.append(teleop_command) + + self._debug_vis() + return torch.cat(self.command_b, dim=1) + + def _debug_vis(self): + if len(self.visualization_pose) > 0: + pose = torch.cat(self.visualization_pose, dim=0).view(-1, 7) + self.pose_marker.visualize( + translations=pose[:, :3], + orientations=pose[:, 3:], + ) + + if len(self.visualization_position) > 0: + position = torch.cat(self.visualization_position, dim=0).view(-1, 3) + self.position_marker.visualize( + translations=position, + ) + + def _world_command_to_robot_command(self, teleop_device: TeleopState, command: torch.Tensor, *args): + robot = self._env.unwrapped.scene["robot"] + command_shape = command.shape[1] + + teleop_device.update_ref(robot) + if teleop_device.attach_scope == "descendants": + teleop_device.update_attach(robot) + + command = command.repeat(self.num_envs, 1) + if command_shape == 3: + command_pos_b, _ = teleop_device.combine_frame_on_ref_b(command[:, :3], None) + command_pos_b, _ = teleop_device.combine_frame_on_attach_b(command_pos_b, None) + + if teleop_device.debug_vis: + pos_w, _ = teleop_device.combine_frame_on_root(robot, command_pos_b, None) + self.visualization_position.append(pos_w) + + command_res = command_pos_b.view(self.num_envs, -1) + else: + command_pos_b, command_quat_b = teleop_device.combine_frame_on_ref_b(command[:, :3], command[:, 3:]) + command_pos_b, command_quat_b = teleop_device.combine_frame_on_attach_b(command_pos_b, command_quat_b) + + if teleop_device.debug_vis: + pos_w, quat_w = teleop_device.combine_frame_on_root(robot, command_pos_b, command_quat_b) + self.visualization_pose.append(torch.cat([pos_w, quat_w], dim=1)) + + command_res = torch.cat([command_pos_b, command_quat_b], dim=1).view(self.num_envs, -1) + + if len(args) > 0: + arg_cat = torch.cat(args, dim=1).repeat(self.num_envs, 1) + command_res = torch.cat([command_res, arg_cat], dim=1) + + return command_res diff --git a/source/uwlab/uwlab/devices/teleop_cfg.py b/source/uwlab/uwlab/devices/teleop_cfg.py new file mode 100644 index 00000000..2a1db4a2 --- /dev/null +++ b/source/uwlab/uwlab/devices/teleop_cfg.py @@ -0,0 +1,36 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import MISSING +from typing import Callable, Literal + +from isaaclab.managers import SceneEntityCfg +from isaaclab.utils import configclass +from uwlab.devices import DeviceBaseTeleopCfg + +from .teleop import Teleop + + +@configclass +class TeleopCfg: + @configclass + class TeleopDevicesCfg: + teleop_interface_cfg: DeviceBaseTeleopCfg = MISSING + + debug_vis: bool = False + + attach_body: SceneEntityCfg = MISSING + + attach_scope: Literal["self", "descendants"] = "self" + + command_type: Literal["position", "pose"] = MISSING + + pose_reference_body: SceneEntityCfg = MISSING + + reference_axis_remap: tuple[str, str, str] = MISSING + + class_type: Callable[..., Teleop] = Teleop + + teleop_devices: dict[str, TeleopDevicesCfg] = {} diff --git a/source/uwlab/uwlab/envs/__init__.py b/source/uwlab/uwlab/envs/__init__.py new file mode 100644 index 00000000..ffd07cc8 --- /dev/null +++ b/source/uwlab/uwlab/envs/__init__.py @@ -0,0 +1,33 @@ +# Copyright (c) 2022-2025, The UW Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-package for environment definitions. + +Environments define the interface between the agent and the simulation. +In the simplest case, the environment provides the agent with the current +observations and executes the actions provided by the agent. However, the +environment can also provide additional information such as the current +reward, done flag, and information about the current episode. + +UW Lab provide Data-Manager-based and Real-Rl environment workflows: + +* **Data-Manager-based**: The IsaacLab Manager Based Rl Environments is designed such that no manager role's is to + hold and share fresh state data and complex data structure. The Command Manager can holds prev_step command + that is used to condition on the policy but calculated data is inherently one step behind if used by + Reward Manager or Event Managers. Data Manager calculates the state immediately after the action so the state + information is always up to date. more, the data doesn't needs to be tensor and is designed to be shared across + managers with data manager providing cfg for the data retrieval. + +* **Real-Rl**: The Real-Rl Environments is designed to be used in real robot applications. The environment uses + universal robot articulation to abstract the real hardware control loop and provides a simple interface to interact + with the rest of uwlab and isaaclab tool kit. This provides seamless experience from simulation to real + +Based on these workflows, there are the following environment classes for single and multi-agent RL: +""" + +from .data_manager_based_rl import DataManagerBasedRLEnv +from .data_manager_based_rl_cfg import DataManagerBasedRLEnvCfg +from .real_rl_env import RealRLEnv +from .real_rl_env_cfg import RealRLEnvCfg diff --git a/source/uwlab/uwlab/envs/data_manager_based_rl.py b/source/uwlab/uwlab/envs/data_manager_based_rl.py new file mode 100644 index 00000000..9f2ad30b --- /dev/null +++ b/source/uwlab/uwlab/envs/data_manager_based_rl.py @@ -0,0 +1,312 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from collections.abc import Sequence + +from isaaclab.envs import ManagerBasedEnvCfg, ManagerBasedRLEnv +from isaaclab.envs.manager_based_env import VecEnvObs +from isaaclab.managers import EventManager +from isaaclab.markers import VisualizationMarkers +from isaaclab.markers.config import FRAME_MARKER_CFG +from isaaclab.utils.timer import Timer +from uwlab.scene import InteractiveScene + +from ..managers.data_manager import DataManager +from .data_manager_based_rl_cfg import DataManagerBasedRLEnvCfg + +VecEnvStepReturn = tuple[VecEnvObs, torch.Tensor, torch.Tensor, torch.Tensor, dict] + + +class DataManagerBasedRLEnv(ManagerBasedRLEnv): + """The superclass for flexible reinforcement learning-based environments. + + This class extending the :class:`ManagerBasedRLEnv` and implements the flexible modules + that makes experimentation and hacking easier while still using manager based style. + Specifically, the two fields, extensions and data manager. + The field self.extension is a dictionary serves as a place to store any data that needs + a centralized access that can not be modularized by managers, this is usually the last + resort when no other optional are easy and it is not recommended to use in any formalized code. + Data Manager in another hand, can be used to provide more structured data structure to be accessed + by rest of managers. In practice, this module should be used carefully and avoid at best + when formalizing the code and only use for advanced special treatment. + + """ + + cfg: DataManagerBasedRLEnvCfg + """Configuration for the environment.""" + + def __init__(self, cfg: DataManagerBasedRLEnvCfg, render_mode: str | None = None, **kwargs): + self.extensions = {} + # there is no good way to override the InteractiveScene class, so we need to patch it + self._interactive_scene_patched_manager_based_env_init(cfg) + self._manager_based_rl_env_init(cfg, render_mode) + # RLTaskEnv.__bases__ = (BaseEnv, gym.Env) + frame_marker_cfg = FRAME_MARKER_CFG.copy() + frame_marker_cfg.markers["frame"].scale = (0.02, 0.02, 0.02) + self.goal_marker = VisualizationMarkers(frame_marker_cfg.replace(prim_path="/Visuals/ee_goal")) + self.symmetry_augmentation_func = None + if hasattr(cfg, "symmetry_augmentation_func"): + self.symmetry_augmentation_func = cfg.symmetry_augmentation_func + + def _interactive_scene_patched_manager_based_env_init(self, cfg: ManagerBasedEnvCfg): + import builtins + + import omni.log + + from isaaclab.envs.ui import ViewportCameraController + from isaaclab.sim import SimulationContext + + # check that the config is valid + cfg.validate() + # store inputs to class + self.cfg = cfg + # initialize internal variables + self._is_closed = False + + # set the seed for the environment + if self.cfg.seed is not None: + self.cfg.seed = self.seed(self.cfg.seed) + else: + omni.log.warn("Seed not set for the environment. The environment creation may not be deterministic.") + + # create a simulation context to control the simulator + if SimulationContext.instance() is None: + # the type-annotation is required to avoid a type-checking error + # since it gets confused with Isaac Sim's SimulationContext class + self.sim: SimulationContext = SimulationContext(self.cfg.sim) + else: + # simulation context should only be created before the environment + # when in extension mode + if not builtins.ISAAC_LAUNCHED_FROM_TERMINAL: + raise RuntimeError("Simulation context already exists. Cannot create a new one.") + self.sim: SimulationContext = SimulationContext.instance() + + # make sure torch is running on the correct device + if "cuda" in self.device: + torch.cuda.set_device(self.device) + + # print useful information + print("[INFO]: Base environment:") + print(f"\tEnvironment device : {self.device}") + print(f"\tEnvironment seed : {self.cfg.seed}") + print(f"\tPhysics step-size : {self.physics_dt}") + print(f"\tRendering step-size : {self.physics_dt * self.cfg.sim.render_interval}") + print(f"\tEnvironment step-size : {self.step_dt}") + + 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 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) + + # counter for simulation steps + self._sim_step_counter = 0 + + # generate scene + with Timer("[INFO]: Time taken for scene creation", "scene_creation"): + self.scene = InteractiveScene(self.cfg.scene) + print("[INFO]: Scene manager: ", self.scene) + + # set up camera viewport controller + # viewport is not available in other rendering modes so the function will throw a warning + # FIXME: This needs to be fixed in the future when we unify the UI functionalities even for + # non-rendering modes. + if self.sim.render_mode >= self.sim.RenderMode.PARTIAL_RENDERING: + self.viewport_camera_controller = ViewportCameraController(self, self.cfg.viewer) + else: + self.viewport_camera_controller = None + + # create event manager + # note: this is needed here (rather than after simulation play) to allow USD-related randomization events + # that must happen before the simulation starts. Example: randomizing mesh scale + self.event_manager = EventManager(self.cfg.events, self) + print("[INFO] Event Manager: ", self.event_manager) + + # apply USD-related randomization events + if "prestartup" in self.event_manager.available_modes: + self.event_manager.apply(mode="prestartup") + + # play the simulator to activate physics handles + # note: this activates the physics simulation view that exposes TensorAPIs + # note: when started in extension mode, first call sim.reset_async() and then initialize the managers + if builtins.ISAAC_LAUNCHED_FROM_TERMINAL is False: + print("[INFO]: Starting the simulation. This may take a few seconds. Please wait...") + with Timer("[INFO]: Time taken for simulation start", "simulation_start"): + self.sim.reset() + # add timeline event to load managers + self.load_managers() + + # extend UI elements + # 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 + self._window = None + + # allocate dictionary to store metrics + self.extras = {} + + # initialize observation buffers + self.obs_buf = {} + + def _manager_based_rl_env_init(self, cfg: DataManagerBasedRLEnvCfg, render_mode): + # store the render mode + self.render_mode = render_mode + + # initialize data and constants + # -- counter for curriculum + self.common_step_counter = 0 + # -- init buffers + self.episode_length_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.long) + # -- set the framerate of the gym video recorder wrapper so that the playback speed of the produced video matches the simulation + self.metadata["render_fps"] = 1 / self.step_dt + + print("[INFO]: Completed setting up the environment...") + + @property + def is_closed(self): + return self._is_closed + + """ + Patches + """ + + def patch_interactive_scene(self): + import sys + + # Make sure we have not imported the original module yet + if "isaaclab.scene" in sys.modules: + del sys.modules["isaaclab.scene"] + + import uwlab.scene + + sys.modules["isaaclab.scene"] = uwlab.scene + + def unpatch_interactive_scene(self): + import sys + + if "isaaclab.scene" in sys.modules: + del sys.modules["isaaclab.scene"] + + """ + Operations - MDP + """ + + def step(self, action: torch.Tensor) -> VecEnvStepReturn: + """Execute one time-step of the environment's dynamics and reset terminated environments. + + Unlike the :class:`ManagerBasedEnvCfg.step` class, the function performs the following operations: + + 1. Process the actions. + 2. Perform physics stepping. + 3. Perform rendering if gui is enabled. + 4. Update the environment counters and compute the rewards and terminations. + 5. Reset the environments that terminated. + 6. Compute the observations. + 7. Return the observations, rewards, resets and extras. + + Args: + action: The actions to apply on the environment. Shape is (num_envs, action_dim). + + Returns: + A tuple containing the observations, rewards, resets (terminated and truncated) and extras. + """ + # process actions + self.action_manager.process_action(action) + + 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() + + for _ in range(self.cfg.decimation): + self._sim_step_counter += 1 + # set actions into buffers + self.action_manager.apply_action() + # set actions into simulator + self.scene.write_data_to_sim() + # simulate + self.sim.step(render=False) + # render between steps only if the GUI or an RTX sensor needs it + # note: we assume the render interval to be the shortest accepted rendering interval. + # If a camera needs rendering at a faster frequency, this will lead to unexpected behavior. + if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering: + self.sim.render() + # update buffers at sim dt + self.scene.update(dt=self.physics_dt) + + if hasattr(self, "data_manager"): + self.data_manager.compute() + # post-step: + # -- update env counters (used for curriculum generation) + self.episode_length_buf += 1 # step in current episode (per env) + self.common_step_counter += 1 # total step (common for all envs) + # -- check terminations + self.reset_buf = self.termination_manager.compute() + self.reset_terminated = self.termination_manager.terminated + self.reset_time_outs = self.termination_manager.time_outs + # -- 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 + if "interval" in self.event_manager.available_modes: + self.event_manager.apply(mode="interval", dt=self.step_dt) + # -- compute observations + # note: done after reset to get the correct observations for reset envs + self.obs_buf = self.observation_manager.compute() + + # return observations, rewards, resets and extras + return self.obs_buf, self.reward_buf, self.reset_terminated, self.reset_time_outs, self.extras + + def load_managers(self): + if getattr(self.cfg, "data", None) is not None: + self.data_manager: DataManager = DataManager(self.cfg.data, self) + print("[INFO] Data Manager: ", self.data_manager) + super().load_managers() + + def _reset_idx(self, env_ids: Sequence[int]): + if getattr(self, "data_manager", None) is not None: + self.data_manager.reset(env_ids) + return super()._reset_idx(env_ids) + + def close(self): + for key, val in self.extensions.items(): + del val + super().close() diff --git a/source/uwlab/uwlab/envs/data_manager_based_rl_cfg.py b/source/uwlab/uwlab/envs/data_manager_based_rl_cfg.py new file mode 100644 index 00000000..df9916e4 --- /dev/null +++ b/source/uwlab/uwlab/envs/data_manager_based_rl_cfg.py @@ -0,0 +1,14 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from isaaclab.envs.manager_based_rl_env_cfg import ManagerBasedRLEnvCfg +from isaaclab.utils import configclass + + +@configclass +class DataManagerBasedRLEnvCfg(ManagerBasedRLEnvCfg): + data: object | None = None diff --git a/source/uwlab/uwlab/envs/diagnosis/__init__.py b/source/uwlab/uwlab/envs/diagnosis/__init__.py new file mode 100644 index 00000000..e9bbfdb9 --- /dev/null +++ b/source/uwlab/uwlab/envs/diagnosis/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .diagnosis_util import * diff --git a/source/uwlab/uwlab/envs/diagnosis/diagnosis.py b/source/uwlab/uwlab/envs/diagnosis/diagnosis.py new file mode 100644 index 00000000..aca32d3a --- /dev/null +++ b/source/uwlab/uwlab/envs/diagnosis/diagnosis.py @@ -0,0 +1,279 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING, Sequence + +from isaaclab.assets import Articulation +from isaaclab.managers import SceneEntityCfg + +if TYPE_CHECKING: + from uwlab.envs import DataManagerBasedRLEnv + + +def get_link_incoming_joint_force( + env: DataManagerBasedRLEnv, + env_ids: Sequence[int] | torch.Tensor | None, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + """ + This function provides force and torque exerted from child link to parent link in world frame + This force/torque is in WorldFrame -> https://nvidia-omniverse.github.io/PhysX/physx/5.4.1/docs/Joints.html + This force/torque is exerted by the joint connecting child link to the parent link -> + https://docs.omniverse.nvidia.com/isaacsim/latest/features/sensors_simulation/isaac_sim_sensors_physics_articulation_force.html + + Return measured joint forces and torques. Shape is (num_articulation, num_joint + 1, 6). Row index 0 is the incoming + joint of the base link. For the last dimension the first 3 values are for forces and the last 3 for torques + """ + asset: Articulation = env.scene[asset_cfg.name] + if env_ids is None: + env_ids = slice(None) + force_from_child_link_to_joints = asset.root_physx_view.get_link_incoming_joint_force().to(env.device)[env_ids] + return force_from_child_link_to_joints + + +def get_dof_projected_joint_forces( + env: DataManagerBasedRLEnv, + env_ids: Sequence[int] | torch.Tensor | None, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + """ + get_measured_joint_efforts specifies the active component (the projection of the joint forces on the + motion direction) of the joint force for all the joints and articulations. + + My guess: get_measured_joint_efforts() provides the actual active efforts acting along the joint axes during the simulation. + It reflects the effective efforts being applied to the joints, considering factors such as: Actuator limitations, + Dynamics of the system, External forces and constraints, Interactions with the environment + + Return dimension: (num_articulations, num_links) + + """ + asset: Articulation = env.scene[asset_cfg.name] + if env_ids is None: + env_ids = slice(None) + projected_joint_forces = asset.root_physx_view.get_dof_projected_joint_forces().to(env.device)[env_ids] + return projected_joint_forces + + +def get_dof_applied_torque( + env: DataManagerBasedRLEnv, + env_ids: Sequence[int] | torch.Tensor | None, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + """ + get the torque that joint drive applies + Return dimension: (num_articulations, num_links) + """ + asset: Articulation = env.scene[asset_cfg.name] + if env_ids is None: + env_ids = slice(None) + applied_torque = asset.data.applied_torque[env_ids] + return applied_torque + + +def get_dof_computed_torque( + env: DataManagerBasedRLEnv, + env_ids: Sequence[int] | torch.Tensor | None, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + asset: Articulation = env.scene[asset_cfg.name] + if env_ids is None: + env_ids = slice(None) + computed_torque = asset.data.computed_torque[env_ids] + return computed_torque + + +def get_dof_target_position( + env: DataManagerBasedRLEnv, + env_ids: Sequence[int] | torch.Tensor | None, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + asset: Articulation = env.scene[asset_cfg.name] + if env_ids is None: + env_ids = slice(None) + target_joint_pos = asset.data.joint_pos_target[env_ids] + return target_joint_pos + + +def get_dof_position( + env: DataManagerBasedRLEnv, + env_ids: Sequence[int] | torch.Tensor | None, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + asset: Articulation = env.scene[asset_cfg.name] + if env_ids is None: + env_ids = slice(None) + joint_pos = asset.data.joint_pos[env_ids] + return joint_pos + + +def get_dof_target_velocity( + env: DataManagerBasedRLEnv, + env_ids: Sequence[int] | torch.Tensor | None, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + asset: Articulation = env.scene[asset_cfg.name] + if env_ids is None: + env_ids = slice(None) + target_joint_vel = asset.data.joint_vel_target[env_ids] + return target_joint_vel + + +def get_dof_velocity( + env: DataManagerBasedRLEnv, + env_ids: Sequence[int] | torch.Tensor | None, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + asset: Articulation = env.scene[asset_cfg.name] + if env_ids is None: + env_ids = slice(None) + joint_vel = asset.data.joint_vel[env_ids] + return joint_vel + + +def get_dof_acceleration( + env: DataManagerBasedRLEnv, + env_ids: Sequence[int] | torch.Tensor | None, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + asset: Articulation = env.scene[asset_cfg.name] + if env_ids is None: + env_ids = slice(None) + joint_acc = asset.data.joint_acc[env_ids] + return joint_acc + + +def get_action_rate( + env: DataManagerBasedRLEnv, + env_ids: Sequence[int] | torch.Tensor | None, +) -> torch.Tensor: + return torch.square(env.action_manager.action[env_ids] - env.action_manager.prev_action[env_ids]) + + +def get_joint_torque_utilization( + env: DataManagerBasedRLEnv, + env_ids: Sequence[int] | torch.Tensor | None, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + asset: Articulation = env.scene[asset_cfg.name] + if env_ids is None: + env_ids = slice(None) + applied_torque = asset.data.applied_torque[env_ids] + torque_max = asset.root_physx_view.get_dof_max_forces().to(env.device)[env_ids] + torque_utilization = torch.abs(applied_torque) / torque_max + return torque_utilization + + +def get_joint_velocity_utilization( + env: DataManagerBasedRLEnv, + env_ids: Sequence[int] | torch.Tensor | None, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + asset: Articulation = env.scene[asset_cfg.name] + if env_ids is None: + env_ids = slice(None) + joint_vel = asset.data.joint_vel[env_ids] + max_vel = asset.root_physx_view.get_dof_max_velocities().to(env.device)[env_ids] + velocity_utilization = torch.abs(joint_vel) / max_vel + return velocity_utilization + + +def get_joint_power( + env: DataManagerBasedRLEnv, + env_ids: Sequence[int] | torch.Tensor | None, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + asset: Articulation = env.scene[asset_cfg.name] + if env_ids is None: + env_ids = slice(None) + joint_vel = asset.data.joint_vel[env_ids] + applied_torque = asset.data.applied_torque[env_ids] + power = torch.abs(applied_torque * joint_vel) + return power + + +def get_joint_mechanical_work( + env: DataManagerBasedRLEnv, + env_ids: Sequence[int] | torch.Tensor | None, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + asset: Articulation = env.scene[asset_cfg.name] + if env_ids is None: + env_ids = slice(None) + joint_pos = asset.data.joint_pos[env_ids] + applied_torque = asset.data.applied_torque[env_ids] + if "prev_joint_pos" in env.extensions: + delta_joint_pos = joint_pos - env.extensions["prev_joint_pos"] + mechanical_work = torch.abs(delta_joint_pos * applied_torque) + env.extensions["prev_joint_pos"] = joint_pos + return mechanical_work + else: + env.extensions["prev_joint_pos"] = joint_pos + return torch.zeros(joint_pos.shape, device=env.device) + + +def effective_torque( + env: DataManagerBasedRLEnv, + env_ids: Sequence[int] | torch.Tensor | None, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + """ + Calculate the ratio of projected joint forces over applied joint forces for each joint. + + Args: + env: The simulation environment. + env_ids: Environment IDs (optional, not used here). + asset_cfg: Asset configuration. + + Returns: + force_ratio: Tensor of the ratio of projected to applied joint forces. + Shape: (num_envs, num_joints) + """ + asset: Articulation = env.scene[asset_cfg.name] + if env_ids is None: + env_ids = slice(None) + # Get applied joint torques (from actuators) + applied_torque = asset.data.applied_torque[env_ids] # Shape: (num_envs, num_joints) + + # Get projected joint forces + projected_joint_forces = asset.root_physx_view.get_dof_projected_joint_forces().to(env.device)[ + env_ids + ] # Shape: (num_envs, num_joints) + + # Handle zero applied torques to avoid division by zero + epsilon = 1e-6 # Small constant + applied_torque[torch.abs(applied_torque) < epsilon] = epsilon + + # Calculate the ratio + force_ratio = projected_joint_forces / applied_torque # Element-wise division + + return force_ratio + + +def get_dof_weight_distribution( + env: DataManagerBasedRLEnv, + env_ids: Sequence[int] | torch.Tensor | None, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + """ + Calculate weight distribution across forelegs and hindlegs based on joint forces, considering the gravity component. + + Args: + env: The simulation environment. + env_ids: Environment IDs (optional, not used here). + asset_cfg: Asset configuration. + + Returns: + weight_distribution: A tensor with weight distribution across forelegs and hindlegs. + Shape: (num_envs, 2) -> [weight_on_forelegs, weight_on_hindlegs] + """ + asset: Articulation = env.scene[asset_cfg.name] + if env_ids is None: + env_ids = slice(None) + force_from_child_link_to_joints = asset.root_physx_view.get_link_incoming_joint_force().to(env.device)[env_ids] + weight_forces = force_from_child_link_to_joints[..., 2] + return weight_forces diff --git a/source/uwlab/uwlab/envs/diagnosis/diagnosis_util.py b/source/uwlab/uwlab/envs/diagnosis/diagnosis_util.py new file mode 100644 index 00000000..e7fcb52d --- /dev/null +++ b/source/uwlab/uwlab/envs/diagnosis/diagnosis_util.py @@ -0,0 +1,56 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch + +import isaaclab.utils.math as math_utils + + +@torch.jit.script +def get_dof_stress_von_mises_optimized( + forces_w: torch.Tensor, + torques_w: torch.Tensor, + bodies_quat: torch.Tensor, + A: torch.Tensor, + I: torch.Tensor, + J: torch.Tensor, + c: float, +) -> torch.Tensor: + # Transform forces and torques to body frame + forces_b = math_utils.quat_rotate_inverse(bodies_quat, forces_w) + torques_b = math_utils.quat_rotate_inverse(bodies_quat, torques_w) + + # Extract force and torque components + F = forces_b # Shape: (..., 3) + M = torques_b # Shape: (..., 3) + + # Precompute ratios + F_over_A = F / A # Shape: (..., 3) + M_c_over_IJ = M * c # Shape: (..., 3) + M_c_over_IJ[..., :2] /= I # M_x and M_y over I + M_c_over_IJ[..., 2] /= J # M_z over J + + # Compute axial stress + sigma_axial = F_over_A[..., 2] # F_z / A + + # Compute bending stress + sigma_bending = M_c_over_IJ[..., 1] + M_c_over_IJ[..., 0] # M_y*c/I + M_x*c/I + + # Compute shear stress due to torsion + tau_torsion = M_c_over_IJ[..., 2] # M_z*c/J + + # Compute shear stress due to shear forces + tau_shear = torch.sqrt(F_over_A[..., 0] ** 2 + F_over_A[..., 1] ** 2) + + # Total equivalent shear stress + tau_eq = torch.sqrt(tau_torsion**2 + tau_shear**2) + + # Total equivalent normal stress + sigma_eq = sigma_axial + sigma_bending + + # Compute von Mises stress + sigma_von_mises = torch.sqrt(sigma_eq**2 + 3.0 * tau_eq**2) + + return sigma_von_mises diff --git a/source/uwlab/uwlab/envs/mdp/__init__.py b/source/uwlab/uwlab/envs/mdp/__init__.py new file mode 100644 index 00000000..85937c3b --- /dev/null +++ b/source/uwlab/uwlab/envs/mdp/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .actions import * +from .commands import * +from .events import * +from .observations import * +from .rewards import * +from .terminations import * # noqa: F401, F403 diff --git a/source/uwlab/uwlab/envs/mdp/actions/__init__.py b/source/uwlab/uwlab/envs/mdp/actions/__init__.py new file mode 100644 index 00000000..a01912b0 --- /dev/null +++ b/source/uwlab/uwlab/envs/mdp/actions/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Various action terms that can be used in the environment.""" + +from .actions_cfg import * +from .default_joint_static_action import * +from .task_space_actions import * diff --git a/source/uwlab/uwlab/envs/mdp/actions/actions_cfg.py b/source/uwlab/uwlab/envs/mdp/actions/actions_cfg.py new file mode 100644 index 00000000..7e9374fd --- /dev/null +++ b/source/uwlab/uwlab/envs/mdp/actions/actions_cfg.py @@ -0,0 +1,98 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from dataclasses import MISSING + +from isaaclab.envs.mdp import JointActionCfg, JointPositionActionCfg +from isaaclab.envs.mdp.actions import DifferentialInverseKinematicsActionCfg +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers.action_manager import ActionTerm +from isaaclab.utils import configclass +from uwlab.envs.mdp.actions import ( + default_joint_static_action, + pca_actions, + task_space_actions, + visualizable_joint_target_position, +) + +## +# Task-space Actions. +## + + +@configclass +class MultiConstraintsDifferentialInverseKinematicsActionCfg(DifferentialInverseKinematicsActionCfg): + """Configuration for inverse differential kinematics action term with multi constraints. + This class amend attr body_name from type:str to type:list[str] reflecting its capability to + received the desired positions, poses from multiple target bodies. This will be particularly + useful for controlling dextrous hand robot with only positions of multiple key frame positions + and poses, and output joint positions that satisfy key frame position/pose constrains + + See :class:`DifferentialInverseKinematicsAction` for more details. + """ + + @configclass + class OffsetCfg: + @configclass + class BodyOffsetCfg: + """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).""" + + pose: dict[str, BodyOffsetCfg] = {} + + body_name: list[str] = MISSING + + class_type: type[ActionTerm] = task_space_actions.MultiConstraintDifferentialInverseKinematicsAction + + body_offset: OffsetCfg | None = None + + task_space_boundary: list[tuple[float, float]] | None = None + + +@configclass +class PCAJointPositionActionCfg(JointPositionActionCfg): + """Configuration for the joint position action term. + + See :class:`JointPositionAction` for more details. + """ + + class_type: type[ActionTerm] = pca_actions.PCAJointPositionAction + + eigenspace_path: str = MISSING + + joint_range: tuple[float, float] | dict[str, tuple[float, float]] = MISSING + + +@configclass +class VisualizableJointTargetPositionCfg(JointActionCfg): + """Joint action term that applies the processed actions to the articulation's joints as position commands.""" + + class_type: type[ActionTerm] = visualizable_joint_target_position.VisualizableJointTargetPosition + + articulation_vis_cfg: SceneEntityCfg = MISSING + + +@configclass +class DefaultJointPositionStaticActionCfg(JointActionCfg): + """Configuration for the joint position action term. + + See :class:`JointPositionAction` for more details. + """ + + class_type: type[ActionTerm] = default_joint_static_action.DefaultJointPositionStaticAction + + use_default_offset: bool = True diff --git a/source/uwlab/uwlab/envs/mdp/actions/default_joint_static_action.py b/source/uwlab/uwlab/envs/mdp/actions/default_joint_static_action.py new file mode 100644 index 00000000..587877c9 --- /dev/null +++ b/source/uwlab/uwlab/envs/mdp/actions/default_joint_static_action.py @@ -0,0 +1,42 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.envs.mdp.actions.joint_actions import JointAction + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + from . import actions_cfg + + +class DefaultJointPositionStaticAction(JointAction): + """Joint action term that applies the processed actions to the articulation's joints as position commands.""" + + cfg: actions_cfg.DefaultJointPositionStaticActionCfg + """The configuration of the action term.""" + + def __init__(self, cfg: actions_cfg.DefaultJointPositionStaticActionCfg, env: ManagerBasedEnv): + # initialize the action term + super().__init__(cfg, env) + # use default joint positions as offset + if cfg.use_default_offset: + self._offset = self._asset.data.default_joint_pos[:, self._joint_ids].clone() + self._default_actions = self._asset.data.default_joint_pos[:, self._joint_ids].clone() + + @property + def action_dim(self) -> int: + return 0 + + def process_actions(self, actions: torch.Tensor): + pass + + def apply_actions(self): + # set position targets + self._asset.set_joint_position_target(self._default_actions, joint_ids=self._joint_ids) diff --git a/source/uwlab/uwlab/envs/mdp/actions/pca_actions.py b/source/uwlab/uwlab/envs/mdp/actions/pca_actions.py new file mode 100644 index 00000000..8b84fbae --- /dev/null +++ b/source/uwlab/uwlab/envs/mdp/actions/pca_actions.py @@ -0,0 +1,57 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import numpy as np +import torch +from io import BytesIO +from typing import TYPE_CHECKING + +import requests + +from isaaclab.envs.mdp.actions import JointPositionAction + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + from . import actions_cfg + + +class PCAJointPositionAction(JointPositionAction): + """Joint action term that applies the processed actions to the articulation's joints as position commands.""" + + cfg: actions_cfg.PCAJointPositionActionCfg + """The configuration of the action term.""" + + def __init__(self, cfg: actions_cfg.PCAJointPositionActionCfg, env: ManagerBasedEnv): + # Download the file + response = requests.get(cfg.eigenspace_path) + response.raise_for_status() # Ensure the download was successful + + # Load the numpy data from the downloaded file + np_data = np.load(BytesIO(response.content)) + # load the the eigen_vector from the file + self.eigenspace = torch.from_numpy(np_data).to(torch.float32) + + # initialize the action term + super().__init__(cfg, env) + + self.eigenspace = self.eigenspace.to(self.device) + self.joint_range = self.cfg.joint_range + + def process_actions(self, actions: torch.Tensor): + # set position targets + self._raw_actions[:] = actions + self._processed_actions = self._raw_actions @ self.eigenspace * self._scale + self._offset + self._processed_actions = torch.clamp(self.processed_actions, min=self.joint_range[0], max=self.joint_range[1]) + + """ + Properties. + """ + + @property + def action_dim(self) -> int: + return self.eigenspace.shape[0] diff --git a/source/uwlab/uwlab/envs/mdp/actions/task_space_actions.py b/source/uwlab/uwlab/envs/mdp/actions/task_space_actions.py new file mode 100644 index 00000000..aabb7c62 --- /dev/null +++ b/source/uwlab/uwlab/envs/mdp/actions/task_space_actions.py @@ -0,0 +1,246 @@ +# Copyright (c) 2024-2025, The UW 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.log + +import isaaclab.utils.math as math_utils +from isaaclab.assets.articulation import Articulation +from isaaclab.managers.action_manager import ActionTerm + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + from . import actions_cfg + + +class MultiConstraintDifferentialInverseKinematicsAction(ActionTerm): + r"""Inverse Kinematics action term. + + This action term performs pre-processing of the raw actions using scaling transformation. + + .. math:: + \text{action} = \text{scaling} \times \text{input action} + \text{joint position} = J^{-} \times \text{action} + + where :math:`\text{scaling}` is the scaling applied to the input action, and :math:`\text{input action}` + is the input action from the user, :math:`J` is the Jacobian over the articulation's actuated joints, + and \text{joint position} is the desired joint position command for the articulation's joints. + """ + + cfg: actions_cfg.MultiConstraintsDifferentialInverseKinematicsActionCfg + """The configuration of the action term.""" + _asset: Articulation + """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).""" + + def __init__(self, cfg: actions_cfg.MultiConstraintsDifferentialInverseKinematicsActionCfg, env: ManagerBasedEnv): + # initialize the action term + super().__init__(cfg, env) + + # 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_joints = len(self._joint_ids) + # parse the body index + body_ids, body_names = self._asset.find_bodies(self.cfg.body_name, preserve_order=True) + # save only the first body index + self._body_idx = body_ids + self._body_name = body_names + # 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_body_idx = [i - 1 for i in self._body_idx] + self._jacobi_joint_ids = self._joint_ids + else: + self._jacobi_body_idx = self._body_idx + self._jacobi_joint_ids = [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 body name for the action term {self.__class__.__name__}: {self._body_name} [{self._body_idx}]" + ) + # Avoid indexing across all joints for efficiency + if self._num_joints == self._asset.num_joints: + self._joint_ids = slice(None) + + # create the differential IK controller + self._ik_controller = self.cfg.controller.class_type( + cfg=self.cfg.controller, num_bodies=len(self._body_idx), 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) + + # save the scale as tensors + self._scale = torch.zeros((self.num_envs, self.action_dim), device=self.device) + self._scale[:] = torch.tensor(self.cfg.scale, device=self.device) + + # convert the fixed offsets to torch tensors of batched shape + if self.cfg.body_offset is not None: + self._offset_pos = torch.zeros((self.num_envs, len(self._body_idx), 3), device=self.device) + self._offset_rot = torch.zeros((self.num_envs, len(self._body_idx), 4), device=self.device) + for body_names, pose in self.cfg.body_offset.pose.items(): + offset_body_ids, offset_body_names = self._asset.find_bodies(body_names, preserve_order=True) + offset_body_ids = [self._body_idx.index(i) for i in offset_body_ids] + self._offset_pos[:, offset_body_ids] = torch.tensor(pose.pos, device=self.device) + self._offset_rot[:, offset_body_ids] = torch.tensor(pose.rot, device=self.device) + else: + self._offset_pos, self._offset_rot = None, None + + self.enable_task_space_boundary = False + if self.cfg.task_space_boundary is not None: + task_space_boundary = torch.tensor(self.cfg.task_space_boundary, device=self.device) + self.enable_task_space_boundary = True + self.min_limits = task_space_boundary[:, 0] + self.max_limits = task_space_boundary[:, 1] + + self.joint_pos_des = torch.zeros(self.num_envs, self.action_dim, device=self.device) + + """ + Properties. + """ + + @property + def action_dim(self) -> int: + return self._ik_controller.action_dim + + @property + def raw_actions(self) -> torch.Tensor: + return self._raw_actions + + @property + def processed_actions(self) -> torch.Tensor: + return self._processed_actions + + @property + def desired_joint_position(self): + return self.joint_pos_des + + @property + def jacobian_w(self) -> torch.Tensor: + return self._asset.root_physx_view.get_jacobians()[:, self._jacobi_body_idx, :, :][:, :, :, self._joint_ids] + + @property + def jacobian_b(self) -> torch.Tensor: + jacobian = self.jacobian_w + B = len(self._jacobi_body_idx) + rot_b = self._asset.data.root_link_quat_w + rot_b_m = math_utils.matrix_from_quat(math_utils.quat_inv(rot_b)) + rot_b_m = rot_b_m.unsqueeze(1).expand(-1, B, -1, -1).reshape(-1, 3, 3) # [N*B, 3, 3] + + jacobian_pos = jacobian[:, :, :3, :].view(-1, 3, self._num_joints) # [N*B, 3, #joints] + jacobian_rot = jacobian[:, :, 3:, :].view(-1, 3, self._num_joints) # [N*B, 3, #joints] + # multiply and reshape back + jacobian[:, :, :3, :] = torch.bmm(rot_b_m, jacobian_pos).view(self.num_envs, B, 3, -1) + jacobian[:, :, 3:, :] = torch.bmm(rot_b_m, jacobian_rot).view(self.num_envs, B, 3, -1) + return jacobian + + """ + Operations. + """ + + def process_actions(self, actions: torch.Tensor): + # store the raw actions + self._raw_actions[:] = actions + self._processed_actions[:] = self.raw_actions * self._scale + # obtain quantities from simulation + ee_pos_curr, ee_quat_curr = self._compute_frame_pose() + # set command into controller + self._ik_controller.set_command(self._processed_actions, ee_pos_curr, ee_quat_curr) + if self.enable_task_space_boundary: + self._ik_controller.ee_pos_des = self._ik_controller.ee_pos_des.clip(self.min_limits, self.max_limits) + + def apply_actions(self): + # obtain quantities from simulation + ee_pos_curr, ee_quat_curr = self._compute_frame_pose() + joint_pos = self._asset.data.joint_pos[:, self._joint_ids] + # compute the delta in joint-space + if ee_quat_curr.norm() != 0: + jacobian = self._compute_frame_jacobian() + self.joint_pos_des = self._ik_controller.compute(ee_pos_curr, ee_quat_curr, jacobian, joint_pos) + else: + self.joint_pos_des = joint_pos.clone() + # set the joint position command + self._asset.set_joint_position_target(self.joint_pos_des, self._joint_ids) + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + self._raw_actions[env_ids] = 0.0 + + """ + Helper functions. + """ + + def _compute_frame_pose(self) -> tuple[torch.Tensor, torch.Tensor]: + """Computes the pose of the target frame in the root frame. + + Returns: + A tuple of the body's position and orientation in the root frame. + """ + # obtain quantities from simulation + num_body_idx = len(self._body_idx) + ee_pose_w = self._asset.data.body_link_state_w[:, self._body_idx, :7].view(-1, 7) + root_pose_w = self._asset.data.root_state_w[:, :7].repeat_interleave(num_body_idx, dim=0) + # compute the pose of the body in the root frame + ee_pos_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] + ) + # account for the offset + if self.cfg.body_offset is not None: + ee_pos_b, ee_quat_b = math_utils.combine_frame_transforms( + ee_pos_b, ee_quat_b, self._offset_pos.view(-1, 3), self._offset_rot.view(-1, 4) + ) + + return ee_pos_b.view(-1, num_body_idx, 3), ee_quat_b.view(-1, num_body_idx, 4) + + def _compute_frame_jacobian(self): + """Computes the geometric Jacobian of the target frame in the 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. + """ + # read the parent jacobian + jacobian = self.jacobian_b + + if self.cfg.body_offset is not None: + # Flatten from (num_envs, num_bodies, 6, num_joints) → (num_envs * num_bodies, 6, num_joints) + jacobian_flat = jacobian.reshape(-1, 6, jacobian.shape[3]) # (N*B, 6, num_joints) + + # Flatten offsets + offset_pos_flat = self._offset_pos.reshape(-1, 3) # (N*B, 3) + offset_rot_flat = self._offset_rot.reshape(-1, 4) # (N*B, 4) + + # 1) Translate part: áč—_link = áč—_ee + ω̇_ee × r_(link←ee) + + # => J_link_lin = J_ee_lin + -[r×]_offset * J_ee_ang + row_lin = jacobian_flat[:, 0:3, :] # (N*B, 3, num_joints) + row_ang = jacobian_flat[:, 3:6, :] # (N*B, 3, num_joints) + + skew = math_utils.skew_symmetric_matrix(offset_pos_flat) # (N*B, 3, 3) + + # áč—_link += -skew(r_offset) ⋅ ω̇_ee + + # We can do bmm: row_lin += bmm(-skew, row_ang) + row_lin += torch.bmm(-skew, row_ang) + + # 2) Rotate part: ω_link = R_offset ⋅ ω_ee + R_offset = math_utils.matrix_from_quat(offset_rot_flat) # (N*B, 3, 3) + row_ang_new = torch.bmm(R_offset, row_ang) + jacobian_flat[:, 3:6, :] = row_ang_new + + # Reshape back + jacobian = jacobian_flat.view(self.num_envs, len(self._body_idx), 6, -1) + return jacobian diff --git a/source/uwlab/uwlab/envs/mdp/actions/visualizable_joint_target_position.py b/source/uwlab/uwlab/envs/mdp/actions/visualizable_joint_target_position.py new file mode 100644 index 00000000..e5390367 --- /dev/null +++ b/source/uwlab/uwlab/envs/mdp/actions/visualizable_joint_target_position.py @@ -0,0 +1,70 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.assets.articulation import Articulation +from isaaclab.envs import ManagerBasedEnv +from isaaclab.managers.action_manager import ActionTerm + +if TYPE_CHECKING: + from . import actions_cfg + + +class VisualizableJointTargetPosition(ActionTerm): + """Joint action term that applies the processed actions to the articulation's joints as position commands.""" + + cfg: actions_cfg.VisualizableJointTargetPositionCfg + """The configuration of the action term.""" + + def __init__(self, cfg: actions_cfg.JointPositionActionCfg, env: ManagerBasedEnv): + super().__init__(cfg, env) + + @property + def action_dim(self) -> int: + return 0 + + @property + def raw_actions(self) -> torch.Tensor: + return torch.tensor([]) + + @property + def processed_actions(self) -> torch.Tensor: + return torch.tensor([]) + + def process_actions(self, actions): + pass + + def apply_actions(self): + pass + + def _set_debug_vis_impl(self, debug_vis: bool): + import isaacsim.core.utils.prims as prim_utils + from pxr import UsdGeom + + if debug_vis: + if not hasattr(self, "vis_articulation"): + if self.cfg.articulation_vis_cfg.name in self._env.scene: + self.vis_articulation: Articulation = self._env.scene[self.cfg.articulation_vis_cfg.name] + prims_paths = prim_utils.find_matching_prim_paths(self.vis_articulation.cfg.prim_path) + prims = [prim_utils.get_prim_at_path(prim) for prim in prims_paths] + for prim in prims: + UsdGeom.Imageable(prim).MakeVisible() + else: + if hasattr(self, "vis_articulation"): + prims_paths = prim_utils.find_matching_prim_paths(self.vis_articulation.cfg.prim_path) + prims = [prim_utils.get_prim_at_path(prim) for prim in prims_paths] + for prim in prims: + UsdGeom.Imageable(prim).MakeInvisible() + + def _debug_vis_callback(self, event): + # update the box marker + self.vis_articulation.write_joint_state_to_sim( + position=self._asset.data.joint_pos_target, + velocity=torch.zeros_like(self._asset.data.joint_pos_target, device=self.device), + ) diff --git a/source/uwlab/uwlab/envs/mdp/commands/__init__.py b/source/uwlab/uwlab/envs/mdp/commands/__init__.py new file mode 100644 index 00000000..1d674664 --- /dev/null +++ b/source/uwlab/uwlab/envs/mdp/commands/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .commands_cfg import CategoricalCommandCfg diff --git a/source/uwlab/uwlab/envs/mdp/commands/categorical_command.py b/source/uwlab/uwlab/envs/mdp/commands/categorical_command.py new file mode 100644 index 00000000..1da43ba5 --- /dev/null +++ b/source/uwlab/uwlab/envs/mdp/commands/categorical_command.py @@ -0,0 +1,133 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-module containing command generators for the velocity-based locomotion task.""" + +from __future__ import annotations + +import torch +from collections.abc import Sequence +from typing import TYPE_CHECKING + +from isaaclab.assets import Articulation +from isaaclab.managers import CommandTerm +from isaaclab.markers import VisualizationMarkers +from isaaclab.markers.config import CUBOID_MARKER_CFG + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + from .commands_cfg import CategoricalCommandCfg + + +class CategoricalCommand(CommandTerm): + r"""Command generator that generates a velocity command in SE(2) from uniform distribution. + + The command comprises of a linear velocity in x and y direction and an angular velocity around + the z-axis. It is given in the robot's base frame. + + If the :attr:`cfg.heading_command` flag is set to True, the angular velocity is computed from the heading + error similar to doing a proportional control on the heading error. The target heading is sampled uniformly + from the provided range. Otherwise, the angular velocity is sampled uniformly from the provided range. + + Mathematically, the angular velocity is computed as follows from the heading command: + + .. math:: + + \omega_z = \frac{1}{2} \text{wrap_to_pi}(\theta_{\text{target}} - \theta_{\text{current}}) + + """ + + cfg: CategoricalCommandCfg + """The configuration of the command generator.""" + + def __init__(self, cfg: CategoricalCommandCfg, env: ManagerBasedEnv): + """Initialize the command generator. + + Args: + cfg: The configuration of the command generator. + env: The environment. + """ + # initialize the base class + super().__init__(cfg, env) # type: ignore + + # obtain the robot asset + # -- robot + self.robot: Articulation = env.scene[cfg.asset_name] + + # crete buffers to store the command + self.category = torch.zeros(self.num_envs, device=self.device, dtype=torch.long) + + self.num_categories = self.cfg.num_category + + def __str__(self) -> str: + """Return a string representation of the command generator.""" + msg = "UniformVelocityCommand:\n" + msg += f"\tCommand dimension: {tuple(self.command.shape[1:])}\n" + msg += f"\tResampling time range: {self.cfg.resampling_time_range}\n" + msg += f"\tNumber of command categories: {self.cfg.num_category}\n" + return msg + + """ + Properties + """ + + @property + def command(self) -> torch.Tensor: + """The desired base velocity command in the base frame. Shape is (num_envs, 3).""" + return self.category + + """ + Implementation specific functions. + """ + + def _update_metrics(self): + """Update the metrics of the command generator.""" + pass + + def _resample_command(self, env_ids: Sequence[int]): + # category + self.category[env_ids] = torch.randint(0, self.num_categories, (len(env_ids),), device=self.device) + + def _update_command(self): + """Post-processes the velocity command. + + This function sets velocity command to zero for standing environments and computes angular + velocity from heading direction if the heading_command flag is set. + """ + pass + + 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, "category_visualizer"): + # -- goal + marker_cfg = CUBOID_MARKER_CFG.copy() # type: ignore + marker_cfg.prim_path = "/Visuals/Command/category" + marker_cfg.markers["cuboid"].scale = (0.1, 0.1, 0.1) + marker_cfg.markers["cuboid"].visual_material.diffuse_color = (0.0, 1.0, 0.0) + self.category_visualizer = VisualizationMarkers(marker_cfg) + # set their visibility to true + self.category_visualizer.set_visibility(True) + else: + if hasattr(self, "category_visualizer"): + self.category_visualizer.set_visibility(False) + + def _debug_vis_callback(self, event): + # check if robot is initialized + # note: this is needed in-case the robot is de-initialized. we can't access the data + if not self.robot.is_initialized: + return + # get marker location + # -- base state + base_pos_w = self.robot.data.root_pos_w.clone() + base_quat_w = self.robot.data.root_quat_w.clone() + base_pos_w[:, 2] += 0.5 + # -- resolve the scales + scale = self.command[:].repeat_interleave(3, 0).view(-1, 3) + + self.category_visualizer.visualize(base_pos_w, base_quat_w, scale) diff --git a/source/uwlab/uwlab/envs/mdp/commands/commands_cfg.py b/source/uwlab/uwlab/envs/mdp/commands/commands_cfg.py new file mode 100644 index 00000000..90f1964c --- /dev/null +++ b/source/uwlab/uwlab/envs/mdp/commands/commands_cfg.py @@ -0,0 +1,23 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import MISSING + +from isaaclab.managers import CommandTermCfg +from isaaclab.utils import configclass + +from .categorical_command import CategoricalCommand + + +@configclass +class CategoricalCommandCfg(CommandTermCfg): + """Configuration for the uniform velocity command generator.""" + + class_type: type = CategoricalCommand + + asset_name: str = MISSING + """Name of the asset in the environment for which the commands are generated.""" + + num_category: int = MISSING diff --git a/source/uwlab/uwlab/envs/mdp/events.py b/source/uwlab/uwlab/envs/mdp/events.py new file mode 100644 index 00000000..304b8761 --- /dev/null +++ b/source/uwlab/uwlab/envs/mdp/events.py @@ -0,0 +1,76 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.assets import Articulation, RigidObject +from isaaclab.managers import SceneEntityCfg + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + from uwlab.envs import DataManagerBasedRLEnv + + +def viewport_follow_robot( + env: DataManagerBasedRLEnv, + env_ids: torch.Tensor | None, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +): + target: RigidObject = env.scene[asset_cfg.name] + target_pos = target.data.root_pos_w.clone() + new_camera_pos = target_pos + torch.tensor([-1.0, 3.0, 2 * target_pos[0, 2]], device=target_pos.device) + direction = target_pos - new_camera_pos + direction[:, 2] -= 0.2 * target_pos[0, 2] + if env.viewport_camera_controller is not None: + env.viewport_camera_controller.update_view_location( + eye=new_camera_pos[0].cpu().tolist(), lookat=target_pos[0].cpu().tolist() + ) + + +def update_joint_target_positions_to_current(env: DataManagerBasedRLEnv, env_ids: torch.Tensor | None, asset_name: str): + asset: Articulation = env.scene[asset_name] + joint_pos_target = asset.data.joint_pos + asset.set_joint_position_target(joint_pos_target) + + +def reset_robot_to_default( + env: DataManagerBasedRLEnv, env_ids: torch.Tensor, robot_cfg: SceneEntityCfg = SceneEntityCfg("robot") +): + """Reset the scene to the default state specified in the scene configuration.""" + robot: Articulation = env.scene[robot_cfg.name] + default_root_state = robot.data.default_root_state[env_ids].clone() + default_root_state[:, 0:3] += env.scene.env_origins[env_ids] + # set into the physics simulation + robot.write_root_state_to_sim(default_root_state, env_ids=env_ids) + # obtain default joint positions + default_joint_pos = robot.data.default_joint_pos[env_ids].clone() + default_joint_vel = robot.data.default_joint_vel[env_ids].clone() + # set into the physics simulation + robot.write_joint_state_to_sim(default_joint_pos, default_joint_vel, env_ids=env_ids) + + +def launch_view_port( + env: ManagerBasedEnv, + env_ids: torch.Tensor, + view_portname: str, + camera_path: str, + viewport_size: tuple[int, int] = (640, 360), + position: tuple[int, int] = (0, 0), +): + if env.sim.has_gui(): + from isaacsim.core.utils.viewports import create_viewport_for_camera, get_viewport_names + + if view_portname not in get_viewport_names(): + create_viewport_for_camera( + viewport_name=view_portname, + camera_prim_path=camera_path, + width=viewport_size[0], + height=viewport_size[1], + position_x=position[0], + position_y=position[1], + ) diff --git a/source/uwlab/uwlab/envs/mdp/observations.py b/source/uwlab/uwlab/envs/mdp/observations.py new file mode 100644 index 00000000..4cc191b6 --- /dev/null +++ b/source/uwlab/uwlab/envs/mdp/observations.py @@ -0,0 +1,32 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch + +from isaaclab.assets import RigidObject +from isaaclab.envs import ManagerBasedRLEnv +from isaaclab.managers import SceneEntityCfg +from isaaclab.utils.math import subtract_frame_transforms + + +def object_position_in_robot_root_frame( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + object_cfg: SceneEntityCfg = SceneEntityCfg("object"), +) -> torch.Tensor: + """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 + object_pos_b, _ = subtract_frame_transforms(robot.data.root_pos_w, robot.data.root_quat_w, object_pos_w) + return object_pos_b + + +def life_spent(env: ManagerBasedRLEnv) -> torch.Tensor: + if hasattr(env, "episode_length_buf"): + life_spent = env.episode_length_buf.float() / env.max_episode_length + else: + life_spent = torch.zeros(env.num_envs, device=env.device, dtype=torch.float) + return life_spent.view(-1, 1) diff --git a/source/uwlab/uwlab/envs/mdp/rewards.py b/source/uwlab/uwlab/envs/mdp/rewards.py new file mode 100644 index 00000000..4df612bf --- /dev/null +++ b/source/uwlab/uwlab/envs/mdp/rewards.py @@ -0,0 +1,196 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.assets import Articulation, RigidObject +from isaaclab.managers import SceneEntityCfg +from isaaclab.sensors import FrameTransformer +from isaaclab.utils.math import combine_frame_transforms, quat_error_magnitude, quat_mul + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +def get_frame1_frame2_distance(frame1, frame2): + frames_distance = torch.norm(frame1.data.target_pos_w[..., 0, :] - frame2.data.target_pos_w[..., 0, :], dim=1) + return frames_distance + + +def get_body1_body2_distance(body1, body2, body1_offset, body2_offset): + bodys_distance = torch.norm((body1.data.root_pos_w + body1_offset) - (body2.data.root_pos_w + body2_offset), dim=1) + return bodys_distance + + +def get_frame1_body2_distance(frame1, body2, body2_offset): + distance = torch.norm(frame1.data.target_pos_w[..., 0, :] - (body2.data.root_pos_w + body2_offset), dim=1) + return distance + + +def reward_body_height_above(env: ManagerBasedRLEnv, minimum_height: float, asset_cfg: SceneEntityCfg) -> torch.Tensor: + """Terminate when the asset's height is below the minimum height. + + Note: + This is currently only supported for flat terrains, i.e. the minimum height is in the world frame. + """ + # extract the used quantities (to enable type-hinting) + asset: RigidObject = env.scene[asset_cfg.name] + return torch.where(asset.data.root_pos_w[:, 2] > minimum_height, 1, 0) + + +def reward_frame1_frame2_distance( + env: ManagerBasedRLEnv, + frame1_cfg: SceneEntityCfg, + frame2_cfg: SceneEntityCfg, +) -> torch.Tensor: + object_frame1: FrameTransformer = env.scene[frame1_cfg.name] + object_frame2: FrameTransformer = env.scene[frame2_cfg.name] + frames_distance = get_frame1_frame2_distance(object_frame1, object_frame2) + return 1 - torch.tanh(frames_distance / 0.1) + + +def reward_body1_body2_distance( + env: ManagerBasedRLEnv, + body1_cfg: SceneEntityCfg, + body2_cfg: SceneEntityCfg, + std: float, + body1_offset: list[float] = [0.0, 0.0, 0.0], + body2_offset: list[float] = [0.0, 0.0, 0.0], +) -> torch.Tensor: + body1: RigidObject = env.scene[body1_cfg.name] + body2: RigidObject = env.scene[body2_cfg.name] + body1_offset_tensor = torch.tensor(body1_offset, device=env.device) + body2_offset_tensor = torch.tensor(body2_offset, device=env.device) + bodys_distance = get_body1_body2_distance(body1, body2, body1_offset_tensor, body2_offset_tensor) + + return 1 - torch.tanh(bodys_distance / std) + + +def reward_body1_frame2_distance( + env: ManagerBasedRLEnv, + body_cfg: SceneEntityCfg, + frame_cfg: SceneEntityCfg, + body_offset: list[float] = [0.0, 0.0, 0.0], +) -> torch.Tensor: + body: RigidObject = env.scene[body_cfg.name] + object_frame: RigidObject = env.scene[frame_cfg.name] + body_offset_tensor = torch.tensor(body_offset, device=env.device) + bodys_distance = get_frame1_body2_distance(object_frame, body, body_offset_tensor) + return 1 - torch.tanh(bodys_distance / 0.1) + + +def reward_body1_body2_within_distance( + env: ManagerBasedRLEnv, + body1_cfg: SceneEntityCfg, + body2_cfg: SceneEntityCfg, + min_distance: float, + body1_offset: list[float] = [0.0, 0.0, 0.0], + body2_offset: list[float] = [0.0, 0.0, 0.0], +) -> torch.Tensor: + body1: RigidObject = env.scene[body1_cfg.name] + body2: RigidObject = env.scene[body2_cfg.name] + body1_offset_tensor = torch.tensor(body1_offset, device=env.device) + body2_offset_tensor = torch.tensor(body2_offset, device=env.device) + bodys_distance = get_body1_body2_distance(body1, body2, body1_offset_tensor, body2_offset_tensor) + reward = torch.where(bodys_distance < min_distance, 1.0, 0.0) + return reward + + +def reward_being_alive( + env: ManagerBasedRLEnv, +) -> torch.Tensor: + return torch.tanh((env.episode_length_buf / env.max_episode_length) / 0.1) + + +def lin_vel_z_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + """Penalize z-axis base linear velocity using L2-kernel.""" + # extract the used quantities (to enable type-hinting) + asset: RigidObject = env.scene[asset_cfg.name] + return torch.where(env.episode_length_buf > 20, torch.square(asset.data.root_lin_vel_b[:, 2]), 0) + + +def l2_norm_joint_position_command_error( + env: ManagerBasedRLEnv, command_name: str, asset_cfg: SceneEntityCfg +) -> torch.Tensor: + asset: Articulation = env.scene[asset_cfg.name] + command = env.command_manager.get_command(command_name) + cur_joint_position = asset.data.joint_pos[:, asset_cfg.joint_ids] + error = torch.norm(command - cur_joint_position, dim=1) + return error + + +def position_command_error(env: ManagerBasedRLEnv, command_name: str, asset_cfg: SceneEntityCfg) -> torch.Tensor: + """Penalize tracking of the position error using L2-norm. + + The function computes the position error between the desired position (from the command) and the + current position of the asset's body (in world frame). The position error is computed as the L2-norm + of the difference between the desired and current positions. + """ + # extract the asset (to enable type hinting) + asset: RigidObject = env.scene[asset_cfg.name] + 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_link_pos_w[:, asset_cfg.body_ids[0], :3] # type: ignore + return torch.norm(curr_pos_w - des_pos_w, dim=1) + + +def position_command_error_tanh( + env: ManagerBasedRLEnv, std: float, command_name: str, asset_cfg: SceneEntityCfg +) -> torch.Tensor: + """Reward tracking of the position using the tanh kernel. + + The function computes the position error between the desired position (from the command) and the + current position of the asset's body (in world frame) and maps it with a tanh kernel. + """ + # extract the asset (to enable type hinting) + asset: RigidObject = env.scene[asset_cfg.name] + 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_link_pos_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) + + +def orientation_command_error_tanh( + env: ManagerBasedRLEnv, std: float, command_name: str, asset_cfg: SceneEntityCfg +) -> torch.Tensor: + """Penalize tracking orientation error using shortest path. + + The function computes the orientation error between the desired orientation (from the command) and the + current orientation of the asset's body (in world frame). The orientation error is computed as the shortest + path between the desired and current orientations. + """ + # extract the asset (to enable type hinting) + asset: RigidObject = env.scene[asset_cfg.name] + 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_link_quat_w[:, asset_cfg.body_ids[0]] # type: ignore + return 1 - torch.tanh(quat_error_magnitude(curr_quat_w, des_quat_w) / std) + + +def orientation_command_error(env: ManagerBasedRLEnv, command_name: str, asset_cfg: SceneEntityCfg) -> torch.Tensor: + """Penalize tracking orientation error using shortest path. + + The function computes the orientation error between the desired orientation (from the command) and the + current orientation of the asset's body (in world frame). The orientation error is computed as the shortest + path between the desired and current orientations. + """ + # extract the asset (to enable type hinting) + asset: RigidObject = env.scene[asset_cfg.name] + 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_link_quat_w[:, asset_cfg.body_ids[0]] # type: ignore + return quat_error_magnitude(curr_quat_w, des_quat_w) diff --git a/source/uwlab/uwlab/envs/mdp/terminations.py b/source/uwlab/uwlab/envs/mdp/terminations.py new file mode 100644 index 00000000..855b4cdd --- /dev/null +++ b/source/uwlab/uwlab/envs/mdp/terminations.py @@ -0,0 +1,37 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Common functions that can be used to activate certain terminations. + +The functions can be passed to the :class:`isaaclab.managers.TerminationTermCfg` object to enable +the termination introduced by the function. +""" + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.assets import Articulation, RigidObject +from isaaclab.managers import SceneEntityCfg + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + +""" +MDP terminations. +""" + + +def invalid_state(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: + """Return true if the RigidBody position reads nan""" + # extract the used quantities (to enable type-hinting) + asset: RigidObject = env.scene[asset_cfg.name] + return torch.isnan(asset.data.body_pos_w).any(dim=-1).any(dim=-1) + + +def abnormal_robot_state(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + robot: Articulation = env.scene[asset_cfg.name] + return (robot.data.joint_vel.abs() > (robot.data.joint_vel_limits * 2)).any(dim=1) diff --git a/source/uwlab/uwlab/envs/real_rl_env.py b/source/uwlab/uwlab/envs/real_rl_env.py new file mode 100644 index 00000000..56036b73 --- /dev/null +++ b/source/uwlab/uwlab/envs/real_rl_env.py @@ -0,0 +1,284 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import math +import numpy as np +import torch +from typing import TYPE_CHECKING, Any, ClassVar, Sequence + +import isaacsim.core.utils.torch as torch_utils + +from isaaclab.envs import ManagerBasedRLEnv +from isaaclab.envs.common import VecEnvObs, VecEnvStepReturn +from isaaclab.managers import ( + ActionManager, + CommandManager, + CurriculumManager, + EventManager, + ObservationManager, + RecorderManager, + RewardManager, + TerminationManager, +) +from isaaclab.utils.timer import Timer + +if TYPE_CHECKING: + from .real_rl_env_cfg import RealRLEnvCfg + + +class RealRLEnv(ManagerBasedRLEnv): + is_vector_env: ClassVar[bool] = True + """Whether the environment is a vectorized environment.""" + + metadata: ClassVar[dict[str, Any]] = {} + """real environment doesn't have metadata""" + + cfg: RealRLEnvCfg + + def __init__(self, cfg: RealRLEnvCfg, **kwargs) -> None: + # cfg.validate() # TODO: uncomment this line when bug resolved + self.cfg = cfg + self.cfg.scene.device = self.cfg.device + self._is_closed = False + + # set the seed for the environment + if self.cfg.seed is not None: + self.cfg.seed = self.seed(self.cfg.seed) + else: + print("Seed not set for the environment. The environment creation may not be deterministic.") + + with Timer("[INFO]: Time taken for context scene creation", "scene_creation"): + self.scene = self.cfg.scene.class_type(self.cfg.scene) + print("[INFO]: Scene manager: ", self.scene) + + # -- counter for sanity check + self.common_step_counter = 0 + # -- init buffers + self.episode_length_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.long) + print("[INFO]: Starting the simulation. This may take a few seconds. Please wait...") + with Timer("[INFO]: Time taken for simulation start", "simulation_start"): + self.scene.start() + self.load_managers() + + self._sim_step_counter = 0 + + # allocate dictionary to store metrics + self.extras = {} + + # initialize observation buffers + self.obs_buf = {} + + """ + Operations - Setup. + """ + + @property + def num_envs(self) -> int: + """The number of instances of the environment that are running.""" + return 1 + + @property + def device(self): + """The device on which the environment is running.""" + return self.cfg.device + + @property + def max_episode_length_s(self) -> float: + """Maximum episode length in seconds.""" + return self.cfg.episode_length_s + + @property + def max_episode_length(self) -> int: + """Maximum episode length in environment steps.""" + return math.ceil(self.max_episode_length_s / self.step_dt) + + @property + def physic_dt(self) -> float: + raise NotImplementedError(f"The property physic_dt is not supported for {self.__class__.__name__}.") + + @property + def step_dt(self) -> float: + return self.cfg.scene.dt + + """ + Operations - Setup. + """ + + def load_managers(self): + # note: this order is important since observation manager needs to know the command and action managers + # and the reward manager needs to know the termination manager + # -- command manager + self.command_manager: CommandManager = CommandManager(self.cfg.commands, self) + print("[INFO] Command Manager: ", self.command_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) + # -- observation manager + self.observation_manager = ObservationManager(self.cfg.observations, self) + print("[INFO] Observation Manager:", self.observation_manager) + # -- event manager + self.event_manager = EventManager(self.cfg.events, self) + print("[INFO] Event Manager: ", self.event_manager) + + # -- termination manager + self.termination_manager = TerminationManager(self.cfg.terminations, self) + print("[INFO] Termination Manager: ", self.termination_manager) + # # -- reward manager + self.reward_manager = RewardManager(self.cfg.rewards, self) + print("[INFO] Reward Manager: ", self.reward_manager) + # -- curriculum manager + self.curriculum_manager = CurriculumManager(self.cfg.curriculum, self) + print("[INFO] Curriculum Manager: ", self.curriculum_manager) + + # setup the action and observation spaces for Gym + self._configure_gym_env_spaces() + + # perform events at the start of the simulation + if "startup" in self.event_manager.available_modes: + self.event_manager.apply(mode="startup") + + def step(self, action: torch.Tensor) -> VecEnvStepReturn: + self.action_manager.process_action(action.to(self.device)) + + self.recorder_manager.record_pre_step() + + # perform physics stepping + for _ in range(self.cfg.decimation): + self._sim_step_counter += 1 + # set actions into buffers + self.action_manager.apply_action() + # set actions into simulator + self.scene.write_data_to_context() + # update buffers at sim dt + self.scene.update(dt=self.physics_dt) + + # post-step: + # -- update env counters (used for curriculum generation) + self.episode_length_buf += 1 # step in current episode (per env) + self.common_step_counter += 1 # total step (common for all envs) + # -- check terminations + self.reset_buf = self.termination_manager.compute() + self.reset_terminated = self.termination_manager.terminated + self.reset_time_outs = self.termination_manager.time_outs + # # -- 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_context() + + # # 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 + if "interval" in self.event_manager.available_modes: + self.event_manager.apply(mode="interval", dt=self.step_dt) + # -- compute observations + # note: done after reset to get the correct observations for reset envs + self.obs_buf = self.observation_manager.compute() + + # return observations, rewards, resets and extras + return self.obs_buf, self.reward_buf, self.reset_terminated, self.reset_time_outs, self.extras + + def render(self, recompute: bool = False) -> np.ndarray | None: + raise NotImplementedError(f"The function render is not supported for {self.__class__.__name__}.") + + def reset( + self, seed: int | None = None, env_ids: Sequence[int] | None = None, options: dict[str, Any] | None = None + ) -> tuple[VecEnvObs, dict]: + if env_ids is None: + env_ids = torch.arange(self.num_envs, dtype=torch.int64, device=self.device) + + # set the seed + if seed is not None: + self.seed(seed) + + # reset state of scene + self._reset_idx(env_ids) + + # update articulation kinematics + self.scene.write_data_to_context() + + # compute observations + self.obs_buf = self.observation_manager.compute() + + self.extras = dict() + + return self.obs_buf, self.extras + + def seed(self, seed: int = -1) -> int: + return torch_utils.set_seed(seed) + + def close(self) -> None: + del self.command_manager + del self.reward_manager + del self.termination_manager + del self.action_manager + del self.observation_manager + del self.event_manager + + def _configure_gym_env_spaces(self): + super()._configure_gym_env_spaces() + + def _reset_idx(self, env_ids: Sequence[int]): + """Reset environments based on specified indices. + + Args: + env_ids: List of environment ids which must be reset + """ + # update the curriculum for environments that need a reset + self.curriculum_manager.compute(env_ids=env_ids) + # reset the internal buffers of the scene elements + self.scene.reset(env_ids) + # apply events such as randomizations for environments that need a reset + if "reset" in self.event_manager.available_modes: + self.event_manager.apply(mode="reset", env_ids=env_ids, global_env_step_count=self.common_step_counter) + + # iterate over all managers and reset them + # this returns a dictionary of information which is stored in the extras + # note: This is order-sensitive! Certain things need be reset before others. + self.extras["log"] = dict() + # -- observation manager + info = self.observation_manager.reset(env_ids) + self.extras["log"].update(info) + # -- action manager + info = self.action_manager.reset(env_ids) + self.extras["log"].update(info) + # -- rewards manager + info = self.reward_manager.reset(env_ids) + self.extras["log"].update(info) + # -- curriculum manager + info = self.curriculum_manager.reset(env_ids) + self.extras["log"].update(info) + # -- command manager + info = self.command_manager.reset(env_ids) + self.extras["log"].update(info) + # -- event manager + info = self.event_manager.reset(env_ids) + self.extras["log"].update(info) + # -- termination manager + info = self.termination_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/uwlab/uwlab/envs/real_rl_env_cfg.py b/source/uwlab/uwlab/envs/real_rl_env_cfg.py new file mode 100644 index 00000000..4e0ee35e --- /dev/null +++ b/source/uwlab/uwlab/envs/real_rl_env_cfg.py @@ -0,0 +1,21 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import MISSING + +from isaaclab.envs import ManagerBasedRLEnvCfg +from isaaclab.utils import configclass +from uwlab.scene import SceneContextCfg + +from .real_rl_env import RealRLEnv + + +@configclass +class RealRLEnvCfg(ManagerBasedRLEnvCfg): + class_type = RealRLEnv + + device = "cpu" + + scene: SceneContextCfg = MISSING diff --git a/source/uwlab/uwlab/genes/__init__.py b/source/uwlab/uwlab/genes/__init__.py new file mode 100644 index 00000000..e050c058 --- /dev/null +++ b/source/uwlab/uwlab/genes/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .gene import * +from .genome import Genome +from .genome_cfg import GenomeCfg diff --git a/source/uwlab/uwlab/genes/gene/__init__.py b/source/uwlab/uwlab/genes/gene/__init__.py new file mode 100644 index 00000000..4216f92b --- /dev/null +++ b/source/uwlab/uwlab/genes/gene/__init__.py @@ -0,0 +1,16 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .gene import FloatGeneOperator, FloatTupleGeneOperator, GeneOperatorBase, IntGeneOperator +from .gene_cfg import ( + FloatGeneCfg, + FloatTupleGeneCfg, + GeneOperatorBaseCfg, + IntGeneCfg, + IntTupleGeneCfg, + StrGeneCfg, + TupleGeneBaseCfg, +) +from .gene_mdp import * diff --git a/source/uwlab/uwlab/genes/gene/gene.py b/source/uwlab/uwlab/genes/gene/gene.py new file mode 100644 index 00000000..6053b964 --- /dev/null +++ b/source/uwlab/uwlab/genes/gene/gene.py @@ -0,0 +1,162 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import numpy as np +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from uwlab.terrains.terrain_generator_cfg import MultiOriginTerrainGeneratorCfg + + from .gene_cfg import FloatGeneCfg, FloatTupleGeneCfg, GeneOperatorBaseCfg, TerrainGeneCfg + + +class GeneOperatorBase: + cfg: GeneOperatorBaseCfg + + def __init__(self, retrive_args, cfg: GeneOperatorBaseCfg, rng: np.random.Generator): + self.retrive_args = retrive_args + self.mutation_args = cfg.mutation_args + self.mutation_func = cfg.mutation_func + self.group = cfg.group + self.rng = rng + + def get(self, source): + raise NotImplementedError + + def set(self, val): + raise NotImplementedError + + def mutate(self, source): + raise NotImplementedError + + def breed(self, this_source, other_source): + raise NotImplementedError + + def traverse_operations(self, src, operation_list, arg_list): + for operation, arg in zip(operation_list, arg_list): + src = operation(src, arg) + return src + + def _set_attr(self, target, key, val): + if isinstance(target, dict): + target[key] = val + else: + setattr(target, key, val) + + +class FloatGeneOperator(GeneOperatorBase): + cfg: FloatGeneCfg + + def __init__(self, retrive_args, cfg: FloatGeneCfg, rng: np.random.Generator): + super().__init__(retrive_args, cfg, rng) + self.fmin = cfg.fmin + self.fmax = cfg.fmax + self.mutation_rate = cfg.mutation_rate + + def get(self, source): + return self.traverse_operations(source, *self.retrive_args[:2]) + + def set(self, source, value): + if value < self.fmin or value > self.fmax: + raise ValueError("you are trying to set a value out of bound") + self._set_func(source, value, *self.retrive_args) + + def mutate(self, source): + val = self.get(source) + new_val = self.mutation_func(self.rng, val, self.mutation_rate, *self.mutation_args) + new_val = np.clip(new_val, self.fmin, self.fmax).item() + self.set(source, new_val) + + def breed(self, this_source, other_source): + this_val = self.get(this_source) + other_val = self.get(other_source) + new_val = (this_val + other_val) / 2 + self.set(this_source, new_val) + + def _set_func(self, src_env, v, ops, args): + self._set_attr(self.traverse_operations(src_env, ops[:-1], args[:-1]), args[-1], float(v)) + + +class IntGeneOperator(FloatGeneOperator): + def set(self, source, value): + if value < self.fmin or value > self.fmax: + raise ValueError("you are trying to set a value out of bound") + self._set_func(source, value, *self.retrive_args) + + def _set_func(self, src_env, v, ops, args): + self._set_attr(self.traverse_operations(src_env, ops[:-1], args[:-1]), args[-1], int(v)) + + +class FloatTupleGeneOperator(GeneOperatorBase): + cfg: FloatTupleGeneCfg + + def __init__(self, retrive_args, cfg: FloatTupleGeneCfg, rng: np.random.Generator): + super().__init__(retrive_args, cfg, rng) + self.mutation_rate = cfg.mutation_rate + self.fmin = cfg.fmin + self.fmax = cfg.fmax + self.element_length = cfg.element_length + self.element_idx: int = cfg.element_idx + + def get(self, source) -> float: + return self.traverse_operations(source, *self.retrive_args[:2])[self.element_idx] + + def set(self, source, value): + self._set_float_tuple_func(source, value, *self.retrive_args) + + def mutate(self, source): + val = self.get(source) + new_val = self.mutation_func(self.rng, val, self.mutation_rate, *self.mutation_args) + new_val = np.clip(new_val, self.fmin[self.element_idx], self.fmax[self.element_idx]).item() + self.set(source, new_val) + + def breed(self, this_source, other_source): + val = self.get(this_source) + other_val = self.get(other_source) + self.set(this_source, (val + other_val) / 2) + + def _set_float_tuple_func(self, src_env, v, ops, args): + # Retrieve the original tuple and convert it to a list to make it mutable + val_list = list(getattr(self.traverse_operations(src_env, ops[:-1], args[:-1]), args[-1])) + # Modify the value at the specified index + val_list[self.element_idx] = v + # Convert the list back to a tuple + val_tuple = tuple(val_list) + # Set the modified tuple back to the source environment + self._set_attr(self.traverse_operations(src_env, ops[:-1], args[:-1]), args[-1], val_tuple) + + +class TerrainGeneOperator(GeneOperatorBase): + cfg: TerrainGeneCfg + + def __init__(self, retrive_args, cfg: TerrainGeneCfg, rng: np.random.Generator): + super().__init__(retrive_args, cfg, rng) + self.mutation_rate = cfg.mutation_rate + + def get(self, source): + return self.traverse_operations(source, *self.retrive_args[:2]) + + def set(self, source, value): + self._set_func(source, value, *self.retrive_args) + + def mutate(self, source): + val = self.get(source) + new_val = self.mutation_func(self.rng, val, self.mutation_rate, *self.mutation_args) + self.set(source, new_val) + + def breed(self, this_source, other_source): + this_val: MultiOriginTerrainGeneratorCfg = self.get(this_source) + other_val: MultiOriginTerrainGeneratorCfg = self.get(other_source) + num_sub_terrains = len(this_val.sub_terrains) + len(other_val.sub_terrains) + width = np.ceil(np.sqrt(num_sub_terrains)).item() + this_val.num_cols = int(width) + this_val.num_rows = int(width) + this_val.sub_terrains.update(other_val.sub_terrains) + self.set(this_source, this_val) + + def _set_func(self, src_env, v, ops, args): + self._set_attr(self.traverse_operations(src_env, ops[:-1], args[:-1]), args[-1], v) diff --git a/source/uwlab/uwlab/genes/gene/gene_cfg.py b/source/uwlab/uwlab/genes/gene/gene_cfg.py new file mode 100644 index 00000000..fc74b817 --- /dev/null +++ b/source/uwlab/uwlab/genes/gene/gene_cfg.py @@ -0,0 +1,88 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from dataclasses import MISSING +from typing import Callable, Literal + +from isaaclab.utils import configclass +from uwlab.genes.gene import gene + + +@configclass +class GeneOperatorBaseCfg: + class_type: type = MISSING # type: ignore + # The group flag indicates which group this gene belongs to. + group: str = "any" + phase: list[Literal["init", "mutate", "breed"]] = MISSING # type: ignore + # The function to be used for mutations, to be defined externally. + mutation_func: Callable = MISSING # type: ignore + # The arguments that supplies mutation_func. + mutation_args: tuple[any, ...] = MISSING # type: ignore + + mutation_rate: float = 1 + + +@configclass +class TupleGeneBaseCfg(GeneOperatorBaseCfg): + element_length: int = MISSING # type: ignore + + element_idx: int = MISSING # type: ignore + + # Defines the type of tuple operation: + # 'descend': Each subsequent tuple element must be less than the previous one. + # 'ascend': Each subsequent tuple element must be greater than the previous one. + # 'equal': All elements in the tuple are the same. + # 'symmetric': The tuple consists of exactly two values, 0th value being the negative of the 1st value. + tuple_type: Literal["descend", "ascend", "equal", "symmetric"] = MISSING # type: ignore + + +@configclass +class TerrainGeneCfg(GeneOperatorBaseCfg): + class_type: type = gene.TerrainGeneOperator + + +@configclass +class FloatGeneCfg(GeneOperatorBaseCfg): + class_type: type = gene.FloatGeneOperator + + fmin: float = -float("inf") + + fmax: float = float("inf") + + +@configclass +class IntGeneCfg(FloatGeneCfg): + class_type: type = gene.IntGeneOperator + + +@configclass +class StrGeneCfg(GeneOperatorBaseCfg): + str_list: list[str] = MISSING # type: ignore + + +@configclass +class StrTupleGeneCfg(TupleGeneBaseCfg): + # Not yet implemented. + pass + + +# Define the configuration class for floating-point tuple gene operations. +@configclass +class FloatTupleGeneCfg(TupleGeneBaseCfg): + class_type: type = gene.FloatTupleGeneOperator + + # Minimum possible values for each element in the tuple, defaulting to negative infinity. + fmin: tuple[float, ...] = tuple(-float("inf") for _ in range(2)) + + # Maximum possible values for each element in the tuple, defaulting to positive infinity. + fmax: tuple[float, ...] = tuple(float("inf") for _ in range(2)) + + +@configclass +class IntTupleGeneCfg(FloatTupleGeneCfg): + # not_yet_implemented + pass diff --git a/source/uwlab/uwlab/genes/gene/gene_mdp.py b/source/uwlab/uwlab/genes/gene/gene_mdp.py new file mode 100644 index 00000000..b7bf3554 --- /dev/null +++ b/source/uwlab/uwlab/genes/gene/gene_mdp.py @@ -0,0 +1,111 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import numpy as np +import torch +from typing import List + +from uwlab.terrains.terrain_generator_cfg import MultiOriginTerrainGeneratorCfg + + +# MUTATION FUNCTIONS +def add_constant(rng: np.random.Generator, val, mutation_rate: float, constant): + sign = -1 if rng.random() < 0.5 else 1 + new_val = val + sign * constant * mutation_rate + return new_val + + +def add_fraction(rng: np.random.Generator, val, mutation_rate: float, fraction): + sign = -1 if rng.random() < 0.5 else 1 + val_offset = sign * fraction * val + new_val = val + val_offset * mutation_rate + return new_val + + +def random_int(rng: np.random.Generator, val, mutation_rate: float, imin, imax): + new_val = rng.integers(imin, imax) + return new_val + + +def random_float(rng: np.random.Generator, val, mutation_rate: float, fmin, fmax): + new_val = rng.random() * (fmax - fmin) + fmin + return new_val + + +def random_selection(rng: np.random.Generator, val, mutation_rate: float, selection_list: list): + rand_int = rng.integers(0, len(selection_list)) + return selection_list[rand_int] + + +def random_dict(rng: np.random.Generator, val, mutation_rate: float, dict: dict): + # select a random pair of key, value from the dictionary + # return as a dictionary + key = random_selection(rng, val, mutation_rate, list(dict.keys())) + value = dict[key] + return {key: value} + + +def mutate_terrain_cfg(rng: np.random.Generator, val, mutation_rate, cfg: MultiOriginTerrainGeneratorCfg): + key = random_selection(rng, val, mutation_rate, list(cfg.sub_terrains.keys())) + value = cfg.sub_terrains[key] + sub_terrain = {key: value} + cfg.sub_terrains = sub_terrain + return cfg + + +# BREEDING FUNCTIONS + + +def breed_terrain_cfg(this_val: MultiOriginTerrainGeneratorCfg, other_val: MultiOriginTerrainGeneratorCfg): + num_sub_terrains = len(this_val.sub_terrains) + len(other_val.sub_terrains) + width = np.ceil(np.sqrt(num_sub_terrains)) + this_val.num_cols = width + this_val.num_rows = width + this_val.sub_terrains.update(other_val.sub_terrains) + + +def value_distribution( + values: List[float], + distribute_to_n_values: int, + value_to_distribute: float | None = None, + equal_distribution: bool = False, +) -> List[float]: + """Redistributes the total sum of values to the top n values based on their initial proportion.""" + if distribute_to_n_values <= 0 or distribute_to_n_values > len(values): + raise ValueError("distribute_to_n_values must be greater than 0 and less than or equal to the length of values") + + # Get indices of the top 'n' values + top_indices = sorted(range(len(values)), key=lambda i: values[i], reverse=True)[:distribute_to_n_values] + top_sum = sum(values[i] for i in top_indices) + if value_to_distribute: + total_value_to_distribute = value_to_distribute + else: + total_value_to_distribute = sum(values) + + # Calculate proportions and distribute the total sum accordingly + proportion = np.zeros(len(values)) + if equal_distribution: + proportion[top_indices] = np.array([1 / distribute_to_n_values for _ in top_indices]) + else: + proportion[top_indices] = np.array([values[i] / top_sum for i in top_indices]) + output = proportion * total_value_to_distribute + return output.tolist() + + +def probability_distribution(vals: List[float], distribute_to_n_values: int) -> List[float]: + """Converts redistributed values into a valid probability distribution using softmax and returns it as a list of floats.""" + # First, redistribute the values + redistributed_vals = value_distribution(vals, distribute_to_n_values) + + # Convert to a torch tensor + logits = torch.tensor(redistributed_vals, dtype=torch.float) + + # Apply softmax to convert logits into probabilities + probabilities = torch.softmax(logits, dim=0) + + # Convert the tensor back to a list of floats + probabilities_list = probabilities.tolist() + + return probabilities_list diff --git a/source/uwlab/uwlab/genes/genome.py b/source/uwlab/uwlab/genes/genome.py new file mode 100644 index 00000000..667dd355 --- /dev/null +++ b/source/uwlab/uwlab/genes/genome.py @@ -0,0 +1,303 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import numpy as np +import re +from typing import TYPE_CHECKING + +from isaaclab.envs import ManagerBasedRLEnvCfg +from uwlab.genes.gene import GeneOperatorBase, GeneOperatorBaseCfg, TupleGeneBaseCfg + +from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg + +if TYPE_CHECKING: + from .genome_cfg import GenomeCfg + + +class Genome: + env_cfg: ManagerBasedRLEnvCfg | None = None + + agent_cfg: RslRlOnPolicyRunnerCfg | None = None + + """ + The Genome class models a genetic structure with specific configuration environments (env_cfg and agent_cfg) + and a genomic_regulatory_profile. It allows for genetic operations like mutation, breedingm and cloning of an + genome. The class also supports the construction of a genetic dictionary where all gene are listed in a dictionary + """ + + def __init__(self, cfg: GenomeCfg): + """Initializes an Genome instance by copying the provided configurations and profile and setting up an empty + genetic dictionary. + cfg (GenomeGeneticCfg): The regulatory profile for gene mutation where rule of mutation is specified. + """ + self.genomic_regulatory_profile = cfg.genomic_mutation_profile + self.genomic_group_processor = cfg.genomic_constraint_profile + self.genetic_dictionary: dict[str, GeneOperatorBase] = {} + self.genetic_groups: dict[str, list] = {} + self.genetic_mutation_phases: dict[str, list[GeneOperatorBase]] = {"init": [], "mutate": [], "breed": []} + self.tuple_genetic_groups: dict[str, list] = {"descend": [], "ascend": [], "equal": [], "symmetric": []} + self.np_random = np.random.default_rng(cfg.seed) + + @property + def my_genetic_manual(self) -> dict[str, GeneOperatorBase]: + return self.genetic_dictionary + + def activate(self, env_cfg, agent_cfg): + """Activates the genome by constructing the genetic modulation linkages. This involves recursively traversing the + environment and regulatory profile, generating a dictionary of genes (genetic_dictionary) that access to all + genetic functions. + + Args: + env_cfg (ManagerBasedRLEnvCfg): The configuration of environment this genome lives in. + agent_cfg (RslRlOnPolicyRunnerCfg): The configuration of agent that trains this genome. + """ + # reset these fields are necessary because the genome may be re-activated, + # which means some of genes may be deleted, and some new genes may be added, the new gene bridge needs + # to be re-constructed + self.env_cfg = env_cfg + self.agent_cfg = agent_cfg + self.genetic_dictionary: dict[str, GeneOperatorBase] = {} + self.genetic_groups: dict[str, list] = {} + self.genetic_mutation_phases: dict[str, list[GeneOperatorBase]] = {"init": [], "mutate": [], "breed": []} + self.tuple_genetic_groups: dict[str, list] = {"descend": [], "ascend": [], "equal": [], "symmetric": []} + + genetic_dictionary = {} + cfgs, keys = self._recursively_construct_genetic_modulation_linkage( + self.env_cfg, self.genomic_regulatory_profile, [] + ) + for cfg, args in zip(cfgs, keys): + cfg.mutation_rate = self.np_random.random() + retrival_func_list = [] + for i in range(len(args)): + if "[]" in args[i][:2]: + retrival_func_list.append(lambda src_dict, key: dict.get(src_dict, key)) + args[i] = args[i][2:] # type: ignore + elif "." in args[i][0]: + retrival_func_list.append(lambda src_class, key: enhanced_attrgetter(key)(src_class)) + args[i] = args[i][1:] # type: ignore + else: + retrival_func_list.append(lambda src_class, key: enhanced_attrgetter(key)(src_class)) + + if isinstance(cfg, TupleGeneBaseCfg): + # process tuple type genes + tuple_gene = [] + for i in range(cfg.element_length): # type: ignore + cfg.element_idx = i + gene_operator = cfg.class_type((retrival_func_list, args), cfg, rng=self.np_random) + gene_identifier = ".".join(args) + f"_{i}" + genetic_dictionary[gene_identifier] = gene_operator + tuple_gene.append(gene_operator) + + # Process Group + if cfg.group not in self.genetic_groups: + self.genetic_groups[cfg.group] = [] + self.genetic_groups[cfg.group].append(gene_operator) + # Process Phase + for phase in cfg.phase: + self.genetic_mutation_phases[phase].append(gene_operator) + self.tuple_genetic_groups[cfg.tuple_type].append(tuple_gene) + else: + # process single number type genes + gene_operator = cfg.class_type((retrival_func_list, args), cfg, rng=self.np_random) + genetic_dictionary[".".join(args)] = gene_operator + # Process Group + if cfg.group not in self.genetic_groups: + self.genetic_groups[cfg.group] = [] + self.genetic_groups[cfg.group].append(gene_operator) + # Process Phase + for phase in cfg.phase: + self.genetic_mutation_phases[phase].append(gene_operator) + self.genetic_dictionary = genetic_dictionary + + def gene_initialize(self): + for gene in self.genetic_mutation_phases["init"]: + gene.mutate(self.env_cfg) + # init functions may modify the env_cfg, some genes may be deleted, + # and genes bridged earlier may be invalid, so we need to re-activate + self.activate(self.env_cfg, self.agent_cfg) + + return self.env_cfg, self.agent_cfg + + def mutate(self): + # Step1: mutation + for gene in self.genetic_mutation_phases["mutate"]: + gene.mutate(self.env_cfg) + + # Step2: Apply rules that make sure float are sanity checked + + # apply rules for those that are tuple + for tuple_gene, tuples in self.tuple_genetic_groups.items(): + if tuple_gene == "equal": + for tuple in tuples: + val = [gene.get(self.env_cfg) for gene in tuple] + val = np.array(val) + new_val = np.average(val).item() + for gene in tuple: + gene.set(self.env_cfg, new_val) + elif tuple_gene == "ascend": + for tuple in tuples: + val = [gene.get(self.env_cfg) for gene in tuple] + val = np.array(val) + new_val = np.sort(val) + for i, gene in enumerate(tuple): + gene.set(self.env_cfg, new_val[i].item()) + elif tuple_gene == "descend": + for tuple in tuples: + val = [gene.get(self.env_cfg) for gene in tuple] + val = np.array(val) + new_val = np.sort(val)[::-1] + for i, gene in enumerate(tuple): + gene.set(self.env_cfg, new_val[i].item()) + elif tuple_gene == "symmetric": + for tuple in tuples: + val = [gene.get(self.env_cfg) for gene in tuple] + val = np.array(val) + avg = np.abs(np.average(val)) + new_val = np.array([-avg, avg]) + for i, gene in enumerate(tuple): + gene.set(self.env_cfg, new_val[i].item()) + + # apply extra customly designed rules + for group_key, gene_list in self.genetic_groups.items(): + if group_key == "any": + continue + val_list = [gene.get(self.env_cfg) for gene in gene_list] + func, arg = self.genomic_group_processor[group_key] + new_val_list = func(val_list, *arg) + for i, gene in enumerate(gene_list): + gene.set(self.env_cfg, new_val_list[i]) + + def breed(self, other_genome: Genome): + for gene in self.genetic_mutation_phases["breed"]: + val = gene.get(self.env_cfg) + try: + other_val = gene.get(other_genome.env_cfg) + # if the gene is not in the other genome, skip + # this error is accepted because the gene may not be in the other genome + except TypeError: + continue + if val is None or other_val is None: + continue + gene.breed(self.env_cfg, other_genome.env_cfg) + + def clone(self): + return Genome(self.env_cfg.copy(), self.agent_cfg.copy(), self.genomic_regulatory_profile.copy()) # type: ignore + + def _recursively_construct_genetic_modulation_linkage( + self, roots, genomic_regulatory_profile, keys + ) -> tuple[list[GeneOperatorBaseCfg], list[str]]: + """A private method that recursively traverses the environment configuration (roots) and the regulatory profile + (genomic_regulatory_profile) to construct the genetic modulation linkages. This method identifies all the relevant genes and + attributes based on the genomic_regulatory_profile and returns the path_to_gene and gene_detail necessary for creating the + genetic dictionary. + + Args: + roots (dict|class): The current level of the environment configuration being traversed. + genomic_regulatory_profile (dict): The current level of the regulatory profile that provides instructions on which attributes/keys to extract. + keys (list): A list of keys (attribute paths) accumulated during recursion. + + Returns: + gene_detail list[(scale, type)]: A list of guides for the identified genes mutation rules. + path_to_gene list[str]: A list of keys indicating the attribute path to retrieve the data in env_cfg. + """ + gene_detail = [] + path_to_gene = [] + profile = genomic_regulatory_profile + if isinstance(profile, GeneOperatorBaseCfg): + gene_detail.append(profile) + path_to_gene.append(keys) + elif isinstance(roots, dict): + for k, v in profile.items(): + if k == ".": + for ext_key, val in roots.items(): + for guide_key in profile["."].keys(): + if enhanced_attrgetter(guide_key)(val) is not None: + sub_guid, sub_keys = self._recursively_construct_genetic_modulation_linkage( + val, profile["."], keys + split_keys(ext_key) + ) + gene_detail.extend(sub_guid) + path_to_gene.extend(sub_keys) + break + else: + if k in roots: + attr = roots[k] + if attr is not None: + sub_guid, sub_keys = self._recursively_construct_genetic_modulation_linkage( + attr, profile[k], keys + split_keys(k) + ) + gene_detail.extend(sub_guid) + path_to_gene.extend(sub_keys) + + elif hasattr(roots, "__dict__"): + for k, v in profile.items(): + if k == ".": + for ext_key in dir(roots): + attr = getattr(roots, ext_key) + for guide_key in profile["."].keys(): + if hasattr(attr, guide_key): + sub_guid, sub_keys = self._recursively_construct_genetic_modulation_linkage( + attr, profile["."], keys + split_keys(ext_key) + ) + gene_detail.extend(sub_guid) + path_to_gene.extend(sub_keys) + break + else: + if enhanced_attrgetter(k)(roots) is not None: + attr = enhanced_attrgetter(k)(roots) + if attr is not None: + sub_guid, sub_keys = self._recursively_construct_genetic_modulation_linkage( + attr, profile[k], keys + split_keys(k) + ) + gene_detail.extend(sub_guid) + path_to_gene.extend(sub_keys) + + else: + gene_detail.append(profile) + path_to_gene.append(keys) + + return gene_detail, path_to_gene + + +def enhanced_attrgetter(attr_string): + def getter(obj): + parts = re.split(r"(\.|\[|\])", attr_string) # Split the string by '.', '[', and ']' + current_obj = obj + # inside_bracket = False + + for part in parts: + if not part or part in ".[]": + continue + try: + if isinstance(current_obj, dict): + current_obj = current_obj[part] # Access dictionary key + else: + current_obj = getattr(current_obj, part) # Access attribute + except (AttributeError, KeyError, TypeError): + return None # Return None if an attribute or key is not found + return current_obj + + return getter + + +def split_keys(attr_string): + parts = re.split(r"(\.|\[|\])", attr_string) # Split the string by '.', '[', and ']' + inside_bracket = False + keys = [] + for part in parts: + if not part or part in ".": + continue + if part == "[": + inside_bracket = True # Set flag when entering a bracket + continue + if part == "]": + inside_bracket = False # Set flag when exiting a bracket + continue + if inside_bracket: + keys.append(f"[]{part}") + else: + keys.append(f".{part}") + return keys diff --git a/source/uwlab/uwlab/genes/genome_cfg.py b/source/uwlab/uwlab/genes/genome_cfg.py new file mode 100644 index 00000000..2b034722 --- /dev/null +++ b/source/uwlab/uwlab/genes/genome_cfg.py @@ -0,0 +1,24 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from dataclasses import MISSING +from typing import Callable + +from isaaclab.utils import configclass + +from .genome import Genome + + +@configclass +class GenomeCfg: + class_type: Callable[..., Genome] = Genome + + genomic_mutation_profile: dict = MISSING # type: ignore + + genomic_constraint_profile: dict = MISSING # type: ignore + + seed: int = 32 diff --git a/source/uwlab/uwlab/managers/__init__.py b/source/uwlab/uwlab/managers/__init__.py new file mode 100644 index 00000000..f4a92432 --- /dev/null +++ b/source/uwlab/uwlab/managers/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-module for environment managers. + +Data manager is written to handle task oriented data structure needed for efficient indexing and operations +""" + +from .data_manager import DataManager, DataTerm +from .manager_term_cfg import DataTermCfg diff --git a/source/uwlab/uwlab/managers/data_manager.py b/source/uwlab/uwlab/managers/data_manager.py new file mode 100644 index 00000000..db835a35 --- /dev/null +++ b/source/uwlab/uwlab/managers/data_manager.py @@ -0,0 +1,385 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Data manager for generating and updating data.""" + +from __future__ import annotations + +import inspect +import torch +import weakref +from abc import abstractmethod +from collections.abc import Sequence +from prettytable import PrettyTable +from typing import TYPE_CHECKING + +import omni.kit.app + +from isaaclab.managers.manager_base import ManagerBase, ManagerTermBase + +from .manager_term_cfg import DataTermCfg + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +class DataTerm(ManagerTermBase): + """The base class for implementing a data term. + + A data term encapsulates a specialized, up-to-date data structure intended for use by multiple managers. + It is designed for scenarios where recalculating computationally expensive data repeatedly is inefficient. + + It is possible to assign a visualization function to the data term + that can be used to visualize the data in the simulator. + """ + + def __init__(self, cfg: DataTermCfg, env: ManagerBasedRLEnv): + """Initialize the data manager class. + + Args: + cfg: The configuration parameters for the data manager. + env: The environment object. + """ + super().__init__(cfg, env) + + # create buffers to store the data + # -- infos that can be used for logging + self.infos = dict() + # add handle for debug visualization (this is set to a valid handle inside set_debug_vis) + self._debug_vis_handle = None + # set initial state of debug visualization + self.set_debug_vis(self.cfg.debug_vis) + + def __del__(self): + """Unsubscribe from the callbacks.""" + if self._debug_vis_handle: + self._debug_vis_handle.unsubscribe() + self._debug_vis_handle = None + + """ + Properties + """ + + @property + @abstractmethod + def data(self) -> torch.Tensor: + """The data tensor. Shape is (num_envs, data_dim).""" + raise NotImplementedError + + @property + def has_debug_vis_implementation(self) -> bool: + """Whether the data generator 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 + + """ + Operations. + """ + + def set_debug_vis(self, debug_vis: bool) -> bool: + """Sets whether to visualize the data. + + Args: + debug_vis: Whether to visualize the data. + + Returns: + Whether the debug visualization was successfully set. False if the data + generator does not support debug visualization. + """ + # 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 + if debug_vis: + # create a subscriber for the post update event if it doesn't exist + if 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: + # remove the subscriber if it exists + if self._debug_vis_handle is not None: + self._debug_vis_handle.unsubscribe() + self._debug_vis_handle = None + # return success + return True + + def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float]: + """Reset the data manager and log infos. + + Args: + env_ids: The list of environment IDs to reset. Defaults to None. + + Returns: + A dictionary containing the information to log under the "{name}" key. + """ + # resolve the environment IDs + if env_ids is None: + env_ids = slice(None) + + # add logging infos + extras = {} + for info_name, info_value in self.infos.items(): + # compute the mean info value + extras[info_name] = torch.mean(info_value[env_ids]).item() + # reset the info value + info_value[env_ids] = 0.0 + + return extras + + def compute(self, dt: float): + """Compute the data. + + Args: + dt: The time step passed since the last call to compute. + """ + # update the infos based on current state + self._update_infos() + # update the data + self._update_data() + + """ + Implementation specific functions. + """ + + @abstractmethod + def _update_infos(self): + """Update the infos based on the current state.""" + raise NotImplementedError + + @abstractmethod + def _update_data(self): + """Update the data based on the current state.""" + raise NotImplementedError + + def _set_debug_vis_impl(self, debug_vis: bool): + """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__}.") + + def _debug_vis_callback(self, event): + """Callback for debug visualization. + + This function calls the visualization objects and sets the data to visualize into them. + """ + raise NotImplementedError(f"Debug visualization is not implemented for {self.__class__.__name__}.") + + +class DataManager(ManagerBase): + """Manager for generating data. + + DataManager is designed to address the challenge of managing computationally expensive + data structures that must be shared across multiple system components. In many scenarios, + these data structures—such as reorganized indices of multi-dimensional tensors—are not easily + derived directly from the base asset data, and recalculating them repeatedly would be inefficient. + + - "Expensive" refers to data structures that require significant computational resources to construct. + Specifically, these are reorganized indices of multi-dimensional tensors that all managers rely on and + cannot be directly derived from the base Isaac Lab asset data without a costly computation process. + + - "Up-to-date" means that the data is current and fresh, reflecting the most recent state of the system + without delay. For example, if the reward and event managers were to use data computed by the + command manager—which performs its computations later—they would end up using information that is + already one step behind, showing that command manager is not a appropriate place to calculate fresh data. + Other managers are not appropriate either because their names and concepts are unrelated. + + Centralizing the data generation within the DataManager ensures that all dependent managers receive + synchronized and immediately relevant data, maintaining optimal system performance. + + The data terms are implemented as classes that inherit from the :class:`DataTerm` class. + Each data manager term should also have a corresponding configuration class that inherits from the + :class:`DataTermCfg` class. + """ + + _env: ManagerBasedRLEnv + """The environment instance.""" + + def __init__(self, cfg: object, env: ManagerBasedRLEnv): + """Initialize the data manager. + + Args: + cfg: The configuration object or dictionary (``dict[str, DataTermCfg]``). + env: The environment instance. + """ + # create buffers to parse and store terms + self._terms: dict[str, DataTerm] = dict() + + # call the base class constructor (this prepares the terms) + super().__init__(cfg, env) + # store the data + self._data = dict() + 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 data manager.""" + msg = f" contains {len(self._terms.values())} active terms.\n" + + # create table for term information + table = PrettyTable() + table.title = "Active Data Terms" + table.field_names = ["Index", "Name", "Type"] + # set alignment of table columns + table.align["Name"] = "l" + # add info on each term + for index, (name, term) in enumerate(self._terms.items()): + table.add_row([index, name, term.__class__.__name__]) + # convert table to string + msg += table.get_string() + msg += "\n" + + return msg + + """ + Properties. + """ + + @property + def active_terms(self) -> list[str]: + """Name of active data terms.""" + return list(self._terms.keys()) + + @property + def has_debug_vis_implementation(self) -> bool: + """Whether the data terms have debug visualization implemented.""" + # check if function raises NotImplementedError + has_debug_vis = False + for term in self._terms.values(): + has_debug_vis |= term.has_debug_vis_implementation + return has_debug_vis + + """ + Operations. + """ + + 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.data[env_idx].cpu().tolist())) + idx += term.data.shape[1] + return terms + + def set_debug_vis(self, debug_vis: bool): + """Sets whether to visualize the data data. + + Args: + debug_vis: Whether to visualize the data data. + + Returns: + Whether the debug visualization was successfully set. False if the data + generator does not support debug visualization. + """ + for term in self._terms.values(): + term.set_debug_vis(debug_vis) + + def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]: + """Reset the data terms and log their infos. + + Args: + env_ids: The list of environment IDs to reset. Defaults to None. + + Returns: + A dictionary containing the information to log under the "Infos/{term_name}/{info_name}" key. + """ + # resolve environment ids + if env_ids is None: + env_ids = slice(None) + # store information + extras = {} + for name, term in self._terms.items(): + # reset the data term + infos = term.reset(env_ids=env_ids) + # compute the mean info value + for info_name, info_value in infos.items(): + extras[f"Infos/{name}/{info_name}"] = info_value + # return logged information + return extras + + def compute(self, dt: float): + """Updates the data. + + This function calls each data term managed by the class. + + Args: + dt: The time-step interval of the environment. + + """ + # iterate over all the data terms + for term in self._terms.values(): + # compute term's value + term.compute(dt) + + def get_data(self, name: str) -> torch.Tensor: + """Returns the data for the specified data term. + + Args: + name: The name of the data term. + + Returns: + The data tensor of the specified data term. + """ + return self._terms[name].data + + def get_term(self, name: str) -> DataTerm: + """Returns the data term with the specified name. + + Args: + name: The name of the data term. + + Returns: + The data term with the specified name. + """ + return self._terms[name] + + """ + Helper functions. + """ + + def _prepare_terms(self): + # check if config is dict already + if isinstance(self.cfg, dict): + cfg_items = self.cfg.items() + else: + cfg_items = self.cfg.__dict__.items() + # iterate over all the terms + for term_name, term_cfg in cfg_items: + # check for non config + if term_cfg is None: + continue + # check for valid config type + if not isinstance(term_cfg, DataTermCfg): + raise TypeError( + f"Configuration for the term '{term_name}' is not of type DataTermCfg." + f" Received: '{type(term_cfg)}'." + ) + # create the action term + term = term_cfg.class_type(term_cfg, self._env) + # sanity check if term is valid type + if not isinstance(term, DataTerm): + raise TypeError(f"Returned object for the term '{term_name}' is not of type DataType.") + # add class to dict + self._terms[term_name] = term diff --git a/source/uwlab/uwlab/managers/manager_term_cfg.py b/source/uwlab/uwlab/managers/manager_term_cfg.py new file mode 100644 index 00000000..8dfcbd8f --- /dev/null +++ b/source/uwlab/uwlab/managers/manager_term_cfg.py @@ -0,0 +1,56 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from collections.abc import Callable +from dataclasses import MISSING + +from isaaclab.managers.manager_term_cfg import ManagerTermBaseCfg +from isaaclab.utils import configclass +from isaaclab.utils.noise import NoiseCfg + + +@configclass +class DataTermCfg(ManagerTermBaseCfg): + """Configuration for an observation term.""" + + func: Callable[..., dict] = MISSING + """The name of the function to be called. + + This function should take the environment object and any other parameters + as input and return the observation signal as torch float tensors of + shape (num_envs, obs_term_dim). + """ + + noise: NoiseCfg | None = None + """The noise to add to the observation. Defaults to None, in which case no noise is added.""" + + clip: tuple[float, float] | None = None + """The clipping range for the observation after adding noise. Defaults to None, + in which case no clipping is applied.""" + + scale: 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`).""" + + history_length: int = 1 + + +@configclass +class DataGroupCfg: + """Configuration for an observation group.""" + + concatenate_terms: bool = True + """Whether to concatenate the observation terms in the group. Defaults to True. + + If true, the observation terms in the group are concatenated along the last dimension. + Otherwise, they are kept separate and returned as a dictionary. + """ + + enable_corruption: bool = False + """Whether to enable corruption for the observation group. Defaults to False. + + If true, the observation terms in the group are corrupted by adding noise (if specified). + Otherwise, no corruption is applied. + """ diff --git a/source/uwlab/uwlab/scene/__init__.py b/source/uwlab/uwlab/scene/__init__.py new file mode 100644 index 00000000..c764cf6c --- /dev/null +++ b/source/uwlab/uwlab/scene/__init__.py @@ -0,0 +1,31 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-package containing an interactive scene definition. + +A scene is a collection of entities (e.g., terrain, articulations, sensors, lights, etc.) that can be added to the +simulation. However, only a subset of these entities are of direct interest for the user to interact with. +For example, the user may want to interact with a robot in the scene, but not with the terrain or the lights. +For this reason, we integrate the different entities into a single class called :class:`InteractiveScene`. + +The interactive scene performs the following tasks: + +1. It parses the configuration class :class:`InteractiveSceneCfg` to create the scene. This configuration class is + inherited by the user to add entities to the scene. +2. It clones the entities based on the number of environments specified by the user. +3. It clubs the entities into different groups based on their type (e.g., articulations, sensors, etc.). +4. It provides a set of methods to unify the common operations on the entities in the scene (e.g., resetting internal + buffers, writing buffers to simulation and updating buffers from simulation). + +The interactive scene can be passed around to different modules in the framework to perform different tasks. +For instance, computing the observations based on the state of the scene, or randomizing the scene, or applying +actions to the scene. All these are handled by different "managers" in the framework. Please refer to the +:mod:`isaaclab.managers` sub-package for more details. +""" + +from .interactive_scene import InteractiveScene +from .interactive_scene_cfg import InteractiveSceneCfg +from .scene_context import SceneContext +from .scene_context_cfg import SceneContextCfg diff --git a/source/uwlab/uwlab/scene/interactive_scene.py b/source/uwlab/uwlab/scene/interactive_scene.py new file mode 100644 index 00000000..c0f774df --- /dev/null +++ b/source/uwlab/uwlab/scene/interactive_scene.py @@ -0,0 +1,643 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch +from collections.abc import Sequence +from typing import Any + +import carb +import omni.usd +from isaacsim.core.cloner import GridCloner +from isaacsim.core.prims import XFormPrim +from pxr import PhysxSchema + +import isaaclab.sim as sim_utils +from isaaclab.assets import ( + Articulation, + ArticulationCfg, + AssetBaseCfg, + DeformableObject, + DeformableObjectCfg, + RigidObject, + RigidObjectCfg, + RigidObjectCollection, + RigidObjectCollectionCfg, +) +from isaaclab.sensors import ContactSensorCfg, FrameTransformerCfg, SensorBase, SensorBaseCfg +from uwlab.terrains import TerrainImporter, TerrainImporterCfg + +from .interactive_scene_cfg import InteractiveSceneCfg + + +class InteractiveScene: + """A scene that contains entities added to the simulation. + + The interactive scene parses the :class:`InteractiveSceneCfg` class to create the scene. + Based on the specified number of environments, it clones the entities and groups them into different + categories (e.g., articulations, sensors, etc.). + + Cloning can be performed in two ways: + + * For tasks where all environments contain the same assets, a more performant cloning paradigm + can be used to allow for faster environment creation. This is specified by the ``replicate_physics`` flag. + + .. code-block:: python + + scene = InteractiveScene(cfg=InteractiveSceneCfg(replicate_physics=True)) + + * For tasks that require having separate assets in the environments, ``replicate_physics`` would have to + be set to False, which will add some costs to the overall startup time. + + .. code-block:: python + + scene = InteractiveScene(cfg=InteractiveSceneCfg(replicate_physics=False)) + + Each entity is registered to scene based on its name in the configuration class. For example, if the user + specifies a robot in the configuration class as follows: + + .. code-block:: python + + from isaaclab.scene import InteractiveSceneCfg + from isaaclab.utils import configclass + + from isaaclab_assets.robots.anymal import ANYMAL_C_CFG + + @configclass + class MySceneCfg(InteractiveSceneCfg): + + robot = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + Then the robot can be accessed from the scene as follows: + + .. code-block:: python + + from isaaclab.scene import InteractiveScene + + # create 128 environments + scene = InteractiveScene(cfg=MySceneCfg(num_envs=128)) + + # access the robot from the scene + robot = scene["robot"] + # access the robot based on its type + robot = scene.articulations["robot"] + + If the :class:`InteractiveSceneCfg` class does not include asset entities, the cloning process + can still be triggered if assets were added to the stage outside of the :class:`InteractiveScene` class: + + .. code-block:: python + + scene = InteractiveScene(cfg=InteractiveSceneCfg(num_envs=128, replicate_physics=True)) + scene.clone_environments() + + .. note:: + It is important to note that the scene only performs common operations on the entities. For example, + resetting the internal buffers, writing the buffers to the simulation and updating the buffers from the + simulation. The scene does not perform any task specific to the entity. For example, it does not apply + actions to the robot or compute observations from the robot. These tasks are handled by different + modules called "managers" in the framework. Please refer to the :mod:`isaaclab.managers` sub-package + for more details. + """ + + def __init__(self, cfg: InteractiveSceneCfg): + """Initializes the scene. + + Args: + cfg: The configuration class for the scene. + """ + # check that the config is valid + cfg.validate() + # store inputs + self.cfg = cfg + # initialize scene elements + self._terrain = None + 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 + self.stage = omni.usd.get_context().get_stage() + # physics scene path + self._physics_scene_path = None + # prepare cloner for environment replication + self.cloner = GridCloner(spacing=self.cfg.env_spacing) + self.cloner.define_base_env(self.env_ns) + self.env_prim_paths = self.cloner.generate_paths(f"{self.env_ns}/env", self.cfg.num_envs) + # create source prim + self.stage.DefinePrim(self.env_prim_paths[0], "Xform") + + # when replicate_physics=False, we assume heterogeneous environments and clone the xforms first. + # this triggers per-object level cloning in the spawner. + if not self.cfg.replicate_physics: + # clone the env xform + env_origins = self.cloner.clone( + source_prim_path=self.env_prim_paths[0], + prim_paths=self.env_prim_paths, + replicate_physics=False, + copy_from_source=True, + enable_env_ids=self.cfg.filter_collisions, # this won't do anything because we are not replicating physics + ) + self._default_env_origins = torch.tensor(env_origins, device=self.device, dtype=torch.float32) + else: + # otherwise, environment origins will be initialized during cloning at the end of environment creation + self._default_env_origins = None + + self._global_prim_paths = list() + if self._is_scene_setup_from_cfg(): + # add entities from config + self._add_entities_from_cfg() + # clone environments on a global scope if environment is homogeneous + if self.cfg.replicate_physics: + self.clone_environments(copy_from_source=False) + # replicate physics if we have more than one environment + # this is done to make scene initialization faster at play time + if self.cfg.replicate_physics and self.cfg.num_envs > 1: + self.cloner.replicate_physics( + source_prim_path=self.env_prim_paths[0], + prim_paths=self.env_prim_paths, + base_env_path=self.env_ns, + root_path=self.env_regex_ns.replace(".*", ""), + enable_env_ids=self.cfg.filter_collisions, + ) + + # since env_ids is only applicable when replicating physics, we have to fallback to the previous method + # to filter collisions if replicate_physics is not enabled + if not self.cfg.replicate_physics and self.cfg.filter_collisions: + self.filter_collisions(self._global_prim_paths) + + def clone_environments(self, copy_from_source: bool = False): + """Creates clones of the environment ``/World/envs/env_0``. + + Args: + copy_from_source: (bool): If set to False, clones inherit from /World/envs/env_0 and mirror its changes. + If True, clones are independent copies of the source prim and won't reflect its changes (start-up time + may increase). Defaults to False. + """ + # check if user spawned different assets in individual environments + # this flag will be None if no multi asset is spawned + carb_settings_iface = carb.settings.get_settings() + has_multi_assets = carb_settings_iface.get("/isaaclab/spawn/multi_assets") + if has_multi_assets and self.cfg.replicate_physics: + omni.log.warn( + "Varying assets might have been spawned under different environments." + " However, the replicate physics flag is enabled in the 'InteractiveScene' configuration." + " This may adversely affect PhysX parsing. We recommend disabling this property." + ) + + # clone the environment + env_origins = self.cloner.clone( + source_prim_path=self.env_prim_paths[0], + prim_paths=self.env_prim_paths, + replicate_physics=self.cfg.replicate_physics, + copy_from_source=copy_from_source, + enable_env_ids=self.cfg.filter_collisions, # this automatically filters collisions between environments + ) + + # since env_ids is only applicable when replicating physics, we have to fallback to the previous method + # to filter collisions if replicate_physics is not enabled + if not self.cfg.replicate_physics and self.cfg.filter_collisions: + omni.log.warn( + "Collision filtering can only be automatically enabled when replicate_physics=True." + " Please call scene.filter_collisions(global_prim_paths) to filter collisions across environments." + ) + + # in case of heterogeneous cloning, the env origins is specified at init + if self._default_env_origins is None: + self._default_env_origins = torch.tensor(env_origins, device=self.device, dtype=torch.float32) + + def filter_collisions(self, global_prim_paths: list[str] | None = None): + """Filter environments collisions. + + Disables collisions between the environments in ``/World/envs/env_.*`` and enables collisions with the prims + in global prim paths (e.g. ground plane). + + Args: + global_prim_paths: A list of global prim paths to enable collisions with. + Defaults to None, in which case no global prim paths are considered. + """ + # validate paths in global prim paths + if global_prim_paths is None: + global_prim_paths = [] + else: + # remove duplicates in paths + global_prim_paths = list(set(global_prim_paths)) + + # set global prim paths list if not previously defined + if len(self._global_prim_paths) < 1: + self._global_prim_paths += global_prim_paths + + # filter collisions within each environment instance + self.cloner.filter_collisions( + self.physics_scene_path, + "/World/collisions", + self.env_prim_paths, + global_paths=self._global_prim_paths, + ) + + def __str__(self) -> str: + """Returns a string representation of the scene.""" + msg = f"\n" + msg += f"\tNumber of environments: {self.cfg.num_envs}\n" + msg += f"\tEnvironment spacing : {self.cfg.env_spacing}\n" + msg += f"\tSource prim name : {self.env_prim_paths[0]}\n" + msg += f"\tGlobal prim paths : {self._global_prim_paths}\n" + msg += f"\tReplicate physics : {self.cfg.replicate_physics}" + return msg + + """ + Properties. + """ + + @property + def physics_scene_path(self) -> str: + """The path to the USD Physics Scene.""" + if self._physics_scene_path is None: + for prim in self.stage.Traverse(): + if prim.HasAPI(PhysxSchema.PhysxSceneAPI): + self._physics_scene_path = prim.GetPrimPath().pathString + omni.log.info(f"Physics scene prim path: {self._physics_scene_path}") + break + if self._physics_scene_path is None: + raise RuntimeError("No physics scene found! Please make sure one exists.") + return self._physics_scene_path + + @property + def physics_dt(self) -> float: + """The physics timestep of the scene.""" + return sim_utils.SimulationContext.instance().get_physics_dt() # pyright: ignore [reportOptionalMemberAccess] + + @property + def device(self) -> str: + """The device on which the scene is created.""" + return sim_utils.SimulationContext.instance().device # pyright: ignore [reportOptionalMemberAccess] + + @property + def env_ns(self) -> str: + """The namespace ``/World/envs`` in which all environments created. + + The environments are present w.r.t. this namespace under "env_{N}" prim, + where N is a natural number. + """ + return "/World/envs" + + @property + def env_regex_ns(self) -> str: + """The namespace ``/World/envs/env_.*`` in which all environments created.""" + return f"{self.env_ns}/env_.*" + + @property + def num_envs(self) -> int: + """The number of environments handled by the scene.""" + return self.cfg.num_envs + + @property + def env_origins(self) -> torch.Tensor: + """The origins of the environments in the scene. Shape is (num_envs, 3).""" + if self._terrain is not None: + return self._terrain.env_origins + else: + return self._default_env_origins + + @property + def terrain(self) -> TerrainImporter | None: + """The terrain in the scene. If None, then the scene has no terrain. + + Note: + We treat terrain separate from :attr:`extras` since terrains define environment origins and are + handled differently from other miscellaneous entities. + """ + return self._terrain + + @property + def articulations(self) -> dict[str, Articulation]: + """A dictionary of articulations in the scene.""" + return self._articulations + + @property + def deformable_objects(self) -> dict[str, DeformableObject]: + """A dictionary of deformable objects in the scene.""" + return self._deformable_objects + + @property + 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.""" + return self._sensors + + @property + def extras(self) -> dict[str, XFormPrim]: + """A dictionary of miscellaneous simulation objects that neither inherit from assets nor sensors. + + The keys are the names of the miscellaneous objects, and the values are the `XFormPrim`_ + of the corresponding prims. + + As an example, lights or other props in the scene that do not have any attributes or properties that you + want to alter at runtime can be added to this dictionary. + + Note: + These are not reset or updated by the scene. They are mainly other prims that are not necessarily + handled by the interactive scene, but are useful to be accessed by the user. + + .. _XFormPrim: https://docs.omniverse.nvidia.com/py/isaacsim/source/isaacsim.core/docs/index.html#isaacsim.core.prims.XFormPrim + + """ + 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_state_w[:, :7].clone() + if is_relative: + asset_state["root_pose"][:, :3] -= self.env_origins + asset_state["root_velocity"] = articulation.data.root_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_state_w[:, :7].clone() + if is_relative: + asset_state["root_pose"][:, :3] -= self.env_origins + asset_state["root_velocity"] = rigid_object.data.root_vel_w.clone() + state["rigid_object"][asset_name] = asset_state + return state + + """ + Operations. + """ + + def reset(self, env_ids: Sequence[int] | None = None): + """Resets the scene entities. + + Args: + env_ids: The indices of the environments to reset. + Defaults to None (all instances). + """ + # -- assets + for articulation in self._articulations.values(): + articulation.reset(env_ids) + for deformable_object in self._deformable_objects.values(): + 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_pose_to_sim(root_pose, env_ids=env_ids) + articulation.write_root_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_pose_to_sim(root_pose, env_ids=env_ids) + rigid_object.write_root_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 + for articulation in self._articulations.values(): + articulation.write_data_to_sim() + for deformable_object in self._deformable_objects.values(): + 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. + + Args: + dt: The amount of time passed from last :meth:`update` call. + """ + # -- assets + for articulation in self._articulations.values(): + articulation.update(dt) + for deformable_object in self._deformable_objects.values(): + 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) + + """ + Operations: Iteration. + """ + + def keys(self) -> list[str]: + """Returns the keys of the scene entities. + + Returns: + The keys of the scene entities. + """ + all_keys = ["terrain"] + for asset_family in [ + self._articulations, + self._deformable_objects, + self._rigid_objects, + self._rigid_object_collections, + self._sensors, + self._extras, + ]: + all_keys += list(asset_family.keys()) + return all_keys + + def __getitem__(self, key: str) -> Any: + """Returns the scene entity with the given key. + + Args: + key: The key of the scene entity. + + Returns: + The scene entity. + """ + # check if it is a terrain + if key == "terrain": + return self._terrain + + all_keys = ["terrain"] + # check if it is in other dictionaries + for asset_family in [ + self._articulations, + self._deformable_objects, + self._rigid_objects, + self._rigid_object_collections, + self._sensors, + self._extras, + ]: + out = asset_family.get(key) + # if found, return + if out is not None: + return out + all_keys += list(asset_family.keys()) + # if not found, raise error + raise KeyError(f"Scene entity with key '{key}' not found. Available Entities: '{all_keys}'") + + """ + Internal methods. + """ + + def _is_scene_setup_from_cfg(self): + return any( + not (asset_name in InteractiveSceneCfg.__dataclass_fields__ or asset_cfg is None) + for asset_name, asset_cfg in self.cfg.__dict__.items() + ) + + def _add_entities_from_cfg(self): + """Add scene entities from the config.""" + # store paths that are in global collision filter + self._global_prim_paths = list() + # parse the entire scene config and resolve regex + for asset_name, asset_cfg in self.cfg.__dict__.items(): + # skip keywords + # note: easier than writing a list of keywords: [num_envs, env_spacing, lazy_sensor_update] + if asset_name in InteractiveSceneCfg.__dataclass_fields__ or asset_cfg is None: + continue + # resolve regex + 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 + asset_cfg.num_envs = self.cfg.num_envs + asset_cfg.env_spacing = self.cfg.env_spacing + self._terrain = asset_cfg.class_type(asset_cfg) + elif isinstance(asset_cfg, ArticulationCfg): + self._articulations[asset_name] = asset_cfg.class_type(asset_cfg) + elif isinstance(asset_cfg, DeformableObjectCfg): + 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): + updated_target_frames = [] + for target_frame in asset_cfg.target_frames: + target_frame.prim_path = target_frame.prim_path.format(ENV_REGEX_NS=self.env_regex_ns) + updated_target_frames.append(target_frame) + asset_cfg.target_frames = updated_target_frames + elif isinstance(asset_cfg, ContactSensorCfg): + updated_filter_prim_paths_expr = [] + for filter_prim_path in asset_cfg.filter_prim_paths_expr: + updated_filter_prim_paths_expr.append(filter_prim_path.format(ENV_REGEX_NS=self.env_regex_ns)) + asset_cfg.filter_prim_paths_expr = updated_filter_prim_paths_expr + + self._sensors[asset_name] = asset_cfg.class_type(asset_cfg) + elif isinstance(asset_cfg, AssetBaseCfg): + # manually spawn asset + if asset_cfg.spawn is not None: + asset_cfg.spawn.func( + asset_cfg.prim_path, + asset_cfg.spawn, + translation=asset_cfg.init_state.pos, + orientation=asset_cfg.init_state.rot, + ) + # store xform prim view corresponding to this asset + # all prims in the scene are Xform prims (i.e. have a transform component) + self._extras[asset_name] = XFormPrim(asset_cfg.prim_path, reset_xform_properties=False) + else: + raise ValueError(f"Unknown asset config type for {asset_name}: {asset_cfg}") + # store global collision paths + if hasattr(asset_cfg, "collision_group") and asset_cfg.collision_group == -1: + asset_paths = sim_utils.find_matching_prim_paths(asset_cfg.prim_path) + self._global_prim_paths += asset_paths diff --git a/source/uwlab/uwlab/scene/interactive_scene_cfg.py b/source/uwlab/uwlab/scene/interactive_scene_cfg.py new file mode 100644 index 00000000..4fe4d664 --- /dev/null +++ b/source/uwlab/uwlab/scene/interactive_scene_cfg.py @@ -0,0 +1,111 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import MISSING + +from isaaclab.utils.configclass import configclass + + +@configclass +class InteractiveSceneCfg: + """Configuration for the interactive scene. + + The users can inherit from this class to add entities to their scene. This is then parsed by the + :class:`InteractiveScene` class to create the scene. + + .. note:: + The adding of entities to the scene is sensitive to the order of the attributes in the configuration. + Please make sure to add the entities in the order you want them to be added to the scene. + The recommended order of specification is terrain, physics-related assets (articulations and rigid bodies), + sensors and non-physics-related assets (lights). + + For example, to add a robot to the scene, the user can create a configuration class as follows: + + .. code-block:: python + + import isaaclab.sim as sim_utils + from isaaclab.assets import AssetBaseCfg + from isaaclab.scene import InteractiveSceneCfg + from isaaclab.sensors.ray_caster import GridPatternCfg, RayCasterCfg + from isaaclab.utils import configclass + + from isaaclab_assets.robots.anymal import ANYMAL_C_CFG + + @configclass + class MySceneCfg(InteractiveSceneCfg): + + # terrain - flat terrain plane + terrain = TerrainImporterCfg( + prim_path="/World/ground", + terrain_type="plane", + ) + + # articulation - robot 1 + robot_1 = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot_1") + # articulation - robot 2 + robot_2 = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot_2") + robot_2.init_state.pos = (0.0, 1.0, 0.6) + + # sensor - ray caster attached to the base of robot 1 that scans the ground + height_scanner = RayCasterCfg( + prim_path="{ENV_REGEX_NS}/Robot_1/base", + offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)), + attach_yaw_only=True, + pattern_cfg=GridPatternCfg(resolution=0.1, size=[1.6, 1.0]), + debug_vis=True, + mesh_prim_paths=["/World/ground"], + ) + + # extras - light + light = AssetBaseCfg( + prim_path="/World/light", + spawn=sim_utils.DistantLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75)), + init_state=AssetBaseCfg.InitialStateCfg(pos=(0.0, 0.0, 500.0)), + ) + + """ + + num_envs: int = MISSING + """Number of environment instances handled by the scene.""" + + env_spacing: float = MISSING + """Spacing between environments. + + This is the default distance between environment origins in the scene. Used only when the + number of environments is greater than one. + """ + + lazy_sensor_update: bool = True + """Whether to update sensors only when they are accessed. Default is True. + + If true, the sensor data is only updated when their attribute ``data`` is accessed. Otherwise, the sensor + data is updated every time sensors are updated. + """ + + replicate_physics: bool = True + """Enable/disable replication of physics schemas when using the Cloner APIs. Default is True. + + If True, the simulation will have the same asset instances (USD prims) in all the cloned environments. + Internally, this ensures optimization in setting up the scene and parsing it via the physics stage parser. + + If False, the simulation allows having separate asset instances (USD prims) in each environment. + This flexibility comes at a cost of slowdowns in setting up and parsing the scene. + + .. note:: + Optimized parsing of certain prim types (such as deformable objects) is not currently supported + by the physics engine. In these cases, this flag needs to be set to False. + """ + + filter_collisions: bool = True + """Enable/disable collision filtering between cloned environments. Default is True. + + If True, collisions will not occur between cloned environments. + + If False, the simulation will generate collisions between environments. + + .. note:: + Collisions can only be filtered automatically in direct workflows when physics replication is enabled. + If ``replicated_physics=False`` and collision filtering is desired, make sure to call ``scene.filter_collisions()``. + """ diff --git a/source/uwlab/uwlab/scene/scene_context.py b/source/uwlab/uwlab/scene/scene_context.py new file mode 100644 index 00000000..3b69095c --- /dev/null +++ b/source/uwlab/uwlab/scene/scene_context.py @@ -0,0 +1,230 @@ +# Copyright (c) 2024-2025, The UW 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, Any + +from isaaclab.sensors import SensorBase, SensorBaseCfg +from uwlab.assets import ArticulationCfg, UniversalArticulation + +if TYPE_CHECKING: + from .scene_context_cfg import SceneContextCfg + + +class SceneContext: + def __init__(self, cfg: SceneContextCfg): + """Initializes the scene. + + Args: + cfg: The configuration class for the scene. + """ + # check that the config is valid + cfg.validate() # type: ignore + # store inputs + self.cfg = cfg + # initialize scene elements + self._terrain = None + self._articulations = dict() + self._deformable_objects = dict() + self._rigid_objects = dict() + self._rigid_object_collections = dict() + self._sensors = dict() + self._extras = dict() + + # physics scene path + self._physics_scene_path = None + + self._global_prim_paths = list() + + if self._is_scene_setup_from_cfg(): + # add entities from config + self._add_entities_from_cfg() + # clone environments on a global scope if environment is homogeneous + + def __str__(self) -> str: + """Returns a string representation of the scene.""" + msg = f"\n" + msg += f"\tNumber of environments: {self.cfg.num_envs}\n" + msg += f"\tGlobal prim paths : {self._global_prim_paths}\n" + return msg + + """ + Properties. + """ + + @property + def physics_dt(self) -> float: + """The physics timestep of the scene.""" + return self.cfg.dt + + @property + def device(self) -> str: + """The device on which the scene is created.""" + return self.cfg.device + + @property + def num_envs(self) -> int: + """The number of environments handled by the scene.""" + return self.cfg.num_envs + + @property + def articulations(self) -> dict[str, UniversalArticulation]: + """A dictionary of articulations in the scene.""" + return self._articulations + + @property + def deformable_objects(self) -> dict[str, any]: + """A dictionary of deformable objects in the scene.""" + return self._deformable_objects + + @property + def rigid_objects(self) -> dict[str, any]: + """A dictionary of rigid objects in the scene.""" + return self._rigid_objects + + @property + def rigid_object_collections(self) -> dict[str, any]: + """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.""" + return self._sensors + + @property + def extras(self) -> dict[str, Any]: + """A dictionary of extra entities in the scene, such as XFormPrimView.""" + return self._extras + + @property + def env_origins(self) -> torch.Tensor: + return torch.zeros((self.num_envs, 3), device=self.device) + + """ + Operations. + """ + + def start(self): + """Starts the scene entities.""" + # -- assets + for articulation in self._articulations.values(): + articulation._initialize_impl(self.device) + # -- sensors + for sensor in self._sensors.values(): + sensor._initialize_impl(self.device) + + def reset(self, env_ids: Sequence[int] | None = None): + """Resets the scene entities. + + Args: + env_ids: The indices of the environments to reset. + Defaults to None (all instances). + """ + # -- assets + for articulation in self._articulations.values(): + articulation.reset(env_ids) + # -- sensors + for sensor in self._sensors.values(): + sensor.reset(env_ids) + + def write_data_to_context(self): + """Writes the data of the scene entities to the simulation.""" + # -- assets + for articulation in self._articulations.values(): + articulation.write_data_to_sim() + + def update(self, dt: float) -> None: + """Update the scene entities. + + Args: + dt: The amount of time passed from last :meth:`update` call. + """ + # -- assets + for articulation in self._articulations.values(): + articulation.update(dt) + # -- sensors + for sensor in self._sensors.values(): + sensor.update(dt, force_recompute=not self.cfg.lazy_sensor_update) + + """ + Operations: Iteration. + """ + + def keys(self) -> list[str]: + """Returns the keys of the scene entities. + + Returns: + The keys of the scene entities. + """ + all_keys = [] + for asset_family in [ + self._articulations, + self._sensors, + self._extras, + ]: + all_keys += list(asset_family.keys()) + return all_keys + + def __getitem__(self, key: str) -> Any: + """Returns the scene entity with the given key. + + Args: + key: The key of the scene entity. + + Returns: + The scene entity. + """ + + all_keys = [] + # check if it is in other dictionaries + for asset_family in [ + self._articulations, + self._sensors, + self._extras, + ]: + out = asset_family.get(key) + # if found, return + if out is not None: + return out + all_keys += list(asset_family.keys()) + # if not found, raise error + raise KeyError(f"Scene entity with key '{key}' not found. Available Entities: '{all_keys}'") + + """ + Internal methods. + """ + + def _is_scene_setup_from_cfg(self): + from .scene_context_cfg import SceneContextCfg + + return any( + not (asset_name in SceneContextCfg.__dataclass_fields__ or asset_cfg is None) + for asset_name, asset_cfg in self.cfg.__dict__.items() + ) + + def _add_entities_from_cfg(self): + """Add scene entities from the config.""" + # store paths that are in global collision filter + self._global_prim_paths = list() + from .scene_context_cfg import SceneContextCfg + + # parse the entire scene config and resolve regex + for asset_name, asset_cfg in self.cfg.__dict__.items(): + # skip keywords + # note: easier than writing a list of keywords: [num_envs, env_spacing, lazy_sensor_update] + if asset_name in SceneContextCfg.__dataclass_fields__ or asset_cfg is None: + continue + # create asset + if isinstance(asset_cfg, ArticulationCfg): + self._articulations[asset_name] = asset_cfg.class_type(asset_cfg) + elif isinstance(asset_cfg, SensorBaseCfg): + self._sensors[asset_name] = asset_cfg.class_type(asset_cfg) + else: + raise ValueError(f"Unknown asset config type for {asset_name}: {asset_cfg}") + # store global collision paths diff --git a/source/uwlab/uwlab/scene/scene_context_cfg.py b/source/uwlab/uwlab/scene/scene_context_cfg.py new file mode 100644 index 00000000..f0dee4b4 --- /dev/null +++ b/source/uwlab/uwlab/scene/scene_context_cfg.py @@ -0,0 +1,32 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from typing import Callable + +from isaaclab.utils.configclass import configclass + +from .scene_context import SceneContext + + +@configclass +class SceneContextCfg: + class_type: Callable[..., SceneContext] = SceneContext + + dt: float = 0.02 + """Time step of the simulation. Default is 0.02.""" + + device: str = "cpu" + """Device to run the simulation on. Default is "cpu".""" + + num_envs: int = 1 + """Number of environment instances handled by the scene.""" + # DO NOT MODIFY NUM_ENVS DEFAULT VALUE, REAL ENVIRONMENT IS UNLIKELY TO HAVE MORE THAN 1 ENVIRONMENT + + lazy_sensor_update: bool = True + """Whether to update sensors only when they are accessed. Default is True. + + If true, the sensor data is only updated when their attribute ``data`` is accessed. Otherwise, the sensor + data is updated every time sensors are updated. + """ diff --git a/source/uwlab/uwlab/sim/converters/__init__.py b/source/uwlab/uwlab/sim/converters/__init__.py new file mode 100644 index 00000000..9f012a3f --- /dev/null +++ b/source/uwlab/uwlab/sim/converters/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .mesh_converter import MeshConverter +from .mesh_converter_cfg import MeshConverterCfg diff --git a/source/uwlab/uwlab/sim/converters/common_material_property_cfg.py b/source/uwlab/uwlab/sim/converters/common_material_property_cfg.py new file mode 100644 index 00000000..c4f07dda --- /dev/null +++ b/source/uwlab/uwlab/sim/converters/common_material_property_cfg.py @@ -0,0 +1,36 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from dataclasses import field + +import isaaclab.sim as sim_utils +from isaaclab.utils import configclass + +from ..spawners.materials import common_materials_cfg as common_materials +from .mesh_converter_cfg import MeshConverterCfg + + +@configclass +class PVCCfg(MeshConverterCfg): + rigid_props: sim_utils.RigidBodyPropertiesCfg = field(default_factory=sim_utils.RigidBodyPropertiesCfg) + collision_props: sim_utils.CollisionPropertiesCfg = field(default_factory=sim_utils.CollisionPropertiesCfg) + mass_props: sim_utils.MassPropertiesCfg = sim_utils.MassPropertiesCfg(density=1380) + visual_material_props: sim_utils.VisualMaterialCfg = field(default_factory=common_materials.PCVVisualMaterialCfg) + physics_material_props: sim_utils.PhysicsMaterialCfg = field( + default_factory=common_materials.PCVPhysicalMaterialCfg + ) + + +@configclass +class SteelCfg(MeshConverterCfg): + rigid_props: sim_utils.RigidBodyPropertiesCfg = field(default_factory=sim_utils.RigidBodyPropertiesCfg) + collision_props: sim_utils.CollisionPropertiesCfg = field(default_factory=sim_utils.CollisionPropertiesCfg) + mass_props: sim_utils.MassPropertiesCfg = sim_utils.MassPropertiesCfg(density=7870) + visual_material_props: sim_utils.VisualMaterialCfg = field(default_factory=common_materials.SteelVisualMaterialCfg) + physics_material_props: sim_utils.PhysicsMaterialCfg = field( + default_factory=common_materials.SteelPhysicalMaterialCfg + ) diff --git a/source/uwlab/uwlab/sim/converters/mesh_converter.py b/source/uwlab/uwlab/sim/converters/mesh_converter.py new file mode 100644 index 00000000..dabb9296 --- /dev/null +++ b/source/uwlab/uwlab/sim/converters/mesh_converter.py @@ -0,0 +1,398 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import asyncio +import os +from typing import List + +import isaacsim.core.utils.prims as prim_utils +import omni +import omni.kit.commands +import omni.usd +from isaacsim.coreutils.extensions import enable_extension +from pxr import Sdf, Usd, UsdGeom, UsdPhysics, UsdShade, UsdUtils + +from isaaclab.sim.converters.asset_converter_base import AssetConverterBase +from isaaclab.sim.schemas import schemas +from isaaclab.sim.utils import clone, export_prim_to_file, get_all_matching_child_prims, safe_set_attribute_on_usd_prim + +from .mesh_converter_cfg import MeshConverterCfg + + +def apply_material_binding(stage: Usd.Stage, prim_path: str, material_path: str) -> None: + """Applies a material to a given prim.""" + prim = stage.GetPrimAtPath(prim_path) + if not prim: + return + + # Ensure that the MaterialBindingAPI is applied + if not prim.HasAPI(UsdShade.MaterialBindingAPI): + UsdShade.MaterialBindingAPI.Apply(prim) + + material_binding_api = UsdShade.MaterialBindingAPI(prim) + material = UsdShade.Material(stage.GetPrimAtPath(material_path)) + + # Unbind any previous materials before binding a new one + material_binding_api.UnbindAllBindings() + material_binding_api.Bind( + material=material, + bindingStrength=UsdShade.Tokens.weakerThanDescendants, + materialPurpose=UsdShade.Tokens.allPurpose, + ) + + +def apply_collision_properties(stage: Usd.Stage, prim_path: str, cfg: MeshConverterCfg) -> None: + """Applies collision properties to a given prim.""" + collision_prim: Usd.Prim = stage.GetPrimAtPath(prim_path) + if not collision_prim: + return + + mesh_collision_api = UsdPhysics.MeshCollisionAPI.Apply(collision_prim) # type: ignore + mesh_collision_api.GetApproximationAttr().Set(cfg.collision_approximation) + # -- Collider properties such as offset, scale, etc. + if cfg.collision_props: + schemas.define_collision_properties(prim_path=collision_prim.GetPath(), cfg=cfg.collision_props, stage=stage) + visibility_attr = collision_prim.GetAttribute("visibility") + if visibility_attr: + visibility_attr.Set("invisible", time=Usd.TimeCode.Default()) + + +def remove_all_other_prims(stage: Usd.Stage, keep_prefix_paths: List[str]) -> None: + """Removes all prims except /World and /World/. + + Args: + stage: The USD stage. + keep_prefix_paths: A list of path prefixes to keep. All prims under these prefixes + (and under /World) will not be removed. + """ + prims_to_remove = [] + root_path = stage.GetDefaultPrim().GetPath() + for prim in stage.Traverse(): + ppath = prim.GetPath() + + if ppath == root_path: + continue + # Check if the prim matches any of the keep prefixes + if not any(ppath.HasPrefix(Sdf.Path(prefix)) for prefix in keep_prefix_paths): + prims_to_remove.append(ppath) + + for ppath in prims_to_remove: + stage.RemovePrim(ppath) + + +@clone +def spawn_preview_surface(prim_path: str, cfg) -> Usd.Prim: + """Create a preview surface prim and override the settings with the given config. + + A preview surface is a physically-based surface that handles simple shaders while supporting + both *specular* and *metallic* workflows. All color inputs are in linear color space (RGB). + For more information, see the `documentation `__. + + The function calls the USD command `CreatePreviewSurfaceMaterialPrim`_ to create the prim. + + .. _CreatePreviewSurfaceMaterialPrim: https://docs.omniverse.nvidia.com/kit/docs/omni.usd/latest/omni.usd.commands/omni.usd.commands.CreatePreviewSurfaceMaterialPrimCommand.html + + .. note:: + This function is decorated with :func:`clone` that resolves prim path into list of paths + if the input prim path is a regex pattern. This is done to support spawning multiple assets + from a single and cloning the USD prim at the given path expression. + + Args: + prim_path: The prim path or pattern to spawn the asset at. If the prim path is a regex pattern, + then the asset is spawned at all the matching prim paths. + cfg: The configuration instance. + + Returns: + The created prim. + + Raises: + ValueError: If a prim already exists at the given path. + """ + # spawn material if it doesn't exist. + if not prim_utils.is_prim_path_valid(prim_path): + omni.kit.commands.execute("CreatePreviewSurfaceMaterialPrim", mtl_path=prim_path, select_new_prim=False) + else: + raise ValueError(f"A prim already exists at path: '{prim_path}'.") + # obtain prim + prim = prim_utils.get_prim_at_path(f"{prim_path}/Shader") + # apply properties + cfg = cfg.to_dict() + del cfg["func"] + for attr_name, attr_value in cfg.items(): + safe_set_attribute_on_usd_prim(prim, f"inputs:{attr_name}", attr_value, camel_case=True) + # return prim + return prim + + +class MeshConverter(AssetConverterBase): + """Converter for a mesh file in OBJ / STL / FBX format to a USD file. + + This class wraps around the `omni.kit.asset_converter`_ extension to provide a lazy implementation + for mesh to USD conversion. It stores the output USD file in an instanceable format since that is + what is typically used in all learning related applications. + + To make the asset instanceable, we must follow a certain structure dictated by how USD scene-graph + instancing and physics work. The rigid body component must be added to each instance and not the + referenced asset (i.e. the prototype prim itself). This is because the rigid body component defines + properties that are specific to each instance and cannot be shared under the referenced asset. For + more information, please check the `documentation `_. + + Due to the above, we follow the following structure: + + * ``{prim_path}`` - The root prim that is an Xform with the rigid body and mass APIs if configured. + * ``{prim_path}/geometry`` - The prim that contains the mesh and optionally the materials if configured. + If instancing is enabled, this prim will be an instanceable reference to the prototype prim. + + .. _omni.kit.asset_converter: https://docs.omniverse.nvidia.com/extensions/latest/ext_asset-converter.html + + .. caution:: + When converting STL files, Z-up convention is assumed, even though this is not the default for many CAD + export programs. Asset orientation convention can either be modified directly in the CAD program's export + process or an offset can be added within the config in Isaac Lab. + + """ + + cfg: MeshConverterCfg + """The configuration instance for mesh to USD conversion.""" + + def __init__(self, cfg: MeshConverterCfg): + """Initializes the class. + + Args: + cfg: The configuration instance for mesh to USD conversion. + """ + super().__init__(cfg=cfg) + + """ + Implementation specific methods. + """ + + def _convert_asset(self, cfg: MeshConverterCfg): + """Generate USD from OBJ, STL or FBX. + + It stores the asset in the following format: + + /file_name (default prim) + |- /mesh_file_name <- Made instanceable if requested + |- /Looks + |- /mesh + + Args: + cfg: The configuration for conversion of mesh to USD. + + Raises: + RuntimeError: If the conversion using the Omniverse asset converter fails. + """ + # resolve mesh name and format + mesh_file_basename, mesh_file_format = os.path.basename(cfg.asset_path).split(".") + mesh_file_format = mesh_file_format.lower() + + # 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}" + ) + ) + # 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 = Usd.Stage.Open(self.usd_path) # type: ignore + # add USD to stage cache + stage_id = UsdUtils.StageCache.Get().Insert(stage) # type: ignore + + stage.DefinePrim(f"/{mesh_file_basename}/visuals", "Xform") + stage.DefinePrim(f"/{mesh_file_basename}/collisions", "Xform") + # Get the default prim (which is the root prim) -- "/{mesh_file_basename}" + xform_prim = stage.GetDefaultPrim() + root_path = f"/{mesh_file_basename}" + root_prim = stage.GetPrimAtPath(f"/{mesh_file_basename}") + visual_prim = stage.GetPrimAtPath(f"{root_prim.GetPath()}/visuals") + collisions_prim = stage.GetPrimAtPath(f"{root_prim.GetPath()}/collisions") + + # Collect meshes + meshes = [] + for child in stage.GetPseudoRoot().GetChildren(): + found_meshes = get_all_matching_child_prims( + child.GetPath(), lambda prim: prim.GetTypeName() == "Mesh", stage=stage + ) + meshes.extend(found_meshes) + if not meshes: + print(f"No meshes found in {self.usd_path}") + + # Process each mesh with visual material, physic material and collision properties + for count, mesh_prim in enumerate(meshes): + # Create visual copy + visual_prim_path = f"{root_path}/visuals/mesh_{count:03d}" + visual_material_path = f"{root_path}/visuals/Looks" + + Sdf.CopySpec( + stage.GetEditTarget().GetLayer(), + str(mesh_prim.GetPath()), + stage.GetEditTarget().GetLayer(), + visual_prim_path, + ) + if cfg.visual_material_props: + visual_cfg = cfg.visual_material_props + visual_cfg.func(prim_path=visual_material_path, cfg=visual_cfg, stage=stage) + apply_material_binding(stage, visual_prim_path, visual_material_path) + else: + Sdf.CopySpec( + stage.GetEditTarget().GetLayer(), + f"{root_path}/temp/Looks", + stage.GetEditTarget().GetLayer(), + visual_material_path, + ) + apply_material_binding(stage, visual_prim_path, visual_material_path + "/DefaultMaterial") + + # Create collision copy + collisions_prim_path = f"{root_path}/collisions/mesh_{count:03d}" + Sdf.CopySpec( + stage.GetEditTarget().GetLayer(), + str(mesh_prim.GetPath()), + stage.GetEditTarget().GetLayer(), + collisions_prim_path, + ) + + # Apply collision properties + if cfg.collision_props: + apply_collision_properties(stage, collisions_prim_path, cfg) + + if cfg.physics_material_props is not None: + physics_material_path = f"{root_path}/collisions/PhysicsMaterial" + cfg.physics_material_props.func( + prim_path=physics_material_path, cfg=cfg.physics_material_props, stage=stage + ) + apply_material_binding(stage, collisions_prim_path, physics_material_path) + + # Remove extraneous prims + remove_all_other_prims(stage, [visual_prim.GetPath().pathString, collisions_prim.GetPath().pathString]) + # Delete the old Xform and make the new Xform the default prim + stage.SetDefaultPrim(xform_prim) + # Handle instanceable + # Create a new Xform prim that will be the prototype prim + instanceable_mesh_path = os.path.join(".", "Props", f"{mesh_file_basename}.usd") + if cfg.make_instanceable: + # Export Xform to a file so we can reference it from all instances + export_prim_to_file( + path=os.path.join(self.usd_dir, instanceable_mesh_path), + source_prim_path=root_path, + stage=stage, + ) + # Delete the original prim that will now be a reference + visual_prim_path = visual_prim.GetPath().pathString + collision_prim_path = collisions_prim.GetPath().pathString + omni.kit.commands.execute("DeletePrims", paths=[visual_prim_path, collision_prim_path], stage=stage) + # Update references to exported Xform and make it instanceable + stage.DefinePrim(visual_prim_path, typeName="Xform") + stage.DefinePrim(collision_prim_path, typeName="Xform") + visual_undef_prim = stage.GetPrimAtPath(visual_prim_path) + collision_undef_prim = stage.GetPrimAtPath(collision_prim_path) + visual_undef_prim_ref: Usd.References = visual_undef_prim.GetReferences() + collision_undef_prim_ref: Usd.References = collision_undef_prim.GetReferences() + visual_undef_prim_ref.AddReference(assetPath=instanceable_mesh_path, primPath=visual_prim_path) # type: ignore + collision_undef_prim_ref.AddReference(assetPath=instanceable_mesh_path, primPath=collision_prim_path) # type: ignore + visual_undef_prim.SetInstanceable(True) + collision_undef_prim.SetInstanceable(True) + + # Apply mass and rigid body properties after everything else + # Properties are applied to the top level prim to avoid the case where all instances of this + # asset unintentionally share the same rigid body properties + # apply mass properties + if cfg.mass_props is not None: + schemas.define_mass_properties(prim_path=xform_prim.GetPath(), cfg=cfg.mass_props, stage=stage) + # apply rigid body properties + if cfg.rigid_props is not None: + schemas.define_rigid_body_properties(prim_path=xform_prim.GetPath(), cfg=cfg.rigid_props, stage=stage) + + # Save changes to USD stage + stage.Save() + if stage_id is not None: + UsdUtils.StageCache.Get().Erase(stage_id) # type: ignore + + """ + Helper methods. + """ + + @staticmethod + async def _convert_mesh_to_usd( + in_file: str, out_file: str, prim_path: str = "/World", 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 + + 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. + + Returns: + 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() + # Set up converter settings + # Don't import/export materials + converter_context.ignore_materials = not load_materials + converter_context.ignore_animations = True + converter_context.ignore_camera = True + converter_context.ignore_light = True + # Merge all meshes into one + converter_context.merge_all_meshes = True + # Sets world units to meters, this will also scale asset if it's centimeters model. + # This does not work right now :(, so we need to scale the mesh manually + converter_context.use_meter_as_world_unit = True + converter_context.baking_scales = True + # Uses double precision for all transform ops. + converter_context.use_double_precision_to_usd_transform_op = True + + # 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) # type: ignore + # 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() # type: ignore + UsdGeom.SetStageUpAxis(temp_stage, UsdGeom.Tokens.z) + UsdGeom.SetStageMetersPerUnit(temp_stage, 1.0) + UsdPhysics.SetStageKilogramsPerUnit(temp_stage, 1.0) # type: ignore + + base_prim = temp_stage.DefinePrim(prim_path, "Xform") + prim = temp_stage.DefinePrim(f"{prim_path}/temp", "Xform") + prim.GetReferences().AddReference(out_file_non_metric) + cache = UsdUtils.StageCache.Get() + cache.Insert(temp_stage) # type: ignore + stage_id = cache.GetId(temp_stage).ToLongInt() # type: ignore + get_metrics_assembler_interface().resolve_stage(stage_id) + temp_stage.SetDefaultPrim(base_prim) + temp_stage.Export(out_file) + return success diff --git a/source/uwlab/uwlab/sim/converters/mesh_converter_cfg.py b/source/uwlab/uwlab/sim/converters/mesh_converter_cfg.py new file mode 100644 index 00000000..c450769a --- /dev/null +++ b/source/uwlab/uwlab/sim/converters/mesh_converter_cfg.py @@ -0,0 +1,60 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.sim.converters.asset_converter_base_cfg import AssetConverterBaseCfg +from isaaclab.sim.schemas import schemas_cfg +from isaaclab.sim.spawners.materials import PhysicsMaterialCfg +from isaaclab.sim.spawners.materials.visual_materials_cfg import VisualMaterialCfg +from isaaclab.utils import configclass + + +@configclass +class MeshConverterCfg(AssetConverterBaseCfg): + """The configuration class for MeshConverter.""" + + visual_material_props: VisualMaterialCfg | None = None + """Visual material properties to apply to the USD. Defaults to None. + + Note: + If None, then default visual properties will be added. + """ + + physics_material_props: PhysicsMaterialCfg | None = None + """Physics material properties to apply to the USD. Defaults to None. + + Note: + If None, then default physics material properties will be added. + """ + + 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 = 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 = None + """Collision properties to apply to the USD. Defaults to None. + + Note: + If None, then no collision properties will be added. + """ + + collision_approximation: str = "convexDecomposition" + """Collision approximation method to use. Defaults to "convexDecomposition". + + Valid options are: + "convexDecomposition", "convexHull", "boundingCube", + "boundingSphere", "meshSimplification", or "none" + + "none" causes no collision mesh to be added. + """ diff --git a/source/uwlab/uwlab/sim/spawners/materials/__init__.py b/source/uwlab/uwlab/sim/spawners/materials/__init__.py new file mode 100644 index 00000000..591d8e62 --- /dev/null +++ b/source/uwlab/uwlab/sim/spawners/materials/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .physics_materials_cfg import * +from .visual_materials_cfg import * diff --git a/source/uwlab/uwlab/sim/spawners/materials/common_materials_cfg.py b/source/uwlab/uwlab/sim/spawners/materials/common_materials_cfg.py new file mode 100644 index 00000000..ae77703d --- /dev/null +++ b/source/uwlab/uwlab/sim/spawners/materials/common_materials_cfg.py @@ -0,0 +1,50 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from isaaclab.utils import configclass + +from .physics_materials_cfg import StageSpecificRigidBodyMaterialCfg +from .visual_materials_cfg import StageSpecificPreviewSurfaceCfg + +""" +PCV MATERIALS +""" + + +@configclass +class PCVVisualMaterialCfg(StageSpecificPreviewSurfaceCfg): + diffuse_color = (0.5, 0.1, 0.1) + roughness = 0.3 + metallic = 0.0 + opacity = 1.0 + + +@configclass +class PCVPhysicalMaterialCfg(StageSpecificRigidBodyMaterialCfg): + static_friction = 0.4 + dynamic_friction = 0.23 + restitution = 0.2 + + +""" +STEEL MATERIALS +""" + + +@configclass +class SteelVisualMaterialCfg(StageSpecificPreviewSurfaceCfg): + diffuse_color = (0.5, 0.5, 0.5) + roughness = 0.3 + metallic = 0.9 + opacity = 1.0 + + +@configclass +class SteelPhysicalMaterialCfg(StageSpecificRigidBodyMaterialCfg): + static_friction = 0.445 + dynamic_friction = 0.375 + restitution = 0.56 diff --git a/source/uwlab/uwlab/sim/spawners/materials/physics_materials.py b/source/uwlab/uwlab/sim/spawners/materials/physics_materials.py new file mode 100644 index 00000000..08a7dffd --- /dev/null +++ b/source/uwlab/uwlab/sim/spawners/materials/physics_materials.py @@ -0,0 +1,136 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import isaacsim.core.utils.prims as prim_utils +from pxr import PhysxSchema, Usd, UsdPhysics, UsdShade + +from isaaclab.sim.utils import clone, safe_set_attribute_on_usd_schema + +if TYPE_CHECKING: + from . import physics_materials_cfg + + +@clone +def stage_specific_spawn_rigid_body_material( + prim_path: str, + cfg: physics_materials_cfg.StageSpecificRigidBodyMaterialCfg, + stage: Usd.Stage, +) -> Usd.Prim: + """Create material with rigid-body physics properties. + + Different from the IsaacLab version, :func:`isaaclab.sim.spawners.materials.spawn_spawn_rigid_body_material` + this function allows users to specify the stage where the material is spawned. + + Rigid body materials are used to define the physical properties to meshes of a rigid body. These + include the friction, restitution, and their respective combination modes. For more information on + rigid body material, please refer to the `documentation on PxMaterial `_. + + .. note:: + This function is decorated with :func:`clone` that resolves prim path into list of paths + if the input prim path is a regex pattern. This is done to support spawning multiple assets + from a single and cloning the USD prim at the given path expression. + + Args: + prim_path: The prim path or pattern to spawn the asset at. If the prim path is a regex pattern, + then the asset is spawned at all the matching prim paths. + cfg: The configuration for the physics material. + + Returns: + The spawned rigid body material prim. + + Raises: + ValueError: When a prim already exists at the specified prim path and is not a material. + """ + # create material prim if no prim exists + if not prim_utils.is_prim_path_valid(prim_path): + _ = UsdShade.Material.Define(stage, prim_path) + + # obtain prim + prim = stage.GetPrimAtPath(prim_path) + # check if prim is a material + if not prim.IsA(UsdShade.Material): + raise ValueError(f"A prim already exists at path: '{prim_path}' but is not a material.") + # retrieve the USD rigid-body api + usd_physics_material_api = UsdPhysics.MaterialAPI(prim) # type: ignore + if not usd_physics_material_api: + usd_physics_material_api = UsdPhysics.MaterialAPI.Apply(prim) # type: ignore + # retrieve the collision api + physx_material_api = PhysxSchema.PhysxMaterialAPI(prim) + if not physx_material_api: + physx_material_api = PhysxSchema.PhysxMaterialAPI.Apply(prim) + + # convert to dict + cfg = cfg.to_dict() # type: ignore + del cfg["func"] # type: ignore + # set into USD API + for attr_name in ["static_friction", "dynamic_friction", "restitution"]: + value = cfg.pop(attr_name, None) # type: ignore + safe_set_attribute_on_usd_schema(usd_physics_material_api, attr_name, value, camel_case=True) + # set into PhysX API + for attr_name, value in cfg.items(): # type: ignore + safe_set_attribute_on_usd_schema(physx_material_api, attr_name, value, camel_case=True) + # return the prim + return prim + + +@clone +def stage_specific_spawn_deformable_body_material( + prim_path: str, + cfg: physics_materials_cfg.StageSpecificDeformableBodyMaterialCfg, + stage: Usd.Stage, +) -> Usd.Prim: + """Create material with deformable-body physics properties. + + Different from the IsaacLab version, :func:`isaaclab.sim.spawners.materials.spawn_spawn_deformable_body_material` + this function allows users to specify the stage where the material is spawned. + + Deformable body materials are used to define the physical properties to meshes of a deformable body. These + include the friction and deformable body properties. For more information on deformable body material, + please refer to the documentation on `PxFEMSoftBodyMaterial`_. + + .. note:: + This function is decorated with :func:`clone` that resolves prim path into list of paths + if the input prim path is a regex pattern. This is done to support spawning multiple assets + from a single and cloning the USD prim at the given path expression. + + Args: + prim_path: The prim path or pattern to spawn the asset at. If the prim path is a regex pattern, + then the asset is spawned at all the matching prim paths. + cfg: The configuration for the physics material. + + Returns: + The spawned deformable body material prim. + + Raises: + ValueError: When a prim already exists at the specified prim path and is not a material. + + .. _PxFEMSoftBodyMaterial: https://nvidia-omniverse.github.io/PhysX/physx/5.4.1/_api_build/structPxFEMSoftBodyMaterialModel.html + """ + # create material prim if no prim exists + if not prim_utils.is_prim_path_valid(prim_path): + _ = UsdShade.Material.Define(stage, prim_path) + + # obtain prim + prim = stage.GetPrimAtPath(prim_path) + # check if prim is a material + if not prim.IsA(UsdShade.Material): + raise ValueError(f"A prim already exists at path: '{prim_path}' but is not a material.") + # retrieve the deformable-body api + physx_deformable_body_material_api = PhysxSchema.PhysxDeformableBodyMaterialAPI(prim) + if not physx_deformable_body_material_api: + physx_deformable_body_material_api = PhysxSchema.PhysxDeformableBodyMaterialAPI.Apply(prim) + + # convert to dict + cfg = cfg.to_dict() # type: ignore + del cfg["func"] # type: ignore + # set into PhysX API + for attr_name, value in cfg.items(): # type: ignore + safe_set_attribute_on_usd_schema(physx_deformable_body_material_api, attr_name, value, camel_case=True) + # return the prim + return prim diff --git a/source/uwlab/uwlab/sim/spawners/materials/physics_materials_cfg.py b/source/uwlab/uwlab/sim/spawners/materials/physics_materials_cfg.py new file mode 100644 index 00000000..c0592beb --- /dev/null +++ b/source/uwlab/uwlab/sim/spawners/materials/physics_materials_cfg.py @@ -0,0 +1,114 @@ +# Copyright (c) 2022-2024, The UW Lab and Isaac Lab Project Developers. +# All rights reserved. + +from collections.abc import Callable +from typing import Literal + +from isaaclab.sim.spawners.materials import PhysicsMaterialCfg +from isaaclab.utils import configclass + +from . import physics_materials + + +@configclass +class StageSpecificRigidBodyMaterialCfg(PhysicsMaterialCfg): + """Physics material parameters for rigid bodies. + + See :meth:`spawn_rigid_body_material` for more information. + + Note: + The default values are the `default values used by PhysX 5 + `__. + """ + + func: Callable = physics_materials.stage_specific_spawn_rigid_body_material + + static_friction: float = 0.5 + """The static friction coefficient. Defaults to 0.5.""" + + dynamic_friction: float = 0.5 + """The dynamic friction coefficient. Defaults to 0.5.""" + + restitution: float = 0.0 + """The restitution coefficient. Defaults to 0.0.""" + + improve_patch_friction: bool = True + """Whether to enable patch friction. Defaults to True.""" + + friction_combine_mode: Literal["average", "min", "multiply", "max"] = "average" + """Determines the way friction will be combined during collisions. Defaults to `"average"`. + + .. attention:: + + When two physics materials with different combine modes collide, the combine mode with the higher + priority will be used. The priority order is provided `here + `__. + """ + + restitution_combine_mode: Literal["average", "min", "multiply", "max"] = "average" + """Determines the way restitution coefficient will be combined during collisions. Defaults to `"average"`. + + .. attention:: + + When two physics materials with different combine modes collide, the combine mode with the higher + priority will be used. The priority order is provided `here + `__. + """ + + compliant_contact_stiffness: float = 0.0 + """Spring stiffness for a compliant contact model using implicit springs. Defaults to 0.0. + + A higher stiffness results in behavior closer to a rigid contact. The compliant contact model is only enabled + if the stiffness is larger than 0. + """ + + compliant_contact_damping: float = 0.0 + """Damping coefficient for a compliant contact model using implicit springs. Defaults to 0.0. + + Irrelevant if compliant contacts are disabled when :obj:`compliant_contact_stiffness` is set to zero and + rigid contacts are active. + """ + + +@configclass +class StageSpecificDeformableBodyMaterialCfg(PhysicsMaterialCfg): + """Physics material parameters for deformable bodies. + + See :meth:`spawn_deformable_body_material` for more information. + + Note: + The default values are the `default values used by PhysX 5 + `__. + """ + + func: Callable = physics_materials.stage_specific_spawn_deformable_body_material + + density: float | None = None + """The material density. Defaults to None, in which case the simulation decides the default density.""" + + dynamic_friction: float = 0.25 + """The dynamic friction. Defaults to 0.25.""" + + youngs_modulus: float = 50000000.0 + """The Young's modulus, which defines the body's stiffness. Defaults to 50000000.0. + + The Young's modulus is a measure of the material's ability to deform under stress. It is measured in Pascals (Pa). + """ + + poissons_ratio: float = 0.45 + """The Poisson's ratio which defines the body's volume preservation. Defaults to 0.45. + + The Poisson's ratio is a measure of the material's ability to expand in the lateral direction when compressed + in the axial direction. It is a dimensionless number between 0 and 0.5. Using a value of 0.5 will make the + material incompressible. + """ + + elasticity_damping: float = 0.005 + """The elasticity damping for the deformable material. Defaults to 0.005.""" + + damping_scale: float = 1.0 + """The damping scale for the deformable material. Defaults to 1.0. + + A scale of 1 corresponds to default damping. A value of 0 will only apply damping to certain motions leading + to special effects that look similar to water filled soft bodies. + """ diff --git a/source/uwlab/uwlab/sim/spawners/materials/visual_materials.py b/source/uwlab/uwlab/sim/spawners/materials/visual_materials.py new file mode 100644 index 00000000..0cd95033 --- /dev/null +++ b/source/uwlab/uwlab/sim/spawners/materials/visual_materials.py @@ -0,0 +1,144 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import isaacsim.core.utils.prims as prim_utils +import omni.kit.commands +from pxr import Gf, Sdf, Usd, UsdShade + +from isaaclab.sim.utils import clone, safe_set_attribute_on_usd_prim +from isaaclab.utils.assets import NVIDIA_NUCLEUS_DIR + +if TYPE_CHECKING: + from . import visual_materials_cfg + + +@clone +def stage_specific_spawn_preview_surface( + prim_path: str, cfg: visual_materials_cfg.StageSpecificPreviewSurfaceCfg, stage: Usd.Stage +) -> Usd.Prim: + """Create a preview surface prim and override the settings with the given config. + + A preview surface is a physically-based surface that handles simple shaders while supporting + both *specular* and *metallic* workflows. All color inputs are in linear color space (RGB). + For more information, see the `documentation `__. + + Different from the IsaacLab version, this function uses pxr libraries to create a preview surface instead of + USD command `CreatePreviewSurfaceMaterialPrim`_ to create the prim in Isaac Lab. + + .. _CreatePreviewSurfaceMaterialPrim: https://docs.omniverse.nvidia.com/kit/docs/omni.usd/latest/omni.usd.commands/omni.usd.commands.CreatePreviewSurfaceMaterialPrimCommand.html + .. PreviewSurfaceMaterial through pxr: https://docs.omniverse.nvidia.com/dev-guide/latest/programmer_ref/usd/materials/create-usdpreviewsurface-material.html + + .. note:: + This function is decorated with :func:`clone` that resolves prim path into list of paths + if the input prim path is a regex pattern. This is done to support spawning multiple assets + from a single and cloning the USD prim at the given path expression. + + Args: + prim_path: The prim path or pattern to spawn the asset at. If the prim path is a regex pattern, + then the asset is spawned at all the matching prim paths. + cfg: The configuration instance. + + Returns: + The created prim. + + Raises: + ValueError: If a prim already exists at the given path. + """ + + # spawn material if it doesn't exist. + if not prim_utils.is_prim_path_valid(prim_path): + preview_material_prim_path = Sdf.Path(prim_path) + shader_prim_path = preview_material_prim_path.AppendPath("Shader") + UsdShade.Material.Define(stage, preview_material_prim_path) + UsdShade.Shader.Define(stage, shader_prim_path) + preview_material_prim = UsdShade.Material(stage.GetPrimAtPath(preview_material_prim_path)) + shader = UsdShade.Shader(stage.GetPrimAtPath(shader_prim_path)) + shader.CreateIdAttr("UsdPreviewSurface", writeSparsely=False) + shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set( + Gf.Vec3f(cfg.diffuse_color), time=Usd.TimeCode.Default() + ) + shader.CreateInput("emissiveColor", Sdf.ValueTypeNames.Color3f).Set( + cfg.emissive_color, time=Usd.TimeCode.Default() + ) + # shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(cfg.opacity) + # shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(cfg.roughness) + # shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(cfg.metallic) + preview_material_prim.CreateSurfaceOutput(renderContext=UsdShade.Tokens.universalRenderContext).ConnectToSource( + shader.ConnectableAPI(), "surface" + ) + + else: + raise ValueError(f"A prim already exists at path: '{prim_path}'.") + # obtain prim + prim = stage.GetPrimAtPath(shader_prim_path) + # apply properties + cfg = cfg.to_dict() # type: ignore + del cfg["func"] # type: ignore + for attr_name in ["opacity", "roughness", "metallic"]: + value = cfg.pop(attr_name, None) # type: ignore + safe_set_attribute_on_usd_prim(prim, f"inputs:{attr_name}", value, camel_case=True) + # return prim + return prim + + +@clone +def stage_specific_spawn_from_mdl_file(prim_path: str, cfg: visual_materials_cfg.PxrMdlFileCfg) -> Usd.Prim: + raise NotImplementedError("This function is not implemented yet.") + + # below is the isaac lab implementation spawn_from_mdl_file, which will not work with custom usd stages. + """Load a material from its MDL file and override the settings with the given config. + + NVIDIA's `Material Definition Language (MDL) `__ + is a language for defining physically-based materials. The MDL file format is a binary format + that can be loaded by Omniverse and other applications such as Adobe Substance Designer. + To learn more about MDL, see the `documentation `_. + + The function calls the USD command `CreateMdlMaterialPrim`_ to create the prim. + + .. _CreateMdlMaterialPrim: https://docs.omniverse.nvidia.com/kit/docs/omni.usd/latest/omni.usd.commands/omni.usd.commands.CreateMdlMaterialPrimCommand.html + + .. note:: + This function is decorated with :func:`clone` that resolves prim path into list of paths + if the input prim path is a regex pattern. This is done to support spawning multiple assets + from a single and cloning the USD prim at the given path expression. + + Args: + prim_path: The prim path or pattern to spawn the asset at. If the prim path is a regex pattern, + then the asset is spawned at all the matching prim paths. + cfg: The configuration instance. + + Returns: + The created prim. + + Raises: + ValueError: If a prim already exists at the given path. + """ + # spawn material if it doesn't exist. + if not prim_utils.is_prim_path_valid(prim_path): + # extract material name from path + material_name = cfg.mdl_path.split("/")[-1].split(".")[0] + omni.kit.commands.execute( + "CreateMdlMaterialPrim", + mtl_url=cfg.mdl_path.format(NVIDIA_NUCLEUS_DIR=NVIDIA_NUCLEUS_DIR), + mtl_name=material_name, + mtl_path=prim_path, + select_new_prim=False, + ) + else: + raise ValueError(f"A prim already exists at path: '{prim_path}'.") + # obtain prim + prim = prim_utils.get_prim_at_path(f"{prim_path}/Shader") + # apply properties + cfg = cfg.to_dict() + del cfg["func"] + del cfg["mdl_path"] + for attr_name, attr_value in cfg.items(): + safe_set_attribute_on_usd_prim(prim, f"inputs:{attr_name}", attr_value, camel_case=False) + # return prim + return prim diff --git a/source/uwlab/uwlab/sim/spawners/materials/visual_materials_cfg.py b/source/uwlab/uwlab/sim/spawners/materials/visual_materials_cfg.py new file mode 100644 index 00000000..d0631e12 --- /dev/null +++ b/source/uwlab/uwlab/sim/spawners/materials/visual_materials_cfg.py @@ -0,0 +1,110 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from collections.abc import Callable +from dataclasses import MISSING + +from isaaclab.sim.spawners.materials.visual_materials_cfg import VisualMaterialCfg +from isaaclab.utils import configclass + +from . import visual_materials + + +@configclass +class StageSpecificPreviewSurfaceCfg(VisualMaterialCfg): + """Configuration parameters for creating a preview surface. + + Different from the IsaacLab version, this uses pxr UsdShadeShader with UsdPreviewSurface attributes + to create a preview surface instead of using the USD command `CreatePreviewSurfaceMaterialPrim`, + which relies on prim created through isaac sim opened usd stage. + + See :meth:`spawn_preview_surface` for more information. + """ + + func: Callable = visual_materials.stage_specific_spawn_preview_surface + + diffuse_color: tuple[float, float, float] = (0.18, 0.18, 0.18) + """The RGB diffusion color. This is the base color of the surface. Defaults to a dark gray.""" + emissive_color: tuple[float, float, float] = (0.0, 0.0, 0.0) + """The RGB emission component of the surface. Defaults to black.""" + roughness: float = 0.5 + """The roughness for specular lobe. Ranges from 0 (smooth) to 1 (rough). Defaults to 0.5.""" + metallic: float = 0.0 + """The metallic component. Ranges from 0 (dielectric) to 1 (metal). Defaults to 0.""" + opacity: float = 1.0 + """The opacity of the surface. Ranges from 0 (transparent) to 1 (opaque). Defaults to 1. + + Note: + Opacity only affects the surface's appearance during interactive rendering. + """ + + +@configclass +class PxrMdlFileCfg(VisualMaterialCfg): + """Configuration parameters for loading an MDL material from a file. + + Different from the IsaacLab version, this uses pxr libraries to create a preview surface instead of + using the USD command `CreateMdlMaterialPrim`, which only creates prim under isaac sim opened usd stage. + + See :meth:`spawn_from_mdl_file` for more information. + """ + + func: Callable = visual_materials.stage_specific_spawn_from_mdl_file + + mdl_path: str = MISSING # type: ignore + """The path to the MDL material. + + NVIDIA Omniverse provides various MDL materials in the NVIDIA Nucleus. + To use these materials, you can set the path of the material in the nucleus directory + using the ``{NVIDIA_NUCLEUS_DIR}`` variable. This is internally resolved to the path of the + NVIDIA Nucleus directory on the host machine through the attribute + :attr:`isaaclab.utils.assets.NVIDIA_NUCLEUS_DIR`. + + For example, to use the "Aluminum_Anodized" material, you can set the path to: + ``{NVIDIA_NUCLEUS_DIR}/Materials/Base/Metals/Aluminum_Anodized.mdl``. + """ + project_uvw: bool | None = None + """Whether to project the UVW coordinates of the material. Defaults to None. + + If None, then the default setting in the MDL material will be used. + """ + albedo_brightness: float | None = None + """Multiplier for the diffuse color of the material. Defaults to None. + + If None, then the default setting in the MDL material will be used. + """ + texture_scale: tuple[float, float] | None = None + """The scale of the texture. Defaults to None. + + If None, then the default setting in the MDL material will be used. + """ + + +@configclass +class GlassMdlCfg(VisualMaterialCfg): + """Configuration parameters for loading a glass MDL material. + + This is a convenience class for loading a glass MDL material. For more information on + glass materials, see the `documentation `__. + + .. note:: + The default values are taken from the glass material in the NVIDIA Nucleus. + """ + + func: Callable = visual_materials.stage_specific_spawn_from_mdl_file + + mdl_path: str = "OmniGlass.mdl" + """The path to the MDL material. Defaults to the glass material in the NVIDIA Nucleus.""" + glass_color: tuple[float, float, float] = (1.0, 1.0, 1.0) + """The RGB color or tint of the glass. Defaults to white.""" + frosting_roughness: float = 0.0 + """The amount of reflectivity of the surface. Ranges from 0 (perfectly clear) to 1 (frosted). + Defaults to 0.""" + thin_walled: bool = False + """Whether to perform thin-walled refraction. Defaults to False.""" + glass_ior: float = 1.491 + """The incidence of refraction to control how much light is bent when passing through the glass. + Defaults to 1.491, which is the IOR of glass. + """ diff --git a/source/uwlab/uwlab/terrains/__init__.py b/source/uwlab/uwlab/terrains/__init__.py new file mode 100644 index 00000000..9f4ebd66 --- /dev/null +++ b/source/uwlab/uwlab/terrains/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .height_field import * # noqa: F401, F403 +from .terrain_generator import TerrainGenerator +from .terrain_generator_cfg import PatchSamplingCfg, SubTerrainBaseCfg, TerrainGeneratorCfg +from .terrain_importer import TerrainImporter +from .terrain_importer_cfg import TerrainImporterCfg +from .trimesh import * # noqa: F401, F403 diff --git a/source/uwlab/uwlab/terrains/config/rough.py b/source/uwlab/uwlab/terrains/config/rough.py new file mode 100644 index 00000000..7efdf088 --- /dev/null +++ b/source/uwlab/uwlab/terrains/config/rough.py @@ -0,0 +1,52 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configuration for custom terrains.""" + +import uwlab.terrains as terrain_gen + +from ..terrain_generator_cfg import TerrainGeneratorCfg + +ROUGH_TERRAINS_CFG = TerrainGeneratorCfg( + size=(8.0, 8.0), + border_width=20.0, + num_rows=10, + num_cols=20, + horizontal_scale=0.1, + vertical_scale=0.005, + slope_threshold=0.75, + use_cache=False, + sub_terrains={ + "pyramid_stairs": terrain_gen.MeshPyramidStairsTerrainCfg( + proportion=0.2, + step_height_range=(0.05, 0.23), + step_width=0.3, + platform_width=3.0, + border_width=1.0, + holes=False, + ), + "pyramid_stairs_inv": terrain_gen.MeshInvertedPyramidStairsTerrainCfg( + proportion=0.2, + step_height_range=(0.05, 0.23), + step_width=0.3, + platform_width=3.0, + border_width=1.0, + holes=False, + ), + "boxes": terrain_gen.MeshRandomGridTerrainCfg( + proportion=0.2, grid_width=0.45, grid_height_range=(0.05, 0.2), platform_width=2.0 + ), + "random_rough": terrain_gen.HfRandomUniformTerrainCfg( + proportion=0.2, noise_range=(0.02, 0.10), noise_step=0.02, border_width=0.25 + ), + "hf_pyramid_slope": terrain_gen.HfPyramidSlopedTerrainCfg( + proportion=0.1, slope_range=(0.0, 0.4), platform_width=2.0, border_width=0.25 + ), + "hf_pyramid_slope_inv": terrain_gen.HfInvertedPyramidSlopedTerrainCfg( + proportion=0.1, slope_range=(0.0, 0.4), platform_width=2.0, border_width=0.25 + ), + }, +) +"""Rough terrains configuration.""" diff --git a/source/uwlab/uwlab/terrains/config/terrain_gen.py b/source/uwlab/uwlab/terrains/config/terrain_gen.py new file mode 100644 index 00000000..424cd0c9 --- /dev/null +++ b/source/uwlab/uwlab/terrains/config/terrain_gen.py @@ -0,0 +1,120 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from typing import Dict + +import isaaclab.terrains as terrain_gen +import uwlab.terrains as uw_terrain_gen +from isaaclab.terrains.terrain_generator_cfg import SubTerrainBaseCfg + +TERRAIN_GEN_SUB_TERRAINS: Dict[str, SubTerrainBaseCfg] = { + "perlin": uw_terrain_gen.CachedTerrainGenCfg( + proportion=1.00, + height=0.25, + levels=3, + task_descriptor="perlin", + include_overhang=False, + ), + "ramp_perlin": uw_terrain_gen.CachedTerrainGenCfg( + proportion=1.00, + height=0.50, + levels=3, + task_descriptor="ramp_perlin", + include_overhang=False, + ), + "wall": uw_terrain_gen.CachedTerrainGenCfg( + proportion=1.00, + height=0.25, + levels=3, + task_descriptor="wall", + include_overhang=False, + ), + "stair_ramp": uw_terrain_gen.CachedTerrainGenCfg( + proportion=1.00, + height=0.25, + levels=3, + task_descriptor="stair_ramp", + include_overhang=False, + ), + "stair_platform": uw_terrain_gen.CachedTerrainGenCfg( + proportion=1.00, + height=0.25, + levels=3, + task_descriptor="stair_platform", + include_overhang=False, + ), + "ramp": uw_terrain_gen.CachedTerrainGenCfg( + proportion=1.00, + height=0.25, + levels=3, + task_descriptor="ramp", + include_overhang=False, + ), + "stair_platform_wall": uw_terrain_gen.CachedTerrainGenCfg( + proportion=1.00, + height=0.25, + levels=3, + task_descriptor="stair_platform_wall", + include_overhang=False, + ), + "perlin_wall": uw_terrain_gen.CachedTerrainGenCfg( + proportion=1.00, + height=0.25, + levels=3, + task_descriptor="perlin_wall", + include_overhang=False, + ), + "box": uw_terrain_gen.CachedTerrainGenCfg( + proportion=1.00, + height=0.25, + levels=3, + task_descriptor="box", + include_overhang=False, + ), + "pyramid_stairs": terrain_gen.MeshPyramidStairsTerrainCfg( + proportion=1.00, + step_height_range=(0.05, 0.07), + step_width=0.3, + platform_width=3.0, + border_width=1.0, + holes=False, + ), + "pyramid_stairs_inv": terrain_gen.MeshInvertedPyramidStairsTerrainCfg( + proportion=1.00, + step_height_range=(0.05, 0.07), + step_width=0.3, + platform_width=3.0, + border_width=1.0, + holes=False, + ), + "boxes": terrain_gen.MeshRandomGridTerrainCfg( + proportion=1.00, grid_width=0.45, grid_height_range=(0.45, 0.57), platform_width=2.0 + ), + "random_rough": terrain_gen.HfRandomUniformTerrainCfg( + proportion=1.00, noise_range=(0.02, 0.04), noise_step=0.02, border_width=0.25 + ), + "hf_pyramid_slope": terrain_gen.HfPyramidSlopedTerrainCfg( + proportion=1.00, slope_range=(0.02, 0.04), platform_width=2.0, border_width=0.25 + ), + "random_grid": terrain_gen.MeshRandomGridTerrainCfg( + proportion=1.00, + platform_width=1.5, + grid_width=0.75, + grid_height_range=(0.025, 0.045), + holes=False, + ), + "discrete_obstacle": terrain_gen.HfDiscreteObstaclesTerrainCfg( + proportion=1.00, + size=(8.0, 8.0), + horizontal_scale=0.1, + vertical_scale=0.005, + border_width=0.0, + num_obstacles=100, + obstacle_height_mode="choice", + obstacle_width_range=(0.25, 0.75), + obstacle_height_range=(1.0, 2.0), + platform_width=1.5, + ), +} diff --git a/source/uwlab/uwlab/terrains/height_field/__init__.py b/source/uwlab/uwlab/terrains/height_field/__init__.py new file mode 100644 index 00000000..2e0f8d2a --- /dev/null +++ b/source/uwlab/uwlab/terrains/height_field/__init__.py @@ -0,0 +1,38 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +""" +This sub-module provides utilities to create different terrains as height fields (HF). + +Height fields are a 2.5D terrain representation that is used in robotics to obtain the +height of the terrain at a given point. This is useful for controls and planning algorithms. + +Each terrain is represented as a 2D numpy array with discretized heights. The shape of the array +is (width, length), where width and length are the number of points along the x and y axis, +respectively. The height of the terrain at a given point is obtained by indexing the array with +the corresponding x and y coordinates. + +.. caution:: + + When working with height field terrains, it is important to remember that the terrain is generated + from a discretized 3D representation. This means that the height of the terrain at a given point + is only an approximation of the real height of the terrain at that point. The discretization + error is proportional to the size of the discretization cells. Therefore, it is important to + choose a discretization size that is small enough for the application. A larger discretization + size will result in a faster simulation, but the terrain will be less accurate. + +""" + +from .hf_terrains_cfg import ( + HfDiscreteObstaclesTerrainCfg, + HfInvertedPyramidSlopedTerrainCfg, + HfInvertedPyramidStairsTerrainCfg, + HfPyramidSlopedTerrainCfg, + HfPyramidStairsTerrainCfg, + HfRandomUniformTerrainCfg, + HfSteppingStonesTerrainCfg, + HfTerrainBaseCfg, + HfWaveTerrainCfg, +) diff --git a/source/uwlab/uwlab/terrains/height_field/hf_terrains.py b/source/uwlab/uwlab/terrains/height_field/hf_terrains.py new file mode 100644 index 00000000..e8791497 --- /dev/null +++ b/source/uwlab/uwlab/terrains/height_field/hf_terrains.py @@ -0,0 +1,436 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Functions to generate height fields for different terrains.""" + +from __future__ import annotations + +import numpy as np +import scipy.interpolate as interpolate +from typing import TYPE_CHECKING + +from .utils import height_field_to_mesh + +if TYPE_CHECKING: + from . import hf_terrains_cfg + + +@height_field_to_mesh +def random_uniform_terrain(difficulty: float, cfg: hf_terrains_cfg.HfRandomUniformTerrainCfg) -> np.ndarray: + """Generate a terrain with height sampled uniformly from a specified range. + + .. image:: ../../_static/terrains/height_field/random_uniform_terrain.jpg + :width: 40% + :align: center + + Note: + The :obj:`difficulty` parameter is ignored for this terrain. + + Args: + difficulty: The difficulty of the terrain. This is a value between 0 and 1. + cfg: The configuration for the terrain. + + Returns: + The height field of the terrain as a 2D numpy array with discretized heights. + The shape of the array is (width, length), where width and length are the number of points + along the x and y axis, respectively. + + Raises: + ValueError: When the downsampled scale is smaller than the horizontal scale. + """ + # check parameters + # -- horizontal scale + if cfg.downsampled_scale is None: + cfg.downsampled_scale = cfg.horizontal_scale + elif cfg.downsampled_scale < cfg.horizontal_scale: + raise ValueError( + "Downsampled scale must be larger than or equal to the horizontal scale:" + f" {cfg.downsampled_scale} < {cfg.horizontal_scale}." + ) + + # switch parameters to discrete units + # -- horizontal scale + width_pixels = int(cfg.size[0] / cfg.horizontal_scale) + length_pixels = int(cfg.size[1] / cfg.horizontal_scale) + # -- downsampled scale + width_downsampled = int(cfg.size[0] / cfg.downsampled_scale) + length_downsampled = int(cfg.size[1] / cfg.downsampled_scale) + # -- height + height_min = int(cfg.noise_range[0] / cfg.vertical_scale) + height_max = int(cfg.noise_range[1] / cfg.vertical_scale) + height_step = int(cfg.noise_step / cfg.vertical_scale) + + # create range of heights possible + height_range = np.arange(height_min, height_max + height_step, height_step) + # sample heights randomly from the range along a grid + height_field_downsampled = np.random.choice(height_range, size=(width_downsampled, length_downsampled)) + # create interpolation function for the sampled heights + x = np.linspace(0, cfg.size[0] * cfg.horizontal_scale, width_downsampled) + y = np.linspace(0, cfg.size[1] * cfg.horizontal_scale, length_downsampled) + func = interpolate.RectBivariateSpline(x, y, height_field_downsampled) + + # interpolate the sampled heights to obtain the height field + x_upsampled = np.linspace(0, cfg.size[0] * cfg.horizontal_scale, width_pixels) + y_upsampled = np.linspace(0, cfg.size[1] * cfg.horizontal_scale, length_pixels) + z_upsampled = func(x_upsampled, y_upsampled) + # round off the interpolated heights to the nearest vertical step + return np.rint(z_upsampled).astype(np.int16) + + +@height_field_to_mesh +def pyramid_sloped_terrain(difficulty: float, cfg: hf_terrains_cfg.HfPyramidSlopedTerrainCfg) -> np.ndarray: + """Generate a terrain with a truncated pyramid structure. + + The terrain is a pyramid-shaped sloped surface with a slope of :obj:`slope` that trims into a flat platform + at the center. The slope is defined as the ratio of the height change along the x axis to the width along the + x axis. For example, a slope of 1.0 means that the height changes by 1 unit for every 1 unit of width. + + If the :obj:`cfg.inverted` flag is set to :obj:`True`, the terrain is inverted such that + the platform is at the bottom. + + .. image:: ../../_static/terrains/height_field/pyramid_sloped_terrain.jpg + :width: 40% + + .. image:: ../../_static/terrains/height_field/inverted_pyramid_sloped_terrain.jpg + :width: 40% + + Args: + difficulty: The difficulty of the terrain. This is a value between 0 and 1. + cfg: The configuration for the terrain. + + Returns: + The height field of the terrain as a 2D numpy array with discretized heights. + The shape of the array is (width, length), where width and length are the number of points + along the x and y axis, respectively. + """ + # resolve terrain configuration + if cfg.inverted: + slope = -cfg.slope_range[0] - difficulty * (cfg.slope_range[1] - cfg.slope_range[0]) + else: + slope = cfg.slope_range[0] + difficulty * (cfg.slope_range[1] - cfg.slope_range[0]) + + # switch parameters to discrete units + # -- horizontal scale + width_pixels = int(cfg.size[0] / cfg.horizontal_scale) + length_pixels = int(cfg.size[1] / cfg.horizontal_scale) + # -- height + # we want the height to be 1/2 of the width since the terrain is a pyramid + height_max = int(slope * cfg.size[0] / 2 / cfg.vertical_scale) + # -- center of the terrain + center_x = int(width_pixels / 2) + center_y = int(length_pixels / 2) + + # create a meshgrid of the terrain + x = np.arange(0, width_pixels) + y = np.arange(0, length_pixels) + xx, yy = np.meshgrid(x, y, sparse=True) + # offset the meshgrid to the center of the terrain + xx = (center_x - np.abs(center_x - xx)) / center_x + yy = (center_y - np.abs(center_y - yy)) / center_y + # reshape the meshgrid to be 2D + xx = xx.reshape(width_pixels, 1) + yy = yy.reshape(1, length_pixels) + # create a sloped surface + hf_raw = np.zeros((width_pixels, length_pixels)) + hf_raw = height_max * xx * yy + + # create a flat platform at the center of the terrain + platform_width = int(cfg.platform_width / cfg.horizontal_scale / 2) + # get the height of the platform at the corner of the platform + x_pf = width_pixels // 2 - platform_width + y_pf = length_pixels // 2 - platform_width + z_pf = hf_raw[x_pf, y_pf] + hf_raw = np.clip(hf_raw, min(0, z_pf), max(0, z_pf)) + + # round off the heights to the nearest vertical step + return np.rint(hf_raw).astype(np.int16) + + +@height_field_to_mesh +def pyramid_stairs_terrain(difficulty: float, cfg: hf_terrains_cfg.HfPyramidStairsTerrainCfg) -> np.ndarray: + """Generate a terrain with a pyramid stair pattern. + + The terrain is a pyramid stair pattern which trims to a flat platform at the center of the terrain. + + If the :obj:`cfg.inverted` flag is set to :obj:`True`, the terrain is inverted such that + the platform is at the bottom. + + .. image:: ../../_static/terrains/height_field/pyramid_stairs_terrain.jpg + :width: 40% + + .. image:: ../../_static/terrains/height_field/inverted_pyramid_stairs_terrain.jpg + :width: 40% + + Args: + difficulty: The difficulty of the terrain. This is a value between 0 and 1. + cfg: The configuration for the terrain. + + Returns: + The height field of the terrain as a 2D numpy array with discretized heights. + The shape of the array is (width, length), where width and length are the number of points + along the x and y axis, respectively. + """ + # resolve terrain configuration + step_height = cfg.step_height_range[0] + difficulty * (cfg.step_height_range[1] - cfg.step_height_range[0]) + if cfg.inverted: + step_height *= -1 + # switch parameters to discrete units + # -- terrain + width_pixels = int(cfg.size[0] / cfg.horizontal_scale) + length_pixels = int(cfg.size[1] / cfg.horizontal_scale) + # -- stairs + step_width = int(cfg.step_width / cfg.horizontal_scale) + step_height = int(step_height / cfg.vertical_scale) + # -- platform + platform_width = int(cfg.platform_width / cfg.horizontal_scale) + + # create a terrain with a flat platform at the center + hf_raw = np.zeros((width_pixels, length_pixels)) + # add the steps + current_step_height = 0 + start_x, start_y = 0, 0 + stop_x, stop_y = width_pixels, length_pixels + while (stop_x - start_x) > platform_width and (stop_y - start_y) > platform_width: + # increment position + # -- x + start_x += step_width + stop_x -= step_width + # -- y + start_y += step_width + stop_y -= step_width + # increment height + current_step_height += step_height + # add the step + hf_raw[start_x:stop_x, start_y:stop_y] = current_step_height + + # round off the heights to the nearest vertical step + return np.rint(hf_raw).astype(np.int16) + + +@height_field_to_mesh +def discrete_obstacles_terrain(difficulty: float, cfg: hf_terrains_cfg.HfDiscreteObstaclesTerrainCfg) -> np.ndarray: + """Generate a terrain with randomly generated obstacles as pillars with positive and negative heights. + + The terrain is a flat platform at the center of the terrain with randomly generated obstacles as pillars + with positive and negative height. The obstacles are randomly generated cuboids with a random width and + height. They are placed randomly on the terrain with a minimum distance of :obj:`cfg.platform_width` + from the center of the terrain. + + .. image:: ../../_static/terrains/height_field/discrete_obstacles_terrain.jpg + :width: 40% + :align: center + + Args: + difficulty: The difficulty of the terrain. This is a value between 0 and 1. + cfg: The configuration for the terrain. + + Returns: + The height field of the terrain as a 2D numpy array with discretized heights. + The shape of the array is (width, length), where width and length are the number of points + along the x and y axis, respectively. + """ + # resolve terrain configuration + obs_height = cfg.obstacle_height_range[0] + difficulty * ( + cfg.obstacle_height_range[1] - cfg.obstacle_height_range[0] + ) + + # switch parameters to discrete units + # -- terrain + width_pixels = int(cfg.size[0] / cfg.horizontal_scale) + length_pixels = int(cfg.size[1] / cfg.horizontal_scale) + # -- obstacles + obs_height = int(obs_height / cfg.vertical_scale) + obs_width_min = int(cfg.obstacle_width_range[0] / cfg.horizontal_scale) + obs_width_max = int(cfg.obstacle_width_range[1] / cfg.horizontal_scale) + # -- center of the terrain + platform_width = int(cfg.platform_width / cfg.horizontal_scale) + + # create discrete ranges for the obstacles + # -- shape + obs_width_range = np.arange(obs_width_min, obs_width_max, 4) + obs_length_range = np.arange(obs_width_min, obs_width_max, 4) + # -- position + obs_x_range = np.arange(0, width_pixels, 4) + obs_y_range = np.arange(0, length_pixels, 4) + + # create a terrain with a flat platform at the center + hf_raw = np.zeros((width_pixels, length_pixels)) + # generate the obstacles + for _ in range(cfg.num_obstacles): + # sample size + if cfg.obstacle_height_mode == "choice": + height = np.random.choice([-obs_height, -obs_height // 2, obs_height // 2, obs_height]) + elif cfg.obstacle_height_mode == "fixed": + height = obs_height + else: + raise ValueError(f"Unknown obstacle height mode '{cfg.obstacle_height_mode}'. Must be 'choice' or 'fixed'.") + width = int(np.random.choice(obs_width_range)) + length = int(np.random.choice(obs_length_range)) + # sample position + x_start = int(np.random.choice(obs_x_range)) + y_start = int(np.random.choice(obs_y_range)) + # clip start position to the terrain + if x_start + width > width_pixels: + x_start = width_pixels - width + if y_start + length > length_pixels: + y_start = length_pixels - length + # add to terrain + hf_raw[x_start : x_start + width, y_start : y_start + length] = height + # clip the terrain to the platform + x1 = (width_pixels - platform_width) // 2 + x2 = (width_pixels + platform_width) // 2 + y1 = (length_pixels - platform_width) // 2 + y2 = (length_pixels + platform_width) // 2 + hf_raw[x1:x2, y1:y2] = 0 + # round off the heights to the nearest vertical step + return np.rint(hf_raw).astype(np.int16) + + +@height_field_to_mesh +def wave_terrain(difficulty: float, cfg: hf_terrains_cfg.HfWaveTerrainCfg) -> np.ndarray: + r"""Generate a terrain with a wave pattern. + + The terrain is a flat platform at the center of the terrain with a wave pattern. The wave pattern + is generated by adding sinusoidal waves based on the number of waves and the amplitude of the waves. + + The height of the terrain at a point :math:`(x, y)` is given by: + + .. math:: + + h(x, y) = A \left(\sin\left(\frac{2 \pi x}{\lambda}\right) + \cos\left(\frac{2 \pi y}{\lambda}\right) \right) + + where :math:`A` is the amplitude of the waves, :math:`\lambda` is the wavelength of the waves. + + .. image:: ../../_static/terrains/height_field/wave_terrain.jpg + :width: 40% + :align: center + + Args: + difficulty: The difficulty of the terrain. This is a value between 0 and 1. + cfg: The configuration for the terrain. + + Returns: + The height field of the terrain as a 2D numpy array with discretized heights. + The shape of the array is (width, length), where width and length are the number of points + along the x and y axis, respectively. + + Raises: + ValueError: When the number of waves is non-positive. + """ + # check number of waves + if cfg.num_waves < 0: + raise ValueError(f"Number of waves must be a positive integer. Got: {cfg.num_waves}.") + + # resolve terrain configuration + amplitude = cfg.amplitude_range[0] + difficulty * (cfg.amplitude_range[1] - cfg.amplitude_range[0]) + # switch parameters to discrete units + # -- terrain + width_pixels = int(cfg.size[0] / cfg.horizontal_scale) + length_pixels = int(cfg.size[1] / cfg.horizontal_scale) + amplitude_pixels = int(0.5 * amplitude / cfg.vertical_scale) + + # compute the wave number: nu = 2 * pi / lambda + wave_length = length_pixels / cfg.num_waves + wave_number = 2 * np.pi / wave_length + # create meshgrid for the terrain + x = np.arange(0, width_pixels) + y = np.arange(0, length_pixels) + xx, yy = np.meshgrid(x, y, sparse=True) + xx = xx.reshape(width_pixels, 1) + yy = yy.reshape(1, length_pixels) + + # create a terrain with a flat platform at the center + hf_raw = np.zeros((width_pixels, length_pixels)) + # add the waves + hf_raw += amplitude_pixels * (np.cos(yy * wave_number) + np.sin(xx * wave_number)) + # round off the heights to the nearest vertical step + return np.rint(hf_raw).astype(np.int16) + + +@height_field_to_mesh +def stepping_stones_terrain(difficulty: float, cfg: hf_terrains_cfg.HfSteppingStonesTerrainCfg) -> np.ndarray: + """Generate a terrain with a stepping stones pattern. + + The terrain is a stepping stones pattern which trims to a flat platform at the center of the terrain. + + .. image:: ../../_static/terrains/height_field/stepping_stones_terrain.jpg + :width: 40% + :align: center + + Args: + difficulty: The difficulty of the terrain. This is a value between 0 and 1. + cfg: The configuration for the terrain. + + Returns: + The height field of the terrain as a 2D numpy array with discretized heights. + The shape of the array is (width, length), where width and length are the number of points + along the x and y axis, respectively. + """ + # resolve terrain configuration + stone_width = cfg.stone_width_range[1] - difficulty * (cfg.stone_width_range[1] - cfg.stone_width_range[0]) + stone_distance = cfg.stone_distance_range[0] + difficulty * ( + cfg.stone_distance_range[1] - cfg.stone_distance_range[0] + ) + + # switch parameters to discrete units + # -- terrain + width_pixels = int(cfg.size[0] / cfg.horizontal_scale) + length_pixels = int(cfg.size[1] / cfg.horizontal_scale) + # -- stones + stone_distance = int(stone_distance / cfg.horizontal_scale) + stone_width = int(stone_width / cfg.horizontal_scale) + stone_height_max = int(cfg.stone_height_max / cfg.vertical_scale) + # -- holes + holes_depth = int(cfg.holes_depth / cfg.vertical_scale) + # -- platform + platform_width = int(cfg.platform_width / cfg.horizontal_scale) + # create range of heights + stone_height_range = np.arange(-stone_height_max - 1, stone_height_max, step=1) + + # create a terrain with a flat platform at the center + hf_raw = np.full((width_pixels, length_pixels), holes_depth) + # add the stones + start_x, start_y = 0, 0 + # -- if the terrain is longer than it is wide then fill the terrain column by column + if length_pixels >= width_pixels: + while start_y < length_pixels: + # ensure that stone stops along y-axis + stop_y = min(length_pixels, start_y + stone_width) + # randomly sample x-position + start_x = np.random.randint(0, stone_width) + stop_x = max(0, start_x - stone_distance) + # fill first stone + hf_raw[0:stop_x, start_y:stop_y] = np.random.choice(stone_height_range) + # fill row with stones + while start_x < width_pixels: + stop_x = min(width_pixels, start_x + stone_width) + hf_raw[start_x:stop_x, start_y:stop_y] = np.random.choice(stone_height_range) + start_x += stone_width + stone_distance + # update y-position + start_y += stone_width + stone_distance + elif width_pixels > length_pixels: + while start_x < width_pixels: + # ensure that stone stops along x-axis + stop_x = min(width_pixels, start_x + stone_width) + # randomly sample y-position + start_y = np.random.randint(0, stone_width) + stop_y = max(0, start_y - stone_distance) + # fill first stone + hf_raw[start_x:stop_x, 0:stop_y] = np.random.choice(stone_height_range) + # fill column with stones + while start_y < length_pixels: + stop_y = min(length_pixels, start_y + stone_width) + hf_raw[start_x:stop_x, start_y:stop_y] = np.random.choice(stone_height_range) + start_y += stone_width + stone_distance + # update x-position + start_x += stone_width + stone_distance + # add the platform in the center + x1 = (width_pixels - platform_width) // 2 + x2 = (width_pixels + platform_width) // 2 + y1 = (length_pixels - platform_width) // 2 + y2 = (length_pixels + platform_width) // 2 + hf_raw[x1:x2, y1:y2] = 0 + # round off the heights to the nearest vertical step + return np.rint(hf_raw).astype(np.int16) diff --git a/source/uwlab/uwlab/terrains/height_field/hf_terrains_cfg.py b/source/uwlab/uwlab/terrains/height_field/hf_terrains_cfg.py new file mode 100644 index 00000000..a05dccc9 --- /dev/null +++ b/source/uwlab/uwlab/terrains/height_field/hf_terrains_cfg.py @@ -0,0 +1,167 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import MISSING + +from isaaclab.utils import configclass + +from ..terrain_generator_cfg import SubTerrainBaseCfg +from . import hf_terrains + + +@configclass +class HfTerrainBaseCfg(SubTerrainBaseCfg): + """The base configuration for height field terrains.""" + + border_width: float = 0.0 + """The width of the border/padding around the terrain (in m). Defaults to 0.0. + + The border width is subtracted from the :obj:`size` of the terrain. If non-zero, it must be + greater than or equal to the :obj:`horizontal scale`. + """ + horizontal_scale: float = 0.1 + """The discretization of the terrain along the x and y axes (in m). Defaults to 0.1.""" + vertical_scale: float = 0.005 + """The discretization of the terrain along the z axis (in m). Defaults to 0.005.""" + slope_threshold: float | None = None + """The slope threshold above which surfaces are made vertical. Defaults to None, + in which case no correction is applied.""" + + +""" +Different height field terrain configurations. +""" + + +@configclass +class HfRandomUniformTerrainCfg(HfTerrainBaseCfg): + """Configuration for a random uniform height field terrain.""" + + function = hf_terrains.random_uniform_terrain + + noise_range: tuple[float, float] = MISSING + """The minimum and maximum height noise (i.e. along z) of the terrain (in m).""" + noise_step: float = MISSING + """The minimum height (in m) change between two points.""" + downsampled_scale: float | None = None + """The distance between two randomly sampled points on the terrain. Defaults to None, + in which case the :obj:`horizontal scale` is used. + + The heights are sampled at this resolution and interpolation is performed for intermediate points. + This must be larger than or equal to the :obj:`horizontal scale`. + """ + + +@configclass +class HfPyramidSlopedTerrainCfg(HfTerrainBaseCfg): + """Configuration for a pyramid sloped height field terrain.""" + + function = hf_terrains.pyramid_sloped_terrain + + slope_range: tuple[float, float] = MISSING + """The slope of the terrain (in radians).""" + platform_width: float = 1.0 + """The width of the square platform at the center of the terrain. Defaults to 1.0.""" + inverted: bool = False + """Whether the pyramid is inverted. Defaults to False. + + If True, the terrain is inverted such that the platform is at the bottom and the slopes are upwards. + """ + + +@configclass +class HfInvertedPyramidSlopedTerrainCfg(HfPyramidSlopedTerrainCfg): + """Configuration for an inverted pyramid sloped height field terrain. + + Note: + This is a subclass of :class:`HfPyramidSlopedTerrainCfg` with :obj:`inverted` set to True. + We make it as a separate class to make it easier to distinguish between the two and match + the naming convention of the other terrains. + """ + + inverted: bool = True + + +@configclass +class HfPyramidStairsTerrainCfg(HfTerrainBaseCfg): + """Configuration for a pyramid stairs height field terrain.""" + + function = hf_terrains.pyramid_stairs_terrain + + step_height_range: tuple[float, float] = MISSING + """The minimum and maximum height of the steps (in m).""" + step_width: float = MISSING + """The width of the steps (in m).""" + platform_width: float = 1.0 + """The width of the square platform at the center of the terrain. Defaults to 1.0.""" + inverted: bool = False + """Whether the pyramid stairs is inverted. Defaults to False. + + If True, the terrain is inverted such that the platform is at the bottom and the stairs are upwards. + """ + + +@configclass +class HfInvertedPyramidStairsTerrainCfg(HfPyramidStairsTerrainCfg): + """Configuration for an inverted pyramid stairs height field terrain. + + Note: + This is a subclass of :class:`HfPyramidStairsTerrainCfg` with :obj:`inverted` set to True. + We make it as a separate class to make it easier to distinguish between the two and match + the naming convention of the other terrains. + """ + + inverted: bool = True + + +@configclass +class HfDiscreteObstaclesTerrainCfg(HfTerrainBaseCfg): + """Configuration for a discrete obstacles height field terrain.""" + + function = hf_terrains.discrete_obstacles_terrain + + obstacle_height_mode: str = "choice" + """The mode to use for the obstacle height. Defaults to "choice". + + The following modes are supported: "choice", "fixed". + """ + obstacle_width_range: tuple[float, float] = MISSING + """The minimum and maximum width of the obstacles (in m).""" + obstacle_height_range: tuple[float, float] = MISSING + """The minimum and maximum height of the obstacles (in m).""" + num_obstacles: int = MISSING + """The number of obstacles to generate.""" + platform_width: float = 1.0 + """The width of the square platform at the center of the terrain. Defaults to 1.0.""" + + +@configclass +class HfWaveTerrainCfg(HfTerrainBaseCfg): + """Configuration for a wave height field terrain.""" + + function = hf_terrains.wave_terrain + + amplitude_range: tuple[float, float] = MISSING + """The minimum and maximum amplitude of the wave (in m).""" + num_waves: int = 1.0 + """The number of waves to generate. Defaults to 1.0.""" + + +@configclass +class HfSteppingStonesTerrainCfg(HfTerrainBaseCfg): + """Configuration for a stepping stones height field terrain.""" + + function = hf_terrains.stepping_stones_terrain + + stone_height_max: float = MISSING + """The maximum height of the stones (in m).""" + stone_width_range: tuple[float, float] = MISSING + """The minimum and maximum width of the stones (in m).""" + stone_distance_range: tuple[float, float] = MISSING + """The minimum and maximum distance between stones (in m).""" + holes_depth: float = -10.0 + """The depth of the holes (negative obstacles). Defaults to -10.0.""" + platform_width: float = 1.0 + """The width of the square platform at the center of the terrain. Defaults to 1.0.""" diff --git a/source/uwlab/uwlab/terrains/height_field/utils.py b/source/uwlab/uwlab/terrains/height_field/utils.py new file mode 100644 index 00000000..19fedb02 --- /dev/null +++ b/source/uwlab/uwlab/terrains/height_field/utils.py @@ -0,0 +1,173 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import copy +import functools +import numpy as np +import trimesh +from collections.abc import Callable +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .hf_terrains_cfg import HfTerrainBaseCfg + + +def height_field_to_mesh(func: Callable) -> Callable: + """Decorator to convert a height field function to a mesh function. + + This decorator converts a height field function to a mesh function by sampling the heights + at a specified resolution and performing interpolation to obtain the intermediate heights. + Additionally, it adds a border around the terrain to avoid artifacts at the edges. + + Args: + func: The height field function to convert. The function should return a 2D numpy array + with the heights of the terrain. + + Returns: + The mesh function. The mesh function returns a tuple containing a list of ``trimesh`` + mesh objects and the origin of the terrain. + """ + + @functools.wraps(func) + def wrapper(difficulty: float, cfg: HfTerrainBaseCfg): + # check valid border width + if cfg.border_width > 0 and cfg.border_width < cfg.horizontal_scale: + raise ValueError( + f"The border width ({cfg.border_width}) must be greater than or equal to the" + f" horizontal scale ({cfg.horizontal_scale})." + ) + # allocate buffer for height field (with border) + width_pixels = int(cfg.size[0] / cfg.horizontal_scale) + 1 + length_pixels = int(cfg.size[1] / cfg.horizontal_scale) + 1 + border_pixels = int(cfg.border_width / cfg.horizontal_scale) + 1 + heights = np.zeros((width_pixels, length_pixels), dtype=np.int16) + # override size of the terrain to account for the border + sub_terrain_size = [width_pixels - 2 * border_pixels, length_pixels - 2 * border_pixels] + sub_terrain_size = [dim * cfg.horizontal_scale for dim in sub_terrain_size] + # update the config + terrain_size = copy.deepcopy(cfg.size) + cfg.size = tuple(sub_terrain_size) + # generate the height field + z_gen = func(difficulty, cfg) + # handle the border for the terrain + heights[border_pixels:-border_pixels, border_pixels:-border_pixels] = z_gen + # set terrain size back to config + cfg.size = terrain_size + + # convert to trimesh + vertices, triangles = convert_height_field_to_mesh( + heights, cfg.horizontal_scale, cfg.vertical_scale, cfg.slope_threshold + ) + mesh = trimesh.Trimesh(vertices=vertices, faces=triangles) + # compute origin + x1 = int((cfg.size[0] * 0.5 - 1) / cfg.horizontal_scale) + x2 = int((cfg.size[0] * 0.5 + 1) / cfg.horizontal_scale) + y1 = int((cfg.size[1] * 0.5 - 1) / cfg.horizontal_scale) + y2 = int((cfg.size[1] * 0.5 + 1) / cfg.horizontal_scale) + origin_z = np.max(heights[x1:x2, y1:y2]) * cfg.vertical_scale + origin = np.array([0.5 * cfg.size[0], 0.5 * cfg.size[1], origin_z]) + # return mesh and origin + return [mesh], origin + + return wrapper + + +def convert_height_field_to_mesh( + height_field: np.ndarray, horizontal_scale: float, vertical_scale: float, slope_threshold: float | None = None +) -> tuple[np.ndarray, np.ndarray]: + """Convert a height-field array to a triangle mesh represented by vertices and triangles. + + This function converts a height-field array to a triangle mesh represented by vertices and triangles. + The height-field array is assumed to be a 2D array of floats, where each element represents the height + of the terrain at that location. The height-field array is assumed to be in the form of a matrix, where + the first dimension represents the x-axis and the second dimension represents the y-axis. + + The function can also correct vertical surfaces above the provide slope threshold. This is helpful to + avoid having long vertical surfaces in the mesh. The correction is done by moving the vertices of the + vertical surfaces to minimum of the two neighboring vertices. + + The correction is done in the following way: + If :math:`\\frac{y_2 - y_1}{x_2 - x_1} > threshold`, then move A to A' (i.e., set :math:`x_1' = x_2`). + This is repeated along all directions. + + .. code-block:: none + + B(x_2,y_2) + /| + / | + / | + (x_1,y_1)A---A'(x_1',y_1) + + Args: + height_field: The input height-field array. + horizontal_scale: The discretization of the terrain along the x and y axis. + vertical_scale: The discretization of the terrain along the z axis. + slope_threshold: The slope threshold above which surfaces are made vertical. + Defaults to None, in which case no correction is applied. + + Returns: + The vertices and triangles of the mesh: + - **vertices** (np.ndarray(float)): Array of shape (num_vertices, 3). + Each row represents the location of each vertex (in m). + - **triangles** (np.ndarray(int)): Array of shape (num_triangles, 3). + Each row represents the indices of the 3 vertices connected by this triangle. + """ + # read height field + num_rows, num_cols = height_field.shape + # create a mesh grid of the height field + y = np.linspace(0, (num_cols - 1) * horizontal_scale, num_cols) + x = np.linspace(0, (num_rows - 1) * horizontal_scale, num_rows) + yy, xx = np.meshgrid(y, x) + # copy height field to avoid modifying the original array + hf = height_field.copy() + + # correct vertical surfaces above the slope threshold + if slope_threshold is not None: + # scale slope threshold based on the horizontal and vertical scale + slope_threshold *= horizontal_scale / vertical_scale + # allocate arrays to store the movement of the vertices + move_x = np.zeros((num_rows, num_cols)) + move_y = np.zeros((num_rows, num_cols)) + move_corners = np.zeros((num_rows, num_cols)) + # move vertices along the x-axis + move_x[: num_rows - 1, :] += hf[1:num_rows, :] - hf[: num_rows - 1, :] > slope_threshold + move_x[1:num_rows, :] -= hf[: num_rows - 1, :] - hf[1:num_rows, :] > slope_threshold + # move vertices along the y-axis + move_y[:, : num_cols - 1] += hf[:, 1:num_cols] - hf[:, : num_cols - 1] > slope_threshold + move_y[:, 1:num_cols] -= hf[:, : num_cols - 1] - hf[:, 1:num_cols] > slope_threshold + # move vertices along the corners + move_corners[: num_rows - 1, : num_cols - 1] += ( + hf[1:num_rows, 1:num_cols] - hf[: num_rows - 1, : num_cols - 1] > slope_threshold + ) + move_corners[1:num_rows, 1:num_cols] -= ( + hf[: num_rows - 1, : num_cols - 1] - hf[1:num_rows, 1:num_cols] > slope_threshold + ) + xx += (move_x + move_corners * (move_x == 0)) * horizontal_scale + yy += (move_y + move_corners * (move_y == 0)) * horizontal_scale + + # create vertices for the mesh + vertices = np.zeros((num_rows * num_cols, 3), dtype=np.float32) + vertices[:, 0] = xx.flatten() + vertices[:, 1] = yy.flatten() + vertices[:, 2] = hf.flatten() * vertical_scale + # create triangles for the mesh + triangles = -np.ones((2 * (num_rows - 1) * (num_cols - 1), 3), dtype=np.uint32) + for i in range(num_rows - 1): + ind0 = np.arange(0, num_cols - 1) + i * num_cols + ind1 = ind0 + 1 + ind2 = ind0 + num_cols + ind3 = ind2 + 1 + start = 2 * i * (num_cols - 1) + stop = start + 2 * (num_cols - 1) + triangles[start:stop:2, 0] = ind0 + triangles[start:stop:2, 1] = ind3 + triangles[start:stop:2, 2] = ind1 + triangles[start + 1 : stop : 2, 0] = ind0 + triangles[start + 1 : stop : 2, 1] = ind2 + triangles[start + 1 : stop : 2, 2] = ind3 + + return vertices, triangles diff --git a/source/uwlab/uwlab/terrains/terrain_generator.py b/source/uwlab/uwlab/terrains/terrain_generator.py new file mode 100644 index 00000000..fbcc0b24 --- /dev/null +++ b/source/uwlab/uwlab/terrains/terrain_generator.py @@ -0,0 +1,386 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import numpy as np +import os +import torch +import trimesh +from typing import TYPE_CHECKING + +import omni.log + +from isaaclab.terrains.height_field import HfTerrainBaseCfg +from isaaclab.terrains.trimesh.utils import make_border +from isaaclab.terrains.utils import color_meshes_by_height +from isaaclab.utils.dict import dict_to_md5_hash +from isaaclab.utils.io import dump_yaml +from isaaclab.utils.timer import Timer +from isaaclab.utils.warp import convert_to_warp_mesh + +if TYPE_CHECKING: + from .terrain_generator_cfg import PatchSamplingCfg, SubTerrainBaseCfg, TerrainGeneratorCfg + + +class TerrainGenerator: + r"""Terrain generator to handle different terrain generation functions. + + The terrains are represented as meshes. These are obtained either from height fields or by using the + `trimesh `__ library. The height field representation is more + flexible, but it is less computationally and memory efficient than the trimesh representation. + + All terrain generation functions take in the argument :obj:`difficulty` which determines the complexity + of the terrain. The difficulty is a number between 0 and 1, where 0 is the easiest and 1 is the hardest. + In most cases, the difficulty is used for linear interpolation between different terrain parameters. + For example, in a pyramid stairs terrain the step height is interpolated between the specified minimum + and maximum step height. + + Each sub-terrain has a corresponding configuration class that can be used to specify the parameters + of the terrain. The configuration classes are inherited from the :class:`SubTerrainBaseCfg` class + which contains the common parameters for all terrains. + + If a curriculum is used, the terrains are generated based on their difficulty parameter. + The difficulty is varied linearly over the number of rows (i.e. along x) with a small random value + added to the difficulty to ensure that the columns with the same sub-terrain type are not exactly + the same. The difficulty parameter for a sub-terrain at a given row is calculated as: + + .. math:: + + \text{difficulty} = \frac{\text{row_id} + \eta}{\text{num_rows}} \times (\text{upper} - \text{lower}) + \text{lower} + + where :math:`\eta\sim\mathcal{U}(0, 1)` is a random perturbation to the difficulty, and + :math:`(\text{lower}, \text{upper})` is the range of the difficulty parameter, specified using the + :attr:`~TerrainGeneratorCfg.difficulty_range` parameter. + + If a curriculum is not used, the terrains are generated randomly. In this case, the difficulty parameter + is randomly sampled from the specified range, given by the :attr:`~TerrainGeneratorCfg.difficulty_range` parameter: + + .. math:: + + \text{difficulty} \sim \mathcal{U}(\text{lower}, \text{upper}) + + If the :attr:`~TerrainGeneratorCfg.flat_patch_sampling` is specified for a sub-terrain, flat patches are sampled + on the terrain. These can be used for spawning robots, targets, etc. The sampled patches are stored + in the :obj:`flat_patches` dictionary. The key specifies the intention of the flat patches and the + value is a tensor containing the flat patches for each sub-terrain. + + If the flag :attr:`~TerrainGeneratorCfg.use_cache` is set to True, the terrains are cached based on their + sub-terrain configurations. This means that if the same sub-terrain configuration is used + multiple times, the terrain is only generated once and then reused. This is useful when + generating complex sub-terrains that take a long time to generate. + + .. attention:: + + The terrain generation has its own seed parameter. This is set using the :attr:`TerrainGeneratorCfg.seed` + parameter. If the seed is not set and the caching is disabled, the terrain generation may not be + completely reproducible. + + """ + + terrain_mesh: trimesh.Trimesh + """A single trimesh.Trimesh object for all the generated sub-terrains.""" + terrain_meshes: list[trimesh.Trimesh] + """List of trimesh.Trimesh objects for all the generated sub-terrains.""" + terrain_origins: np.ndarray + """The origin of each sub-terrain. Shape is (num_rows, num_cols, 3).""" + flat_patches: dict[str, torch.Tensor] + """A dictionary of sampled valid (flat) patches for each sub-terrain. + + The dictionary keys are the names of the flat patch sampling configurations. This maps to a + tensor containing the flat patches for each sub-terrain. The shape of the tensor is + (num_rows, num_cols, num_patches, 3). + + For instance, the key "root_spawn" maps to a tensor containing the flat patches for spawning an asset. + Similarly, the key "target_spawn" maps to a tensor containing the flat patches for setting targets. + """ + + def __init__(self, cfg: TerrainGeneratorCfg, device: str = "cpu"): + """Initialize the terrain generator. + + Args: + cfg: Configuration for the terrain generator. + device: The device to use for the flat patches tensor. + """ + # check inputs + if len(cfg.sub_terrains) == 0: + raise ValueError("No sub-terrains specified! Please add at least one sub-terrain.") + # store inputs + self.cfg = cfg + self.device = device + + # set common values to all sub-terrains config + for sub_cfg in self.cfg.sub_terrains.values(): + # size of all terrains + sub_cfg.size = self.cfg.size + # params for height field terrains + if isinstance(sub_cfg, HfTerrainBaseCfg): + sub_cfg.horizontal_scale = self.cfg.horizontal_scale + sub_cfg.vertical_scale = self.cfg.vertical_scale + sub_cfg.slope_threshold = self.cfg.slope_threshold + + # throw a warning if the cache is enabled but the seed is not set + if self.cfg.use_cache and self.cfg.seed is None: + omni.log.warn( + "Cache is enabled but the seed is not set. The terrain generation will not be reproducible." + " Please set the seed in the terrain generator configuration to make the generation reproducible." + ) + + # if the seed is not set, we assume there is a global seed set and use that. + # this ensures that the terrain is reproducible if the seed is set at the beginning of the program. + if self.cfg.seed is not None: + seed = self.cfg.seed + else: + seed = np.random.get_state()[1][0] + # set the seed for reproducibility + # note: we create a new random number generator to avoid affecting the global state + # in the other places where random numbers are used. + self.np_rng = np.random.default_rng(seed) + + # buffer for storing valid patches + self.flat_patches = {} + # create a list of all sub-terrains + self.terrain_meshes = list() + self.terrain_origins = np.zeros((self.cfg.num_rows, self.cfg.num_cols, 3)) + + # parse configuration and add sub-terrains + # create terrains based on curriculum or randomly + if self.cfg.curriculum: + with Timer("[INFO] Generating terrains based on curriculum took"): + self._generate_curriculum_terrains() + else: + with Timer("[INFO] Generating terrains randomly took"): + self._generate_random_terrains() + # add a border around the terrains + self._add_terrain_border() + # combine all the sub-terrains into a single mesh + self.terrain_mesh = trimesh.util.concatenate(self.terrain_meshes) + + # color the terrain mesh + if self.cfg.color_scheme == "height": + self.terrain_mesh = color_meshes_by_height(self.terrain_mesh) + elif self.cfg.color_scheme == "random": + self.terrain_mesh.visual.vertex_colors = self.np_rng.choice( + range(256), size=(len(self.terrain_mesh.vertices), 4) + ) + elif self.cfg.color_scheme == "none": + pass + else: + raise ValueError(f"Invalid color scheme: {self.cfg.color_scheme}.") + + # offset the entire terrain and origins so that it is centered + # -- terrain mesh + transform = np.eye(4) + transform[:2, -1] = -self.cfg.size[0] * self.cfg.num_rows * 0.5, -self.cfg.size[1] * self.cfg.num_cols * 0.5 + self.terrain_mesh.apply_transform(transform) + # -- terrain origins + self.terrain_origins += transform[:3, -1] + # -- valid patches + terrain_origins_torch = torch.tensor(self.terrain_origins, dtype=torch.float, device=self.device).unsqueeze(2) + for name, value in self.flat_patches.items(): + self.flat_patches[name] = value + terrain_origins_torch + + def __str__(self): + """Return a string representation of the terrain generator.""" + msg = "Terrain Generator:" + msg += f"\n\tSeed: {self.cfg.seed}" + msg += f"\n\tNumber of rows: {self.cfg.num_rows}" + msg += f"\n\tNumber of columns: {self.cfg.num_cols}" + msg += f"\n\tSub-terrain size: {self.cfg.size}" + msg += f"\n\tSub-terrain types: {list(self.cfg.sub_terrains.keys())}" + msg += f"\n\tCurriculum: {self.cfg.curriculum}" + msg += f"\n\tDifficulty range: {self.cfg.difficulty_range}" + msg += f"\n\tColor scheme: {self.cfg.color_scheme}" + msg += f"\n\tUse cache: {self.cfg.use_cache}" + if self.cfg.use_cache: + msg += f"\n\tCache directory: {self.cfg.cache_dir}" + + return msg + + """ + Terrain generator functions. + """ + + def _generate_random_terrains(self): + """Add terrains based on randomly sampled difficulty parameter.""" + # normalize the proportions of the sub-terrains + proportions = np.array([sub_cfg.proportion for sub_cfg in self.cfg.sub_terrains.values()]) + proportions /= np.sum(proportions) + # create a list of all terrain configs + sub_terrains_cfgs = list(self.cfg.sub_terrains.values()) + + # randomly sample sub-terrains + for index in range(self.cfg.num_rows * self.cfg.num_cols): + # coordinate index of the sub-terrain + (sub_row, sub_col) = np.unravel_index(index, (self.cfg.num_rows, self.cfg.num_cols)) + # randomly sample terrain index + sub_index = self.np_rng.choice(len(proportions), p=proportions) + # randomly sample difficulty parameter + difficulty = self.np_rng.uniform(*self.cfg.difficulty_range) + # generate terrain + mesh, origin = self._get_terrain_mesh(difficulty, sub_terrains_cfgs[sub_index]) + # add to sub-terrains + self._add_sub_terrain(mesh, origin, sub_row, sub_col, sub_terrains_cfgs[sub_index]) + + def _generate_curriculum_terrains(self): + """Add terrains based on the difficulty parameter.""" + # normalize the proportions of the sub-terrains + proportions = np.array([sub_cfg.proportion for sub_cfg in self.cfg.sub_terrains.values()]) + proportions /= np.sum(proportions) + + # find the sub-terrain index for each column + # we generate the terrains based on their proportion (not randomly sampled) + sub_indices = [] + for index in range(self.cfg.num_cols): + sub_index = np.min(np.where(index / self.cfg.num_cols + 0.001 < np.cumsum(proportions))[0]) + sub_indices.append(sub_index) + sub_indices = np.array(sub_indices, dtype=np.int32) + # create a list of all terrain configs + sub_terrains_cfgs = list(self.cfg.sub_terrains.values()) + + # curriculum-based sub-terrains + for sub_col in range(self.cfg.num_cols): + for sub_row in range(self.cfg.num_rows): + # vary the difficulty parameter linearly over the number of rows + # note: based on the proportion, multiple columns can have the same sub-terrain type. + # Thus to increase the diversity along the rows, we add a small random value to the difficulty. + # This ensures that the terrains are not exactly the same. For example, if the + # the row index is 2 and the number of rows is 10, the nominal difficulty is 0.2. + # We add a small random value to the difficulty to make it between 0.2 and 0.3. + lower, upper = self.cfg.difficulty_range + difficulty = (sub_row + self.np_rng.uniform()) / self.cfg.num_rows + difficulty = lower + (upper - lower) * difficulty + # generate terrain + mesh, origin = self._get_terrain_mesh(difficulty, sub_terrains_cfgs[sub_indices[sub_col]]) + # add to sub-terrains + self._add_sub_terrain(mesh, origin, sub_row, sub_col, sub_terrains_cfgs[sub_indices[sub_col]]) + + """ + Internal helper functions. + """ + + def _add_terrain_border(self): + """Add a surrounding border over all the sub-terrains into the terrain meshes.""" + # border parameters + border_size = ( + self.cfg.num_rows * self.cfg.size[0] + 2 * self.cfg.border_width, + self.cfg.num_cols * self.cfg.size[1] + 2 * self.cfg.border_width, + ) + inner_size = (self.cfg.num_rows * self.cfg.size[0], self.cfg.num_cols * self.cfg.size[1]) + border_center = ( + self.cfg.num_rows * self.cfg.size[0] / 2, + self.cfg.num_cols * self.cfg.size[1] / 2, + -self.cfg.border_height / 2, + ) + # border mesh + border_meshes = make_border(border_size, inner_size, height=self.cfg.border_height, position=border_center) + border = trimesh.util.concatenate(border_meshes) + # update the faces to have minimal triangles + selector = ~(np.asarray(border.triangles)[:, :, 2] < -0.1).any(1) + border.update_faces(selector) + # add the border to the list of meshes + self.terrain_meshes.append(border) + + def _add_sub_terrain( + self, mesh: trimesh.Trimesh, origin: np.ndarray, row: int, col: int, sub_terrain_cfg: SubTerrainBaseCfg + ): + """Add input sub-terrain to the list of sub-terrains. + + This function adds the input sub-terrain mesh to the list of sub-terrains and updates the origin + of the sub-terrain in the list of origins. It also samples flat patches if specified. + + Args: + mesh: The mesh of the sub-terrain. + origin: The origin of the sub-terrain. + row: The row index of the sub-terrain. + col: The column index of the sub-terrain. + """ + # sample flat patches if specified + if sub_terrain_cfg.patch_sampling is not None: + omni.log.info(f"Sampling flat patches for sub-terrain at (row, col): ({row}, {col})") + # convert the mesh to warp mesh + wp_mesh = convert_to_warp_mesh(mesh.vertices, mesh.faces, device=self.device) + # sample flat patches based on each patch configuration for that sub-terrain + for name, patch_cfg in sub_terrain_cfg.patch_sampling.items(): + patch_cfg: PatchSamplingCfg + # create the flat patches tensor (if not already created) + if name not in self.flat_patches: + self.flat_patches[name] = torch.zeros( + (self.cfg.num_rows, self.cfg.num_cols, patch_cfg.num_patches, 3), device=self.device + ) + # add the flat patches to the tensor + self.flat_patches[name][row, col] = patch_cfg.func( + wp_mesh=wp_mesh, + origin=origin, + cfg=patch_cfg, + ) + + # transform the mesh to the correct position + transform = np.eye(4) + transform[0:2, -1] = (row + 0.5) * self.cfg.size[0], (col + 0.5) * self.cfg.size[1] + mesh.apply_transform(transform) + # add mesh to the list + self.terrain_meshes.append(mesh) + # add origin to the list + self.terrain_origins[row, col] = origin + transform[:3, -1] + + def _get_terrain_mesh(self, difficulty: float, cfg: SubTerrainBaseCfg) -> tuple[trimesh.Trimesh, np.ndarray]: + """Generate a sub-terrain mesh based on the input difficulty parameter. + + If caching is enabled, the sub-terrain is cached and loaded from the cache if it exists. + The cache is stored in the cache directory specified in the configuration. + + .. Note: + This function centers the 2D center of the mesh and its specified origin such that the + 2D center becomes :math:`(0, 0)` instead of :math:`(size[0] / 2, size[1] / 2). + + Args: + difficulty: The difficulty parameter. + cfg: The configuration of the sub-terrain. + + Returns: + The sub-terrain mesh and origin. + """ + # copy the configuration + cfg = cfg.copy() + # add other parameters to the sub-terrain configuration + cfg.difficulty = float(difficulty) + cfg.seed = self.cfg.seed + # generate hash for the sub-terrain + sub_terrain_hash = dict_to_md5_hash(cfg.to_dict()) + # generate the file name + sub_terrain_cache_dir = os.path.join(self.cfg.cache_dir, sub_terrain_hash) + sub_terrain_obj_filename = os.path.join(sub_terrain_cache_dir, "mesh.obj") + sub_terrain_csv_filename = os.path.join(sub_terrain_cache_dir, "origin.csv") + sub_terrain_meta_filename = os.path.join(sub_terrain_cache_dir, "cfg.yaml") + + # check if hash exists - if true, load the mesh and origin and return + if self.cfg.use_cache and os.path.exists(sub_terrain_obj_filename): + # load existing mesh + mesh = trimesh.load_mesh(sub_terrain_obj_filename, process=False) + origin = np.loadtxt(sub_terrain_csv_filename, delimiter=",") + # return the generated mesh + return mesh, origin + + # generate the terrain + meshes, origin = cfg.function(difficulty, cfg) + mesh = trimesh.util.concatenate(meshes) + # offset mesh such that they are in their center + transform = np.eye(4) + transform[0:2, -1] = -cfg.size[0] * 0.5, -cfg.size[1] * 0.5 + mesh.apply_transform(transform) + # change origin to be in the center of the sub-terrain + origin += transform[0:3, -1] + + # if caching is enabled, save the mesh and origin + if self.cfg.use_cache: + # create the cache directory + os.makedirs(sub_terrain_cache_dir, exist_ok=True) + # save the data + mesh.export(sub_terrain_obj_filename) + np.savetxt(sub_terrain_csv_filename, origin, delimiter=",", header="x,y,z") + dump_yaml(sub_terrain_meta_filename, cfg) + # return the generated mesh + return mesh, origin diff --git a/source/uwlab/uwlab/terrains/terrain_generator_cfg.py b/source/uwlab/uwlab/terrains/terrain_generator_cfg.py new file mode 100644 index 00000000..c96b0c5e --- /dev/null +++ b/source/uwlab/uwlab/terrains/terrain_generator_cfg.py @@ -0,0 +1,166 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +""" +Configuration classes defining the different terrains available. Each configuration class must +inherit from ``isaaclab.terrains.terrains_cfg.TerrainConfig`` and define the following attributes: + +- ``name``: Name of the terrain. This is used for the prim name in the USD stage. +- ``function``: Function to generate the terrain. This function must take as input the terrain difficulty + and the configuration parameters and return a `tuple with the `trimesh`` mesh object and terrain origin. +""" + +from __future__ import annotations + +import numpy as np +import trimesh +from collections.abc import Callable +from dataclasses import MISSING +from typing import Literal + +from isaaclab.utils import configclass + +from .terrain_generator import TerrainGenerator +from .utils.patch_sampling_cfg import PatchSamplingCfg + + +@configclass +class SubTerrainBaseCfg: + """Base class for terrain configurations. + + All the sub-terrain configurations must inherit from this class. + + The :attr:`size` attribute is the size of the generated sub-terrain. Based on this, the terrain must + extend from :math:`(0, 0)` to :math:`(size[0], size[1])`. + """ + + function: Callable[[float, SubTerrainBaseCfg], tuple[list[trimesh.Trimesh], np.ndarray]] = MISSING + """Function to generate the terrain. + + This function must take as input the terrain difficulty and the configuration parameters and + return a tuple with a list of ``trimesh`` mesh objects and the terrain origin. + """ + + proportion: float = 1.0 + """Proportion of the terrain to generate. Defaults to 1.0. + + This is used to generate a mix of terrains. The proportion corresponds to the probability of sampling + the particular terrain. For example, if there are two terrains, A and B, with proportions 0.3 and 0.7, + respectively, then the probability of sampling terrain A is 0.3 and the probability of sampling terrain B + is 0.7. + """ + + 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:`~isaaclab.terrains.TerrainImporterCfg` is used, this parameter gets overridden by + :attr:`isaaclab.scene.TerrainImporterCfg.size` attribute. + """ + + patch_sampling: dict[str, PatchSamplingCfg] | None = None + """Dictionary of configurations for sampling flat patches on the sub-terrain. Defaults to None, + in which case no flat patch sampling is performed. + + The keys correspond to the name of the flat patch sampling configuration and the values are the + corresponding configurations. + """ + + +@configclass +class TerrainGeneratorCfg: + """Configuration for the terrain generator.""" + + class_type: Callable = TerrainGenerator + + seed: int | None = None + """The seed for the random number generator. Defaults to None, in which case the seed from the + current NumPy's random state is used. + + When the seed is set, the random number generator is initialized with the given seed. This ensures + that the generated terrains are deterministic across different runs. If the seed is not set, the + seed from the current NumPy's random state is used. This assumes that the seed is set elsewhere in + the code. + """ + + curriculum: bool = False + """Whether to use the curriculum mode. Defaults to False. + + If True, the terrains are generated based on their difficulty parameter. Otherwise, + they are randomly generated. + """ + + size: tuple[float, float] = MISSING + """The width (along x) and length (along y) of each sub-terrain (in m). + + Note: + This value is passed on to all the sub-terrain configurations. + """ + + border_width: float = 0.0 + """The width of the border around the terrain (in m). Defaults to 0.0.""" + + border_height: float = 1.0 + """The height of the border around the terrain (in m). Defaults to 1.0.""" + + num_rows: int = 1 + """Number of rows of sub-terrains to generate. Defaults to 1.""" + + num_cols: int = 1 + """Number of columns of sub-terrains to generate. Defaults to 1.""" + + color_scheme: Literal["height", "random", "none"] = "none" + """Color scheme to use for the terrain. Defaults to "none". + + The available color schemes are: + + - "height": Color based on the height of the terrain. + - "random": Random color scheme. + - "none": No color scheme. + """ + + horizontal_scale: float = 0.1 + """The discretization of the terrain along the x and y axes (in m). Defaults to 0.1. + + This value is passed on to all the height field sub-terrain configurations. + """ + + vertical_scale: float = 0.005 + """The discretization of the terrain along the z axis (in m). Defaults to 0.005. + + This value is passed on to all the height field sub-terrain configurations. + """ + + slope_threshold: float | None = 0.75 + """The slope threshold above which surfaces are made vertical. Defaults to 0.75. + + If None no correction is applied. + + This value is passed on to all the height field sub-terrain configurations. + """ + + sub_terrains: dict[str, SubTerrainBaseCfg] = MISSING + """Dictionary of sub-terrain configurations. + + The keys correspond to the name of the sub-terrain configuration and the values are the corresponding + configurations. + """ + + difficulty_range: tuple[float, float] = (0.0, 1.0) + """The range of difficulty values for the sub-terrains. Defaults to (0.0, 1.0). + + If curriculum is enabled, the terrains will be generated based on this range in ascending order + of difficulty. Otherwise, the terrains will be generated based on this range in a random order. + """ + + use_cache: bool = False + """Whether to load the sub-terrain from cache if it exists. Defaults to True. + + If enabled, the generated terrains are stored in the cache directory. When generating terrains, the cache + is checked to see if the terrain already exists. If it does, the terrain is loaded from the cache. Otherwise, + the terrain is generated and stored in the cache. Caching can be used to speed up terrain generation. + """ + + cache_dir: str = "/tmp/isaaclab/terrains" + """The directory where the terrain cache is stored. Defaults to "/tmp/isaaclab/terrains".""" diff --git a/source/uwlab/uwlab/terrains/terrain_importer.py b/source/uwlab/uwlab/terrains/terrain_importer.py new file mode 100644 index 00000000..b6049cc9 --- /dev/null +++ b/source/uwlab/uwlab/terrains/terrain_importer.py @@ -0,0 +1,365 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import numpy as np +import torch +import trimesh +from typing import TYPE_CHECKING + +import warp +from pxr import UsdGeom + +import isaaclab.sim as sim_utils +from isaaclab.markers import VisualizationMarkers +from isaaclab.markers.config import FRAME_MARKER_CFG +from isaaclab.terrains.trimesh.utils import make_plane +from isaaclab.terrains.utils import create_prim_from_mesh +from isaaclab.utils.warp import convert_to_warp_mesh +from uwlab.terrains.terrain_generator import TerrainGenerator + +if TYPE_CHECKING: + from .terrain_importer_cfg import TerrainImporterCfg + + +class TerrainImporter: + r"""A class to handle terrain meshes and import them into the simulator. + + We assume that a terrain mesh comprises of sub-terrains that are arranged in a grid with + rows ``num_rows`` and columns ``num_cols``. The terrain origins are the positions of the sub-terrains + where the robot should be spawned. + + Based on the configuration, the terrain importer handles computing the environment origins from the sub-terrain + origins. In a typical setup, the number of sub-terrains (:math:`num\_rows \times num\_cols`) is smaller than + the number of environments (:math:`num\_envs`). In this case, the environment origins are computed by + sampling the sub-terrain origins. + + If a curriculum is used, it is possible to update the environment origins to terrain origins that correspond + to a harder difficulty. This is done by calling :func:`update_terrain_levels`. The idea comes from game-based + curriculum. For example, in a game, the player starts with easy levels and progresses to harder levels. + """ + + meshes: dict[str, trimesh.Trimesh] + """A dictionary containing the names of the meshes and their keys.""" + warp_meshes: dict[str, warp.Mesh] + """A dictionary containing the names of the warp meshes and their keys.""" + terrain_origins: torch.Tensor | None + """The origins of the sub-terrains in the added terrain mesh. Shape is (num_rows, num_cols, 3). + + If None, then it is assumed no sub-terrains exist. The environment origins are computed in a grid. + """ + env_origins: torch.Tensor + """The origins of the environments. Shape is (num_envs, 3).""" + + def __init__(self, cfg: TerrainImporterCfg): + """Initialize the terrain importer. + + Args: + cfg: The configuration for the terrain importer. + + Raises: + ValueError: If input terrain type is not supported. + ValueError: If terrain type is 'generator' and no configuration provided for ``terrain_generator``. + 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 + + # create a dict of meshes + self.meshes = dict() + self.warp_meshes = dict() + self.env_origins = None + self.terrain_origins = None + # private variables + self._terrain_flat_patches = dict() + + # auto-import the terrain based on the config + if self.cfg.terrain_type == "generator": + # check config is provided + if self.cfg.terrain_generator is None: + raise ValueError("Input terrain type is 'generator' but no value provided for 'terrain_generator'.") + # generate the terrain + terrain_generator: TerrainGenerator = self.cfg.terrain_generator.class_type( + cfg=self.cfg.terrain_generator, device=self.device + ) + self.import_mesh("terrain", terrain_generator.terrain_mesh) + # configure the terrain origins based on the terrain generator + self.configure_env_origins(terrain_generator.terrain_origins) + # refer to the flat patches + self._terrain_flat_patches = terrain_generator.flat_patches + elif self.cfg.terrain_type == "usd": + # check if config is provided + if self.cfg.usd_path is None: + raise ValueError("Input terrain type is 'usd' but no value provided for 'usd_path'.") + # import the terrain + self.import_usd("terrain", self.cfg.usd_path) + # configure the origins in a grid + self.configure_env_origins() + elif self.cfg.terrain_type == "plane": + # load the plane + self.import_ground_plane("terrain") + # configure the origins in a grid + self.configure_env_origins() + else: + raise ValueError(f"Terrain type '{self.cfg.terrain_type}' not available.") + + # set initial state of debug visualization + self.set_debug_vis(self.cfg.debug_vis) + + """ + Properties. + """ + + @property + def has_debug_vis_implementation(self) -> bool: + """Whether the terrain importer has a debug visualization implemented. + + This always returns True. + """ + return True + + @property + def flat_patches(self) -> dict[str, torch.Tensor]: + """A dictionary containing the sampled valid (flat) patches for the terrain. + + This is only available if the terrain type is 'generator'. For other terrain types, this feature + is not available and the function returns an empty dictionary. + + Please refer to the :attr:`TerrainGenerator.flat_patches` for more information. + """ + return self._terrain_flat_patches + + """ + Operations - Visibility. + """ + + def set_debug_vis(self, debug_vis: bool) -> bool: + """Set the debug visualization of the terrain importer. + + Args: + debug_vis: Whether to visualize the terrain origins. + + Returns: + Whether the debug visualization was successfully set. False if the terrain + importer does not support debug visualization. + + Raises: + RuntimeError: If terrain origins are not configured. + """ + # create a marker if necessary + if debug_vis: + if not hasattr(self, "origin_visualizer"): + self.origin_visualizer = VisualizationMarkers( + cfg=FRAME_MARKER_CFG.replace(prim_path="/Visuals/TerrainOrigin") + ) + if self.terrain_origins is not None: + self.origin_visualizer.visualize(self.terrain_origins.reshape(-1, 3)) + elif self.env_origins is not None: + self.origin_visualizer.visualize(self.env_origins.reshape(-1, 3)) + else: + raise RuntimeError("Terrain origins are not configured.") + # set visibility + self.origin_visualizer.set_visibility(True) + else: + if hasattr(self, "origin_visualizer"): + self.origin_visualizer.set_visibility(False) + # report success + return True + + """ + Operations - Import. + """ + + def import_ground_plane(self, key: str, size: tuple[float, float] = (2.0e6, 2.0e6)): + """Add a plane to the terrain importer. + + Args: + key: The key to store the mesh. + size: The size of the plane. Defaults to (2.0e6, 2.0e6). + + Raises: + ValueError: If a terrain with the same key already exists. + """ + # check if key exists + if key in self.meshes: + raise ValueError(f"Mesh with key {key} already exists. Existing keys: {self.meshes.keys()}.") + # create a plane + mesh = make_plane(size, height=0.0, center_zero=True) + # store the mesh + self.meshes[key] = mesh + # create a warp mesh + device = "cuda" if "cuda" in self.device else "cpu" + self.warp_meshes[key] = convert_to_warp_mesh(mesh.vertices, mesh.faces, device=device) + + # get the mesh + ground_plane_cfg = sim_utils.GroundPlaneCfg(physics_material=self.cfg.physics_material, size=size) + ground_plane_cfg.func(self.cfg.prim_path, ground_plane_cfg) + + def import_mesh(self, key: str, mesh: trimesh.Trimesh): + """Import a mesh into the simulator. + + The mesh is imported into the simulator under the prim path ``cfg.prim_path/{key}``. The created path + contains the mesh as a :class:`pxr.UsdGeom` instance along with visual or physics material prims. + + Args: + key: The key to store the mesh. + mesh: The mesh to import. + + Raises: + ValueError: If a terrain with the same key already exists. + """ + # check if key exists + if key in self.meshes: + raise ValueError(f"Mesh with key {key} already exists. Existing keys: {self.meshes.keys()}.") + # store the mesh + self.meshes[key] = mesh + # create a warp mesh + device = "cuda" if "cuda" in self.device else "cpu" + self.warp_meshes[key] = convert_to_warp_mesh(mesh.vertices, mesh.faces, device=device) + + # get the mesh + mesh = self.meshes[key] + mesh_prim_path = self.cfg.prim_path + f"/{key}" + # import the mesh + create_prim_from_mesh( + mesh_prim_path, + mesh, + visual_material=self.cfg.visual_material, + physics_material=self.cfg.physics_material, + ) + + def import_usd(self, key: str, usd_path: str): + """Import a mesh from a USD file. + + We assume that the USD file contains a single mesh. If the USD file contains multiple meshes, then + the first mesh is used. The function mainly helps in registering the mesh into the warp meshes + and the meshes dictionary. + + Note: + We do not apply any material properties to the mesh. The material properties should + be defined in the USD file. + + Args: + key: The key to store the mesh. + usd_path: The path to the USD file. + + Raises: + ValueError: If a terrain with the same key already exists. + """ + # add mesh to the dict + if key in self.meshes: + raise ValueError(f"Mesh with key {key} already exists. Existing keys: {self.meshes.keys()}.") + # add the prim path + cfg = sim_utils.UsdFileCfg(usd_path=usd_path) + cfg.func(self.cfg.prim_path + f"/{key}", cfg) + + # traverse the prim and get the collision mesh + # THINK: Should the user specify the collision mesh? + mesh_prim = sim_utils.get_first_matching_child_prim( + self.cfg.prim_path + f"/{key}", lambda prim: prim.GetTypeName() == "Mesh" + ) + # check if the mesh is valid + if mesh_prim is None: + raise ValueError(f"Could not find any collision mesh in {usd_path}. Please check asset.") + # cast into UsdGeomMesh + mesh_prim = UsdGeom.Mesh(mesh_prim) + # store the mesh + vertices = np.asarray(mesh_prim.GetPointsAttr().Get()) + faces = np.asarray(mesh_prim.GetFaceVertexIndicesAttr().Get()).reshape(-1, 3) + self.meshes[key] = trimesh.Trimesh(vertices=vertices, faces=faces) + # create a warp mesh + device = "cuda" if "cuda" in self.device else "cpu" + self.warp_meshes[key] = convert_to_warp_mesh(vertices, faces, device=device) + + """ + Operations - Origins. + """ + + def configure_env_origins(self, origins: np.ndarray | None = None): + """Configure the origins of the environments based on the added terrain. + + Args: + origins: The origins of the sub-terrains. Shape is (num_rows, num_cols, 3). + """ + # decide whether to compute origins in a grid or based on curriculum + if origins is not None: + # convert to numpy + if isinstance(origins, np.ndarray): + origins = torch.from_numpy(origins) + # store the origins + self.terrain_origins = origins.to(self.device, dtype=torch.float) + # compute environment origins + self.env_origins = self._compute_env_origins_curriculum(self.cfg.num_envs, self.terrain_origins) + else: + self.terrain_origins = None + # check if env spacing is valid + if self.cfg.env_spacing is None: + raise ValueError("Environment spacing must be specified for configuring grid-like origins.") + # compute environment origins + self.env_origins = self._compute_env_origins_grid(self.cfg.num_envs, self.cfg.env_spacing) + + def update_env_origins(self, env_ids: torch.Tensor, move_up: torch.Tensor, move_down: torch.Tensor): + """Update the environment origins based on the terrain levels.""" + # check if grid-like spawning + if self.terrain_origins is None: + return + # update terrain level for the envs + self.terrain_levels[env_ids] += 1 * move_up - 1 * move_down + # robots that solve the last level are sent to a random one + # the minimum level is zero + self.terrain_levels[env_ids] = torch.where( + self.terrain_levels[env_ids] >= self.max_terrain_level, + torch.randint_like(self.terrain_levels[env_ids], self.max_terrain_level), + torch.clip(self.terrain_levels[env_ids], 0), + ) + # update the env origins + self.env_origins[env_ids] = self.terrain_origins[self.terrain_levels[env_ids], self.terrain_types[env_ids]] + + """ + Internal helpers. + """ + + def _compute_env_origins_curriculum(self, num_envs: int, origins: torch.Tensor) -> torch.Tensor: + """Compute the origins of the environments defined by the sub-terrains origins.""" + # extract number of rows and cols + num_rows, num_cols = origins.shape[:2] + # maximum initial level possible for the terrains + if self.cfg.max_init_terrain_level is None: + max_init_level = num_rows - 1 + else: + max_init_level = min(self.cfg.max_init_terrain_level, num_rows - 1) + # store maximum terrain level possible + self.max_terrain_level = num_rows + # define all terrain levels and types available + self.terrain_levels = torch.randint(0, max_init_level + 1, (num_envs,), device=self.device) + self.terrain_types = torch.div( + torch.arange(num_envs, device=self.device), + (num_envs / num_cols), + rounding_mode="floor", + ).to(torch.long) + # create tensor based on number of environments + env_origins = torch.zeros(num_envs, 3, device=self.device) + env_origins[:] = origins[self.terrain_levels, self.terrain_types] + return env_origins + + def _compute_env_origins_grid(self, num_envs: int, env_spacing: float) -> torch.Tensor: + """Compute the origins of the environments in a grid based on configured spacing.""" + # create tensor based on number of environments + env_origins = torch.zeros(num_envs, 3, device=self.device) + # create a grid of origins + num_rows = np.ceil(num_envs / int(np.sqrt(num_envs))) + num_cols = np.ceil(num_envs / num_rows) + ii, jj = torch.meshgrid( + torch.arange(num_rows, device=self.device), torch.arange(num_cols, device=self.device), indexing="ij" + ) + env_origins[:, 0] = -(ii.flatten()[:num_envs] - (num_rows - 1) / 2) * env_spacing + env_origins[:, 1] = (jj.flatten()[:num_envs] - (num_cols - 1) / 2) * env_spacing + env_origins[:, 2] = 0.0 + return env_origins diff --git a/source/uwlab/uwlab/terrains/terrain_importer_cfg.py b/source/uwlab/uwlab/terrains/terrain_importer_cfg.py new file mode 100644 index 00000000..f77ab871 --- /dev/null +++ b/source/uwlab/uwlab/terrains/terrain_importer_cfg.py @@ -0,0 +1,103 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from dataclasses import MISSING +from typing import TYPE_CHECKING, Literal + +import isaaclab.sim as sim_utils +from isaaclab.utils import configclass + +from .terrain_importer import TerrainImporter + +if TYPE_CHECKING: + from .terrain_generator_cfg import TerrainGeneratorCfg + + +@configclass +class TerrainImporterCfg: + """Configuration for the terrain manager.""" + + class_type: type = TerrainImporter + """The class to use for the terrain importer. + + Defaults to :class:`isaaclab.terrains.terrain_importer.TerrainImporter`. + """ + + collision_group: int = -1 + """The collision group of the terrain. Defaults to -1.""" + + prim_path: str = MISSING + """The absolute path of the USD terrain prim. + + All sub-terrains are imported relative to this prim path. + """ + + num_envs: int = 1 + """The number of environment origins to consider. Defaults to 1. + + In case, the :class:`~isaaclab.scene.InteractiveSceneCfg` is used, this parameter gets overridden by + :attr:`isaaclab.scene.InteractiveSceneCfg.num_envs` attribute. + """ + + terrain_type: Literal["generator", "plane", "usd"] = "generator" + """The type of terrain to generate. Defaults to "generator". + + Available options are "plane", "usd", and "generator". + """ + + terrain_generator: TerrainGeneratorCfg | None = None + """The terrain generator configuration. + + Only used if ``terrain_type`` is set to "generator". + """ + + usd_path: str | None = None + """The path to the USD file containing the terrain. + + Only used if ``terrain_type`` is set to "usd". + """ + + env_spacing: float | None = None + """The spacing between environment origins when defined in a grid. Defaults to None. + + Note: + This parameter is used only when the ``terrain_type`` is ``"plane"`` or ``"usd"``. + """ + + visual_material: sim_utils.VisualMaterialCfg | None = sim_utils.PreviewSurfaceCfg( + diffuse_color=(0.065, 0.0725, 0.080) + ) + """The visual material of the terrain. Defaults to a dark gray color material. + + The material is created at the path: ``{prim_path}/visualMaterial``. If `None`, then no material is created. + + .. note:: + This parameter is used only when the ``terrain_type`` is ``"generator"``. + """ + + physics_material: sim_utils.RigidBodyMaterialCfg = sim_utils.RigidBodyMaterialCfg() + """The physics material of the terrain. Defaults to a default physics material. + + The material is created at the path: ``{prim_path}/physicsMaterial``. + + .. note:: + This parameter is used only when the ``terrain_type`` is ``"generator"`` or ``"plane"``. + """ + + max_init_terrain_level: int | None = None + """The maximum initial terrain level for defining environment origins. Defaults to None. + + The terrain levels are specified by the number of rows in the grid arrangement of + sub-terrains. If None, then the initial terrain level is set to the maximum + terrain level available (``num_rows - 1``). + + Note: + This parameter is used only when sub-terrain origins are defined. + """ + + debug_vis: bool = False + """Whether to enable visualization of terrain origins for the terrain. Defaults to False.""" diff --git a/source/uwlab/uwlab/terrains/trimesh/__init__.py b/source/uwlab/uwlab/terrains/trimesh/__init__.py new file mode 100644 index 00000000..630ed9c6 --- /dev/null +++ b/source/uwlab/uwlab/terrains/trimesh/__init__.py @@ -0,0 +1,40 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +""" +This sub-module provides methods to create different terrains using the ``trimesh`` library. + +In contrast to the height-field representation, the trimesh representation does not +create arbitrarily small triangles. Instead, the terrain is represented as a single +tri-mesh primitive. Thus, this representation is more computationally and memory +efficient than the height-field representation, but it is not as flexible. +""" + +from .basic_mesh_terrains_cfg import ( + MeshBoxTerrainCfg, + MeshFloatingRingTerrainCfg, + MeshGapTerrainCfg, + MeshInvertedPyramidStairsTerrainCfg, + MeshPitTerrainCfg, + MeshPlaneTerrainCfg, + MeshPyramidStairsTerrainCfg, + MeshRailsTerrainCfg, + MeshRandomGridTerrainCfg, + MeshRepeatedBoxesTerrainCfg, + MeshRepeatedCylindersTerrainCfg, + MeshRepeatedPyramidsTerrainCfg, + MeshStarTerrainCfg, +) +from .mesh_terrains_cfg import ( + CachedTerrainGenCfg, + MeshBalanceBeamsTerrainCfg, + MeshDiversityBoxTerrainCfg, + MeshObjTerrainCfg, + MeshPassageTerrainCfg, + MeshSteppingBeamsTerrainCfg, + MeshStonesEverywhereTerrainCfg, + MeshStructuredTerrainCfg, + TerrainGenCfg, +) diff --git a/source/uwlab/uwlab/terrains/trimesh/basic_mesh_terrains.py b/source/uwlab/uwlab/terrains/trimesh/basic_mesh_terrains.py new file mode 100644 index 00000000..6abbd817 --- /dev/null +++ b/source/uwlab/uwlab/terrains/trimesh/basic_mesh_terrains.py @@ -0,0 +1,859 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Functions to generate different terrains using the ``trimesh`` library.""" + +from __future__ import annotations + +import numpy as np +import scipy.spatial.transform as tf +import torch +import trimesh +from typing import TYPE_CHECKING + +from .utils import * # noqa: F401, F403 +from .utils import make_border, make_plane + +if TYPE_CHECKING: + from . import basic_mesh_terrains_cfg + + +def flat_terrain( + difficulty: float, cfg: basic_mesh_terrains_cfg.MeshPlaneTerrainCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray]: + """Generate a flat terrain as a plane. + + .. image:: ../../_static/terrains/trimesh/flat_terrain.jpg + :width: 45% + :align: center + + Note: + The :obj:`difficulty` parameter is ignored for this terrain. + + Args: + difficulty: The difficulty of the terrain. This is a value between 0 and 1. + cfg: The configuration for the terrain. + + Returns: + A tuple containing the tri-mesh of the terrain and the origin of the terrain (in m). + """ + # compute the position of the terrain + origin = (cfg.size[0] / 2.0, cfg.size[1] / 2.0, 0.0) + # compute the vertices of the terrain + plane_mesh = make_plane(cfg.size, 0.0, center_zero=False) + # return the tri-mesh and the position + return [plane_mesh], np.array(origin) + + +def pyramid_stairs_terrain( + difficulty: float, cfg: basic_mesh_terrains_cfg.MeshPyramidStairsTerrainCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray]: + """Generate a terrain with a pyramid stair pattern. + + The terrain is a pyramid stair pattern which trims to a flat platform at the center of the terrain. + + If :obj:`cfg.holes` is True, the terrain will have pyramid stairs of length or width + :obj:`cfg.platform_width` (depending on the direction) with no steps in the remaining area. Additionally, + no border will be added. + + .. image:: ../../_static/terrains/trimesh/pyramid_stairs_terrain.jpg + :width: 45% + + .. image:: ../../_static/terrains/trimesh/pyramid_stairs_terrain_with_holes.jpg + :width: 45% + + Args: + difficulty: The difficulty of the terrain. This is a value between 0 and 1. + cfg: The configuration for the terrain. + + Returns: + A tuple containing the tri-mesh of the terrain and the origin of the terrain (in m). + """ + # resolve the terrain configuration + step_height = cfg.step_height_range[0] + difficulty * (cfg.step_height_range[1] - cfg.step_height_range[0]) + + # compute number of steps in x and y direction + num_steps_x = (cfg.size[0] - 2 * cfg.border_width - cfg.platform_width) // (2 * cfg.step_width) + 1 + num_steps_y = (cfg.size[1] - 2 * cfg.border_width - cfg.platform_width) // (2 * cfg.step_width) + 1 + # we take the minimum number of steps in x and y direction + num_steps = int(min(num_steps_x, num_steps_y)) + + # initialize list of meshes + meshes_list = list() + + # generate the border if needed + if cfg.border_width > 0.0 and not cfg.holes: + # obtain a list of meshes for the border + border_center = [0.5 * cfg.size[0], 0.5 * cfg.size[1], -step_height / 2] + border_inner_size = (cfg.size[0] - 2 * cfg.border_width, cfg.size[1] - 2 * cfg.border_width) + make_borders = make_border(cfg.size, border_inner_size, step_height, border_center) + # add the border meshes to the list of meshes + meshes_list += make_borders + + # generate the terrain + # -- compute the position of the center of the terrain + terrain_center = [0.5 * cfg.size[0], 0.5 * cfg.size[1], 0.0] + terrain_size = (cfg.size[0] - 2 * cfg.border_width, cfg.size[1] - 2 * cfg.border_width) + # -- generate the stair pattern + for k in range(num_steps): + # check if we need to add holes around the steps + if cfg.holes: + box_size = (cfg.platform_width, cfg.platform_width) + else: + box_size = (terrain_size[0] - 2 * k * cfg.step_width, terrain_size[1] - 2 * k * cfg.step_width) + # compute the quantities of the box + # -- location + box_z = terrain_center[2] + k * step_height / 2.0 + box_offset = (k + 0.5) * cfg.step_width + # -- dimensions + box_height = (k + 2) * step_height + # generate the boxes + # top/bottom + box_dims = (box_size[0], cfg.step_width, box_height) + # -- top + box_pos = (terrain_center[0], terrain_center[1] + terrain_size[1] / 2.0 - box_offset, box_z) + box_top = trimesh.creation.box(box_dims, trimesh.transformations.translation_matrix(box_pos)) + # -- bottom + box_pos = (terrain_center[0], terrain_center[1] - terrain_size[1] / 2.0 + box_offset, box_z) + box_bottom = trimesh.creation.box(box_dims, trimesh.transformations.translation_matrix(box_pos)) + # right/left + if cfg.holes: + box_dims = (cfg.step_width, box_size[1], box_height) + else: + box_dims = (cfg.step_width, box_size[1] - 2 * cfg.step_width, box_height) + # -- right + box_pos = (terrain_center[0] + terrain_size[0] / 2.0 - box_offset, terrain_center[1], box_z) + box_right = trimesh.creation.box(box_dims, trimesh.transformations.translation_matrix(box_pos)) + # -- left + box_pos = (terrain_center[0] - terrain_size[0] / 2.0 + box_offset, terrain_center[1], box_z) + box_left = trimesh.creation.box(box_dims, trimesh.transformations.translation_matrix(box_pos)) + # add the boxes to the list of meshes + meshes_list += [box_top, box_bottom, box_right, box_left] + + # generate final box for the middle of the terrain + box_dims = ( + terrain_size[0] - 2 * num_steps * cfg.step_width, + terrain_size[1] - 2 * num_steps * cfg.step_width, + (num_steps + 2) * step_height, + ) + box_pos = (terrain_center[0], terrain_center[1], terrain_center[2] + num_steps * step_height / 2) + box_middle = trimesh.creation.box(box_dims, trimesh.transformations.translation_matrix(box_pos)) + meshes_list.append(box_middle) + # origin of the terrain + origin = np.array([terrain_center[0], terrain_center[1], (num_steps + 1) * step_height]) + + return meshes_list, origin + + +def inverted_pyramid_stairs_terrain( + difficulty: float, cfg: basic_mesh_terrains_cfg.MeshInvertedPyramidStairsTerrainCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray]: + """Generate a terrain with a inverted pyramid stair pattern. + + The terrain is an inverted pyramid stair pattern which trims to a flat platform at the center of the terrain. + + If :obj:`cfg.holes` is True, the terrain will have pyramid stairs of length or width + :obj:`cfg.platform_width` (depending on the direction) with no steps in the remaining area. Additionally, + no border will be added. + + .. image:: ../../_static/terrains/trimesh/inverted_pyramid_stairs_terrain.jpg + :width: 45% + + .. image:: ../../_static/terrains/trimesh/inverted_pyramid_stairs_terrain_with_holes.jpg + :width: 45% + + Args: + difficulty: The difficulty of the terrain. This is a value between 0 and 1. + cfg: The configuration for the terrain. + + Returns: + A tuple containing the tri-mesh of the terrain and the origin of the terrain (in m). + """ + # resolve the terrain configuration + step_height = cfg.step_height_range[0] + difficulty * (cfg.step_height_range[1] - cfg.step_height_range[0]) + + # compute number of steps in x and y direction + num_steps_x = (cfg.size[0] - 2 * cfg.border_width - cfg.platform_width) // (2 * cfg.step_width) + 1 + num_steps_y = (cfg.size[1] - 2 * cfg.border_width - cfg.platform_width) // (2 * cfg.step_width) + 1 + # we take the minimum number of steps in x and y direction + num_steps = int(min(num_steps_x, num_steps_y)) + # total height of the terrain + total_height = (num_steps + 1) * step_height + + # initialize list of meshes + meshes_list = list() + + # generate the border if needed + if cfg.border_width > 0.0 and not cfg.holes: + # obtain a list of meshes for the border + border_center = [0.5 * cfg.size[0], 0.5 * cfg.size[1], -0.5 * step_height] + border_inner_size = (cfg.size[0] - 2 * cfg.border_width, cfg.size[1] - 2 * cfg.border_width) + make_borders = make_border(cfg.size, border_inner_size, step_height, border_center) + # add the border meshes to the list of meshes + meshes_list += make_borders + # generate the terrain + # -- compute the position of the center of the terrain + terrain_center = [0.5 * cfg.size[0], 0.5 * cfg.size[1], 0.0] + terrain_size = (cfg.size[0] - 2 * cfg.border_width, cfg.size[1] - 2 * cfg.border_width) + # -- generate the stair pattern + for k in range(num_steps): + # check if we need to add holes around the steps + if cfg.holes: + box_size = (cfg.platform_width, cfg.platform_width) + else: + box_size = (terrain_size[0] - 2 * k * cfg.step_width, terrain_size[1] - 2 * k * cfg.step_width) + # compute the quantities of the box + # -- location + box_z = terrain_center[2] - total_height / 2 - (k + 1) * step_height / 2.0 + box_offset = (k + 0.5) * cfg.step_width + # -- dimensions + box_height = total_height - (k + 1) * step_height + # generate the boxes + # top/bottom + box_dims = (box_size[0], cfg.step_width, box_height) + # -- top + box_pos = (terrain_center[0], terrain_center[1] + terrain_size[1] / 2.0 - box_offset, box_z) + box_top = trimesh.creation.box(box_dims, trimesh.transformations.translation_matrix(box_pos)) + # -- bottom + box_pos = (terrain_center[0], terrain_center[1] - terrain_size[1] / 2.0 + box_offset, box_z) + box_bottom = trimesh.creation.box(box_dims, trimesh.transformations.translation_matrix(box_pos)) + # right/left + if cfg.holes: + box_dims = (cfg.step_width, box_size[1], box_height) + else: + box_dims = (cfg.step_width, box_size[1] - 2 * cfg.step_width, box_height) + # -- right + box_pos = (terrain_center[0] + terrain_size[0] / 2.0 - box_offset, terrain_center[1], box_z) + box_right = trimesh.creation.box(box_dims, trimesh.transformations.translation_matrix(box_pos)) + # -- left + box_pos = (terrain_center[0] - terrain_size[0] / 2.0 + box_offset, terrain_center[1], box_z) + box_left = trimesh.creation.box(box_dims, trimesh.transformations.translation_matrix(box_pos)) + # add the boxes to the list of meshes + meshes_list += [box_top, box_bottom, box_right, box_left] + # generate final box for the middle of the terrain + box_dims = ( + terrain_size[0] - 2 * num_steps * cfg.step_width, + terrain_size[1] - 2 * num_steps * cfg.step_width, + step_height, + ) + box_pos = (terrain_center[0], terrain_center[1], terrain_center[2] - total_height - step_height / 2) + box_middle = trimesh.creation.box(box_dims, trimesh.transformations.translation_matrix(box_pos)) + meshes_list.append(box_middle) + # origin of the terrain + origin = np.array([terrain_center[0], terrain_center[1], -(num_steps + 1) * step_height]) + + return meshes_list, origin + + +def random_grid_terrain( + difficulty: float, cfg: basic_mesh_terrains_cfg.MeshRandomGridTerrainCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray]: + """Generate a terrain with cells of random heights and fixed width. + + The terrain is generated in the x-y plane and has a height of 1.0. It is then divided into a grid of the + specified size :obj:`cfg.grid_width`. Each grid cell is then randomly shifted in the z-direction by a value uniformly + sampled between :obj:`cfg.grid_height_range`. At the center of the terrain, a platform of the specified width + :obj:`cfg.platform_width` is generated. + + If :obj:`cfg.holes` is True, the terrain will have randomized grid cells only along the plane extending + from the platform (like a plus sign). The remaining area remains empty and no border will be added. + + .. image:: ../../_static/terrains/trimesh/random_grid_terrain.jpg + :width: 45% + + .. image:: ../../_static/terrains/trimesh/random_grid_terrain_with_holes.jpg + :width: 45% + + Args: + difficulty: The difficulty of the terrain. This is a value between 0 and 1. + cfg: The configuration for the terrain. + + Returns: + A tuple containing the tri-mesh of the terrain and the origin of the terrain (in m). + + Raises: + ValueError: If the terrain is not square. This method only supports square terrains. + RuntimeError: If the grid width is large such that the border width is negative. + """ + # check to ensure square terrain + if cfg.size[0] != cfg.size[1]: + raise ValueError(f"The terrain must be square. Received size: {cfg.size}.") + # resolve the terrain configuration + grid_height = cfg.grid_height_range[0] + difficulty * (cfg.grid_height_range[1] - cfg.grid_height_range[0]) + + # initialize list of meshes + meshes_list = list() + # compute the number of boxes in each direction + num_boxes_x = int(cfg.size[0] / cfg.grid_width) + num_boxes_y = int(cfg.size[1] / cfg.grid_width) + # constant parameters + terrain_height = 1.0 + device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + + # generate the border + border_width = cfg.size[0] - min(num_boxes_x, num_boxes_y) * cfg.grid_width + if border_width > 0: + # compute parameters for the border + border_center = (0.5 * cfg.size[0], 0.5 * cfg.size[1], -terrain_height / 2) + border_inner_size = (cfg.size[0] - border_width, cfg.size[1] - border_width) + # create border meshes + make_borders = make_border(cfg.size, border_inner_size, terrain_height, border_center) + meshes_list += make_borders + else: + raise RuntimeError("Border width must be greater than 0! Adjust the parameter 'cfg.grid_width'.") + + # create a template grid of terrain height + grid_dim = [cfg.grid_width, cfg.grid_width, terrain_height] + grid_position = [0.5 * cfg.grid_width, 0.5 * cfg.grid_width, -terrain_height / 2] + template_box = trimesh.creation.box(grid_dim, trimesh.transformations.translation_matrix(grid_position)) + # extract vertices and faces of the box to create a template + template_vertices = template_box.vertices # (8, 3) + template_faces = template_box.faces + + # repeat the template box vertices to span the terrain (num_boxes_x * num_boxes_y, 8, 3) + vertices = torch.tensor(template_vertices, device=device).repeat(num_boxes_x * num_boxes_y, 1, 1) + # create a meshgrid to offset the vertices + x = torch.arange(0, num_boxes_x, device=device) + y = torch.arange(0, num_boxes_y, device=device) + xx, yy = torch.meshgrid(x, y, indexing="ij") + xx = xx.flatten().view(-1, 1) + yy = yy.flatten().view(-1, 1) + xx_yy = torch.cat((xx, yy), dim=1) + # offset the vertices + offsets = cfg.grid_width * xx_yy + border_width / 2 + vertices[:, :, :2] += offsets.unsqueeze(1) + # mask the vertices to create holes, s.t. only grids along the x and y axis are present + if cfg.holes: + # -- x-axis + mask_x = torch.logical_and( + (vertices[:, :, 0] > (cfg.size[0] - border_width - cfg.platform_width) / 2).all(dim=1), + (vertices[:, :, 0] < (cfg.size[0] + border_width + cfg.platform_width) / 2).all(dim=1), + ) + vertices_x = vertices[mask_x] + # -- y-axis + mask_y = torch.logical_and( + (vertices[:, :, 1] > (cfg.size[1] - border_width - cfg.platform_width) / 2).all(dim=1), + (vertices[:, :, 1] < (cfg.size[1] + border_width + cfg.platform_width) / 2).all(dim=1), + ) + vertices_y = vertices[mask_y] + # -- combine these vertices + vertices = torch.cat((vertices_x, vertices_y)) + # add noise to the vertices to have a random height over each grid cell + num_boxes = len(vertices) + # create noise for the z-axis + h_noise = torch.zeros((num_boxes, 3), device=device) + h_noise[:, 2].uniform_(-grid_height, grid_height) + # reshape noise to match the vertices (num_boxes, 4, 3) + # only the top vertices of the box are affected + vertices_noise = torch.zeros((num_boxes, 4, 3), device=device) + vertices_noise += h_noise.unsqueeze(1) + # add height only to the top vertices of the box + vertices[vertices[:, :, 2] == 0] += vertices_noise.view(-1, 3) + # move to numpy + vertices = vertices.reshape(-1, 3).cpu().numpy() + + # create faces for boxes (num_boxes, 12, 3). Each box has 6 faces, each face has 2 triangles. + faces = torch.tensor(template_faces, device=device).repeat(num_boxes, 1, 1) + face_offsets = torch.arange(0, num_boxes, device=device).unsqueeze(1).repeat(1, 12) * 8 + faces += face_offsets.unsqueeze(2) + # move to numpy + faces = faces.view(-1, 3).cpu().numpy() + # convert to trimesh + grid_mesh = trimesh.Trimesh(vertices=vertices, faces=faces) + meshes_list.append(grid_mesh) + + # add a platform in the center of the terrain that is accessible from all sides + dim = (cfg.platform_width, cfg.platform_width, terrain_height + grid_height) + pos = (0.5 * cfg.size[0], 0.5 * cfg.size[1], -terrain_height / 2 + grid_height / 2) + box_platform = trimesh.creation.box(dim, trimesh.transformations.translation_matrix(pos)) + meshes_list.append(box_platform) + + # specify the origin of the terrain + origin = np.array([0.5 * cfg.size[0], 0.5 * cfg.size[1], grid_height]) + + return meshes_list, origin + + +def rails_terrain( + difficulty: float, cfg: basic_mesh_terrains_cfg.MeshRailsTerrainCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray]: + """Generate a terrain with box rails as extrusions. + + The terrain contains two sets of box rails created as extrusions. The first set (inner rails) is extruded from + the platform at the center of the terrain, and the second set is extruded between the first set of rails + and the terrain border. Each set of rails is extruded to the same height. + + .. image:: ../../_static/terrains/trimesh/rails_terrain.jpg + :width: 40% + :align: center + + Args: + difficulty: The difficulty of the terrain. this is a value between 0 and 1. + cfg: The configuration for the terrain. + + Returns: + A tuple containing the tri-mesh of the terrain and the origin of the terrain (in m). + """ + # resolve the terrain configuration + rail_height = cfg.rail_height_range[1] - difficulty * (cfg.rail_height_range[1] - cfg.rail_height_range[0]) + + # initialize list of meshes + meshes_list = list() + # extract quantities + rail_1_thickness, rail_2_thickness = cfg.rail_thickness_range + rail_center = (0.5 * cfg.size[0], 0.5 * cfg.size[1], rail_height * 0.5) + # constants for terrain generation + terrain_height = 1.0 + rail_2_ratio = 0.6 + + # generate first set of rails + rail_1_inner_size = (cfg.platform_width, cfg.platform_width) + rail_1_outer_size = (cfg.platform_width + 2.0 * rail_1_thickness, cfg.platform_width + 2.0 * rail_1_thickness) + meshes_list += make_border(rail_1_outer_size, rail_1_inner_size, rail_height, rail_center) + # generate second set of rails + rail_2_inner_x = cfg.platform_width + (cfg.size[0] - cfg.platform_width) * rail_2_ratio + rail_2_inner_y = cfg.platform_width + (cfg.size[1] - cfg.platform_width) * rail_2_ratio + rail_2_inner_size = (rail_2_inner_x, rail_2_inner_y) + rail_2_outer_size = (rail_2_inner_x + 2.0 * rail_2_thickness, rail_2_inner_y + 2.0 * rail_2_thickness) + meshes_list += make_border(rail_2_outer_size, rail_2_inner_size, rail_height, rail_center) + # generate the ground + dim = (cfg.size[0], cfg.size[1], terrain_height) + pos = (0.5 * cfg.size[0], 0.5 * cfg.size[1], -terrain_height / 2) + ground_meshes = trimesh.creation.box(dim, trimesh.transformations.translation_matrix(pos)) + meshes_list.append(ground_meshes) + + # specify the origin of the terrain + origin = np.array([pos[0], pos[1], 0.0]) + + return meshes_list, origin + + +def pit_terrain( + difficulty: float, cfg: basic_mesh_terrains_cfg.MeshPitTerrainCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray]: + """Generate a terrain with a pit with levels (stairs) leading out of the pit. + + The terrain contains a platform at the center and a staircase leading out of the pit. + The staircase is a series of steps that are aligned along the x- and y- axis. The steps are + created by extruding a ring along the x- and y- axis. If :obj:`is_double_pit` is True, the pit + contains two levels. + + .. image:: ../../_static/terrains/trimesh/pit_terrain.jpg + :width: 40% + + .. image:: ../../_static/terrains/trimesh/pit_terrain_with_two_levels.jpg + :width: 40% + + Args: + difficulty: The difficulty of the terrain. This is a value between 0 and 1. + cfg: The configuration for the terrain. + + Returns: + A tuple containing the tri-mesh of the terrain and the origin of the terrain (in m). + """ + # resolve the terrain configuration + pit_depth = cfg.pit_depth_range[0] + difficulty * (cfg.pit_depth_range[1] - cfg.pit_depth_range[0]) + + # initialize list of meshes + meshes_list = list() + # extract quantities + inner_pit_size = (cfg.platform_width, cfg.platform_width) + total_depth = pit_depth + # constants for terrain generation + terrain_height = 1.0 + ring_2_ratio = 0.6 + + # if the pit is double, the inner ring is smaller to fit the second level + if cfg.double_pit: + # increase the total height of the pit + total_depth *= 2.0 + # reduce the size of the inner ring + inner_pit_x = cfg.platform_width + (cfg.size[0] - cfg.platform_width) * ring_2_ratio + inner_pit_y = cfg.platform_width + (cfg.size[1] - cfg.platform_width) * ring_2_ratio + inner_pit_size = (inner_pit_x, inner_pit_y) + + # generate the pit (outer ring) + pit_center = [0.5 * cfg.size[0], 0.5 * cfg.size[1], -total_depth * 0.5] + meshes_list += make_border(cfg.size, inner_pit_size, total_depth, pit_center) + # generate the second level of the pit (inner ring) + if cfg.double_pit: + pit_center[2] = -total_depth + meshes_list += make_border(inner_pit_size, (cfg.platform_width, cfg.platform_width), total_depth, pit_center) + # generate the ground + dim = (cfg.size[0], cfg.size[1], terrain_height) + pos = (0.5 * cfg.size[0], 0.5 * cfg.size[1], -total_depth - terrain_height / 2) + ground_meshes = trimesh.creation.box(dim, trimesh.transformations.translation_matrix(pos)) + meshes_list.append(ground_meshes) + + # specify the origin of the terrain + origin = np.array([pos[0], pos[1], -total_depth]) + + return meshes_list, origin + + +def box_terrain( + difficulty: float, cfg: basic_mesh_terrains_cfg.MeshBoxTerrainCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray]: + """Generate a terrain with boxes (similar to a pyramid). + + The terrain has a ground with boxes on top of it that are stacked on top of each other. + The boxes are created by extruding a rectangle along the z-axis. If :obj:`double_box` is True, + then two boxes of height :obj:`box_height` are stacked on top of each other. + + .. image:: ../../_static/terrains/trimesh/box_terrain.jpg + :width: 40% + + .. image:: ../../_static/terrains/trimesh/box_terrain_with_two_boxes.jpg + :width: 40% + + Args: + difficulty: The difficulty of the terrain. This is a value between 0 and 1. + cfg: The configuration for the terrain. + + Returns: + A tuple containing the tri-mesh of the terrain and the origin of the terrain (in m). + """ + # resolve the terrain configuration + box_height = cfg.box_height_range[0] + difficulty * (cfg.box_height_range[1] - cfg.box_height_range[0]) + + # initialize list of meshes + meshes_list = list() + # extract quantities + total_height = box_height + if cfg.double_box: + total_height *= 2.0 + # constants for terrain generation + terrain_height = 1.0 + box_2_ratio = 0.6 + + # Generate the top box + dim = (cfg.platform_width, cfg.platform_width, terrain_height + total_height) + pos = (0.5 * cfg.size[0], 0.5 * cfg.size[1], (total_height - terrain_height) / 2) + box_mesh = trimesh.creation.box(dim, trimesh.transformations.translation_matrix(pos)) + meshes_list.append(box_mesh) + # Generate the lower box + if cfg.double_box: + # calculate the size of the lower box + outer_box_x = cfg.platform_width + (cfg.size[0] - cfg.platform_width) * box_2_ratio + outer_box_y = cfg.platform_width + (cfg.size[1] - cfg.platform_width) * box_2_ratio + # create the lower box + dim = (outer_box_x, outer_box_y, terrain_height + total_height / 2) + pos = (0.5 * cfg.size[0], 0.5 * cfg.size[1], (total_height - terrain_height) / 2 - total_height / 4) + box_mesh = trimesh.creation.box(dim, trimesh.transformations.translation_matrix(pos)) + meshes_list.append(box_mesh) + # Generate the ground + pos = (0.5 * cfg.size[0], 0.5 * cfg.size[1], -terrain_height / 2) + dim = (cfg.size[0], cfg.size[1], terrain_height) + ground_mesh = trimesh.creation.box(dim, trimesh.transformations.translation_matrix(pos)) + meshes_list.append(ground_mesh) + + # specify the origin of the terrain + origin = np.array([pos[0], pos[1], total_height]) + + return meshes_list, origin + + +def gap_terrain( + difficulty: float, cfg: basic_mesh_terrains_cfg.MeshGapTerrainCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray]: + """Generate a terrain with a gap around the platform. + + The terrain has a ground with a platform in the middle. The platform is surrounded by a gap + of width :obj:`gap_width` on all sides. + + .. image:: ../../_static/terrains/trimesh/gap_terrain.jpg + :width: 40% + :align: center + + Args: + difficulty: The difficulty of the terrain. This is a value between 0 and 1. + cfg: The configuration for the terrain. + + Returns: + A tuple containing the tri-mesh of the terrain and the origin of the terrain (in m). + """ + # resolve the terrain configuration + gap_width = cfg.gap_width_range[0] + difficulty * (cfg.gap_width_range[1] - cfg.gap_width_range[0]) + + # initialize list of meshes + meshes_list = list() + # constants for terrain generation + terrain_height = 1.0 + terrain_center = (0.5 * cfg.size[0], 0.5 * cfg.size[1], -terrain_height / 2) + + # Generate the outer ring + inner_size = (cfg.platform_width + 2 * gap_width, cfg.platform_width + 2 * gap_width) + meshes_list += make_border(cfg.size, inner_size, terrain_height, terrain_center) + # Generate the inner box + box_dim = (cfg.platform_width, cfg.platform_width, terrain_height) + box = trimesh.creation.box(box_dim, trimesh.transformations.translation_matrix(terrain_center)) + meshes_list.append(box) + + # specify the origin of the terrain + origin = np.array([terrain_center[0], terrain_center[1], 0.0]) + + return meshes_list, origin + + +def floating_ring_terrain( + difficulty: float, cfg: basic_mesh_terrains_cfg.MeshFloatingRingTerrainCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray]: + """Generate a terrain with a floating square ring. + + The terrain has a ground with a floating ring in the middle. The ring extends from the center from + :obj:`platform_width` to :obj:`platform_width` + :obj:`ring_width` in the x and y directions. + The thickness of the ring is :obj:`ring_thickness` and the height of the ring from the terrain + is :obj:`ring_height`. + + .. image:: ../../_static/terrains/trimesh/floating_ring_terrain.jpg + :width: 40% + :align: center + + Args: + difficulty: The difficulty of the terrain. This is a value between 0 and 1. + cfg: The configuration for the terrain. + + Returns: + A tuple containing the tri-mesh of the terrain and the origin of the terrain (in m). + """ + # resolve the terrain configuration + ring_height = cfg.ring_height_range[1] - difficulty * (cfg.ring_height_range[1] - cfg.ring_height_range[0]) + ring_width = cfg.ring_width_range[0] + difficulty * (cfg.ring_width_range[1] - cfg.ring_width_range[0]) + + # initialize list of meshes + meshes_list = list() + # constants for terrain generation + terrain_height = 1.0 + + # Generate the floating ring + ring_center = (0.5 * cfg.size[0], 0.5 * cfg.size[1], ring_height + 0.5 * cfg.ring_thickness) + ring_outer_size = (cfg.platform_width + 2 * ring_width, cfg.platform_width + 2 * ring_width) + ring_inner_size = (cfg.platform_width, cfg.platform_width) + meshes_list += make_border(ring_outer_size, ring_inner_size, cfg.ring_thickness, ring_center) + # Generate the ground + dim = (cfg.size[0], cfg.size[1], terrain_height) + pos = (0.5 * cfg.size[0], 0.5 * cfg.size[1], -terrain_height / 2) + ground = trimesh.creation.box(dim, trimesh.transformations.translation_matrix(pos)) + meshes_list.append(ground) + + # specify the origin of the terrain + origin = np.asarray([pos[0], pos[1], 0.0]) + + return meshes_list, origin + + +def star_terrain( + difficulty: float, cfg: basic_mesh_terrains_cfg.MeshStarTerrainCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray]: + """Generate a terrain with a star. + + The terrain has a ground with a cylinder in the middle. The star is made of :obj:`num_bars` bars + with a width of :obj:`bar_width` and a height of :obj:`bar_height`. The bars are evenly + spaced around the cylinder and connect to the peripheral of the terrain. + + .. image:: ../../_static/terrains/trimesh/star_terrain.jpg + :width: 40% + :align: center + + Args: + difficulty: The difficulty of the terrain. This is a value between 0 and 1. + cfg: The configuration for the terrain. + + Returns: + A tuple containing the tri-mesh of the terrain and the origin of the terrain (in m). + + Raises: + ValueError: If :obj:`num_bars` is less than 2. + """ + # check the number of bars + if cfg.num_bars < 2: + raise ValueError(f"The number of bars in the star must be greater than 2. Received: {cfg.num_bars}") + + # resolve the terrain configuration + bar_height = cfg.bar_height_range[0] + difficulty * (cfg.bar_height_range[1] - cfg.bar_height_range[0]) + bar_width = cfg.bar_width_range[1] - difficulty * (cfg.bar_width_range[1] - cfg.bar_width_range[0]) + + # initialize list of meshes + meshes_list = list() + # Generate a platform in the middle + platform_center = (0.5 * cfg.size[0], 0.5 * cfg.size[1], -bar_height / 2) + platform_transform = trimesh.transformations.translation_matrix(platform_center) + platform = trimesh.creation.cylinder( + cfg.platform_width * 0.5, bar_height, sections=2 * cfg.num_bars, transform=platform_transform + ) + meshes_list.append(platform) + # Generate bars to connect the platform to the terrain + transform = np.eye(4) + transform[:3, -1] = np.asarray(platform_center) + yaw = 0.0 + for _ in range(cfg.num_bars): + # compute the length of the bar based on the yaw + # length changes since the bar is connected to a square border + bar_length = cfg.size[0] + if yaw < 0.25 * np.pi: + bar_length /= np.math.cos(yaw) + elif yaw < 0.75 * np.pi: + bar_length /= np.math.sin(yaw) + else: + bar_length /= np.math.cos(np.pi - yaw) + # compute the transform of the bar + transform[0:3, 0:3] = tf.Rotation.from_euler("z", yaw).as_matrix() + # add the bar to the mesh + dim = [bar_length - bar_width, bar_width, bar_height] + bar = trimesh.creation.box(dim, transform) + meshes_list.append(bar) + # increment the yaw + yaw += np.pi / cfg.num_bars + # Generate the exterior border + inner_size = (cfg.size[0] - 2 * bar_width, cfg.size[1] - 2 * bar_width) + meshes_list += make_border(cfg.size, inner_size, bar_height, platform_center) + # Generate the ground + ground = make_plane(cfg.size, -bar_height, center_zero=False) + meshes_list.append(ground) + # specify the origin of the terrain + origin = np.asarray([0.5 * cfg.size[0], 0.5 * cfg.size[1], 0.0]) + + return meshes_list, origin + + +def repeated_objects_terrain( + difficulty: float, cfg: basic_mesh_terrains_cfg.MeshRepeatedObjectsTerrainCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray]: + """Generate a terrain with a set of repeated objects. + + The terrain has a ground with a platform in the middle. The objects are randomly placed on the + terrain s.t. they do not overlap with the platform. + + Depending on the object type, the objects are generated with different parameters. The objects + The types of objects that can be generated are: ``"cylinder"``, ``"box"``, ``"cone"``. + + The object parameters are specified in the configuration as curriculum parameters. The difficulty + is used to linearly interpolate between the minimum and maximum values of the parameters. + + .. image:: ../../_static/terrains/trimesh/repeated_objects_cylinder_terrain.jpg + :width: 30% + + .. image:: ../../_static/terrains/trimesh/repeated_objects_box_terrain.jpg + :width: 30% + + .. image:: ../../_static/terrains/trimesh/repeated_objects_pyramid_terrain.jpg + :width: 30% + + Args: + difficulty: The difficulty of the terrain. This is a value between 0 and 1. + cfg: The configuration for the terrain. + + Returns: + A tuple containing the tri-mesh of the terrain and the origin of the terrain (in m). + + Raises: + ValueError: If the object type is not supported. It must be either a string or a callable. + """ + # import the object functions -- this is done here to avoid circular imports + from .basic_mesh_terrains_cfg import ( + MeshRepeatedBoxesTerrainCfg, + MeshRepeatedCylindersTerrainCfg, + MeshRepeatedPyramidsTerrainCfg, + ) + + # if object type is a string, get the function: make_{object_type} + if isinstance(cfg.object_type, str): + object_func = globals().get(f"make_{cfg.object_type}") + else: + object_func = cfg.object_type + if not callable(object_func): + raise ValueError(f"The attribute 'object_type' must be a string or a callable. Received: {object_func}") + + # Resolve the terrain configuration + # -- pass parameters to make calling simpler + cp_0 = cfg.object_params_start + cp_1 = cfg.object_params_end + # -- common parameters + num_objects = cp_0.num_objects + int(difficulty * (cp_1.num_objects - cp_0.num_objects)) + height = cp_0.height + difficulty * (cp_1.height - cp_0.height) + # -- object specific parameters + # note: SIM114 requires duplicated logical blocks under a single body. + if isinstance(cfg, MeshRepeatedBoxesTerrainCfg): + cp_0: MeshRepeatedBoxesTerrainCfg.ObjectCfg + cp_1: MeshRepeatedBoxesTerrainCfg.ObjectCfg + object_kwargs = { + "length": cp_0.size[0] + difficulty * (cp_1.size[0] - cp_0.size[0]), + "width": cp_0.size[1] + difficulty * (cp_1.size[1] - cp_0.size[1]), + "max_yx_angle": cp_0.max_yx_angle + difficulty * (cp_1.max_yx_angle - cp_0.max_yx_angle), + "degrees": cp_0.degrees, + } + elif isinstance(cfg, MeshRepeatedPyramidsTerrainCfg): # noqa: SIM114 + cp_0: MeshRepeatedPyramidsTerrainCfg.ObjectCfg + cp_1: MeshRepeatedPyramidsTerrainCfg.ObjectCfg + object_kwargs = { + "radius": cp_0.radius + difficulty * (cp_1.radius - cp_0.radius), + "max_yx_angle": cp_0.max_yx_angle + difficulty * (cp_1.max_yx_angle - cp_0.max_yx_angle), + "degrees": cp_0.degrees, + } + elif isinstance(cfg, MeshRepeatedCylindersTerrainCfg): # noqa: SIM114 + cp_0: MeshRepeatedCylindersTerrainCfg.ObjectCfg + cp_1: MeshRepeatedCylindersTerrainCfg.ObjectCfg + object_kwargs = { + "radius": cp_0.radius + difficulty * (cp_1.radius - cp_0.radius), + "max_yx_angle": cp_0.max_yx_angle + difficulty * (cp_1.max_yx_angle - cp_0.max_yx_angle), + "degrees": cp_0.degrees, + } + else: + raise ValueError(f"Unknown terrain configuration: {cfg}") + # constants for the terrain + platform_clearance = 0.1 + + # initialize list of meshes + meshes_list = list() + # compute quantities + origin = np.asarray((0.5 * cfg.size[0], 0.5 * cfg.size[1], 0.5 * height)) + platform_corners = np.asarray( + [ + [origin[0] - cfg.platform_width / 2, origin[1] - cfg.platform_width / 2], + [origin[0] + cfg.platform_width / 2, origin[1] + cfg.platform_width / 2], + ] + ) + platform_corners[0, :] *= 1 - platform_clearance + platform_corners[1, :] *= 1 + platform_clearance + # sample valid center for objects + object_centers = np.zeros((num_objects, 3)) + # use a mask to track invalid objects that still require sampling + mask_objects_left = np.ones((num_objects,), dtype=bool) + # loop until no objects are left to sample + while np.any(mask_objects_left): + # only sample the centers of the remaining invalid objects + num_objects_left = mask_objects_left.sum() + object_centers[mask_objects_left, 0] = np.random.uniform(0, cfg.size[0], num_objects_left) + object_centers[mask_objects_left, 1] = np.random.uniform(0, cfg.size[1], num_objects_left) + # filter out the centers that are on the platform + is_within_platform_x = np.logical_and( + object_centers[mask_objects_left, 0] >= platform_corners[0, 0], + object_centers[mask_objects_left, 0] <= platform_corners[1, 0], + ) + is_within_platform_y = np.logical_and( + object_centers[mask_objects_left, 1] >= platform_corners[0, 1], + object_centers[mask_objects_left, 1] <= platform_corners[1, 1], + ) + # update the mask to track the validity of the objects sampled in this iteration + mask_objects_left[mask_objects_left] = np.logical_and(is_within_platform_x, is_within_platform_y) + + # generate obstacles (but keep platform clean) + for index in range(len(object_centers)): + # randomize the height of the object + ob_height = height + np.random.uniform(-cfg.max_height_noise, cfg.max_height_noise) + if ob_height > 0.0: + object_mesh = object_func(center=object_centers[index], height=ob_height, **object_kwargs) + meshes_list.append(object_mesh) + + # generate a ground plane for the terrain + ground_plane = make_plane(cfg.size, height=0.0, center_zero=False) + meshes_list.append(ground_plane) + # generate a platform in the middle + dim = (cfg.platform_width, cfg.platform_width, 0.5 * height) + pos = (0.5 * cfg.size[0], 0.5 * cfg.size[1], 0.25 * height) + platform = trimesh.creation.box(dim, trimesh.transformations.translation_matrix(pos)) + meshes_list.append(platform) + + return meshes_list, origin diff --git a/source/uwlab/uwlab/terrains/trimesh/basic_mesh_terrains_cfg.py b/source/uwlab/uwlab/terrains/trimesh/basic_mesh_terrains_cfg.py new file mode 100644 index 00000000..bba95af6 --- /dev/null +++ b/source/uwlab/uwlab/terrains/trimesh/basic_mesh_terrains_cfg.py @@ -0,0 +1,269 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import MISSING +from typing import Literal + +import isaaclab.terrains.trimesh.utils as mesh_utils_terrains +from isaaclab.utils import configclass + +from ..terrain_generator_cfg import SubTerrainBaseCfg +from . import basic_mesh_terrains as mesh_terrains + +""" +Different trimesh terrain configurations. +""" + + +@configclass +class MeshPlaneTerrainCfg(SubTerrainBaseCfg): + """Configuration for a plane mesh terrain.""" + + function = mesh_terrains.flat_terrain + + +@configclass +class MeshPyramidStairsTerrainCfg(SubTerrainBaseCfg): + """Configuration for a pyramid stair mesh terrain.""" + + function = mesh_terrains.pyramid_stairs_terrain + + border_width: float = 0.0 + """The width of the border around the terrain (in m). Defaults to 0.0. + + The border is a flat terrain with the same height as the terrain. + """ + step_height_range: tuple[float, float] = MISSING + """The minimum and maximum height of the steps (in m).""" + step_width: float = MISSING + """The width of the steps (in m).""" + platform_width: float = 1.0 + """The width of the square platform at the center of the terrain. Defaults to 1.0.""" + holes: bool = False + """If True, the terrain will have holes in the steps. Defaults to False. + + If :obj:`holes` is True, the terrain will have pyramid stairs of length or width + :obj:`platform_width` (depending on the direction) with no steps in the remaining area. Additionally, + no border will be added. + """ + + +@configclass +class MeshInvertedPyramidStairsTerrainCfg(MeshPyramidStairsTerrainCfg): + """Configuration for an inverted pyramid stair mesh terrain. + + Note: + This is the same as :class:`MeshPyramidStairsTerrainCfg` except that the steps are inverted. + """ + + function = mesh_terrains.inverted_pyramid_stairs_terrain + + +@configclass +class MeshRandomGridTerrainCfg(SubTerrainBaseCfg): + """Configuration for a random grid mesh terrain.""" + + function = mesh_terrains.random_grid_terrain + + grid_width: float = MISSING + """The width of the grid cells (in m).""" + grid_height_range: tuple[float, float] = MISSING + """The minimum and maximum height of the grid cells (in m).""" + platform_width: float = 1.0 + """The width of the square platform at the center of the terrain. Defaults to 1.0.""" + holes: bool = False + """If True, the terrain will have holes in the steps. Defaults to False. + + If :obj:`holes` is True, the terrain will have randomized grid cells only along the plane extending + from the platform (like a plus sign). The remaining area remains empty and no border will be added. + """ + + +@configclass +class MeshRailsTerrainCfg(SubTerrainBaseCfg): + """Configuration for a terrain with box rails as extrusions.""" + + function = mesh_terrains.rails_terrain + + rail_thickness_range: tuple[float, float] = MISSING + """The thickness of the inner and outer rails (in m).""" + rail_height_range: tuple[float, float] = MISSING + """The minimum and maximum height of the rails (in m).""" + platform_width: float = 1.0 + """The width of the square platform at the center of the terrain. Defaults to 1.0.""" + + +@configclass +class MeshPitTerrainCfg(SubTerrainBaseCfg): + """Configuration for a terrain with a pit that leads out of the pit.""" + + function = mesh_terrains.pit_terrain + + pit_depth_range: tuple[float, float] = MISSING + """The minimum and maximum height of the pit (in m).""" + platform_width: float = 1.0 + """The width of the square platform at the center of the terrain. Defaults to 1.0.""" + double_pit: bool = False + """If True, the pit contains two levels of stairs. Defaults to False.""" + + +@configclass +class MeshBoxTerrainCfg(SubTerrainBaseCfg): + """Configuration for a terrain with boxes (similar to a pyramid).""" + + function = mesh_terrains.box_terrain + + box_height_range: tuple[float, float] = MISSING + """The minimum and maximum height of the box (in m).""" + platform_width: float = 1.0 + """The width of the square platform at the center of the terrain. Defaults to 1.0.""" + double_box: bool = False + """If True, the pit contains two levels of stairs/boxes. Defaults to False.""" + + +@configclass +class MeshGapTerrainCfg(SubTerrainBaseCfg): + """Configuration for a terrain with a gap around the platform.""" + + function = mesh_terrains.gap_terrain + + gap_width_range: tuple[float, float] = MISSING + """The minimum and maximum width of the gap (in m).""" + platform_width: float = 1.0 + """The width of the square platform at the center of the terrain. Defaults to 1.0.""" + + +@configclass +class MeshFloatingRingTerrainCfg(SubTerrainBaseCfg): + """Configuration for a terrain with a floating ring around the center.""" + + function = mesh_terrains.floating_ring_terrain + + ring_width_range: tuple[float, float] = MISSING + """The minimum and maximum width of the ring (in m).""" + ring_height_range: tuple[float, float] = MISSING + """The minimum and maximum height of the ring (in m).""" + ring_thickness: float = MISSING + """The thickness (along z) of the ring (in m).""" + platform_width: float = 1.0 + """The width of the square platform at the center of the terrain. Defaults to 1.0.""" + + +@configclass +class MeshStarTerrainCfg(SubTerrainBaseCfg): + """Configuration for a terrain with a star pattern.""" + + function = mesh_terrains.star_terrain + + num_bars: int = MISSING + """The number of bars per-side the star. Must be greater than 2.""" + bar_width_range: tuple[float, float] = MISSING + """The minimum and maximum width of the bars in the star (in m).""" + bar_height_range: tuple[float, float] = MISSING + """The minimum and maximum height of the bars in the star (in m).""" + platform_width: float = 1.0 + """The width of the cylindrical platform at the center of the terrain. Defaults to 1.0.""" + + +@configclass +class MeshRepeatedObjectsTerrainCfg(SubTerrainBaseCfg): + """Base configuration for a terrain with repeated objects.""" + + @configclass + class ObjectCfg: + """Configuration of repeated objects.""" + + num_objects: int = MISSING + """The number of objects to add to the terrain.""" + height: float = MISSING + """The height (along z) of the object (in m).""" + + function = mesh_terrains.repeated_objects_terrain + + object_type: Literal["cylinder", "box", "cone"] | callable = MISSING + """The type of object to generate. + + The type can be a string or a callable. If it is a string, the function will look for a function called + ``make_{object_type}`` in the current module scope. If it is a callable, the function will + use the callable to generate the object. + """ + object_params_start: ObjectCfg = MISSING + """The object curriculum parameters at the start of the curriculum.""" + object_params_end: ObjectCfg = MISSING + """The object curriculum parameters at the end of the curriculum.""" + + max_height_noise: float = 0.0 + """The maximum amount of noise to add to the height of the objects (in m). Defaults to 0.0.""" + platform_width: float = 1.0 + """The width of the cylindrical platform at the center of the terrain. Defaults to 1.0.""" + + +@configclass +class MeshRepeatedPyramidsTerrainCfg(MeshRepeatedObjectsTerrainCfg): + """Configuration for a terrain with repeated pyramids.""" + + @configclass + class ObjectCfg(MeshRepeatedObjectsTerrainCfg.ObjectCfg): + """Configuration for a curriculum of repeated pyramids.""" + + radius: float = MISSING + """The radius of the pyramids (in m).""" + max_yx_angle: float = 0.0 + """The maximum angle along the y and x axis. Defaults to 0.0.""" + degrees: bool = True + """Whether the angle is in degrees. Defaults to True.""" + + object_type = mesh_utils_terrains.make_cone + + object_params_start: ObjectCfg = MISSING + """The object curriculum parameters at the start of the curriculum.""" + object_params_end: ObjectCfg = MISSING + """The object curriculum parameters at the end of the curriculum.""" + + +@configclass +class MeshRepeatedBoxesTerrainCfg(MeshRepeatedObjectsTerrainCfg): + """Configuration for a terrain with repeated boxes.""" + + @configclass + class ObjectCfg(MeshRepeatedObjectsTerrainCfg.ObjectCfg): + """Configuration for repeated boxes.""" + + size: tuple[float, float] = MISSING + """The width (along x) and length (along y) of the box (in m).""" + max_yx_angle: float = 0.0 + """The maximum angle along the y and x axis. Defaults to 0.0.""" + degrees: bool = True + """Whether the angle is in degrees. Defaults to True.""" + + object_type = mesh_utils_terrains.make_box + + object_params_start: ObjectCfg = MISSING + """The box curriculum parameters at the start of the curriculum.""" + object_params_end: ObjectCfg = MISSING + """The box curriculum parameters at the end of the curriculum.""" + + +@configclass +class MeshRepeatedCylindersTerrainCfg(MeshRepeatedObjectsTerrainCfg): + """Configuration for a terrain with repeated cylinders.""" + + @configclass + class ObjectCfg(MeshRepeatedObjectsTerrainCfg.ObjectCfg): + """Configuration for repeated cylinder.""" + + radius: float = MISSING + """The radius of the pyramids (in m).""" + max_yx_angle: float = 0.0 + """The maximum angle along the y and x axis. Defaults to 0.0.""" + degrees: bool = True + """Whether the angle is in degrees. Defaults to True.""" + + object_type = mesh_utils_terrains.make_cylinder + + object_params_start: ObjectCfg = MISSING + """The box curriculum parameters at the start of the curriculum.""" + object_params_end: ObjectCfg = MISSING + """The box curriculum parameters at the end of the curriculum.""" diff --git a/source/uwlab/uwlab/terrains/trimesh/mesh_terrains.py b/source/uwlab/uwlab/terrains/trimesh/mesh_terrains.py new file mode 100644 index 00000000..d9791909 --- /dev/null +++ b/source/uwlab/uwlab/terrains/trimesh/mesh_terrains.py @@ -0,0 +1,645 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Functions to generate different terrains using the ``trimesh`` library.""" + +from __future__ import annotations + +import io +import numpy as np +import os +import random +import subprocess +import torch +import trimesh +import yaml +from scipy.spatial.transform import Rotation as R +from typing import TYPE_CHECKING + +import requests + +from uwlab_assets import UWLAB_CLOUD_ASSETS_DIR + +from isaaclab.terrains.trimesh.mesh_terrains import inverted_pyramid_stairs_terrain, pyramid_stairs_terrain +from isaaclab.terrains.trimesh.mesh_terrains_cfg import MeshInvertedPyramidStairsTerrainCfg, MeshPyramidStairsTerrainCfg +from uwlab.terrains.trimesh.utils import make_border, make_plane + +if TYPE_CHECKING: + from . import mesh_terrains_cfg + + +def obj_terrain( + difficulty: float, cfg: mesh_terrains_cfg.MeshObjTerrainCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray, np.ndarray] | tuple[list[trimesh.Trimesh], np.ndarray]: + mesh: trimesh.Trimesh = trimesh.load(cfg.obj_path) # type: ignore + mesh: trimesh.Trimesh = trimesh.load(cfg.obj_path) # type: ignore + xy_scale = cfg.size / (mesh.bounds[1] - mesh.bounds[0])[:2] + # set the height scale to the average between length and width scale to preserve as much original shap as possible + height_scale = (xy_scale[0] + xy_scale[1]) / 2 + xyz_scale = np.array([*xy_scale, height_scale]) + mesh.apply_scale(xyz_scale) + translation = -mesh.bounds[0] + mesh.apply_translation(translation) + + extend = mesh.bounds[1] - mesh.bounds[0] + origin = (*((extend[:2]) / 2), mesh.bounds[1][2] / 2) + + if isinstance(cfg.spawn_origin_path, str): + spawning_option = np.load(cfg.spawn_origin_path, allow_pickle=True) + spawning_option *= xyz_scale + spawning_option += translation + # insert the center of the terrain as the first indices + # the rest of the indices represents the spawning locations + return [mesh], np.insert(spawning_option, 0, origin, axis=0) + else: + return [mesh], np.array(origin) + + +def terrain_gen( + difficulty: float, cfg: mesh_terrains_cfg.TerrainGenCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray, np.ndarray] | tuple[list[trimesh.Trimesh], np.ndarray]: + terrain_yaml = { + "terrain": { + "shape": [int(cfg.size[0] / 2), int(cfg.size[1] / 2)], + "height": cfg.height, + "levels": cfg.levels, + "include_overhang": cfg.include_overhang, + "all_terrain_styles": ["stair", "ramp", "box", "platform", "random_box", "perlin", "wall"], + "terrain_styles": cfg.terrain_styles, + } + } + terrain_style = "_".join(cfg.terrain_styles) + os.makedirs(os.path.dirname(cfg.yaml_path), exist_ok=True) + yaml_file_path = cfg.yaml_path.replace(".yaml", f"_{terrain_style}.yaml") + with open(yaml_file_path, "w") as file: + yaml.dump(terrain_yaml, file, default_flow_style=False) + + mesh_origin_dir = os.path.dirname(cfg.obj_path) + mesh_dir = os.path.dirname(mesh_origin_dir) + # Prepare the command and arguments for the subprocess + command = [ + "python", + cfg.python_script, + "--input_path", + yaml_file_path, + "--enable_sdf", + "--mesh_dir", + mesh_dir, + "--mesh_name", + f"{terrain_style}", + ] + + # Invoke the subprocess and run the other script + try: + result = subprocess.run(command, check=True, capture_output=True) + print("Subprocess completed successfully!") + print("Output:", result.stdout.decode()) + print("Errors:", result.stderr.decode()) + except subprocess.CalledProcessError as e: + print(f"Subprocess failed with error: {e}") + print(f"Subprocess output: {e.output.decode()}") + print(f"Subprocess stderr: {e.stderr.decode()}") + + return obj_terrain(difficulty, cfg) + + +def cached_terrain_gen( + difficulty: float, cfg: mesh_terrains_cfg.CachedTerrainGenCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray, np.ndarray] | tuple[list[trimesh.Trimesh], np.ndarray]: + terrain_type = cfg.task_descriptor + level = cfg.levels + height = cfg.height + overhang = "overhang_yes" if cfg.include_overhang else "overhang_no" + mesh_id = "mesh_0" + root_path = f"{UWLAB_CLOUD_ASSETS_DIR}/dataset/terrains/dataset/generated_terrain/{terrain_type}/shape_8/height_{height}/level_{level}/{overhang}/{mesh_id}" + + terrain_mesh_path = os.path.join(root_path, "mesh_terrain.obj") + spawnfile_path = os.path.join(root_path, "spawnable_locations.npy") + + mesh: trimesh.Trimesh = load_mesh(terrain_mesh_path) + xy_scale = cfg.size / (mesh.bounds[1] - mesh.bounds[0])[:2] + # set the height scale to the average between length and width scale to preserve as much original shap as possible + height_scale = (xy_scale[0] + xy_scale[1]) / 2 + xyz_scale = np.array([*xy_scale, height_scale]) + mesh.apply_scale(xyz_scale) + translation = -mesh.bounds[0] + mesh.apply_translation(translation) + + extend = mesh.bounds[1] - mesh.bounds[0] + origin = (*((extend[:2]) / 2), mesh.bounds[1][2] / 2) + + if isinstance(spawnfile_path, str): + spawning_option = load_numpy(spawnfile_path) + spawning_option *= xyz_scale + spawning_option += translation + # insert the center of the terrain as the first indices + # the rest of the indices represents the spawning locations + return [mesh], np.insert(spawning_option, 0, origin, axis=0) + else: + return [mesh], np.array(origin) + + +def load_mesh(terrain_mesh_path: str) -> trimesh.Trimesh: + """Load a mesh from a URL or a local file.""" + if terrain_mesh_path.startswith("http"): + # Load from URL + response = requests.get(terrain_mesh_path) + if response.status_code == 200: + mesh = trimesh.load(io.BytesIO(response.content), file_type="obj") + return mesh # type: ignore + return mesh # type: ignore + else: + raise Exception(f"Failed to load mesh from {terrain_mesh_path}") + else: + # Load from local path + return trimesh.load(terrain_mesh_path) # type: ignore + + return trimesh.load(terrain_mesh_path) # type: ignore + + +def load_numpy(spawnfile_path: str) -> np.ndarray: + """Load a NumPy array from a URL or a local file.""" + if spawnfile_path.startswith("http"): + # Load from URL + response = requests.get(spawnfile_path) + if response.status_code == 200: + data = np.load(io.BytesIO(response.content), allow_pickle=True) + return data + else: + raise Exception(f"Failed to load NumPy file from {spawnfile_path}") + else: + # Load from local path + return np.load(spawnfile_path, allow_pickle=True) + + +def stones_everywhere_terrain( + difficulty: float, cfg: mesh_terrains_cfg.MeshStonesEverywhereTerrainCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray]: + # check to ensure square terrain + assert cfg.size[0] == cfg.size[1], "The terrain should be square" + + # resolve the terrain configuration based on the difficulty + gap_width = cfg.w_gap[0] + difficulty * (cfg.w_gap[1] - cfg.w_gap[0]) + stone_width = cfg.w_stone[0] - difficulty * (cfg.w_stone[0] - cfg.w_stone[1]) + s_max = cfg.s_max[0] + difficulty * (cfg.s_max[1] - cfg.s_max[0]) + h_max = cfg.h_max[0] + difficulty * (cfg.h_max[1] - cfg.h_max[0]) + + # initialize list of meshes + meshes_list = list() + + # compute the number of stones in x and y directions + num_stones_axis = int(cfg.size[0] / (gap_width + stone_width)) + + # constants + terrain_height = -cfg.holes_depth + device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + + # generate the border + border_width = cfg.size[0] - num_stones_axis * (gap_width + stone_width) + if border_width > 0: + border_center = (0.5 * cfg.size[0], 0.5 * cfg.size[1], -terrain_height / 2) + border_inner_size = (cfg.size[0] - border_width, cfg.size[1] - border_width) + # create border meshes + make_borders = make_border(cfg.size, border_inner_size, terrain_height, border_center) + meshes_list += make_borders + # create a template grid of the terrain height + grid_dim = [stone_width, stone_width, terrain_height] + grid_position = [0.5 * (stone_width + gap_width), 0.5 * (stone_width + gap_width), -terrain_height / 2] + template_box = trimesh.creation.box(grid_dim, trimesh.transformations.translation_matrix(grid_position)) + # extract vertices and faces + template_vertices = template_box.vertices # (8, 3) + template_faces = template_box.faces + + # repeat the template box vertices to space the terrain(num_boxes_axis**2, 8, 3) + vertices = torch.tensor(template_vertices, device=device).repeat(num_stones_axis**2, 1, 1) + # create a meshgrid to offset the vertices + x = torch.arange(0, num_stones_axis, device=device) + y = torch.arange(0, num_stones_axis, device=device) + xx, yy = torch.meshgrid(x, y, indexing="ij") + xx = xx.flatten().view(-1, 1) + yy = yy.flatten().view(-1, 1) + xx_yy = torch.cat((xx, yy), dim=1) + # offset the vertices + offsets = ( + (stone_width + gap_width) * xx_yy + + border_width / 2 + + (2 * torch.rand(*xx_yy.shape, device=xx_yy.device) - 1) * s_max + ) + vertices[:, :, :2] += offsets.unsqueeze(1) + + # add noise on height + num_boxes = len(vertices) + h_noise = torch.zeros((num_boxes, 3), device=device) + h_noise[:, 2].uniform_(-h_max, h_max) + # reshape noise to match the vertices (num_boxes, 4, 3) + # only top vertices are affected + vertices_noise = torch.zeros((num_boxes, 4, 3), device=device) + vertices_noise += h_noise.unsqueeze(1) + # add height only to the top vertices of the box + vertices[vertices[:, :, 2] == 0] += vertices_noise.view(-1, 3) + # move to numpy + vertices = vertices.reshape(-1, 3).cpu().numpy() + + # create faces for boxes(num_boxes, 12, 3), each box has 6 faces, each face has 2 triangles + faces = torch.tensor(template_faces, device=device).repeat(num_boxes, 1, 1) + face_offsets = torch.arange(0, num_boxes, device=device).unsqueeze(1).repeat(1, 12) * 8 + faces += face_offsets.unsqueeze(2) + faces = faces.view(-1, 3).cpu().numpy() + + # convert to trimesh + grid_mesh = trimesh.Trimesh(vertices=vertices, faces=faces) + meshes_list.append(grid_mesh) + + # add a platform in the center of the terrain that is accessible from all sides + dim = (cfg.platform_width, cfg.platform_width, terrain_height + h_max) + pos = (0.5 * cfg.size[0], 0.5 * cfg.size[1], -terrain_height / 2 + h_max / 2) + box_platform = trimesh.creation.box(dim, trimesh.transformations.translation_matrix(pos)) + meshes_list.append(box_platform) + + # specify the origin of the terrain + origin = np.array([0.5 * cfg.size[0], 0.5 * cfg.size[1], h_max]) + + return meshes_list, origin + + +def balance_beams_terrain( + difficulty: float, cfg: mesh_terrains_cfg.MeshBalanceBeamsTerrainCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray]: + # check to ensure square terrain + assert cfg.size[0] == cfg.size[1], "The terrain should be square" + + stone_width = cfg.w_stone[0] - difficulty * (cfg.w_stone[0] - cfg.w_stone[1]) + h_offset = cfg.h_offset[0] + difficulty * (cfg.h_offset[1] - cfg.h_offset[0]) + mid_gap = (cfg.mid_gap + stone_width) * (1 - difficulty) + + meshes_list = list() + num_stones = int(((cfg.size[0] - 0.25 - cfg.platform_width) / 2 - 1) / stone_width) + + terrain_height = 1 + device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + + border_width = (cfg.size[1] - cfg.platform_width) / 2 - 1 - num_stones * stone_width + if border_width > 0: + border_center = (0.5 * cfg.size[0], 0.5 * cfg.size[1], -terrain_height / 2) + border_inner_size = (cfg.size[0] - border_width - 2, cfg.size[1] - border_width - 2) + # create border meshes + make_borders = make_border(cfg.size, border_inner_size, terrain_height, border_center) + meshes_list += make_borders + + grid_dim = [stone_width, stone_width, terrain_height] + grid_position = [0.5 * stone_width, 0.5 * stone_width, -0.5 * terrain_height] + template_box = trimesh.creation.box(grid_dim, trimesh.transformations.translation_matrix(grid_position)) + # extract vertices and faces + template_vertices = template_box.vertices # (8, 3) + template_faces = template_box.faces + # repeat the template box vertices to space the terrain(num_stones, 8, 3) + vertices = torch.tensor(template_vertices, device=device).repeat(num_stones, 1, 1) + index = torch.arange(0, num_stones, device=device) + vertices[:, :, 0] += cfg.size[0] / 2 + cfg.platform_width / 2 + vertices[:, :, 0] += (index * stone_width).unsqueeze(-1) + vertices[(index % 2) == 0, :, 1] += cfg.size[1] / 2 - mid_gap / 2 - stone_width / 2 + vertices[(index % 2) == 1, :, 1] += cfg.size[1] / 2 + mid_gap / 2 - stone_width / 2 + + num_boxes = len(vertices) + h_noise = torch.zeros((num_boxes, 3), device=device) + h_noise[:, 2].uniform_(-h_offset, h_offset) + # reshape noise to match the vertices (num_boxes, 4, 3) + # only top vertices are affected + vertices_noise = torch.zeros((num_boxes, 4, 3), device=device) + vertices_noise += h_noise.unsqueeze(1) + # add height only to the top vertices of the box + vertices[vertices[:, :, 2] == 0] += vertices_noise.view(-1, 3) + # move to numpy + vertices = vertices.reshape(-1, 3).cpu().numpy() + + # create faces for boxes(num_boxes, 12, 3), each box has 6 faces, each face has 2 triangles + faces = torch.tensor(template_faces, device=device).repeat(num_boxes, 1, 1) + face_offsets = torch.arange(0, num_boxes, device=device).unsqueeze(1).repeat(1, 12) * 8 + faces += face_offsets.unsqueeze(2) + faces = faces.view(-1, 3).cpu().numpy() + + # convert to trimesh + grid_mesh = trimesh.Trimesh(vertices=vertices, faces=faces) + meshes_list.append(grid_mesh) + + # add a platform in the center of the terrain that is accessible from all sides + dim = (cfg.platform_width, cfg.platform_width, terrain_height) + pos = (0.5 * cfg.size[0], 0.5 * cfg.size[1], -terrain_height / 2) + box_platform = trimesh.creation.box(dim, trimesh.transformations.translation_matrix(pos)) + meshes_list.append(box_platform) + + # specify the origin of the terrain + origin = np.array([0.5 * cfg.size[0], 0.5 * cfg.size[1], 0]) + + return meshes_list, origin + + +def stepping_beams_terrain( + difficulty: float, cfg: mesh_terrains_cfg.MeshSteppingBeamsTerrainCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray]: + stone_width = cfg.w_stone[0] - difficulty * (cfg.w_stone[0] - cfg.w_stone[1]) + h_offset = cfg.h_offset[0] + difficulty * (cfg.h_offset[1] - cfg.h_offset[0]) + gap_width = cfg.gap[0] + difficulty * (cfg.gap[1] - cfg.gap[0]) + yaw = cfg.yaw[0] + difficulty * (cfg.yaw[1] - cfg.yaw[0]) + assert cfg.yaw[0] < cfg.yaw[1], "The yaw range should be in ascending order(0 means no yaw)" + low_stone_l = cfg.l_stone[0] + high_stone_l = cfg.l_stone[1] + + meshes_list = list() + num_stones = int(((cfg.size[0] - cfg.platform_width) / 2) / (gap_width + stone_width)) + + terrain_height = 1 + + border_width = (cfg.size[1] - cfg.platform_width) / 2 - num_stones * (stone_width + gap_width) + + if border_width > 0: + border_center = (0.5 * cfg.size[0], 0.5 * cfg.size[1], -terrain_height / 2) + border_inner_size = (cfg.size[0] - border_width - 2, cfg.size[1] - border_width - 2) + # create border meshes + make_borders = make_border(cfg.size, border_inner_size, terrain_height, border_center) + meshes_list += make_borders + # calculate the center of all the stones + # add random noise to the center + # add noise to the height + # create the stones + for i in range(num_stones): + transform = np.eye(4) + grid_dim = [ + stone_width, + low_stone_l + random.uniform(0, high_stone_l - low_stone_l), + terrain_height + random.uniform(-h_offset, h_offset), + ] + center = [ + cfg.size[0] / 2 + + cfg.platform_width / 2 + + (i + 1) * gap_width + + (i + 0.5) * stone_width + + random.uniform(-0.25, 0.25) * gap_width, + cfg.size[1] / 2 + random.uniform(-0.1, 0.1) * grid_dim[1], + -terrain_height / 2, + ] + transform[0:3, -1] = np.asarray(center) + # create rotation matrix + transform[0:3, 0:3] = R.from_euler("z", random.uniform(-yaw, yaw), degrees=True).as_matrix() + meshes_list.append(trimesh.creation.box(grid_dim, transform)) + # add a platform in the center of the terrain that is accessible from all sides + dim = (cfg.platform_width, cfg.platform_width, terrain_height) + pos = (0.5 * cfg.size[0], 0.5 * cfg.size[1], -terrain_height / 2) + box_platform = trimesh.creation.box(dim, trimesh.transformations.translation_matrix(pos)) + meshes_list.append(box_platform) + + # specify the origin of the terrain + origin = np.array([0.5 * cfg.size[0], 0.5 * cfg.size[1], 0]) + + return meshes_list, origin + + +def box_terrain( + difficulty: float, cfg: mesh_terrains_cfg.MeshDiversityBoxTerrainCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray]: + # + box_width = cfg.box_width_range[1] - difficulty * (cfg.box_width_range[1] - cfg.box_width_range[0]) + box_length = cfg.box_length_range[1] - difficulty * (cfg.box_length_range[1] - cfg.box_length_range[0]) + box_height = cfg.box_height_range[0] + difficulty * (cfg.box_height_range[1] - cfg.box_height_range[0]) + meshes_list = [] + terrain_height = 1.0 + middle_height = 0.0 + # check if box_gap_range is a tuple + if isinstance(cfg.box_gap_range, tuple): + # Task of jumping over neighboring boxes + gap_width = cfg.box_gap_range[0] + difficulty * (cfg.box_gap_range[1] - cfg.box_gap_range[0]) + # generate the box at the origin + box_dim = (box_width, box_length, box_height + terrain_height) + pos = (cfg.size[0] / 2, cfg.size[1] / 2, -terrain_height / 2 + box_height / 2) + box = trimesh.creation.box(box_dim, trimesh.transformations.translation_matrix(pos)) + meshes_list.append(box) + # generate the neighboring boxes + box_dim = (box_width, box_length, box_height + terrain_height) + offset_x = box_width / 2 + box_width / 2 + gap_width + pos = (cfg.size[0] / 2 + offset_x, cfg.size[1] / 2, -terrain_height / 2 + box_height / 2) + box = trimesh.creation.box(box_dim, trimesh.transformations.translation_matrix(pos)) + meshes_list.append(box) + middle_height = box_height + elif cfg.box_gap_range is None: + # Task for climbing up or down boxes + if cfg.up_or_down == "up": + # for climbing up + box_dim = (box_width, box_length, box_height + terrain_height) + offset_x = box_width + pos = (cfg.size[0] / 2 + offset_x, cfg.size[1] / 2, -terrain_height / 2 + box_height / 2) + box = trimesh.creation.box(box_dim, trimesh.transformations.translation_matrix(pos)) + meshes_list.append(box) + middle_height = 0.0 + elif cfg.up_or_down == "down": + # for climbing down + box_dim = (box_width, box_length, box_height + terrain_height) + pos = (cfg.size[0] / 2, cfg.size[1] / 2, -terrain_height / 2 + box_height / 2) + box = trimesh.creation.box(box_dim, trimesh.transformations.translation_matrix(pos)) + meshes_list.append(box) + middle_height = box_height + else: + raise ValueError("up_or_down should be either 'up' or 'down'") + else: + raise ValueError("box_gap_range should be a tuple or None") + + # generate the ground + pos = (cfg.size[0] / 2, cfg.size[1] / 2, -terrain_height / 2) + dim = (cfg.size[0], cfg.size[1], terrain_height) + ground = trimesh.creation.box(dim, trimesh.transformations.translation_matrix(pos)) + meshes_list.append(ground) + + # specify the origin of the terrain + origin = np.array([cfg.size[0] / 2, cfg.size[1] / 2, middle_height]) + + return meshes_list, origin + + +def passage_terrain( + difficulty: float, cfg: mesh_terrains_cfg.MeshPassageTerrainCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray]: + if isinstance(cfg.passage_width, tuple): + width = cfg.passage_width[1] - difficulty * (cfg.passage_width[1] - cfg.passage_width[0]) + elif isinstance(cfg.passage_width, float): + width = cfg.passage_width + else: + raise ValueError("passage_width should be a tuple or a float") + if isinstance(cfg.passage_length, tuple): + length = cfg.passage_length[0] + difficulty * (cfg.passage_length[1] - cfg.passage_length[0]) + elif isinstance(cfg.passage_length, float): + length = cfg.passage_length + else: + raise ValueError("passage_length should be a tuple or a float") + if isinstance(cfg.passage_height, tuple): + height = cfg.passage_height[1] - difficulty * (cfg.passage_height[1] - cfg.passage_height[0]) + elif isinstance(cfg.passage_height, float): + height = cfg.passage_height + else: + raise ValueError("passage_height should be a tuple or a float") + # generate the passage + meshes_list = [] + terrain_height = 1.0 + offset_x = 1.0 + # four legs of the passage + dim = (0.05 + np.random.uniform(0.0, 0.1), 0.05 + np.random.uniform(0.0, 0.1), terrain_height + height) + pos1 = (offset_x + cfg.size[0] / 2 - length / 2, cfg.size[1] / 2 - width / 2, -terrain_height / 2 + height / 2) + box1 = trimesh.creation.box(dim, trimesh.transformations.translation_matrix(pos1)) + meshes_list.append(box1) + pos2 = (offset_x + cfg.size[0] / 2 - length / 2, cfg.size[1] / 2 + width / 2, -terrain_height / 2 + height / 2) + box2 = trimesh.creation.box(dim, trimesh.transformations.translation_matrix(pos2)) + meshes_list.append(box2) + pos3 = (offset_x + cfg.size[0] / 2 + length / 2, cfg.size[1] / 2 - width / 2, -terrain_height / 2 + height / 2) + box3 = trimesh.creation.box(dim, trimesh.transformations.translation_matrix(pos3)) + meshes_list.append(box3) + pos4 = (offset_x + cfg.size[0] / 2 + length / 2, cfg.size[1] / 2 + width / 2, -terrain_height / 2 + height / 2) + box4 = trimesh.creation.box(dim, trimesh.transformations.translation_matrix(pos4)) + meshes_list.append(box4) + # top of the passage + dim = (length + dim[0], width + dim[1], 0.05 + np.random.uniform(0, 0.1)) + pos = (offset_x + cfg.size[0] / 2, cfg.size[1] / 2, dim[2] / 2 + height) + top = trimesh.creation.box(dim, trimesh.transformations.translation_matrix(pos)) + meshes_list.append(top) + # ground + pos = (cfg.size[0] / 2, cfg.size[1] / 2, -terrain_height / 2) + dim = (cfg.size[0], cfg.size[1], terrain_height) + ground = trimesh.creation.box(dim, trimesh.transformations.translation_matrix(pos)) + meshes_list.append(ground) + + # specify the origin of the terrain + origin = np.array([cfg.size[0] / 2 - 1.0, cfg.size[1] / 2, 0.0]) + + return meshes_list, origin + + +def structured_terrain( + difficulty: float, cfg: mesh_terrains_cfg.MeshStructuredTerrainCfg +) -> tuple[list[trimesh.Trimesh], np.ndarray]: + mesh_list = [] + terrain = cfg.terrain_type + # generate the terrain + if terrain == "obstacles": + origin = np.array([cfg.size[0] / 2, cfg.size[1] / 2, 0.0]) + for i in range(12): + if i < 8: + length = random.uniform(0.2, 2.0) + width = random.uniform(0.2, 2.0) + height = random.uniform(0.08, 0.25) + else: + length = random.uniform(0.2, 1.0) + width = random.uniform(0.2, 1.0) + height = 3.0 + center = ( + cfg.size[0] / 2 + random.uniform(1, cfg.size[0] / 2) * (-1) ** (random.randint(1, 2)), + cfg.size[1] / 2 + random.uniform(1, cfg.size[0] / 2) * (-1) ** (random.randint(1, 2)), + height / 2, + ) + transform = np.eye(4) + transform[0:3, -1] = np.asarray(center) + # create the box + dims = (length, width, height) + mesh = trimesh.creation.box(dims, transform=transform) + mesh_list.append(mesh) + # add walls + if random.uniform(0, 1) > 0.1: + center_pts = [(0, 0, 0), (cfg.size[0], 0, 0), (0, cfg.size[1], 0), (cfg.size[0], cfg.size[1], 0)] + for i, center in enumerate(center_pts): + if random.uniform(0, 1) > 0.5: + continue + length = cfg.size[0] * random.uniform(0.2, 0.4) + width = cfg.size[1] * random.uniform(0.2, 0.4) + height = 6.0 + transform = np.eye(4) + c = (center[0] + (-1) ** i * length / 2, center[1] + (-1) ** (i // 2) * width / 2, center[2]) + transform[0:3, -1] = np.asarray(c) + # create the box + dims = (length, width, height) + mesh = trimesh.creation.box(dims, transform=transform) + mesh_list.append(mesh) + # add plane + ground_plane = make_plane(cfg.size, height=0.0, center_zero=False) + mesh_list.append(ground_plane) + + elif terrain == "stairs": + step_width = random.uniform(0.2, 0.5) + _mesh_list, origin = pyramid_stairs_terrain( + difficulty, + MeshPyramidStairsTerrainCfg( + size=cfg.size, + border_width=1.0, + step_height_range=(0.08, 0.20), + step_width=step_width, + platform_width=2.0, + ), + ) + mesh_list += _mesh_list + # add walls + if random.uniform(0, 1) > 0.05: + center_pts = [(0, 0, 0), (cfg.size[0], 0, 0), (0, cfg.size[1], 0), (cfg.size[0], cfg.size[1], 0)] + for i, center in enumerate(center_pts): + if random.uniform(0, 1) > 0.75: + continue + length = cfg.size[0] * random.uniform(0.3, 0.4) + width = cfg.size[1] * random.uniform(0.3, 0.4) + height = 6.0 + transform = np.eye(4) + c = (center[0] + (-1) ** i * length / 2, center[1] + (-1) ** (i // 2) * width / 2, center[2]) + transform[0:3, -1] = np.asarray(c) + # create the box + dims = (length, width, height) + mesh = trimesh.creation.box(dims, transform=transform) + mesh_list.append(mesh) + elif terrain == "inverted_stairs": + step_width = random.uniform(0.2, 0.5) + # inverted prymaid + _mesh_list, origin = inverted_pyramid_stairs_terrain( + difficulty, + MeshInvertedPyramidStairsTerrainCfg( + size=cfg.size, + border_width=1.0, + step_height_range=(0.08, 0.20), + step_width=step_width, + platform_width=2.0, + ), + ) + mesh_list += _mesh_list + # add walls + if random.uniform(0, 1) > 0.05: + center_pts = [(0, 0, 0), (cfg.size[0], 0, 0), (0, cfg.size[1], 0), (cfg.size[0], cfg.size[1], 0)] + for i, center in enumerate(center_pts): + if random.uniform(0, 1) > 0.75: + continue + length = cfg.size[0] * random.uniform(0.3, 0.4) + width = cfg.size[1] * random.uniform(0.3, 0.4) + height = 6.0 + transform = np.eye(4) + c = (center[0] + (-1) ** i * length / 2, center[1] + (-1) ** (i // 2) * width / 2, center[2]) + transform[0:3, -1] = np.asarray(c) + # create the box + dims = (length, width, height) + mesh = trimesh.creation.box(dims, transform=transform) + mesh_list.append(mesh) + elif terrain == "walls": + origin = np.array([cfg.size[0] / 2, cfg.size[1] / 2, 0.0]) + # add walls + center_pts = [(0, 0, 0), (cfg.size[0], 0, 0), (0, cfg.size[1], 0), (cfg.size[0], cfg.size[1], 0)] + for i, center in enumerate(center_pts): + if random.uniform(0, 1) > 0.75: + continue + length = cfg.size[0] * random.uniform(0.3, 0.4) + width = cfg.size[1] * random.uniform(0.3, 0.4) + height = 6.0 + transform = np.eye(4) + c = (center[0] + (-1) ** i * length / 2, center[1] + (-1) ** (i // 2) * width / 2, center[2]) + transform[0:3, -1] = np.asarray(c) + # create the box + dims = (length, width, height) + mesh = trimesh.creation.box(dims, transform=transform) + mesh_list.append(mesh) + # add plane + ground_plane = make_plane(cfg.size, height=0.0, center_zero=False) + mesh_list.append(ground_plane) + else: + raise ValueError(f"terrain_type {terrain} is not supported") + # update the origin in a free space + return mesh_list, origin diff --git a/source/uwlab/uwlab/terrains/trimesh/mesh_terrains_cfg.py b/source/uwlab/uwlab/terrains/trimesh/mesh_terrains_cfg.py new file mode 100644 index 00000000..6067a0a0 --- /dev/null +++ b/source/uwlab/uwlab/terrains/trimesh/mesh_terrains_cfg.py @@ -0,0 +1,188 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import MISSING +from typing import Literal + +import uwlab.terrains.trimesh.mesh_terrains as mesh_terrains +from isaaclab.utils import configclass + +from ..terrain_generator_cfg import SubTerrainBaseCfg + +""" +Different trimesh terrain configurations. +""" + + +@configclass +class MeshObjTerrainCfg(SubTerrainBaseCfg): + """Configuration for a plane mesh terrain.""" + + function = mesh_terrains.obj_terrain + + obj_path: str = MISSING + + spawn_origin_path: str = MISSING + + +@configclass +class CachedTerrainGenCfg(MeshObjTerrainCfg): + """Configuration for a plane mesh terrain.""" + + function = mesh_terrains.cached_terrain_gen + + height: float = MISSING + + levels: float = MISSING + + include_overhang: bool = MISSING + + task_descriptor: str = MISSING + + +@configclass +class TerrainGenCfg(MeshObjTerrainCfg): + """Configuration for a plane mesh terrain.""" + + function = mesh_terrains.terrain_gen + + height: float = MISSING + + levels: float = MISSING + + include_overhang: bool = MISSING + + terrain_styles: list = MISSING + + yaml_path: str = (MISSING,) + + spawn_origin_path: str = MISSING + + python_script: str = MISSING + + +@configclass +class MeshStonesEverywhereTerrainCfg(SubTerrainBaseCfg): + """ + A terrain with stones everywhere + """ + + function = mesh_terrains.stones_everywhere_terrain + + # stone gap width + w_gap: tuple[float, float] = MISSING + + # grid square stone size (width) + w_stone: tuple[float, float] = MISSING + + # the maximum shift, both x and y shift is uniformly sample from [-s_max, s_max] + s_max: tuple[float, float] = MISSING + + # the maximum height, the height is uniformly sample from [-hmax, h_max], default height is 1.0 m + h_max: tuple[float, float] = MISSING + + # holes depth + holes_depth: float = MISSING + + # the platform width + platform_width: float = MISSING + + +@configclass +class MeshBalanceBeamsTerrainCfg(SubTerrainBaseCfg): + """ + A terrain with balance-beams + """ + + # balance beams terrain function + function = mesh_terrains.balance_beams_terrain + + # the platform width + platform_width: float = MISSING + + # the height offset + h_offset: tuple[float, float] = MISSING + + # stone width + w_stone: tuple[float, float] = MISSING + + # the gap between two beams + mid_gap: float = MISSING + + +@configclass +class MeshSteppingBeamsTerrainCfg(SubTerrainBaseCfg): + """ + A terrain with stepping-beams + """ + + # stepping beams terrain function + function = mesh_terrains.stepping_beams_terrain + + # the platform width + platform_width: float = MISSING + + # the height offset + h_offset: tuple[float, float] = MISSING + + # stone width + w_stone: tuple[float, float] = MISSING + + # length of the stepping beams + l_stone: tuple[float, float] = MISSING + + # the gap between two beams + gap: tuple[float, float] = MISSING + + # the yaw angle of the stepping beams + yaw: tuple[float, float] = MISSING + + +@configclass +class MeshDiversityBoxTerrainCfg(SubTerrainBaseCfg): + """ + A terrain with boxes for anymal parkour + """ + + function = mesh_terrains.box_terrain + + # the box width range + box_width_range: tuple[float, float] = MISSING + # the box length range + box_length_range: tuple[float, float] = MISSING + # the box height range + box_height_range: tuple[float, float] = MISSING + + # the gap between two boxes + box_gap_range: tuple[float, float] = None # type: ignore + + # flag for climbing up (box is set at the origin ) or climb down (box is set near the origin) + up_or_down: str = None # type: ignore + + +@configclass +class MeshPassageTerrainCfg(SubTerrainBaseCfg): + """ + A terrain with passage + """ + + function = mesh_terrains.passage_terrain + + # the passage width (y dir) + passage_width: float | tuple[float, float] = MISSING + + # the passage height + passage_height: float | tuple[float, float] = MISSING + + # the passage length (x dir) + passage_length: float | tuple[float, float] = MISSING + + +@configclass +class MeshStructuredTerrainCfg(SubTerrainBaseCfg): + """Configuration for a structured terrain.""" + + function = mesh_terrains.structured_terrain + terrain_type: Literal["stairs", "inverted_stairs", "obstacles", "walls"] = MISSING diff --git a/source/uwlab/uwlab/terrains/trimesh/utils.py b/source/uwlab/uwlab/terrains/trimesh/utils.py new file mode 100644 index 00000000..085a88f4 --- /dev/null +++ b/source/uwlab/uwlab/terrains/trimesh/utils.py @@ -0,0 +1,194 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import numpy as np +import scipy.spatial.transform as tf +import trimesh + +""" +Primitive functions to generate meshes. +""" + + +def make_plane(size: tuple[float, float], height: float, center_zero: bool = True) -> trimesh.Trimesh: + """Generate a plane mesh. + + If :obj:`center_zero` is True, the origin is at center of the plane mesh i.e. the mesh extends from + :math:`(-size[0] / 2, -size[1] / 2, 0)` to :math:`(size[0] / 2, size[1] / 2, height)`. + Otherwise, the origin is :math:`(size[0] / 2, size[1] / 2)` and the mesh extends from + :math:`(0, 0, 0)` to :math:`(size[0], size[1], height)`. + + Args: + size: The length (along x) and width (along y) of the terrain (in m). + height: The height of the plane (in m). + center_zero: Whether the 2D origin of the plane is set to the center of mesh. + Defaults to True. + + Returns: + A trimesh.Trimesh objects for the plane. + """ + # compute the vertices of the terrain + x0 = [size[0], size[1], height] + x1 = [size[0], 0.0, height] + x2 = [0.0, size[1], height] + x3 = [0.0, 0.0, height] + # generate the tri-mesh with two triangles + vertices = np.array([x0, x1, x2, x3]) + faces = np.array([[1, 0, 2], [2, 3, 1]]) + plane_mesh = trimesh.Trimesh(vertices=vertices, faces=faces) + # center the plane at the origin + if center_zero: + plane_mesh.apply_translation(-np.array([size[0] / 2.0, size[1] / 2.0, 0.0])) + # return the tri-mesh and the position + return plane_mesh + + +def make_border( + size: tuple[float, float], inner_size: tuple[float, float], height: float, position: tuple[float, float, float] +) -> list[trimesh.Trimesh]: + """Generate meshes for a rectangular border with a hole in the middle. + + .. code:: text + + +---------------------+ + |#####################| + |##+---------------+##| + |##| |##| + |##| |##| length + |##| |##| (y-axis) + |##| |##| + |##+---------------+##| + |#####################| + +---------------------+ + width (x-axis) + + Args: + size: The length (along x) and width (along y) of the terrain (in m). + inner_size: The inner length (along x) and width (along y) of the hole (in m). + height: The height of the border (in m). + position: The center of the border (in m). + + Returns: + A list of trimesh.Trimesh objects that represent the border. + """ + # compute thickness of the border + thickness_x = (size[0] - inner_size[0]) / 2.0 + thickness_y = (size[1] - inner_size[1]) / 2.0 + # generate tri-meshes for the border + # top/bottom border + box_dims = (size[0], thickness_y, height) + # -- top + box_pos = (position[0], position[1] + inner_size[1] / 2.0 + thickness_y / 2.0, position[2]) + box_mesh_top = trimesh.creation.box(box_dims, trimesh.transformations.translation_matrix(box_pos)) + # -- bottom + box_pos = (position[0], position[1] - inner_size[1] / 2.0 - thickness_y / 2.0, position[2]) + box_mesh_bottom = trimesh.creation.box(box_dims, trimesh.transformations.translation_matrix(box_pos)) + # left/right border + box_dims = (thickness_x, inner_size[1], height) + # -- left + box_pos = (position[0] - inner_size[0] / 2.0 - thickness_x / 2.0, position[1], position[2]) + box_mesh_left = trimesh.creation.box(box_dims, trimesh.transformations.translation_matrix(box_pos)) + # -- right + box_pos = (position[0] + inner_size[0] / 2.0 + thickness_x / 2.0, position[1], position[2]) + box_mesh_right = trimesh.creation.box(box_dims, trimesh.transformations.translation_matrix(box_pos)) + # return the tri-meshes + return [box_mesh_left, box_mesh_right, box_mesh_top, box_mesh_bottom] + + +def make_box( + length: float, + width: float, + height: float, + center: tuple[float, float, float], + max_yx_angle: float = 0, + degrees: bool = True, +) -> trimesh.Trimesh: + """Generate a box mesh with a random orientation. + + Args: + length: The length (along x) of the box (in m). + width: The width (along y) of the box (in m). + height: The height of the cylinder (in m). + center: The center of the cylinder (in m). + max_yx_angle: The maximum angle along the y and x axis. Defaults to 0. + degrees: Whether the angle is in degrees. Defaults to True. + + Returns: + A trimesh.Trimesh object for the cylinder. + """ + # create a pose for the cylinder + transform = np.eye(4) + transform[0:3, -1] = np.asarray(center) + # -- create a random rotation + euler_zyx = tf.Rotation.random().as_euler("zyx") # returns rotation of shape (3,) + # -- cap the rotation along the y and x axis + if degrees: + max_yx_angle = max_yx_angle / 180.0 + euler_zyx[1:] *= max_yx_angle + # -- apply the rotation + transform[0:3, 0:3] = tf.Rotation.from_euler("zyx", euler_zyx).as_matrix() + # create the box + dims = (length, width, height) + return trimesh.creation.box(dims, transform=transform) + + +def make_cylinder( + radius: float, height: float, center: tuple[float, float, float], max_yx_angle: float = 0, degrees: bool = True +) -> trimesh.Trimesh: + """Generate a cylinder mesh with a random orientation. + + Args: + radius: The radius of the cylinder (in m). + height: The height of the cylinder (in m). + center: The center of the cylinder (in m). + max_yx_angle: The maximum angle along the y and x axis. Defaults to 0. + degrees: Whether the angle is in degrees. Defaults to True. + + Returns: + A trimesh.Trimesh object for the cylinder. + """ + # create a pose for the cylinder + transform = np.eye(4) + transform[0:3, -1] = np.asarray(center) + # -- create a random rotation + euler_zyx = tf.Rotation.random().as_euler("zyx") # returns rotation of shape (3,) + # -- cap the rotation along the y and x axis + if degrees: + max_yx_angle = max_yx_angle / 180.0 + euler_zyx[1:] *= max_yx_angle + # -- apply the rotation + transform[0:3, 0:3] = tf.Rotation.from_euler("zyx", euler_zyx).as_matrix() + # create the cylinder + return trimesh.creation.cylinder(radius, height, sections=np.random.randint(4, 6), transform=transform) + + +def make_cone( + radius: float, height: float, center: tuple[float, float, float], max_yx_angle: float = 0, degrees: bool = True +) -> trimesh.Trimesh: + """Generate a cone mesh with a random orientation. + + Args: + radius: The radius of the cone (in m). + height: The height of the cone (in m). + center: The center of the cone (in m). + max_yx_angle: The maximum angle along the y and x axis. Defaults to 0. + degrees: Whether the angle is in degrees. Defaults to True. + + Returns: + A trimesh.Trimesh object for the cone. + """ + # create a pose for the cylinder + transform = np.eye(4) + transform[0:3, -1] = np.asarray(center) + # -- create a random rotation + euler_zyx = tf.Rotation.random().as_euler("zyx") # returns rotation of shape (3,) + # -- cap the rotation along the y and x axis + if degrees: + max_yx_angle = max_yx_angle / 180.0 + euler_zyx[1:] *= max_yx_angle + # -- apply the rotation + transform[0:3, 0:3] = tf.Rotation.from_euler("zyx", euler_zyx).as_matrix() + # create the cone + return trimesh.creation.cone(radius, height, sections=np.random.randint(4, 6), transform=transform) diff --git a/source/uwlab/uwlab/terrains/utils/__init__.py b/source/uwlab/uwlab/terrains/utils/__init__.py new file mode 100644 index 00000000..beaa54bd --- /dev/null +++ b/source/uwlab/uwlab/terrains/utils/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .patch_sampling_cfg import * diff --git a/source/uwlab/uwlab/terrains/utils/patch_sampling.py b/source/uwlab/uwlab/terrains/utils/patch_sampling.py new file mode 100644 index 00000000..878b0378 --- /dev/null +++ b/source/uwlab/uwlab/terrains/utils/patch_sampling.py @@ -0,0 +1,490 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import numpy as np +import torch +from typing import TYPE_CHECKING + +import warp as wp # Warp (https://github.com/NVIDIA/warp) + +from isaaclab.utils.warp import raycast_mesh + +if TYPE_CHECKING: + from . import patch_sampling_cfg as patch_cfg + + +def uniform_sample_multiple_ranges( + ranges: list[tuple[float, float]], + sample_size: int, + device: torch.device, +) -> torch.Tensor: + """ + Sample `sample_size` values from the provided list of (min, max) ranges. + Each sampled value is drawn from one interval chosen uniformly at random. + """ + if not ranges: + raise ValueError("`ranges` cannot be empty") + + # Number of intervals + num_intervals = len(ranges) + + # Randomly select which interval each sample should come from + interval_indices = torch.randint(low=0, high=num_intervals, size=(sample_size,), device=device) + + # Prepare the output buffer + samples = torch.empty(sample_size, device=device, dtype=torch.float32) + + # For each interval, sample the required number of values + for i in range(num_intervals): + mask = interval_indices == i + count_i = mask.sum() + if count_i > 0: + low, high = ranges[i] + samples[mask] = torch.empty(count_i, device=device).uniform_(low, high) + + return samples + + +def find_piecewise_range_flat_patches( + wp_mesh: wp.Mesh, + origin: np.ndarray | torch.Tensor | tuple[float, float, float], + cfg: patch_cfg.PieceWiseRangeFlatPatchSamplingCfg, +) -> torch.Tensor: + """ + Finds flat patches of a given radius in the input mesh, but now supports + multiple intervals for x, y, and z to give you more control over where to sample. + + Args: + wp_mesh: The Warp mesh to find patches in. + num_patches: The desired number of patches to find. + patch_radius: The radii used to form patches (float or list of floats). + origin: The origin defining the center of the search space (in mesh frame). + x_range: A list of (min, max) intervals for X sampling (in mesh frame). + y_range: A list of (min, max) intervals for Y sampling (in mesh frame). + z_range: A list of (min, max) intervals for Z filtering (in mesh frame). + max_height_diff: The maximum allowed distance between lowest and highest + points on a patch to consider it valid. + max_iterations: The maximum number of rejection-sampling iterations. + + Returns: + A (num_patches, 3) torch.Tensor containing the valid flat patch centers, + in the mesh frame, offset so that the origin is subtracted at the end. + + Raises: + RuntimeError: If the function cannot find valid patches within max_iterations. + """ + device = wp.device_to_torch(wp_mesh.device) + + # Handle patch_radius input + if isinstance(cfg.patch_radius, float): + patch_radius = [cfg.patch_radius] + + # Convert origin to a torch tensor (on the correct device) + if isinstance(origin, np.ndarray): + origin = torch.from_numpy(origin).float().to(device) + elif isinstance(origin, torch.Tensor): + origin = origin.float().to(device) + else: + origin = torch.tensor(origin, dtype=torch.float, device=device) + + # --- 1) Clip each interval to the bounding box of the mesh and shift by origin + # Convert mesh points to numpy for min/max + mesh_pts = wp_mesh.points.numpy() + mesh_xmin, mesh_xmax = mesh_pts[:, 0].min(), mesh_pts[:, 0].max() + mesh_ymin, mesh_ymax = mesh_pts[:, 1].min(), mesh_pts[:, 1].max() + + x_range = [cfg.x_ranges] if isinstance(cfg.x_ranges, tuple) else cfg.x_ranges + y_range = [cfg.y_ranges] if isinstance(cfg.y_ranges, tuple) else cfg.y_ranges + z_range = [cfg.z_ranges] if isinstance(cfg.z_ranges, tuple) else cfg.z_ranges + + # For x-ranges + x_ranges_clipped = [] + for low, high in x_range: + new_low = max(low + origin[0].item(), mesh_xmin) + new_high = min(high + origin[0].item(), mesh_xmax) + if new_low < new_high: + x_ranges_clipped.append((new_low, new_high)) + + # For y-ranges + y_ranges_clipped = [] + for low, high in y_range: + new_low = max(low + origin[1].item(), mesh_ymin) + new_high = min(high + origin[1].item(), mesh_ymax) + if new_low < new_high: + y_ranges_clipped.append((new_low, new_high)) + + # For z-ranges, we won't clip by mesh bounding box (optional), + # but we shift them by the origin's Z: + z_ranges_shifted = [] + for low, high in z_range: + new_low = low + origin[2].item() + new_high = high + origin[2].item() + z_ranges_shifted.append((new_low, new_high)) + + if not x_ranges_clipped: + raise ValueError("No valid x-ranges remain after clipping to bounding box.") + if not y_ranges_clipped: + raise ValueError("No valid y-ranges remain after clipping to bounding box.") + if not z_ranges_shifted: + raise ValueError("z_ranges cannot be empty.") + + # --- 2) Create a ring of points around (0, 0) in the XY plane to query patch validity + angle = torch.linspace(0, 2 * np.pi, 10, device=device) + query_x = [] + query_y = [] + for radius in patch_radius: + query_x.append(radius * torch.cos(angle)) + query_y.append(radius * torch.sin(angle)) + query_x = torch.cat(query_x).unsqueeze(1) # (num_radii*10, 1) + query_y = torch.cat(query_y).unsqueeze(1) # (num_radii*10, 1) + # shape: (num_radii*10, 3) + query_points = torch.cat([query_x, query_y, torch.zeros_like(query_x)], dim=-1) + + # Buffers to keep track of invalid patches + points_ids = torch.arange(cfg.num_patches, device=device) + flat_patches = torch.zeros(cfg.num_patches, 3, device=device) + + # --- 3) Rejection sampling + iter_count = 0 + while len(points_ids) > 0 and iter_count < cfg.max_iterations: + # (A) Sample X and Y from the multiple intervals + pos_x = uniform_sample_multiple_ranges(x_ranges_clipped, len(points_ids), device) + pos_y = uniform_sample_multiple_ranges(y_ranges_clipped, len(points_ids), device) + + # Store the new (x, y) + flat_patches[points_ids, 0] = pos_x + flat_patches[points_ids, 1] = pos_y + + # (B) Raycast from above (z=100, say) straight down + # Build the 3D query points for each patch + # shape after unsqueeze: (n_ids, 1, 3) + (query_points) => (n_ids, num_radii*10, 3) + points = flat_patches[points_ids].unsqueeze(1) + query_points + # start from 'far above' in Z + points[..., 2] = 100.0 + + # direction is straight down + dirs = torch.zeros_like(points) + dirs[..., 2] = -1.0 + + # Flatten for raycasting + ray_hits = raycast_mesh(points.view(-1, 3), dirs.view(-1, 3), wp_mesh)[0] + # Reshape back to (n_ids, num_radii*10, 3) + heights = ray_hits.view(points.shape)[..., 2] + + # We'll set the patch center's final Z as the last set of ring hits + # so that e.g. flat_patches[:, 2] is the Z of the center ring point + flat_patches[points_ids, 2] = heights[..., -1] + + # (C) Check validity: + # 1) The patch ring must lie entirely within at least one z-range interval + # We'll check each ring point's Z to see if it's within ANY of the z_ranges. + # If the ring fails in all intervals, it's invalid. + z_ok_mask = torch.zeros(len(points_ids), dtype=torch.bool, device=device) + for zlow, zhigh in z_ranges_shifted: + in_this_range = (heights >= zlow) & (heights <= zhigh) + # We only say "ok" if *all* ring points are within the range + # for that interval: + fully_in_this_interval = in_this_range.all(dim=1) # shape: (len(points_ids)) + z_ok_mask |= fully_in_this_interval + + # 2) Height difference check + # For all ring points, difference between min and max must be <= max_height_diff + height_diff = heights.max(dim=1)[0] - heights.min(dim=1)[0] + + # Final "not valid" condition + not_valid = (~z_ok_mask) | (height_diff > cfg.max_height_diff) + + # Filter out the invalid patch IDs + points_ids = points_ids[not_valid] + + iter_count += 1 + + # If we still have leftover invalid patches, raise an error + if len(points_ids) > 0: + raise RuntimeError( + "Failed to find valid patches within the maximum number of iterations!\n" + f" Iterations: {iter_count}\n" + f" Still invalid patches: {len(points_ids)}\n" + " Consider adjusting your ranges or max_height_diff." + ) + + # Return the flat patches, subtracting the origin to keep consistency + # with the original function's behavior of returning in "mesh frame minus origin". + return flat_patches - origin + + +def find_flat_patches( + wp_mesh: wp.Mesh, + origin: np.ndarray | torch.Tensor | tuple[float, float, float], + cfg: patch_cfg.FlatPatchSamplingCfg, +) -> torch.Tensor: + """Finds flat patches of given radius in the input mesh. + + The function finds flat patches of given radius based on the search space defined by the input ranges. + The search space is characterized by origin in the mesh frame, and the x, y, and z ranges. The x and y + ranges are used to sample points in the 2D region around the origin, and the z range is used to filter + patches based on the height of the points. + + The function performs rejection sampling to find the patches based on the following steps: + + 1. Sample patch locations in the 2D region around the origin. + 2. Define a ring of points around each patch location to query the height of the points using ray-casting. + 3. Reject patches that are outside the z range or have a height difference that is too large. + 4. Keep sampling until all patches are valid. + + Args: + wp_mesh: The warp mesh to find patches in. + num_patches: The desired number of patches to find. + patch_radius: The radii used to form patches. If a list is provided, multiple patch sizes are checked. + This is useful to deal with holes or other artifacts in the mesh. + origin: The origin defining the center of the search space. This is specified in the mesh frame. + x_range: The range of X coordinates to sample from. + y_range: The range of Y coordinates to sample from. + z_range: The range of valid Z coordinates used for filtering patches. + max_height_diff: The maximum allowable distance between the lowest and highest points + on a patch to consider it as valid. If the difference is greater than this value, + the patch is rejected. + + Returns: + A tensor of shape (num_patches, 3) containing the flat patches. The patches are defined in the mesh frame. + + Raises: + RuntimeError: If the function fails to find valid patches. This can happen if the input parameters + are not suitable for finding valid patches and maximum number of iterations is reached. + """ + # set device to warp mesh device + device = wp.device_to_torch(wp_mesh.device) + + # resolve inputs to consistent type + # -- patch radii + patch_radius = [cfg.patch_radius] if isinstance(cfg.patch_radius, float) else cfg.patch_radius + + # -- origin + if isinstance(origin, np.ndarray): + origin = torch.from_numpy(origin).to(torch.float).to(device) + elif isinstance(origin, torch.Tensor): + origin = origin.to(device) + else: + origin = torch.tensor(origin, dtype=torch.float, device=device) + + # create ranges for the x and y coordinates around the origin. + # The provided ranges are bounded by the mesh's bounding box. + x_range = ( + max(cfg.x_range[0] + origin[0].item(), wp_mesh.points.numpy()[:, 0].min()), + min(cfg.x_range[1] + origin[0].item(), wp_mesh.points.numpy()[:, 0].max()), + ) + y_range = ( + max(cfg.y_range[0] + origin[1].item(), wp_mesh.points.numpy()[:, 1].min()), + min(cfg.y_range[1] + origin[1].item(), wp_mesh.points.numpy()[:, 1].max()), + ) + z_range = ( + cfg.z_range[0] + origin[2].item(), + cfg.z_range[1] + origin[2].item(), + ) + + # create a circle of points around (0, 0) to query validity of the patches + # the ring of points is uniformly distributed around the circle + angle = torch.linspace(0, 2 * np.pi, 10, device=device) + query_x = [] + query_y = [] + for radius in patch_radius: + query_x.append(radius * torch.cos(angle)) + query_y.append(radius * torch.sin(angle)) + query_x = torch.cat(query_x).unsqueeze(1) # dim: (num_radii * 10, 1) + query_y = torch.cat(query_y).unsqueeze(1) # dim: (num_radii * 10, 1) + # dim: (num_radii * 10, 3) + query_points = torch.cat([query_x, query_y, torch.zeros_like(query_x)], dim=-1) + + # create buffers + # -- a buffer to store indices of points that are not valid + points_ids = torch.arange(cfg.num_patches, device=device) + # -- a buffer to store the flat patches locations + flat_patches = torch.zeros(cfg.num_patches, 3, device=device) + + # sample points and raycast to find the height. + # 1. Reject points that are outside the z_range or have a height difference that is too large. + # 2. Keep sampling until all points are valid. + iter_count = 0 + while len(points_ids) > 0 and iter_count < 10000: + # sample points in the 2D region around the origin + pos_x = torch.empty(len(points_ids), device=device).uniform_(*x_range) + pos_y = torch.empty(len(points_ids), device=device).uniform_(*y_range) + flat_patches[points_ids, :2] = torch.stack([pos_x, pos_y], dim=-1) + + # define the query points to check validity of the patch + # dim: (num_patches, num_radii * 10, 3) + points = flat_patches[points_ids].unsqueeze(1) + query_points + points[..., 2] = 100.0 + # ray-cast direction is downwards + dirs = torch.zeros_like(points) + dirs[..., 2] = -1.0 + + # ray-cast to find the height of the patches + ray_hits = raycast_mesh(points.view(-1, 3), dirs.view(-1, 3), wp_mesh)[0] + heights = ray_hits.view(points.shape)[..., 2] + # set the height of the patches + # note: for invalid patches, they would be overwritten in the next iteration + # so it's safe to set the height to the last value + flat_patches[points_ids, 2] = heights[..., -1] + + # check validity + # -- height is within the z range + not_valid = torch.any(torch.logical_or(heights < z_range[0], heights > z_range[1]), dim=1) + # -- height difference is within the max height difference + not_valid = torch.logical_or(not_valid, (heights.max(dim=1)[0] - heights.min(dim=1)[0]) > cfg.max_height_diff) + + # remove invalid patches indices + points_ids = points_ids[not_valid] + # increment count + iter_count += 1 + + # check all patches are valid + if len(points_ids) > 0: + raise RuntimeError( + "Failed to find valid patches! Please check the input parameters." + f"\n\tMaximum number of iterations reached: {iter_count}" + f"\n\tNumber of invalid patches: {len(points_ids)}" + f"\n\tMaximum height difference: {cfg.max_height_diff}" + ) + + # return the flat patches (in the mesh frame) + return flat_patches - origin + + +def find_flat_patches_by_radius( + wp_mesh: wp.Mesh, + origin: np.ndarray | torch.Tensor | tuple[float, float, float], + cfg: patch_cfg.FlatPatchSamplingByRadiusCfg, +) -> torch.Tensor: + """Finds flat patches of given radius in the input mesh by sampling patch + centers in a circular region around `origin`. + + Instead of taking x_range, y_range, this function takes radius_range (min, max) + and uniformly samples: + - radius in [radius_range[0], radius_range[1]] + - angle in [0, 2*pi] + Then, (x, y) = radius * cos(angle), radius * sin(angle), around `origin`. + + The function uses rejection sampling to ensure patches are valid according to: + 1. The patch ring is fully within the z_range. + 2. The ring’s height difference is no greater than max_height_diff. + + Args: + wp_mesh: The Warp mesh to find patches in. + origin: The origin defining the center of the circular region for patch sampling. + Specified in the mesh frame. + cfg: A configuration object with the following attributes: + - num_patches (int): Number of patches to find. + - patch_radius (float | list[float]): Single or multiple radii for the validation ring. + - radius_range (tuple[float, float]): The min/max radius used for sampling patch centers. + - z_range (tuple[float, float]): The min/max Z used for validating patches. + - max_height_diff (float): The maximum allowed height difference across patch ring. + - max_iterations (int): The maximum number of iterations for rejection sampling. + + Returns: + A torch.Tensor of shape (num_patches, 3) containing the patch centers in the mesh frame, + offset so that `origin` is subtracted at the end. + + Raises: + RuntimeError: If the function fails to find valid patches within `max_iterations`. + """ + device = wp.device_to_torch(wp_mesh.device) + + # -- handle patch_radius input + if isinstance(cfg.patch_radius, float): + patch_radius = [cfg.patch_radius] + else: + patch_radius = cfg.patch_radius + + # -- resolve the origin to a torch tensor (on the correct device) + if isinstance(origin, np.ndarray): + origin = torch.from_numpy(origin).float().to(device) + elif isinstance(origin, torch.Tensor): + origin = origin.float().to(device) + else: + origin = torch.tensor(origin, dtype=torch.float, device=device) + + # -- expand z_range by origin + z_range_shifted = (cfg.z_range[0] + origin[2].item(), cfg.z_range[1] + origin[2].item()) + + # -- create ring (circle) of points around (0, 0) to test patch "flatness" + # We'll sample 10 angles around the circle for each candidate patch radius + angle = torch.linspace(0, 2 * np.pi, 10, device=device) # shape: (10,) + ring_x = [] + ring_y = [] + for radius in patch_radius: + ring_x.append(radius * torch.cos(angle)) # shape: (10,) + ring_y.append(radius * torch.sin(angle)) # shape: (10,) + + ring_x = torch.cat(ring_x).unsqueeze(1) # shape: (num_radii * 10, 1) + ring_y = torch.cat(ring_y).unsqueeze(1) # shape: (num_radii * 10, 1) + # final ring of shape: (num_radii * 10, 3) + ring_points = torch.cat([ring_x, ring_y, torch.zeros_like(ring_x)], dim=-1) + + # -- Prepare arrays for sampling + # We'll keep track of which patch indices are "still invalid" and require more sampling + patch_ids = torch.arange(cfg.num_patches, device=device) + flat_patches = torch.zeros((cfg.num_patches, 3), device=device) + + # -- Rejection sampling + iteration = 0 + while len(patch_ids) > 0 and iteration < cfg.max_iterations: + # (1) Sample radius in [r_min, r_max] + r_min, r_max = cfg.radius_range + cur_radius = torch.empty(len(patch_ids), device=device).uniform_(r_min, r_max) + # (2) Sample angle in [0, 2*pi] + cur_angle = torch.empty(len(patch_ids), device=device).uniform_(0, 2 * np.pi) + + # Convert polar -> cartesian + pos_x = cur_radius * torch.cos(cur_angle) + pos_y = cur_radius * torch.sin(cur_angle) + + # Store new (x, y) in our patch buffer + flat_patches[patch_ids, 0] = pos_x + origin[0] + flat_patches[patch_ids, 1] = pos_y + origin[1] + + # Raycast ring points from "far above" in Z + # shape for ring offset: (len(patch_ids), num_radii * 10, 3) + ring_in_world = flat_patches[patch_ids].unsqueeze(1) + ring_points + ring_in_world[..., 2] = 100.0 # set starting Z for the ray + # directions: straight down + dirs = torch.zeros_like(ring_in_world) + dirs[..., 2] = -1.0 + + # Flatten for raycasting + ray_hits = raycast_mesh(ring_in_world.view(-1, 3), dirs.view(-1, 3), wp_mesh)[0] # shape: (N, 3) + + # Reshape back to (len(patch_ids), num_radii * 10, 3) + ring_hits_3d = ray_hits.view(ring_in_world.shape) + # We'll define the patch center's final Z from the last ring point's Z + flat_patches[patch_ids, 2] = ring_hits_3d[..., -1, 2] + + # (3) Check validity: + # -- all ring points must lie fully in the z_range + heights = ring_hits_3d[..., 2] # shape: (len(patch_ids), num_radii*10) + out_of_range = (heights < z_range_shifted[0]) | (heights > z_range_shifted[1]) + # -- height difference must be within max_height_diff + height_diff = heights.max(dim=1)[0] - heights.min(dim=1)[0] + # "not valid" if out of range or if height diff is too large + not_valid = out_of_range.any(dim=1) | (height_diff > cfg.max_height_diff) + + # Keep only the patches that are invalid for another iteration + patch_ids = patch_ids[not_valid] + iteration += 1 + + if len(patch_ids) > 0: + # We ran out of iterations but still have invalid patches + raise RuntimeError( + f"Failed to find valid patches within {cfg.max_iterations} iterations.\n" + f"Still invalid patches: {len(patch_ids)}.\n" + "Consider relaxing your constraints or increasing max_iterations." + ) + + # Return patch centers in the "mesh frame minus origin" + # (i.e., subtract origin so that final returned coords are consistent + # with the original code's behavior) + return flat_patches - origin diff --git a/source/uwlab/uwlab/terrains/utils/patch_sampling_cfg.py b/source/uwlab/uwlab/terrains/utils/patch_sampling_cfg.py new file mode 100644 index 00000000..ae34ff93 --- /dev/null +++ b/source/uwlab/uwlab/terrains/utils/patch_sampling_cfg.py @@ -0,0 +1,99 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from dataclasses import MISSING +from typing import Callable + +from isaaclab.utils import configclass + +from . import patch_sampling as sampling_functions + + +@configclass +class PatchSamplingCfg: + """Configuration for sampling patches on the sub-terrain.""" + + func: Callable = MISSING + """The function to use for sampling patches.""" + + num_patches: int = MISSING + """Number of patches to sample.""" + + +@configclass +class FlatPatchSamplingCfg(PatchSamplingCfg): + func: Callable = sampling_functions.find_flat_patches + """The function to use for sampling patches.""" + + patch_radius: float | list[float] = MISSING + """Radius of the patches. + + A list of radii can be provided to check for patches of different sizes. This is useful to deal with + cases where the terrain may have holes or obstacles in some areas. + """ + + x_range: tuple[float, float] = (-1e6, 1e6) + """The range of x-coordinates to sample from. Defaults to (-1e6, 1e6). + + This range is internally clamped to the size of the terrain mesh. + """ + + y_range: tuple[float, float] = (-1e6, 1e6) + """The range of y-coordinates to sample from. Defaults to (-1e6, 1e6). + + This range is internally clamped to the size of the terrain mesh. + """ + + z_range: tuple[float, float] = (-1e6, 1e6) + """Allowed range of z-coordinates for the sampled patch. Defaults to (-1e6, 1e6).""" + + max_height_diff: float = MISSING + """Maximum allowed height difference between the highest and lowest points on the patch.""" + + +@configclass +class PieceWiseRangeFlatPatchSamplingCfg(PatchSamplingCfg): + """Configuration for sampling flat patches on the sub-terrain with piece-wise ranges.""" + + func: Callable = sampling_functions.find_piecewise_range_flat_patches + """The function to use for sampling patches with piece wise ranges.""" + + patch_radius: float | list[float] = MISSING + """Radius of the patches. + + A list of radii can be provided to check for patches of different sizes. This is useful to deal with + cases where the terrain may have holes or obstacles in some areas. + """ + + x_ranges: list[tuple[float, float]] | tuple[float, float] = (-1e6, 1e6) + """The list of (min, max) intervals for X sampling (in mesh frame).""" + + y_ranges: list[tuple[float, float]] | tuple[float, float] = (-1e6, 1e6) + """The list of (min, max) intervals for Y sampling (in mesh frame).""" + + z_ranges: list[tuple[float, float]] | tuple[float, float] = (-1e6, 1e6) + """The list of (min, max) intervals for Z filtering (in mesh frame).""" + + max_height_diff: float = MISSING + """Maximum allowed height difference between the highest and lowest points on the patch.""" + + max_iterations: int = 100 + + +@configclass +class FlatPatchSamplingByRadiusCfg(PatchSamplingCfg): + func: Callable = sampling_functions.find_flat_patches_by_radius + + patch_radius: float | list[float] = MISSING + + radius_range: tuple[float, float] = MISSING + + z_range: tuple[float, float] = (-1e6, 1e6) + + max_height_diff: float = MISSING + + max_iterations: int = 100 diff --git a/source/uwlab/uwlab/utils/__init__.py b/source/uwlab/uwlab/utils/__init__.py new file mode 100644 index 00000000..ac860d6d --- /dev/null +++ b/source/uwlab/uwlab/utils/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause diff --git a/source/uwlab/uwlab/utils/io/__init__.py b/source/uwlab/uwlab/utils/io/__init__.py new file mode 100644 index 00000000..b3dfe19d --- /dev/null +++ b/source/uwlab/uwlab/utils/io/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .encode import encode_dict_to_yaml, encode_pickle diff --git a/source/uwlab/uwlab/utils/io/encode.py b/source/uwlab/uwlab/utils/io/encode.py new file mode 100644 index 00000000..473ff250 --- /dev/null +++ b/source/uwlab/uwlab/utils/io/encode.py @@ -0,0 +1,91 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Utilities for file I/O with yaml.""" + +import io +import logging +import pickle +import yaml +from typing import Any + +from isaaclab.utils import class_to_dict + +logging.basicConfig( + level=logging.INFO, # Set to DEBUG for more detailed logs + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", +) +logger = logging.getLogger(__name__) + + +def encode_pickle(data: Any) -> bytes: + """Saves data into a pickle file safely. + + Args: + filename: The path to save the file at. + data: The data to save. + """ + try: + # Serialize the dictionary to a bytes buffer using pickle + buffer = io.BytesIO() + pickle.dump(data, buffer) + buffer.seek(0) # Reset buffer pointer to the beginning + + serialized_data = buffer.read() + logger.info("Data successfully serialized to pickle format.") + return serialized_data + + except pickle.PicklingError as pe: + logger.error(f"PicklingError: Failed to serialize data: {pe}") + raise + + except Exception as e: + logger.error(f"Unexpected error during pickle serialization: {e}") + raise + + +def encode_dict_to_yaml(data: dict | object, sort_keys: bool = False) -> bytes: + """Serializes data into a YAML-formatted byte string. + + Args: + data (dict | object): The data to serialize. Can be a dictionary or any serializable object. + sort_keys (bool, optional): Whether to sort the keys in the output YAML. Defaults to False. + + Returns: + bytes: The serialized YAML byte string. + + Raises: + TypeError: If the data provided is not serializable to a dictionary. + yaml.YAMLError: If an error occurs during YAML serialization. + Exception: For any other exceptions that occur during serialization. + """ + try: + # Convert data into a dictionary if it's not already one + if not isinstance(data, dict): + data = class_to_dict(data) # Assumes class_to_dict is defined elsewhere + logger.debug("Converted object to dictionary for YAML serialization.") + + # Serialize the dictionary to a YAML-formatted string + yaml_str = yaml.dump(data, sort_keys=sort_keys) + logger.info("Data successfully serialized to YAML format.") + + # Encode the YAML string to bytes (UTF-8) + yaml_bytes = yaml_str.encode("utf-8") + logger.debug("YAML string encoded to bytes successfully.") + + return yaml_bytes + + except TypeError as type_error: + logger.error(f"TypeError: Data provided is not serializable to a dictionary: {type_error}") + raise + + except yaml.YAMLError as yaml_error: + logger.error(f"YAMLError: Failed to serialize data to YAML: {yaml_error}") + raise + + except Exception as e: + logger.error(f"Unexpected error during YAML serialization: {e}") + raise diff --git a/source/uwlab/uwlab/utils/math.py b/source/uwlab/uwlab/utils/math.py new file mode 100644 index 00000000..b0ba1ba5 --- /dev/null +++ b/source/uwlab/uwlab/utils/math.py @@ -0,0 +1,76 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn.functional + + +def create_axis_remap_function(forward: str = "x", left: str = "y", up: str = "z", device: str = "cpu"): + """Creates a function to remap and reorient the axes of input tensors. + + This function generates a new function that remaps the axes of input tensors + according to the specified forward, left, and up axes. The resulting + function can be used to reorder and reorient both positional and rotational data. + + Args: + forward: The axis that should be mapped to the primary axis (e.g., "x" or "-x"). + left: The axis that should be mapped to the secondary axis (e.g., "y" or "-y"). + up: The axis that should be mapped to the tertiary axis (e.g., "z" or "-z"). + device: The device on which the resulting tensors will be processed (e.g., "cpu" or "cuda"). + + Returns: + A function that takes two tensors as inputs: + - positions (torch.Tensor): A positional tensor of shape (N, 3). + - rotations (torch.Tensor): A rotational tensor of shape (N, 3). + + The output function returns a tuple containing: + - new_positions (torch.Tensor): The remapped positional tensor of shape (N, 3). + - new_rotations (torch.Tensor): The remapped rotational tensor of shape (N, 3). + + Example: + .. code-block:: python + + remap_fn = create_axis_remap_function(forward="z", left="-x", up="y") + new_positions, new_rotations = remap_fn(positions, rotations) + """ + # Define the mapping from axis labels to indices + axis_to_index = {"x": 0, "y": 1, "z": 2} + + # Create a mapping for the new axis order + new_axis_order = [forward, left, up] + indices = [] + signs = [] + + for axis in new_axis_order: + sign = 1 + if axis.startswith("-"): + sign = -1 + axis = axis[1:] + + index = axis_to_index[axis] + indices.append(index) + signs.append(sign) + + signs = torch.tensor(signs, device=device) + + def remap_positions_and_rotations(positions: torch.Tensor | None, rotations: torch.Tensor | None) -> tuple: + """Remaps the positions and rotations tensors according to the specified axis order. + + Args: + positions: Input positional tensor of shape (N, 3). + rotations: Input rotational tensor of shape (N, 3). + + Returns: + A tuple containing the remapped positional and rotational tensors, both of shape (N, 3). + if respective input is not None + """ + # Apply sign first, then reorder the axes + new_positions = (positions * signs)[:, indices] if positions is not None else None + new_rotations = (rotations * signs)[:, indices] if rotations is not None else None + return new_positions, new_rotations + + return remap_positions_and_rotations diff --git a/source/uwlab/uwlab/utils/noise/__init__.py b/source/uwlab/uwlab/utils/noise/__init__.py new file mode 100644 index 00000000..8de4351a --- /dev/null +++ b/source/uwlab/uwlab/utils/noise/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .noise_cfg import NoiseModelGroupCfg, OutlierNoiseCfg diff --git a/source/uwlab/uwlab/utils/noise/noise_cfg.py b/source/uwlab/uwlab/utils/noise/noise_cfg.py new file mode 100644 index 00000000..d62ea5d9 --- /dev/null +++ b/source/uwlab/uwlab/utils/noise/noise_cfg.py @@ -0,0 +1,41 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch +from dataclasses import MISSING + +from isaaclab.utils import configclass +from isaaclab.utils.noise import NoiseCfg, NoiseModelCfg + +from . import noise_model + + +@configclass +class OutlierNoiseCfg(NoiseCfg): + """Configuration for an outlier noise term.""" + + func = noise_model.outlier_noise + + probability: torch.Tensor | float = 0.1 + """The probability an outlier occurs, default to 0.1.""" + + noise_cfg: NoiseCfg = MISSING + """The noise function to use for outliers.""" + + +@configclass +class NoiseModelGroupCfg(NoiseModelCfg): + """Configuration for groups of different noise models across environments.""" + + class_type: type = noise_model.NoiseModelGroup + + noise_cfg = None + + @configclass + class NoiseModelMember: + noise_cfg_dict: dict[str, NoiseCfg] = {} + proportion: float = 1.0 + + noise_model_groups: dict[str, NoiseModelMember] = {} diff --git a/source/uwlab/uwlab/utils/noise/noise_model.py b/source/uwlab/uwlab/utils/noise/noise_model.py new file mode 100644 index 00000000..eeb9d8a7 --- /dev/null +++ b/source/uwlab/uwlab/utils/noise/noise_model.py @@ -0,0 +1,88 @@ +# Copyright (c) 2024-2025, The UW 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 + +from isaaclab.utils.noise import NoiseModel + +if TYPE_CHECKING: + from . import noise_cfg + + +def outlier_noise(data: torch.Tensor, cfg: noise_cfg.OutlierNoiseCfg) -> torch.Tensor: + """Applies a gaussian noise to a given data set. + + Args: + data: The unmodified data set to apply noise to. + cfg: The configuration parameters for gaussian noise. + + Returns: + The data modified by the noise parameters provided. + """ + + # fix tensor device for mean on first call and update config parameters + if isinstance(cfg.probability, torch.Tensor): + cfg.probability = cfg.probability.to(data.device) + + # generate a mask for the outliers + mask = torch.rand_like(data) < cfg.probability + + data[mask] = cfg.noise_cfg.func(data[mask], cfg.noise_cfg) + + return data + + +class NoiseModelGroup(NoiseModel): + def __init__(self, noise_model_cfg: noise_cfg.NoiseModelGroupCfg, num_envs: int, device: str): + # initialize parent class + super().__init__(noise_model_cfg, num_envs, device) + # store noise groups + self.proportions = torch.tensor( + [noise_group.proportion for noise_group in noise_model_cfg.noise_model_groups.values()], device=device + ) + self.proportions /= torch.sum(self.proportions) + self.noise_group_assignment = torch.zeros(num_envs, device=device, dtype=torch.int64) + + self.noise_model_list: list[list[noise_cfg.NoiseCfg]] = [] + + for noise_group in noise_model_cfg.noise_model_groups.values(): + noise_models = list(noise_group.noise_cfg_dict.values()) + self.noise_model_list.append(noise_models) + + def reset(self, env_ids: Sequence[int] | None = None): + """Reset the noise model. + + This method resets resample the corresponding model group for the env_ids base on the proportion. + + Args: + env_ids: The environment ids to reset the noise model for. Defaults to None, + in which case all environments are considered. + """ + # resolve the environment ids + if env_ids is None: + env_ids = torch.arange(self._num_envs, device=self._device) + + self.noise_group_assignment[env_ids] = torch.multinomial(self.proportions, len(env_ids), replacement=True) + + def apply(self, data: torch.Tensor) -> torch.Tensor: + """Apply bias noise to the data. + + Args: + data: The data to apply the noise to. Shape is (num_envs, ...). + + Returns: + The data with the noise applied. Shape is the same as the input data. + """ + for i in range(len(self.noise_model_list)): + noise_models = self.noise_model_list[i] + noise_group = self.noise_group_assignment == i + for noise_model in noise_models: + data[noise_group] = noise_model.func(data[noise_group], noise_model) + + return data diff --git a/source/uwlab_assets/config/extension.toml b/source/uwlab_assets/config/extension.toml new file mode 100644 index 00000000..1cd133e5 --- /dev/null +++ b/source/uwlab_assets/config/extension.toml @@ -0,0 +1,18 @@ +[package] +# Semantic Versioning is used: https://semver.org/ +version = "0.5.1" + +# Description +title = "UW Lab Assets" +description="Extension containing configuration instances of different assets and sensors" +readme = "docs/README.md" +repository = "https://github.com/UW-Lab/UWLab" +category = "robotics" +keywords = ["isaaclab", "robotics", "rl", "il", "learning"] + +[dependencies] +"uwlab" = {} + +# Main python module this extension provides. +[[python.module]] +name = "uwlab_assets" diff --git a/source/uwlab_assets/data/.gitkeep b/source/uwlab_assets/data/.gitkeep new file mode 100644 index 00000000..3b41fe31 --- /dev/null +++ b/source/uwlab_assets/data/.gitkeep @@ -0,0 +1,6 @@ +For Isaac Lab, we primarily store assets on the Omniverse Nucleus server. However, at times, it may be +needed to store the assets locally (for debugging purposes). In such cases, this directory can be +used for temporary hosting of assets. + +Inside the `data` directory, we recommend following the same structure as our Nucleus directory +`Isaac/IsaacLab`. Please check the extension's README for further details. diff --git a/source/uwlab_assets/docs/CHANGELOG.rst b/source/uwlab_assets/docs/CHANGELOG.rst new file mode 100644 index 00000000..75d9f746 --- /dev/null +++ b/source/uwlab_assets/docs/CHANGELOG.rst @@ -0,0 +1,123 @@ +Changelog +--------- + +0.5.1 (2024-10-20) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* add teleoperation config to leap xarm + + +0.5.0 (2024-10-20) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* added custom franka setup at :folder:`uwlab_asset.franka` + +0.4.1 (2024-09-01) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Fixed the reversed y axis in Tycho Teleop Cfg at :const:`uwlab_asset.tycho.action.TELEOP_CFG` + +0.4.0 (2024-09-01) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Introduced Teleop Cfg in robot tycho and franka + +0.3.3 (2024-08-24) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* fixed the device inconsistency in :func:`uwlab_asset.tycho.mdp.termination:terminate_extremely_bad_posture` + where the device is hardcoded as "cuda" instead of "env.device" + +0.3.2 (2024-08-19) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* renamed HEBI_ORIBIT_ARTICULATION, and HEBI_CUSTOM_ARTICULATION to HEBI_ARTICULATION + at :file:`uwlab_asset.tycho.tycho.py` since they are all same. + + +0.3.1 (2024-08-19) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* fixed problem where the order of tycho gripper joint action idex and body joint pos are reversed + :class:`uwlab_asset.tycho.actions.IkdeltaAction` and :class:`uwlab_asset.tycho.actions.IkabsoluteAction` + +0.3.0 (2024-07-29) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* updated dependency and meta information to isaac sim 4.1.0 + + +0.2.0 (2024-07-29) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^^^ + +* Created new folder storing :class:`uwlab_asset.unitree` extensions + + + +0.1.3 (2024-07-27) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Bug fix at :const:`uwlab_assets.robots.leap.actions.LEAP_JOINT_POSITION` + and :const:`uwlab_assets.robots.leap.actions.LEAP_JOINT_EFFORT` because + previous version did not include all joint name. it used to be + ``joint_names=["j.*"]`` now becomes ``joint_names=["w.*", "j.*"]`` + + + + +0.1.2 (2024-07-27) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Created new folder storing ``uwlab_asset.anymal`` +* Created new folder storing ``uwlab_asset.leap`` + + +0.1.1 (2024-07-26) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Created new folder storing ``uwlab_asset.tycho`` + + +0.1.0 (2024-07-25) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Created new folder storing ``uwlab_asset`` diff --git a/source/uwlab_assets/docs/README.md b/source/uwlab_assets/docs/README.md new file mode 100644 index 00000000..67512f0a --- /dev/null +++ b/source/uwlab_assets/docs/README.md @@ -0,0 +1,41 @@ +# Isaac Lab: Assets for Robots and Objects + +This extension contains configurations for various assets and sensors. The configuration instances are +used to spawn and configure the instances in the simulation. They are passed to their corresponding +classes during construction. + +## Organizing custom assets + +For Isaac Lab, we primarily store assets on the Omniverse Nucleus server. However, at times, it may be +needed to store the assets locally (for debugging purposes). In such cases, the extension's `data` +directory can be used for temporary hosting of assets. + +Inside the `data` directory, we recommend following the same structure as our Nucleus directory +`Isaac/IsaacLab`. This helps us later to move these assets to the Nucleus server seamlessly. + +The recommended directory structure inside `data` is as follows: + +- **`Robots//`**: The USD files should be inside `` directory with + the name of the robot. +- **`Props//`**: The USD files should be inside `` directory with the name + of the prop. This includes mounts, objects and markers. +- **`ActuatorNets/`**: The actuator networks should inside ``**: The policy should be JIT/ONNX compiled with the name `policy.pt`. It should also + contain the parameters used for training the checkpoint. This is to ensure reproducibility. +- **`Test/`**: The asset used for unit testing purposes. + +## Referring to the assets in your code + +You can use the following snippet to refer to the assets: + +```python + +from isaaclab_assets import ISAACLAB_ASSETS_DATA_DIR + + +# ANYmal-C +ANYMAL_C_USD_PATH = f"{ISAACLAB_ASSETS_DATA_DIR}/Robots/ANYbotics/ANYmal-C/anymal_c.usd" +# ANYmal-D +ANYMAL_D_USD_PATH = f"{ISAACLAB_ASSETS_DATA_DIR}/Robots/ANYbotics/ANYmal-D/anymal_d.usd" +``` diff --git a/source/uwlab_assets/pyproject.toml b/source/uwlab_assets/pyproject.toml new file mode 100644 index 00000000..d90ac353 --- /dev/null +++ b/source/uwlab_assets/pyproject.toml @@ -0,0 +1,3 @@ +[build-system] +requires = ["setuptools", "wheel", "toml"] +build-backend = "setuptools.build_meta" diff --git a/source/uwlab_assets/setup.py b/source/uwlab_assets/setup.py new file mode 100644 index 00000000..20419906 --- /dev/null +++ b/source/uwlab_assets/setup.py @@ -0,0 +1,36 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Installation script for the 'uwlab_assets' python package.""" + +import os +import toml + +from setuptools import setup + +# Obtain the extension data from the extension.toml file +EXTENSION_PATH = os.path.dirname(os.path.realpath(__file__)) +# Read the extension.toml file +EXTENSION_TOML_DATA = toml.load(os.path.join(EXTENSION_PATH, "config", "extension.toml")) +# Installation operation +setup( + name="uwlab_assets", + author="UW and Isaac Lab Project Developers", + maintainer="UW and Isaac Lab Project Developers", + url=EXTENSION_TOML_DATA["package"]["repository"], + version=EXTENSION_TOML_DATA["package"]["version"], + description=EXTENSION_TOML_DATA["package"]["description"], + keywords=EXTENSION_TOML_DATA["package"]["keywords"], + license="BSD-3-Clause", + include_package_data=True, + python_requires=">=3.10", + packages=["uwlab_assets"], + classifiers=[ + "Natural Language :: English", + "Programming Language :: Python :: 3.10", + "Isaac Sim :: 4.5.0", + ], + zip_safe=False, +) diff --git a/source/uwlab_assets/test/test_valid_configs.py b/source/uwlab_assets/test/test_valid_configs.py new file mode 100644 index 00000000..c2fb237b --- /dev/null +++ b/source/uwlab_assets/test/test_valid_configs.py @@ -0,0 +1,70 @@ +# Copyright (c) 2024-2025, The UW 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 isaaclab.app import AppLauncher, run_tests + +# launch the simulator +app_launcher = AppLauncher(headless=True) +simulation_app = app_launcher.app + + +"""Rest everything follows.""" + +import unittest + +import isaaclab_assets as lab_assets # noqa: F401 + +from isaaclab.assets import AssetBase, AssetBaseCfg +from isaaclab.sim import build_simulation_context + + +class TestValidEntitiesConfigs(unittest.TestCase): + """Test cases for all registered entities configurations.""" + + @classmethod + def setUpClass(cls): + # load all registered entities configurations from the module + cls.registered_entities: dict[str, AssetBaseCfg] = {} + # inspect all classes from the module + for obj_name in dir(lab_assets): + obj = getattr(lab_assets, obj_name) + # store all registered entities configurations + if isinstance(obj, AssetBaseCfg): + cls.registered_entities[obj_name] = obj + # print all existing entities names + print(">>> All registered entities:", list(cls.registered_entities.keys())) + + """ + Test fixtures. + """ + + def test_asset_configs(self): + """Check all registered asset configurations.""" + # iterate over all registered assets + for asset_name, entity_cfg in self.registered_entities.items(): + for device in ("cuda:0", "cpu"): + with self.subTest(asset_name=asset_name, device=device): + with build_simulation_context(device=device, auto_add_lighting=True) as sim: + # print the asset name + print(f">>> Testing entity {asset_name} on device {device}") + # name the prim path + entity_cfg.prim_path = "/World/asset" + # create the asset / sensors + entity: AssetBase = entity_cfg.class_type(entity_cfg) # type: ignore + + # play the sim + sim.reset() + + # check asset is initialized successfully + self.assertTrue(entity.is_initialized) + + +if __name__ == "__main__": + run_tests() diff --git a/source/uwlab_assets/uwlab_assets/__init__.py b/source/uwlab_assets/uwlab_assets/__init__.py new file mode 100644 index 00000000..896f6a62 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Package containing asset and sensor configurations.""" + +import os +import toml + +# Conveniences to other module directories via relative paths +UWLAB_ASSETS_EXT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "../")) +"""Path to the extension source directory.""" +UWLAB_ASSETS_DATA_DIR = os.path.join(UWLAB_ASSETS_EXT_DIR, "data") +"""Path to the extension data directory.""" +UWLAB_ASSETS_METADATA = toml.load(os.path.join(UWLAB_ASSETS_EXT_DIR, "config", "extension.toml")) +"""Extension metadata dictionary parsed from the extension.toml file.""" +UWLAB_CLOUD_ASSETS_DIR = "https://uwlab-assets.s3.us-west-004.backblazeb2.com" +# Configure the module-level variables +__version__ = UWLAB_ASSETS_METADATA["package"]["version"] diff --git a/source/uwlab_assets/uwlab_assets/props/workbench/workbench_conversion_cfg.py b/source/uwlab_assets/uwlab_assets/props/workbench/workbench_conversion_cfg.py new file mode 100644 index 00000000..19a563c2 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/props/workbench/workbench_conversion_cfg.py @@ -0,0 +1,38 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from typing import List + +from uwlab.sim.converters import MeshConverterCfg +from uwlab.sim.converters.common_material_property_cfg import PVCCfg, SteelCfg +from uwlab.sim.spawners.materials import common_materials_cfg as common_materials + +BLOCK = PVCCfg( + asset_path="datasets/workbench/block.stl", + usd_dir="datasets/workbench", + usd_file_name="block.usda", + collision_approximation="convexDecomposition", + force_usd_conversion=True, + visual_material_props=common_materials.PCVVisualMaterialCfg(diffuse_color=(0.1, 0.8, 0.1)), +) + +BOX = PVCCfg( + asset_path="datasets/workbench/box.stl", + usd_dir="datasets/workbench", + usd_file_name="box.usda", + collision_approximation="convexDecomposition", + force_usd_conversion=True, +) + +SHELF = SteelCfg( + asset_path="datasets/workbench/shelf.stl", + usd_dir="datasets/workbench", + usd_file_name="shelf.usda", + collision_approximation="convexDecomposition", + force_usd_conversion=True, + visual_material_props=common_materials.SteelVisualMaterialCfg(diffuse_color=(0.1, 0.1, 0.4)), +) + +WORKBENCH_CONVERSION_CFG: List[MeshConverterCfg] = [BLOCK, BOX, SHELF] diff --git a/source/uwlab_assets/uwlab_assets/robots/articulation_drive/ur_driver.py b/source/uwlab_assets/uwlab_assets/robots/articulation_drive/ur_driver.py new file mode 100644 index 00000000..19943474 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/articulation_drive/ur_driver.py @@ -0,0 +1,106 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import numpy as np +import torch +from typing import TYPE_CHECKING + +import pymodbus.client as ModbusClient +import urx +from pymodbus.constants import Endian +from pymodbus.framer import Framer +from pymodbus.payload import BinaryPayloadBuilder + +from uwlab.assets.articulation.articulation_drive import ArticulationDrive + +if TYPE_CHECKING: + from .ur_driver_cfg import URDriverCfg + + +class URDriver(ArticulationDrive): + def __init__(self, cfg: URDriverCfg, data_indices: slice = slice(None)): + self.device = torch.device("cpu") + self.cfg = cfg + # self.work_space_limit = cfg.work_space_limit + self.data_idx = data_indices + + self.current_pos = torch.zeros(1, 6, device=self.device) + self.current_vel = torch.zeros(1, 6, device=self.device) + self.current_eff = torch.zeros(1, 6, device=self.device) + + @property + def ordered_joint_names(self): + return [ + "shoulder_pan_joint", + "shoulder_lift_joint", + "elbow_joint", + "wrist_1_joint", + "wrist_2_joint", + "wrist_3_joint", + ] + + def _prepare(self): + # Initialize urx connection + self.arm = urx.Robot(self.cfg.ip, use_rt=True) + + # Initialize Modbus Client + self.modbus_client = ModbusClient.ModbusTcpClient(self.cfg.ip, port=self.cfg.port, framer=Framer.SOCKET) + self.modbus_client.connect() + + def sendModbusValues(self, values): + # Values will be divided by 100 in URScript + values = np.array(values) * 100 + builder = BinaryPayloadBuilder(byteorder=Endian.BIG, wordorder=Endian.BIG) + # Loop through each pose value and write it to a register + for i in range(6): + builder.reset() + builder.add_16bit_int(int(values[i])) + payload = builder.to_registers() + self.modbus_client.write_register(128 + i, payload[0]) + + def write_dof_targets(self, pos_target: torch.Tensor, vel_target: torch.Tensor, eff_target: torch.Tensor): + # Non-blocking motion + + self.sendModbusValues(pos_target[0].tolist()) + + def read_dof_states(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Blocking call to get_joint_states, storing the data in local torch Tensors.""" + joint_pos = np.array(self.arm.getj()) + pos = torch.tensor(joint_pos, device=self.device) + self.current_pos[:] = pos + return pos, self.current_vel, self.current_eff + + def set_dof_stiffnesses(self, stiffnesses): + pass + + def set_dof_armatures(self, armatures): + pass + + def set_dof_frictions(self, frictions): + pass + + def set_dof_dampings(self, dampings): + pass + + def set_dof_limits(self, limits): + pass + + +# uncomment below code to run the worker +# if __name__ == "__main__": +# # Create the worker +# class Cfg: +# ip = "192.168.1.2" +# port = 602 +# driver = URDriver(cfg=Cfg()) +# driver._prepare() +# pos, vel, eff = driver.read_dof_states() +# print(pos, vel, eff) +# driver.write_dof_targets(pos, vel, eff) + +# while True: +# time.sleep(1) diff --git a/source/uwlab_assets/uwlab_assets/robots/articulation_drive/ur_driver_cfg.py b/source/uwlab_assets/uwlab_assets/robots/articulation_drive/ur_driver_cfg.py new file mode 100644 index 00000000..d3b326f5 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/articulation_drive/ur_driver_cfg.py @@ -0,0 +1,23 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from dataclasses import MISSING +from typing import Callable + +from isaaclab.utils import configclass +from uwlab.assets.articulation.articulation_drive import ArticulationDriveCfg + +from .ur_driver import URDriver + + +@configclass +class URDriverCfg(ArticulationDriveCfg): + class_type: Callable[..., URDriver] = URDriver + + ip: str = MISSING # type: ignore + + port: int = MISSING # type: ignore diff --git a/source/uwlab_assets/uwlab_assets/robots/franka/__init__.py b/source/uwlab_assets/uwlab_assets/robots/franka/__init__.py new file mode 100644 index 00000000..63b4a547 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/franka/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .action import * +from .teleop import * diff --git a/source/uwlab_assets/uwlab_assets/robots/franka/action.py b/source/uwlab_assets/uwlab_assets/robots/franka/action.py new file mode 100644 index 00000000..7ba83850 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/franka/action.py @@ -0,0 +1,60 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.controllers.differential_ik_cfg import DifferentialIKControllerCfg +from isaaclab.envs.mdp.actions.actions_cfg import ( + BinaryJointPositionActionCfg, + DifferentialInverseKinematicsActionCfg, + JointPositionActionCfg, +) +from isaaclab.utils import configclass + +JOINT_POSITION: JointPositionActionCfg = JointPositionActionCfg( + asset_name="robot", + joint_names=["panda_joint.*", "panda_finger_joint.*"], + scale=0.1, +) + +JOINT_IKDELTA: DifferentialInverseKinematicsActionCfg = 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.1034), rot=(1.0, 0.0, 0, 0)), +) + +JOINT_IKABSOLUTE: DifferentialInverseKinematicsActionCfg = DifferentialInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["panda_joint.*"], + body_name="panda_hand", + controller=DifferentialIKControllerCfg(command_type="pose", use_relative_mode=False, ik_method="dls"), + scale=1, + body_offset=DifferentialInverseKinematicsActionCfg.OffsetCfg(pos=(0.0, 0.0, 0.1034), rot=(1.0, 0.0, 0, 0)), +) + +BINARY_GRIPPER = BinaryJointPositionActionCfg( + asset_name="robot", + joint_names=["panda_finger.*"], + open_command_expr={"panda_finger_.*": 0.04}, + close_command_expr={"panda_finger_.*": 0.0}, +) + + +@configclass +class JointPositionAction: + jointpos = JOINT_POSITION + + +@configclass +class IkDeltaAction: + body_joint_pos = JOINT_IKDELTA + gripper_joint_pos = BINARY_GRIPPER + + +@configclass +class IkAbsoluteAction: + body_joint_pos = JOINT_IKABSOLUTE + gripper_joint_pos = BINARY_GRIPPER diff --git a/source/uwlab_assets/uwlab_assets/robots/franka/teleop.py b/source/uwlab_assets/uwlab_assets/robots/franka/teleop.py new file mode 100644 index 00000000..fbb45352 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/franka/teleop.py @@ -0,0 +1,29 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.managers import SceneEntityCfg +from isaaclab.utils import configclass +from uwlab.devices import KeyboardCfg, TeleopCfg + + +@configclass +class FrankaTeleopCfg: + keyboard: TeleopCfg = TeleopCfg( + teleop_devices={ + "device1": TeleopCfg.TeleopDevicesCfg( + attach_body=SceneEntityCfg("robot", body_names="panda_hand"), + attach_scope="self", + pose_reference_body=SceneEntityCfg("robot", body_names="panda_link0"), + reference_axis_remap=("x", "y", "z"), + command_type="pose", + debug_vis=True, + teleop_interface_cfg=KeyboardCfg( + pos_sensitivity=0.01, + rot_sensitivity=0.01, + enable_gripper_command=True, + ), + ), + } + ) diff --git a/source/uwlab_assets/uwlab_assets/robots/leap/__init__.py b/source/uwlab_assets/uwlab_assets/robots/leap/__init__.py new file mode 100644 index 00000000..979df227 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/leap/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .actions import * +from .leap import FRAME_EE, IMPLICIT_LEAP, IMPLICIT_LEAP6D diff --git a/source/uwlab_assets/uwlab_assets/robots/leap/actions.py b/source/uwlab_assets/uwlab_assets/robots/leap/actions.py new file mode 100644 index 00000000..63b56e00 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/leap/actions.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +from isaaclab.envs.mdp.actions.actions_cfg import JointEffortActionCfg, JointPositionActionCfg +from isaaclab.utils import configclass +from uwlab.controllers.differential_ik_cfg import MultiConstraintDifferentialIKControllerCfg +from uwlab.envs.mdp.actions.actions_cfg import MultiConstraintsDifferentialInverseKinematicsActionCfg + +""" +LEAP ACTIONS +""" + +LEAP_JOINT_POSITION: JointPositionActionCfg = JointPositionActionCfg( + asset_name="robot", joint_names=["w.*", "j.*"], scale=0.1 +) + +LEAP_JOINT_EFFORT: JointEffortActionCfg = JointEffortActionCfg( + asset_name="robot", joint_names=["w.*", "j.*"], scale=0.1 +) + +LEAP_MC_IKABSOLUTE = MultiConstraintsDifferentialInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["w.*", "j.*"], + body_name=["wrist", "pip", "pip_2", "pip_3", "thumb_fingertip", "tip", "tip_2", "tip_3"], + controller=MultiConstraintDifferentialIKControllerCfg( + command_type="position", use_relative_mode=False, ik_method="dls" + ), + scale=1, +) + +LEAP_MC_IKDELTA = MultiConstraintsDifferentialInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["w.*", "j.*"], + body_name=["wrist", "pip", "pip_2", "pip_3", "thumb_fingertip", "tip", "tip_2", "tip_3"], + controller=MultiConstraintDifferentialIKControllerCfg( + command_type="position", use_relative_mode=True, ik_method="dls" + ), + scale=0.1, +) + + +@configclass +class LeapMcIkAbsoluteAction: + joint_pos = LEAP_MC_IKABSOLUTE + + +@configclass +class LeapMcIkDeltaAction: + joint_pos = LEAP_MC_IKDELTA + + +@configclass +class LeapJointPositionAction: + joint_pos = LEAP_JOINT_POSITION + + +@configclass +class LeapJointEffortAction: + joint_pos = LEAP_JOINT_EFFORT diff --git a/source/uwlab_assets/uwlab_assets/robots/leap/articulation_drive/__init__.py b/source/uwlab_assets/uwlab_assets/robots/leap/articulation_drive/__init__.py new file mode 100644 index 00000000..80d7de2f --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/leap/articulation_drive/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .dynamixel_driver import DynamixelDriver +from .dynamixel_driver_cfg import DynamixelDriverCfg diff --git a/source/uwlab_assets/uwlab_assets/robots/leap/articulation_drive/dynamixel_client.py b/source/uwlab_assets/uwlab_assets/robots/leap/articulation_drive/dynamixel_client.py new file mode 100644 index 00000000..c7414bf0 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/leap/articulation_drive/dynamixel_client.py @@ -0,0 +1,583 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# This is based off of the dynamixel SDK +import atexit +import logging +import numpy as np +import time +from typing import Optional, Sequence, Tuple, Union + +PROTOCOL_VERSION = 2.0 + +# The following addresses assume XH motors. +ADDR_TORQUE_ENABLE = 64 +ADDR_GOAL_POSITION = 116 +ADDR_PRESENT_POSITION = 132 +ADDR_PRESENT_VELOCITY = 128 +ADDR_PRESENT_CURRENT = 126 +ADDR_PRESENT_POS_VEL_CUR = 126 + +# Data Byte Length +LEN_PRESENT_POSITION = 4 +LEN_PRESENT_VELOCITY = 4 +LEN_PRESENT_CURRENT = 2 +LEN_PRESENT_POS_VEL_CUR = 10 +LEN_GOAL_POSITION = 4 + +DEFAULT_POS_SCALE = 2.0 * np.pi / 4096 # 0.088 degrees +# See http://emanual.robotis.com/docs/en/dxl/x/xh430-v210/#goal-velocity +DEFAULT_VEL_SCALE = 0.229 * 2.0 * np.pi / 60.0 # 0.229 rpm +DEFAULT_CUR_SCALE = 1.34 + + +def dynamixel_cleanup_handler(): + """Cleanup function to ensure Dynamixels are disconnected properly.""" + open_clients = list(DynamixelClient.OPEN_CLIENTS) + for open_client in open_clients: + if open_client.port_handler.is_using: + logging.warning("Forcing client to close.") + open_client.port_handler.is_using = False + open_client.disconnect() + + +def signed_to_unsigned(value: int, size: int) -> int: + """Converts the given value to its unsigned representation.""" + if value < 0: + bit_size = 8 * size + max_value = (1 << bit_size) - 1 + value = max_value + value + return value + + +def unsigned_to_signed(value: int, size: int) -> int: + """Converts the given value from its unsigned representation.""" + bit_size = 8 * size + if (value & (1 << (bit_size - 1))) != 0: + value = -((1 << bit_size) - value) + return value + + +class DynamixelClient: + """Client for communicating with Dynamixel motors. + NOTE: This only supports Protocol 2. + """ + + # The currently open clients. + OPEN_CLIENTS = set() + + def __init__( + self, + motor_ids: Sequence[int], + port: str = "/dev/ttyUSB0", + baudrate: int = 1000000, + lazy_connect: bool = False, + pos_scale: Optional[float] = None, + vel_scale: Optional[float] = None, + cur_scale: Optional[float] = None, + ): + """Initializes a new client. + Args: + motor_ids: All motor IDs being used by the client. + port: The Dynamixel device to talk to. e.g. + - Linux: /dev/ttyUSB0 + - Mac: /dev/tty.usbserial-* + - Windows: COM1 + baudrate: The Dynamixel baudrate to communicate with. + lazy_connect: If True, automatically connects when calling a method + that requires a connection, if not already connected. + pos_scale: The scaling factor for the positions. This is + motor-dependent. If not provided, uses the default scale. + vel_scale: The scaling factor for the velocities. This is + motor-dependent. If not provided uses the default scale. + cur_scale: The scaling factor for the currents. This is + motor-dependent. If not provided uses the default scale. + """ + import dynamixel_sdk + + self.dxl = dynamixel_sdk + + self.motor_ids = list(motor_ids) + self.port_name = port + self.baudrate = baudrate + self.lazy_connect = lazy_connect + + self.port_handler = self.dxl.PortHandler(port) + self.packet_handler = self.dxl.PacketHandler(PROTOCOL_VERSION) + + self._pos_vel_cur_reader = DynamixelPosVelCurReader( + self, + self.motor_ids, + pos_scale=pos_scale if pos_scale is not None else DEFAULT_POS_SCALE, + vel_scale=vel_scale if vel_scale is not None else DEFAULT_VEL_SCALE, + cur_scale=cur_scale if cur_scale is not None else DEFAULT_CUR_SCALE, + ) + self._pos_reader = DynamixelPosReader( + self, + self.motor_ids, + pos_scale=pos_scale if pos_scale is not None else DEFAULT_POS_SCALE, + vel_scale=vel_scale if vel_scale is not None else DEFAULT_VEL_SCALE, + cur_scale=cur_scale if cur_scale is not None else DEFAULT_CUR_SCALE, + ) + self._vel_reader = DynamixelVelReader( + self, + self.motor_ids, + pos_scale=pos_scale if pos_scale is not None else DEFAULT_POS_SCALE, + vel_scale=vel_scale if vel_scale is not None else DEFAULT_VEL_SCALE, + cur_scale=cur_scale if cur_scale is not None else DEFAULT_CUR_SCALE, + ) + self._cur_reader = DynamixelCurReader( + self, + self.motor_ids, + pos_scale=pos_scale if pos_scale is not None else DEFAULT_POS_SCALE, + vel_scale=vel_scale if vel_scale is not None else DEFAULT_VEL_SCALE, + cur_scale=cur_scale if cur_scale is not None else DEFAULT_CUR_SCALE, + ) + self._sync_writers = {} + + self.OPEN_CLIENTS.add(self) + + @property + def is_connected(self) -> bool: + return self.port_handler.is_open + + def connect(self): + """Connects to the Dynamixel motors. + NOTE: This should be called after all DynamixelClients on the same + process are created. + """ + assert not self.is_connected, "Client is already connected." + + if self.port_handler.openPort(): + logging.info("Succeeded to open port: %s", self.port_name) + else: + raise OSError( + ( + "Failed to open port at {} (Check that the device is powered on and connected to your computer)." + ).format(self.port_name) + ) + + if self.port_handler.setBaudRate(self.baudrate): + logging.info("Succeeded to set baudrate to %d", self.baudrate) + else: + raise OSError( + ("Failed to set the baudrate to {} (Ensure that the device was configured for this baudrate).").format( + self.baudrate + ) + ) + + # Start with all motors enabled. NO, I want to set settings before enabled + # self.set_torque_enabled(self.motor_ids, True) + + def disconnect(self): + """Disconnects from the Dynamixel device.""" + if not self.is_connected: + return + if self.port_handler.is_using: + logging.error("Port handler in use; cannot disconnect.") + return + # Ensure motors are disabled at the end. + self.set_torque_enabled(self.motor_ids, False, retries=0) + self.port_handler.closePort() + if self in self.OPEN_CLIENTS: + self.OPEN_CLIENTS.remove(self) + + def set_torque_enabled( + self, motor_ids: Sequence[int], enabled: bool, retries: int = -1, retry_interval: float = 0.25 + ): + """Sets whether torque is enabled for the motors. + Args: + motor_ids: The motor IDs to configure. + enabled: Whether to engage or disengage the motors. + retries: The number of times to retry. If this is <0, will retry + forever. + retry_interval: The number of seconds to wait between retries. + """ + remaining_ids = list(motor_ids) + while remaining_ids: + remaining_ids = self.write_byte( + remaining_ids, + int(enabled), + ADDR_TORQUE_ENABLE, + ) + if remaining_ids: + logging.error( + "Could not set torque %s for IDs: %s", "enabled" if enabled else "disabled", str(remaining_ids) + ) + if retries == 0: + break + time.sleep(retry_interval) + retries -= 1 + + def read_pos_vel_cur(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Returns the current positions and velocities.""" + return self._pos_vel_cur_reader.read() + + def read_pos(self) -> np.ndarray: + """Returns the current positions and velocities.""" + return self._pos_reader.read() + + def read_vel(self) -> np.ndarray: + """Returns the current positions and velocities.""" + return self._vel_reader.read() + + def read_cur(self) -> np.ndarray: + """Returns the current positions and velocities.""" + return self._cur_reader.read() + + def write_desired_pos(self, motor_ids: Sequence[int], positions: np.ndarray): + """Writes the given desired positions. + Args: + motor_ids: The motor IDs to write to. + positions: The joint angles in radians to write. + """ + assert len(motor_ids) == len(positions) + + # Convert to Dynamixel position space. + positions = positions / self._pos_vel_cur_reader.pos_scale + self.sync_write(motor_ids, positions, ADDR_GOAL_POSITION, LEN_GOAL_POSITION) + + def write_byte( + self, + motor_ids: Sequence[int], + value: int, + address: int, + ) -> Sequence[int]: + """Writes a value to the motors. + Args: + motor_ids: The motor IDs to write to. + value: The value to write to the control table. + address: The control table address to write to. + Returns: + A list of IDs that were unsuccessful. + """ + self.check_connected() + errored_ids = [] + for motor_id in motor_ids: + comm_result, dxl_error = self.packet_handler.write1ByteTxRx(self.port_handler, motor_id, address, value) + success = self.handle_packet_result(comm_result, dxl_error, motor_id, context="write_byte") + if not success: + errored_ids.append(motor_id) + return errored_ids + + def sync_write(self, motor_ids: Sequence[int], values: Sequence[Union[int, float]], address: int, size: int): + """Writes values to a group of motors. + Args: + motor_ids: The motor IDs to write to. + values: The values to write. + address: The control table address to write to. + size: The size of the control table value being written to. + """ + self.check_connected() + key = (address, size) + if key not in self._sync_writers: + self._sync_writers[key] = self.dxl.GroupSyncWrite(self.port_handler, self.packet_handler, address, size) + sync_writer = self._sync_writers[key] + + errored_ids = [] + for motor_id, desired_pos in zip(motor_ids, values): + value = signed_to_unsigned(int(desired_pos), size=size) + value = value.to_bytes(size, byteorder="little") + success = sync_writer.addParam(motor_id, value) + if not success: + errored_ids.append(motor_id) + + if errored_ids: + logging.error("Sync write failed for: %s", str(errored_ids)) + + comm_result = sync_writer.txPacket() + self.handle_packet_result(comm_result, context="sync_write") + + sync_writer.clearParam() + + def check_connected(self): + """Ensures the robot is connected.""" + if self.lazy_connect and not self.is_connected: + self.connect() + if not self.is_connected: + raise OSError("Must call connect() first.") + + def handle_packet_result( + self, + comm_result: int, + dxl_error: Optional[int] = None, + dxl_id: Optional[int] = None, + context: Optional[str] = None, + ): + """Handles the result from a communication request.""" + error_message = None + if comm_result != self.dxl.COMM_SUCCESS: + error_message = self.packet_handler.getTxRxResult(comm_result) + elif dxl_error is not None: + error_message = self.packet_handler.getRxPacketError(dxl_error) + if error_message: + if dxl_id is not None: + error_message = f"[Motor ID: {dxl_id}] {error_message}" + if context is not None: + error_message = f"> {context}: {error_message}" + # logging.error(error_message) + return False + return True + + def convert_to_unsigned(self, value: int, size: int) -> int: + """Converts the given value to its unsigned representation.""" + if value < 0: + max_value = (1 << (8 * size)) - 1 + value = max_value + value + return value + + def __enter__(self): + """Enables use as a context manager.""" + if not self.is_connected: + self.connect() + return self + + def __exit__(self, *args): + """Enables use as a context manager.""" + self.disconnect() + + def __del__(self): + """Automatically disconnect on destruction.""" + self.disconnect() + + +class DynamixelReader: + """Reads data from Dynamixel motors. + This wraps a GroupBulkRead from the DynamixelSDK. + """ + + def __init__(self, client: DynamixelClient, motor_ids: Sequence[int], address: int, size: int): + """Initializes a new reader.""" + self.client = client + self.motor_ids = motor_ids + self.address = address + self.size = size + self._initialize_data() + + self.operation = self.client.dxl.GroupBulkRead(client.port_handler, client.packet_handler) + + for motor_id in motor_ids: + success = self.operation.addParam(motor_id, address, size) + if not success: + raise OSError(f"[Motor ID: {motor_id}] Could not add parameter to bulk read.") + + def read(self): + """Reads data from the motors.""" + self.client.check_connected() + success = False + while not success: + comm_result = self.operation.txRxPacket() + success = self.client.handle_packet_result(comm_result, context="read") + + # If we failed, send a copy of the previous data. + if not success: + return self._get_data() + + errored_ids = [] + for i, motor_id in enumerate(self.motor_ids): + # Check if the data is available. + available = self.operation.isAvailable(motor_id, self.address, self.size) + if not available: + errored_ids.append(motor_id) + continue + + self._update_data(i, motor_id) + + if errored_ids: + logging.error("Bulk read data is unavailable for: %s", str(errored_ids)) + + return self._get_data() + + def _initialize_data(self): + """Initializes the cached data.""" + self._data = np.zeros(len(self.motor_ids), dtype=np.float32) + + def _update_data(self, index: int, motor_id: int): + """Updates the data index for the given motor ID.""" + self._data[index] = self.operation.getData(motor_id, self.address, self.size) + + def _get_data(self): + """Returns a copy of the data.""" + return self._data.copy() + + +class DynamixelPosVelCurReader(DynamixelReader): + """Reads positions and velocities.""" + + def __init__( + self, + client: DynamixelClient, + motor_ids: Sequence[int], + pos_scale: float = 1.0, + vel_scale: float = 1.0, + cur_scale: float = 1.0, + ): + super().__init__( + client, + motor_ids, + address=ADDR_PRESENT_POS_VEL_CUR, + size=LEN_PRESENT_POS_VEL_CUR, + ) + self.pos_scale = pos_scale + self.vel_scale = vel_scale + self.cur_scale = cur_scale + + def _initialize_data(self): + """Initializes the cached data.""" + self._pos_data = np.zeros(len(self.motor_ids), dtype=np.float32) + self._vel_data = np.zeros(len(self.motor_ids), dtype=np.float32) + self._cur_data = np.zeros(len(self.motor_ids), dtype=np.float32) + + def _update_data(self, index: int, motor_id: int): + """Updates the data index for the given motor ID.""" + cur = self.operation.getData(motor_id, ADDR_PRESENT_CURRENT, LEN_PRESENT_CURRENT) + vel = self.operation.getData(motor_id, ADDR_PRESENT_VELOCITY, LEN_PRESENT_VELOCITY) + pos = self.operation.getData(motor_id, ADDR_PRESENT_POSITION, LEN_PRESENT_POSITION) + cur = unsigned_to_signed(cur, size=2) + vel = unsigned_to_signed(vel, size=4) + pos = unsigned_to_signed(pos, size=4) + self._pos_data[index] = float(pos) * self.pos_scale + self._vel_data[index] = float(vel) * self.vel_scale + self._cur_data[index] = float(cur) * self.cur_scale + + def _get_data(self): + """Returns a copy of the data.""" + return (self._pos_data.copy(), self._vel_data.copy(), self._cur_data.copy()) + + +class DynamixelPosReader(DynamixelReader): + """Reads positions and velocities.""" + + def __init__( + self, + client: DynamixelClient, + motor_ids: Sequence[int], + pos_scale: float = 1.0, + vel_scale: float = 1.0, + cur_scale: float = 1.0, + ): + super().__init__( + client, + motor_ids, + address=ADDR_PRESENT_POS_VEL_CUR, + size=LEN_PRESENT_POS_VEL_CUR, + ) + self.pos_scale = pos_scale + + def _initialize_data(self): + """Initializes the cached data.""" + self._pos_data = np.zeros(len(self.motor_ids), dtype=np.float32) + + def _update_data(self, index: int, motor_id: int): + """Updates the data index for the given motor ID.""" + pos = self.operation.getData(motor_id, ADDR_PRESENT_POSITION, LEN_PRESENT_POSITION) + pos = unsigned_to_signed(pos, size=4) + self._pos_data[index] = float(pos) * self.pos_scale + + def _get_data(self): + """Returns a copy of the data.""" + return self._pos_data.copy() + + +class DynamixelVelReader(DynamixelReader): + """Reads positions and velocities.""" + + def __init__( + self, + client: DynamixelClient, + motor_ids: Sequence[int], + pos_scale: float = 1.0, + vel_scale: float = 1.0, + cur_scale: float = 1.0, + ): + super().__init__( + client, + motor_ids, + address=ADDR_PRESENT_POS_VEL_CUR, + size=LEN_PRESENT_POS_VEL_CUR, + ) + self.pos_scale = pos_scale + self.vel_scale = vel_scale + self.cur_scale = cur_scale + + def _initialize_data(self): + """Initializes the cached data.""" + self._vel_data = np.zeros(len(self.motor_ids), dtype=np.float32) + + def _update_data(self, index: int, motor_id: int): + """Updates the data index for the given motor ID.""" + vel = self.operation.getData(motor_id, ADDR_PRESENT_VELOCITY, LEN_PRESENT_VELOCITY) + vel = unsigned_to_signed(vel, size=4) + self._vel_data[index] = float(vel) * self.vel_scale + + def _get_data(self): + """Returns a copy of the data.""" + return self._vel_data.copy() + + +class DynamixelCurReader(DynamixelReader): + """Reads positions and velocities.""" + + def __init__( + self, + client: DynamixelClient, + motor_ids: Sequence[int], + pos_scale: float = 1.0, + vel_scale: float = 1.0, + cur_scale: float = 1.0, + ): + super().__init__( + client, + motor_ids, + address=ADDR_PRESENT_POS_VEL_CUR, + size=LEN_PRESENT_POS_VEL_CUR, + ) + self.cur_scale = cur_scale + + def _initialize_data(self): + """Initializes the cached data.""" + self._cur_data = np.zeros(len(self.motor_ids), dtype=np.float32) + + def _update_data(self, index: int, motor_id: int): + """Updates the data index for the given motor ID.""" + cur = self.operation.getData(motor_id, ADDR_PRESENT_CURRENT, LEN_PRESENT_CURRENT) + cur = unsigned_to_signed(cur, size=2) + self._cur_data[index] = float(cur) * self.cur_scale + + def _get_data(self): + """Returns a copy of the data.""" + return self._cur_data.copy() + + +# Register global cleanup function. +atexit.register(dynamixel_cleanup_handler) + +if __name__ == "__main__": + import argparse + import itertools + + parser = argparse.ArgumentParser() + parser.add_argument("-m", "--motors", required=True, help="Comma-separated list of motor IDs.") + parser.add_argument("-d", "--device", default="/dev/ttyUSB0", help="The Dynamixel device to connect to.") + parser.add_argument("-b", "--baud", default=1000000, help="The baudrate to connect with.") + parsed_args = parser.parse_args() + + motors = [int(motor) for motor in parsed_args.motors.split(",")] + + way_points = [np.zeros(len(motors)), np.full(len(motors), np.pi)] + + with DynamixelClient(motors, parsed_args.device, parsed_args.baud) as dxl_client: + for step in itertools.count(): + if step > 0 and step % 50 == 0: + way_point = way_points[(step // 100) % len(way_points)] + print(f"Writing: {way_point.tolist()}") + dxl_client.write_desired_pos(motors, way_point) + read_start = time.time() + pos_now, vel_now, cur_now = dxl_client.read_pos_vel_cur() + if step % 5 == 0: + print(f"[{step}] Frequency: {1.0 / (time.time() - read_start):.2f} Hz") + print(f"> Pos: {pos_now.tolist()}") + print(f"> Vel: {vel_now.tolist()}") + print(f"> Cur: {cur_now.tolist()}") diff --git a/source/uwlab_assets/uwlab_assets/robots/leap/articulation_drive/dynamixel_driver.py b/source/uwlab_assets/uwlab_assets/robots/leap/articulation_drive/dynamixel_driver.py new file mode 100644 index 00000000..051089ef --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/leap/articulation_drive/dynamixel_driver.py @@ -0,0 +1,82 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from uwlab.assets.articulation.articulation_drive import ArticulationDrive + +if TYPE_CHECKING: + from .dynamixel_driver_cfg import DynamixelDriverCfg + + +class DynamixelDriver(ArticulationDrive): + def __init__(self, cfg: DynamixelDriverCfg, data_indices: slice = slice(None)): + self.device = cfg.device + assert self.device == "cpu", "Dynamixel driver only supports CPU mode" + self.cfg = cfg + self.data_idx = data_indices + hand_kI = cfg.hand_kI + hand_curr_lim = cfg.hand_curr_lim + self.offset = 3.14159 + + # Initialize client in CPU-only environment + self.joint_idx = torch.arange(16, device=self.device) + self.joint_integral = torch.ones(len(self.joint_idx), device=self.device) * hand_kI + self.magic = torch.ones(len(self.joint_idx), device=self.device) * 5 + self.joint_curr_lim = torch.ones(len(self.joint_idx), device=self.device) * hand_curr_lim + self._prepare() + + @property + def ordered_joint_names(self): + return ["j" + str(i) for i in range(16)] + + def close(self): + self.dxl_client.disconnect() + + def read_dof_states(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + pos = torch.tensor((self.dxl_client.read_pos() - self.offset), device=self.device).view(1, -1) + vel = torch.tensor(self.dxl_client.read_vel(), device=self.device).view(1, -1) + eff = torch.zeros_like(pos, device=self.device) + return pos, vel, eff + + def write_dof_targets(self, pos_target: torch.Tensor, vel_target: torch.Tensor, eff_target: torch.Tensor): + pos_target = pos_target[0] + self.offset + self.dxl_client.write_desired_pos(self.joint_idx.tolist(), pos_target.numpy()) + + def set_dof_stiffnesses(self, stiffnesses): + self.dxl_client.sync_write(self.joint_idx.tolist(), stiffnesses[0].tolist(), 84, 2) # Pgain stiffness + + def set_dof_armatures(self, armatures): + pass + + def set_dof_frictions(self, frictions): + pass + + def set_dof_dampings(self, dampings): + self.dxl_client.sync_write(self.joint_idx.tolist(), dampings[0].tolist(), 80, 2) + + def set_dof_limits(self, limits): + pass + + def _prepare(self): + from .dynamixel_client import DynamixelClient + + self.dxl_client = DynamixelClient( + port=self.cfg.port, + motor_ids=[i for i in range(16)], + baudrate=4000000, + ) + self.dxl_client.connect() + self.dxl_client.sync_write(self.joint_idx.tolist(), self.joint_integral.tolist(), 82, 2) # Igain + self.dxl_client.sync_write(self.joint_idx.tolist(), self.magic.tolist(), 11, 1) + self.dxl_client.set_torque_enabled(self.joint_idx.tolist(), True) + + # Max at current (in unit 1ma) so don't overheat and grip too hard #500 normal or #350 for lite + self.dxl_client.sync_write(self.joint_idx.tolist(), self.joint_curr_lim.tolist(), 102, 2) + + print("Dynamixel driver initialized") diff --git a/source/uwlab_assets/uwlab_assets/robots/leap/articulation_drive/dynamixel_driver_cfg.py b/source/uwlab_assets/uwlab_assets/robots/leap/articulation_drive/dynamixel_driver_cfg.py new file mode 100644 index 00000000..7c2fa96e --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/leap/articulation_drive/dynamixel_driver_cfg.py @@ -0,0 +1,25 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from dataclasses import MISSING +from typing import Callable + +from isaaclab.utils import configclass +from uwlab.assets.articulation.articulation_drive import ArticulationDriveCfg + +from .dynamixel_driver import DynamixelDriver + + +@configclass +class DynamixelDriverCfg(ArticulationDriveCfg): + class_type: Callable[..., DynamixelDriver] = DynamixelDriver + + port: str = MISSING # type: ignore + + hand_kI: float = 0.0 # type: ignore + + hand_curr_lim: float = 350 # type: ignore diff --git a/source/uwlab_assets/uwlab_assets/robots/leap/articulation_drive/installation.md b/source/uwlab_assets/uwlab_assets/robots/leap/articulation_drive/installation.md new file mode 100644 index 00000000..baa4cf34 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/leap/articulation_drive/installation.md @@ -0,0 +1,23 @@ +# Leap Hand Hardware Installation Notes + +## Installation Steps + +### 1. Install Required Packages +To control the hand, install the necessary Python package: +```bash +pip install dynamixel-sdk +``` + +### 2. Install Additional Packages for Teleoperation with Rokoko +If you need to teleoperate the hand with Rokoko, install these additional packages +```bash +pip install lz4 pyrealsense2==2.53.1.4623 +sudo apt install librealsense2=2.53.1-0~realsense0.8251 +sudo apt install librealsense2-gl=2.53.1-0~realsense0.8251 +sudo apt install librealsense2-net=2.53.1-0~realsense0.8251 +sudo apt install librealsense2-utils=2.53.1-0~realsense0.8251 +``` +and configure the following firewall settings: +```bash +sudo ufw allow from your_mac_ip/24 +``` diff --git a/source/uwlab_assets/uwlab_assets/robots/leap/leap.py b/source/uwlab_assets/uwlab_assets/robots/leap/leap.py new file mode 100644 index 00000000..2bf5f0e8 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/leap/leap.py @@ -0,0 +1,117 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from uwlab_assets import UWLAB_CLOUD_ASSETS_DIR + +import isaaclab.sim as sim_utils +from isaaclab.actuators import ImplicitActuatorCfg +from isaaclab.assets.articulation import ArticulationCfg +from isaaclab.markers.config import FRAME_MARKER_CFG +from isaaclab.sensors import FrameTransformerCfg +from isaaclab.sensors.frame_transformer.frame_transformer_cfg import OffsetCfg + +## +# Configuration +## + +# fmt: off +LEAP_DEFAULT_JOINT_POS = {".*": 0.0} + +LEAP6D_DEFAULT_JOINT_POS = {".*": 0.0} +# fmt: on + + +LEAP_ARTICULATION = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{UWLAB_CLOUD_ASSETS_DIR}/Robots/LeapHand/leap_hand.usd", + activate_contact_sensors=False, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + retain_accelerations=True, + max_depenetration_velocity=5.0, + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=True, solver_position_iteration_count=8, solver_velocity_iteration_count=0 + ), + joint_drive_props=sim_utils.JointDrivePropertiesCfg(drive_type="force"), + ), + init_state=ArticulationCfg.InitialStateCfg(pos=(0, 0, 0), rot=(1, 0, 0, 0), joint_pos=LEAP_DEFAULT_JOINT_POS), + soft_joint_pos_limit_factor=1, +) + +LEAP6D_ARTICULATION = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{UWLAB_CLOUD_ASSETS_DIR}/Robots/LeapHand/leap_hand_6d.usd", + activate_contact_sensors=False, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=True, + max_depenetration_velocity=5.0, + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=True, solver_position_iteration_count=1, solver_velocity_iteration_count=0 + ), + ), + init_state=ArticulationCfg.InitialStateCfg(pos=(0, 0, 0), rot=(1, 0, 0, 0), joint_pos=LEAP6D_DEFAULT_JOINT_POS), + soft_joint_pos_limit_factor=1, +) + +IMPLICIT_LEAP = LEAP_ARTICULATION.copy() +IMPLICIT_LEAP.actuators = { + "j": ImplicitActuatorCfg( + joint_names_expr=["j.*"], + stiffness=200.0, + damping=30.0, + armature=0.001, + friction=0.2, + velocity_limit=8.48, + effort_limit=0.95, + ), +} + +IMPLICIT_LEAP6D = LEAP6D_ARTICULATION.copy() # type: ignore +IMPLICIT_LEAP6D.actuators = { + "w": ImplicitActuatorCfg( + joint_names_expr=["w.*"], + stiffness=200.0, + damping=50.0, + armature=0.001, + friction=0.2, + velocity_limit=1, + effort_limit=50, + ), + "j": ImplicitActuatorCfg( + joint_names_expr=["j.*"], + stiffness=200.0, + damping=30.0, + armature=0.001, + friction=0.2, + velocity_limit=8.48, + effort_limit=0.95, + ), +} + + +""" +FRAMES +""" +marker_cfg = FRAME_MARKER_CFG.copy() # type: ignore +marker_cfg.markers["frame"].scale = (0.01, 0.01, 0.01) +marker_cfg.prim_path = "/Visuals/FrameTransformer" + +FRAME_EE = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/link_base", + debug_vis=False, + visualizer_cfg=marker_cfg, + target_frames=[ + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/palm_lower", + name="ee", + offset=OffsetCfg( + pos=(-0.028, -0.04, -0.07), + rot=(1.0, 0.0, 0.0, 0.0), + ), + ), + ], +) diff --git a/source/uwlab_assets/uwlab_assets/robots/leap/mdp/__init__.py b/source/uwlab_assets/uwlab_assets/robots/leap/mdp/__init__.py new file mode 100644 index 00000000..fe7e5ea5 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/leap/mdp/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .actions import * diff --git a/source/uwlab_assets/uwlab_assets/robots/leap/mdp/actions/__init__.py b/source/uwlab_assets/uwlab_assets/robots/leap/mdp/actions/__init__.py new file mode 100644 index 00000000..ddb1b62f --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/leap/mdp/actions/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .actions_cfg import LeapJointPositionActionCorrectionCfg diff --git a/source/uwlab_assets/uwlab_assets/robots/leap/mdp/actions/actions.py b/source/uwlab_assets/uwlab_assets/robots/leap/mdp/actions/actions.py new file mode 100644 index 00000000..9b32c98a --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/leap/mdp/actions/actions.py @@ -0,0 +1,53 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from isaaclab.envs import ManagerBasedEnv +from isaaclab.envs.mdp.actions import JointAction + +if TYPE_CHECKING: + from . import actions_cfg + + +class LeapJointPositionActionCorrection(JointAction): + """Joint action term that applies the processed actions to the articulation's joints as position commands.""" + + cfg: actions_cfg.LeapJointPositionActionCorrectionCfg + """The configuration of the action term.""" + + def __init__(self, cfg: actions_cfg.LeapJointPositionActionCorrectionCfg, env: ManagerBasedEnv): + # initialize the action term + super().__init__(cfg, env) + # 0 open, 1.57 bent + self.knuckle_idx, knuckle_names = self._asset.find_joints(["j1", "j5", "j9"], preserve_order=True) + # -0.34 straight, +-1.04 bent + self.tip_idx, tip_names = self._asset.find_joints(["j3", "j7", "j11"], preserve_order=True) + # -0.4886 straight, +2.02 bent + self.dip_idx, dip_names = self._asset.find_joints(["j2", "j6", "j10"], preserve_order=True) + # 0 straight, +-1.04 bent + self.mpc_idx, mpc_names = self._asset.find_joints(["j0", "j4", "j8"], preserve_order=True) + + @property + def action_dim(self) -> int: + return 0 + + def process_actions(self, actions): + pass + + def apply_actions(self): + target_pos = self._asset.data.joint_pos_target.clone() + knuckle_bending_pos = target_pos[:, self.knuckle_idx] + target_pos[:, self.mpc_idx] -= (knuckle_bending_pos / 1.00).clamp(0, 1.00) * target_pos[:, self.mpc_idx] + + target_pos[:, self.dip_idx] = target_pos[:, self.dip_idx].clip(min=0.0) + # print(target_pos[:, self.dip_idx]) + dip_bending_pos = target_pos[:, self.dip_idx] + tip_max = dip_bending_pos.maximum(knuckle_bending_pos) + tip_min = dip_bending_pos.minimum(knuckle_bending_pos) / 2 + target_pos[:, self.tip_idx] = target_pos[:, self.tip_idx].maximum(tip_min).minimum(tip_max) + self._asset.set_joint_position_target(target_pos, self._joint_ids) diff --git a/source/uwlab_assets/uwlab_assets/robots/leap/mdp/actions/actions_cfg.py b/source/uwlab_assets/uwlab_assets/robots/leap/mdp/actions/actions_cfg.py new file mode 100644 index 00000000..0efb3c9c --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/leap/mdp/actions/actions_cfg.py @@ -0,0 +1,26 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from isaaclab.envs.mdp.actions import JointPositionActionCfg +from isaaclab.managers.action_manager import ActionTerm +from isaaclab.utils import configclass + +from .actions import LeapJointPositionActionCorrection + + +@configclass +class LeapJointPositionActionCorrectionCfg(JointPositionActionCfg): + """Configuration for inverse differential kinematics action term with multi constraints. + This class amend attr body_name from type:str to type:list[str] reflecting its capability to + received the desired positions, poses from multiple target bodies. This will be particularly + useful for controlling dextrous hand robot with only positions of multiple key frame positions + and poses, and output joint positions that satisfy key frame position/pose constrains + + See :class:`DifferentialInverseKinematicsAction` for more details. + """ + + class_type: type[ActionTerm] = LeapJointPositionActionCorrection diff --git a/source/uwlab_assets/uwlab_assets/robots/leap/teleop.py b/source/uwlab_assets/uwlab_assets/robots/leap/teleop.py new file mode 100644 index 00000000..8402d899 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/leap/teleop.py @@ -0,0 +1,89 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.managers import SceneEntityCfg +from isaaclab.utils import configclass +from uwlab.devices import KeyboardCfg, RealsenseT265Cfg, RokokoGlovesCfg, TeleopCfg + + +@configclass +class LeapTeleopCfg: + keyboard_rokokoglove: TeleopCfg = TeleopCfg( + teleop_devices={ + "device1": TeleopCfg.TeleopDevicesCfg( + attach_body=SceneEntityCfg("robot", body_names="wrist"), + attach_scope="self", + pose_reference_body=SceneEntityCfg("robot", body_names="link_base"), + reference_axis_remap=("x", "y", "z"), + command_type="pose", + debug_vis=True, + teleop_interface_cfg=KeyboardCfg( + pos_sensitivity=0.01, + rot_sensitivity=0.01, + enable_gripper_command=False, + ), + ), + "device2": TeleopCfg.TeleopDevicesCfg( + attach_body=SceneEntityCfg("robot", body_names="wrist"), + attach_scope="descendants", + pose_reference_body=SceneEntityCfg("robot", body_names="link_base"), + reference_axis_remap=("x", "y", "z"), + command_type="position", + debug_vis=True, + teleop_interface_cfg=RokokoGlovesCfg( + UDP_IP="0.0.0.0", + UDP_PORT=14043, + scale=1.55, + thumb_scale=0.9, + right_hand_track=[ + "rightIndexMedial", + "rightMiddleMedial", + "rightRingMedial", + "rightIndexTip", + "rightMiddleTip", + "rightRingTip", + "rightThumbTip", + ], + ), + ), + } + ) + + realsense_rokokoglove: TeleopCfg = TeleopCfg( + teleop_devices={ + "device1": TeleopCfg.TeleopDevicesCfg( + attach_body=SceneEntityCfg("robot", body_names="wrist"), + attach_scope="self", + pose_reference_body=SceneEntityCfg("robot", body_names="link_base"), + reference_axis_remap=("x", "y", "z"), + command_type="pose", + debug_vis=True, + teleop_interface_cfg=RealsenseT265Cfg(), + ), + "device2": TeleopCfg.TeleopDevicesCfg( + attach_body=SceneEntityCfg("robot", body_names="wrist"), + attach_scope="descendants", + pose_reference_body=SceneEntityCfg("robot", body_names="link_base"), + reference_axis_remap=("x", "y", "z"), + command_type="position", + debug_vis=True, + teleop_interface_cfg=RokokoGlovesCfg( + UDP_IP="0.0.0.0", + UDP_PORT=14043, + scale=1.55, + thumb_scale=0.9, + right_hand_track=[ + "rightIndexMedial", + "rightMiddleMedial", + "rightRingMedial", + "rightIndexTip", + "rightMiddleTip", + "rightRingTip", + "rightThumbTip", + ], + ), + ), + } + ) diff --git a/source/uwlab_assets/uwlab_assets/robots/robotiq_gripper/__init__.py b/source/uwlab_assets/uwlab_assets/robots/robotiq_gripper/__init__.py new file mode 100644 index 00000000..2c1b17f9 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/robotiq_gripper/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .articulation_drive import * diff --git a/source/uwlab_assets/uwlab_assets/robots/robotiq_gripper/articulation_drive/__init__.py b/source/uwlab_assets/uwlab_assets/robots/robotiq_gripper/articulation_drive/__init__.py new file mode 100644 index 00000000..ac860d6d --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/robotiq_gripper/articulation_drive/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause diff --git a/source/uwlab_assets/uwlab_assets/robots/robotiq_gripper/articulation_drive/robotiq_driver.py b/source/uwlab_assets/uwlab_assets/robots/robotiq_gripper/articulation_drive/robotiq_driver.py new file mode 100644 index 00000000..0cd926f4 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/robotiq_gripper/articulation_drive/robotiq_driver.py @@ -0,0 +1,84 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from uwlab.assets.articulation.articulation_drive import ArticulationDrive + +if TYPE_CHECKING: + from .robotiq_driver_cfg import RobotiqDriverCfg + + +class RobotiqDriver(ArticulationDrive): + def __init__(self, cfg: RobotiqDriverCfg, data_indices: slice = slice(None)): + self.device = torch.device("cpu") + self.cfg = cfg + # self.work_space_limit = cfg.work_space_limit + self.data_idx = data_indices + self.num_joint = len(self.ordered_joint_names) + self.current_pos = torch.zeros(1, self.num_joint, device=self.device) + self.current_vel = torch.zeros(1, self.num_joint, device=self.device) + self.current_eff = torch.zeros(1, self.num_joint, device=self.device) + + @property + def ordered_joint_names(self): + return ( + [ + "finger_joint", + "right_outer_knuckle_joint", + # "left_outer_finger_joint", # urdf does not have this joint + "right_outer_finger_joint", + "left_inner_finger_joint", + "right_inner_finger_joint", + "left_inner_finger_knuckle_joint", + "right_inner_finger_knuckle_joint", + ], + ) + + def _prepare(self): + # Initialize urx connection + pass + + def write_dof_targets(self, pos_target: torch.Tensor, vel_target: torch.Tensor, eff_target: torch.Tensor): + # Non-blocking motion + pass + + def read_dof_states(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Blocking call to get_joint_states, storing the data in local torch Tensors.""" + return self.current_pos, self.current_vel, self.current_eff + + def set_dof_stiffnesses(self, stiffnesses): + pass + + def set_dof_armatures(self, armatures): + pass + + def set_dof_frictions(self, frictions): + pass + + def set_dof_dampings(self, dampings): + pass + + def set_dof_limits(self, limits): + pass + + +# uncomment below code to run the worker +# if __name__ == "__main__": +# # Create the worker +# class Cfg: +# ip = "192.168.1.2" +# port = 602 +# driver = URDriver(cfg=Cfg()) +# driver._prepare() +# pos, vel, eff = driver.read_dof_states() +# print(pos, vel, eff) +# driver.write_dof_targets(pos, vel, eff) + +# while True: +# time.sleep(1) diff --git a/source/uwlab_assets/uwlab_assets/robots/robotiq_gripper/articulation_drive/robotiq_driver_cfg.py b/source/uwlab_assets/uwlab_assets/robots/robotiq_gripper/articulation_drive/robotiq_driver_cfg.py new file mode 100644 index 00000000..246967a5 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/robotiq_gripper/articulation_drive/robotiq_driver_cfg.py @@ -0,0 +1,18 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from typing import Callable + +from isaaclab.utils import configclass +from uwlab.assets.articulation.articulation_drive import ArticulationDriveCfg + +from .robotiq_driver import RobotiqDriver + + +@configclass +class RobotiqDriverCfg(ArticulationDriveCfg): + class_type: Callable[..., RobotiqDriver] = RobotiqDriver diff --git a/source/uwlab_assets/uwlab_assets/robots/spot/__init__.py b/source/uwlab_assets/uwlab_assets/robots/spot/__init__.py new file mode 100644 index 00000000..1a705042 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/spot/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .actions import * +from .arm_spot import SPOT_WITH_ARM_CFG +from .spot import SPOT_CFG diff --git a/source/uwlab_assets/uwlab_assets/robots/spot/actions.py b/source/uwlab_assets/uwlab_assets/robots/spot/actions.py new file mode 100644 index 00000000..8f3bf375 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/spot/actions.py @@ -0,0 +1,16 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from isaaclab.envs.mdp.actions.actions_cfg import JointPositionActionCfg + +SPOT_JOINT_POSITION: JointPositionActionCfg = JointPositionActionCfg( + asset_name="robot", joint_names=[".*"], scale=0.2, use_default_offset=True +) + +ARM_DEFAULT_JOINT_POSITION: JointPositionActionCfg = JointPositionActionCfg( + asset_name="arm", joint_names=[".*"], scale=0.2, use_default_offset=True +) diff --git a/source/uwlab_assets/uwlab_assets/robots/spot/arm_spot.py b/source/uwlab_assets/uwlab_assets/robots/spot/arm_spot.py new file mode 100644 index 00000000..b6d12e73 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/spot/arm_spot.py @@ -0,0 +1,194 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configuration for the Boston Dynamics robot. + +The following configuration parameters are available: + +* :obj:`SPOT_CFG`: The Spot robot with delay PD and remote PD actuators. +""" + +from uwlab_assets import UWLAB_CLOUD_ASSETS_DIR + +import isaaclab.sim as sim_utils +from isaaclab.actuators import DelayedPDActuatorCfg, ImplicitActuatorCfg, RemotizedPDActuatorCfg +from isaaclab.assets.articulation import ArticulationCfg + +# Note: This data was collected by the Boston Dynamics AI Institute. +joint_parameter_lookup = [ + [-2.792900, -24.776718, 37.165077], + [-2.767442, -26.290108, 39.435162], + [-2.741984, -27.793369, 41.690054], + [-2.716526, -29.285997, 43.928996], + [-2.691068, -30.767536, 46.151304], + [-2.665610, -32.237423, 48.356134], + [-2.640152, -33.695168, 50.542751], + [-2.614694, -35.140221, 52.710331], + [-2.589236, -36.572052, 54.858078], + [-2.563778, -37.990086, 56.985128], + [-2.538320, -39.393730, 59.090595], + [-2.512862, -40.782406, 61.173609], + [-2.487404, -42.155487, 63.233231], + [-2.461946, -43.512371, 65.268557], + [-2.436488, -44.852371, 67.278557], + [-2.411030, -46.174873, 69.262310], + [-2.385572, -47.479156, 71.218735], + [-2.360114, -48.764549, 73.146824], + [-2.334656, -50.030334, 75.045502], + [-2.309198, -51.275761, 76.913641], + [-2.283740, -52.500103, 78.750154], + [-2.258282, -53.702587, 80.553881], + [-2.232824, -54.882442, 82.323664], + [-2.207366, -56.038860, 84.058290], + [-2.181908, -57.171028, 85.756542], + [-2.156450, -58.278133, 87.417200], + [-2.130992, -59.359314, 89.038971], + [-2.105534, -60.413738, 90.620607], + [-2.080076, -61.440529, 92.160793], + [-2.054618, -62.438812, 93.658218], + [-2.029160, -63.407692, 95.111538], + [-2.003702, -64.346268, 96.519402], + [-1.978244, -65.253670, 97.880505], + [-1.952786, -66.128944, 99.193417], + [-1.927328, -66.971176, 100.456764], + [-1.901870, -67.779457, 101.669186], + [-1.876412, -68.552864, 102.829296], + [-1.850954, -69.290451, 103.935677], + [-1.825496, -69.991325, 104.986988], + [-1.800038, -70.654541, 105.981812], + [-1.774580, -71.279190, 106.918785], + [-1.749122, -71.864319, 107.796478], + [-1.723664, -72.409088, 108.613632], + [-1.698206, -72.912567, 109.368851], + [-1.672748, -73.373871, 110.060806], + [-1.647290, -73.792130, 110.688194], + [-1.621832, -74.166512, 111.249767], + [-1.596374, -74.496147, 111.744221], + [-1.570916, -74.780251, 112.170376], + [-1.545458, -75.017998, 112.526997], + [-1.520000, -75.208656, 112.812984], + [-1.494542, -75.351448, 113.027172], + [-1.469084, -75.445686, 113.168530], + [-1.443626, -75.490677, 113.236015], + [-1.418168, -75.485771, 113.228657], + [-1.392710, -75.430344, 113.145515], + [-1.367252, -75.323830, 112.985744], + [-1.341794, -75.165688, 112.748531], + [-1.316336, -74.955406, 112.433109], + [-1.290878, -74.692551, 112.038826], + [-1.265420, -74.376694, 111.565041], + [-1.239962, -74.007477, 111.011215], + [-1.214504, -73.584579, 110.376869], + [-1.189046, -73.107742, 109.661613], + [-1.163588, -72.576752, 108.865128], + [-1.138130, -71.991455, 107.987183], + [-1.112672, -71.351707, 107.027561], + [-1.087214, -70.657486, 105.986229], + [-1.061756, -69.908813, 104.863220], + [-1.036298, -69.105721, 103.658581], + [-1.010840, -68.248337, 102.372505], + [-0.985382, -67.336861, 101.005291], + [-0.959924, -66.371513, 99.557270], + [-0.934466, -65.352615, 98.028923], + [-0.909008, -64.280533, 96.420799], + [-0.883550, -63.155693, 94.733540], + [-0.858092, -61.978588, 92.967882], + [-0.832634, -60.749775, 91.124662], + [-0.807176, -59.469845, 89.204767], + [-0.781718, -58.139503, 87.209255], + [-0.756260, -56.759487, 85.139231], + [-0.730802, -55.330616, 82.995924], + [-0.705344, -53.853729, 80.780594], + [-0.679886, -52.329796, 78.494694], + [-0.654428, -50.759762, 76.139643], + [-0.628970, -49.144699, 73.717049], + [-0.603512, -47.485737, 71.228605], + [-0.578054, -45.784004, 68.676006], + [-0.552596, -44.040764, 66.061146], + [-0.527138, -42.257267, 63.385900], + [-0.501680, -40.434883, 60.652325], + [-0.476222, -38.574947, 57.862421], + [-0.450764, -36.678982, 55.018473], + [-0.425306, -34.748432, 52.122648], + [-0.399848, -32.784836, 49.177254], + [-0.374390, -30.789810, 46.184715], + [-0.348932, -28.764952, 43.147428], + [-0.323474, -26.711969, 40.067954], + [-0.298016, -24.632576, 36.948864], + [-0.272558, -22.528547, 33.792821], + [-0.247100, -20.401667, 30.602500], +] +"""The lookup table for the knee joint parameters of the Boston Dynamics Spot robot. + +This table describes the relationship between the joint angle (rad), the transmission ratio (in/out), +and the output torque (N*m). It is used to interpolate the output torque based on the joint angle. +""" + +## +# Configuration +## + + +SPOT_WITH_ARM_CFG = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{UWLAB_CLOUD_ASSETS_DIR}/Robots/BostonDynamic/SpotWithArm/spot_with_arm.usd", + activate_contact_sensors=True, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + retain_accelerations=False, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=1000.0, + max_depenetration_velocity=1.0, + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=True, solver_position_iteration_count=4, solver_velocity_iteration_count=0 + ), + ), + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.0, 0.0, 0.5), + joint_pos={ + "[fh]l_hx": 0.1, # all left hip_x + "[fh]r_hx": -0.1, # all right hip_x + "f[rl]_hy": 0.9, # front hip_y + "h[rl]_hy": 1.1, # hind hip_y + ".*_kn": -1.5, # all knees + "arm0_sh1": -3.14, # arm shoulder + "arm0_el0": 3.14, # arm elbow + "arm0_el1": 0.0, # arm elbow roll + "arm0_wr0": 0.0, # arm wrist pitch + "arm0_wr1": 0.0, # arm wrist roll + "arm0_f1x": 0.0, # gripper position + }, + joint_vel={".*": 0.0}, + ), + actuators={ + "spot_hip": DelayedPDActuatorCfg( + joint_names_expr=[".*_h[xy]"], + effort_limit=45.0, + stiffness=60.0, + damping=1.5, + min_delay=0, # physics time steps (min: 2.0*0=0.0ms) + max_delay=4, # physics time steps (max: 2.0*4=8.0ms) + ), + "spot_knee": RemotizedPDActuatorCfg( + joint_names_expr=[".*_kn"], + joint_parameter_lookup=joint_parameter_lookup, + effort_limit=None, # torque limits are handled based experimental data (`RemotizedPDActuatorCfg.data`) + stiffness=60.0, + damping=1.5, + min_delay=0, # physics time steps (min: 2.0*0=0.0ms) + max_delay=4, # physics time steps (max: 2.0*4=8.0ms) + ), + "spot_arm": ImplicitActuatorCfg( + joint_names_expr=["arm0.*"], + effort_limit=None, + stiffness=60.0, + damping=1.5, + ), + }, +) +"""Configuration for the Boston Dynamics Spot robot.""" diff --git a/source/uwlab_assets/uwlab_assets/robots/spot/spot.py b/source/uwlab_assets/uwlab_assets/robots/spot/spot.py new file mode 100644 index 00000000..27be7264 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/spot/spot.py @@ -0,0 +1,182 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configuration for the Boston Dynamics robot. + +The following configuration parameters are available: + +* :obj:`SPOT_CFG`: The Spot robot with delay PD and remote PD actuators. +""" + +from uwlab_assets import UWLAB_CLOUD_ASSETS_DIR + +import isaaclab.sim as sim_utils +from isaaclab.actuators import DelayedPDActuatorCfg, RemotizedPDActuatorCfg +from isaaclab.assets.articulation import ArticulationCfg + +# Note: This data was collected by the Boston Dynamics AI Institute. +joint_parameter_lookup = [ + [-2.792900, -24.776718, 37.165077], + [-2.767442, -26.290108, 39.435162], + [-2.741984, -27.793369, 41.690054], + [-2.716526, -29.285997, 43.928996], + [-2.691068, -30.767536, 46.151304], + [-2.665610, -32.237423, 48.356134], + [-2.640152, -33.695168, 50.542751], + [-2.614694, -35.140221, 52.710331], + [-2.589236, -36.572052, 54.858078], + [-2.563778, -37.990086, 56.985128], + [-2.538320, -39.393730, 59.090595], + [-2.512862, -40.782406, 61.173609], + [-2.487404, -42.155487, 63.233231], + [-2.461946, -43.512371, 65.268557], + [-2.436488, -44.852371, 67.278557], + [-2.411030, -46.174873, 69.262310], + [-2.385572, -47.479156, 71.218735], + [-2.360114, -48.764549, 73.146824], + [-2.334656, -50.030334, 75.045502], + [-2.309198, -51.275761, 76.913641], + [-2.283740, -52.500103, 78.750154], + [-2.258282, -53.702587, 80.553881], + [-2.232824, -54.882442, 82.323664], + [-2.207366, -56.038860, 84.058290], + [-2.181908, -57.171028, 85.756542], + [-2.156450, -58.278133, 87.417200], + [-2.130992, -59.359314, 89.038971], + [-2.105534, -60.413738, 90.620607], + [-2.080076, -61.440529, 92.160793], + [-2.054618, -62.438812, 93.658218], + [-2.029160, -63.407692, 95.111538], + [-2.003702, -64.346268, 96.519402], + [-1.978244, -65.253670, 97.880505], + [-1.952786, -66.128944, 99.193417], + [-1.927328, -66.971176, 100.456764], + [-1.901870, -67.779457, 101.669186], + [-1.876412, -68.552864, 102.829296], + [-1.850954, -69.290451, 103.935677], + [-1.825496, -69.991325, 104.986988], + [-1.800038, -70.654541, 105.981812], + [-1.774580, -71.279190, 106.918785], + [-1.749122, -71.864319, 107.796478], + [-1.723664, -72.409088, 108.613632], + [-1.698206, -72.912567, 109.368851], + [-1.672748, -73.373871, 110.060806], + [-1.647290, -73.792130, 110.688194], + [-1.621832, -74.166512, 111.249767], + [-1.596374, -74.496147, 111.744221], + [-1.570916, -74.780251, 112.170376], + [-1.545458, -75.017998, 112.526997], + [-1.520000, -75.208656, 112.812984], + [-1.494542, -75.351448, 113.027172], + [-1.469084, -75.445686, 113.168530], + [-1.443626, -75.490677, 113.236015], + [-1.418168, -75.485771, 113.228657], + [-1.392710, -75.430344, 113.145515], + [-1.367252, -75.323830, 112.985744], + [-1.341794, -75.165688, 112.748531], + [-1.316336, -74.955406, 112.433109], + [-1.290878, -74.692551, 112.038826], + [-1.265420, -74.376694, 111.565041], + [-1.239962, -74.007477, 111.011215], + [-1.214504, -73.584579, 110.376869], + [-1.189046, -73.107742, 109.661613], + [-1.163588, -72.576752, 108.865128], + [-1.138130, -71.991455, 107.987183], + [-1.112672, -71.351707, 107.027561], + [-1.087214, -70.657486, 105.986229], + [-1.061756, -69.908813, 104.863220], + [-1.036298, -69.105721, 103.658581], + [-1.010840, -68.248337, 102.372505], + [-0.985382, -67.336861, 101.005291], + [-0.959924, -66.371513, 99.557270], + [-0.934466, -65.352615, 98.028923], + [-0.909008, -64.280533, 96.420799], + [-0.883550, -63.155693, 94.733540], + [-0.858092, -61.978588, 92.967882], + [-0.832634, -60.749775, 91.124662], + [-0.807176, -59.469845, 89.204767], + [-0.781718, -58.139503, 87.209255], + [-0.756260, -56.759487, 85.139231], + [-0.730802, -55.330616, 82.995924], + [-0.705344, -53.853729, 80.780594], + [-0.679886, -52.329796, 78.494694], + [-0.654428, -50.759762, 76.139643], + [-0.628970, -49.144699, 73.717049], + [-0.603512, -47.485737, 71.228605], + [-0.578054, -45.784004, 68.676006], + [-0.552596, -44.040764, 66.061146], + [-0.527138, -42.257267, 63.385900], + [-0.501680, -40.434883, 60.652325], + [-0.476222, -38.574947, 57.862421], + [-0.450764, -36.678982, 55.018473], + [-0.425306, -34.748432, 52.122648], + [-0.399848, -32.784836, 49.177254], + [-0.374390, -30.789810, 46.184715], + [-0.348932, -28.764952, 43.147428], + [-0.323474, -26.711969, 40.067954], + [-0.298016, -24.632576, 36.948864], + [-0.272558, -22.528547, 33.792821], + [-0.247100, -20.401667, 30.602500], +] +"""The lookup table for the knee joint parameters of the Boston Dynamics Spot robot. + +This table describes the relationship between the joint angle (rad), the transmission ratio (in/out), +and the output torque (N*m). It is used to interpolate the output torque based on the joint angle. +""" + +## +# Configuration +## + + +SPOT_CFG = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{UWLAB_CLOUD_ASSETS_DIR}/Robots/BostonDynamic/Spot/spot.usd", + activate_contact_sensors=True, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + retain_accelerations=False, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=1000.0, + max_depenetration_velocity=1.0, + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=True, solver_position_iteration_count=4, solver_velocity_iteration_count=0 + ), + ), + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.0, 0.0, 0.5), + joint_pos={ + "[fh]l_hx": 0.1, # all left hip_x + "[fh]r_hx": -0.1, # all right hip_x + "f[rl]_hy": 0.9, # front hip_y + "h[rl]_hy": 1.1, # hind hip_y + ".*_kn": -1.5, # all knees + }, + joint_vel={".*": 0.0}, + ), + actuators={ + "spot_hip": DelayedPDActuatorCfg( + joint_names_expr=[".*_h[xy]"], + effort_limit=45.0, + stiffness=60.0, + damping=1.5, + min_delay=0, # physics time steps (min: 2.0*0=0.0ms) + max_delay=4, # physics time steps (max: 2.0*4=8.0ms) + ), + "spot_knee": RemotizedPDActuatorCfg( + joint_names_expr=[".*_kn"], + joint_parameter_lookup=joint_parameter_lookup, + effort_limit=None, # torque limits are handled based experimental data (`RemotizedPDActuatorCfg.data`) + stiffness=60.0, + damping=1.5, + min_delay=0, # physics time steps (min: 2.0*0=0.0ms) + max_delay=4, # physics time steps (max: 2.0*4=8.0ms) + ), + }, +) +"""Configuration for the Boston Dynamics Spot robot.""" diff --git a/source/uwlab_assets/uwlab_assets/robots/tycho/__init__.py b/source/uwlab_assets/uwlab_assets/robots/tycho/__init__.py new file mode 100644 index 00000000..e6cf9cb8 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/tycho/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .actions import * +from .tycho import * diff --git a/source/uwlab_assets/uwlab_assets/robots/tycho/actions.py b/source/uwlab_assets/uwlab_assets/robots/tycho/actions.py new file mode 100644 index 00000000..04c82746 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/tycho/actions.py @@ -0,0 +1,63 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.controllers.differential_ik_cfg import DifferentialIKControllerCfg +from isaaclab.envs.mdp.actions.actions_cfg import ( + BinaryJointPositionActionCfg, + DifferentialInverseKinematicsActionCfg, + JointPositionActionCfg, +) +from isaaclab.utils import configclass + +JOINT_POSITION: JointPositionActionCfg = JointPositionActionCfg( + asset_name="robot", + joint_names=["HEBI_(base|elbow|shoulder|wrist|chopstick).*"], + scale=0.1, +) + + +IKDELTA: DifferentialInverseKinematicsActionCfg = DifferentialInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["HEBI_(base|elbow|shoulder|wrist|chopstick).*"], + body_name="static_chop_tip", + controller=DifferentialIKControllerCfg(command_type="pose", use_relative_mode=True, ik_method="dls"), + scale=0.05, + body_offset=DifferentialInverseKinematicsActionCfg.OffsetCfg(pos=(0.0, 0.0, 0.0), rot=(1.0, 0.0, 0, 0)), +) + + +IKABSOLUTE: DifferentialInverseKinematicsActionCfg = DifferentialInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["HEBI_(base|shoulder|elbow|wrist).*"], + body_name="static_chop_tip", # Do not work if this is not end_effector + controller=DifferentialIKControllerCfg(command_type="pose", use_relative_mode=False, ik_method="dls"), + scale=1, + body_offset=DifferentialInverseKinematicsActionCfg.OffsetCfg(pos=(0.0, 0.0, 0.0), rot=(1, 0, 0, 0)), +) + + +BINARY_GRIPPER = BinaryJointPositionActionCfg( + asset_name="robot", + joint_names=["HEBI_chopstick_X5_1"], + open_command_expr={"HEBI_chopstick_X5_1": -0.175}, + close_command_expr={"HEBI_chopstick_X5_1": -0.646}, +) + + +@configclass +class IkdeltaAction: + body_joint_pos = IKDELTA + gripper_joint_pos = BINARY_GRIPPER + + +@configclass +class IkabsoluteAction: + body_joint_pos = IKABSOLUTE + gripper_joint_pos = BINARY_GRIPPER + + +@configclass +class JointPositionAction: + jointpos = JOINT_POSITION diff --git a/source/uwlab_assets/uwlab_assets/robots/tycho/mdp/__init__.py b/source/uwlab_assets/uwlab_assets/robots/tycho/mdp/__init__.py new file mode 100644 index 00000000..400a45b7 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/tycho/mdp/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .observations import * +from .rewards import * +from .terminations import * diff --git a/source/uwlab_assets/uwlab_assets/robots/tycho/mdp/observations.py b/source/uwlab_assets/uwlab_assets/robots/tycho/mdp/observations.py new file mode 100644 index 00000000..d9c68cc1 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/tycho/mdp/observations.py @@ -0,0 +1,108 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch + +from isaaclab.assets import RigidObject +from isaaclab.envs import ManagerBasedRLEnv +from isaaclab.managers import SceneEntityCfg +from isaaclab.sensors import FrameTransformer +from isaaclab.utils import convert_dict_to_backend +from isaaclab.utils.math import subtract_frame_transforms +from uwlab.envs import DataManagerBasedRLEnv + + +def object_position_in_robot_root_frame( + env: ManagerBasedRLEnv, + offset: list[float] = [0.0, 0.0, 0.0], + robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + object_cfg: SceneEntityCfg = SceneEntityCfg("object"), +) -> torch.Tensor: + """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].clone() + object_pos_w[:] += torch.tensor(offset, device=env.device) + object_pos_b, _ = subtract_frame_transforms( + robot.data.root_state_w[:, :3], robot.data.root_state_w[:, 3:7], object_pos_w + ) + return object_pos_b + + +def position_in_robot_root_frame(env: DataManagerBasedRLEnv, position_b: str) -> torch.Tensor: + """The position of the object in the robot's root frame.""" + return env.data_manager.get_active_term("data", position_b) + + +def object_frame_position_in_robot_root_frame( + env: DataManagerBasedRLEnv, + robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + object_frame_cfg: SceneEntityCfg = SceneEntityCfg("object_frame"), +) -> torch.Tensor: + """The position of the object in the robot's root frame.""" + robot: RigidObject = env.scene[robot_cfg.name] + object_frame: FrameTransformer = env.scene[object_frame_cfg.name] + object_frame_pos_w = object_frame.data.target_pos_w[..., 0, :] + object_frame_pos_b, _ = subtract_frame_transforms( + robot.data.root_state_w[:, :3], robot.data.root_state_w[:, 3:7], object_frame_pos_w + ) + return object_frame_pos_b + + +def end_effector_speed(env: DataManagerBasedRLEnv, end_effector_speed_str: str) -> torch.Tensor: + """The position of the object in the robot's root frame.""" + return env.data_manager.get_active_term("data", end_effector_speed_str) + + +def end_effector_pose_in_robot_root_frame( + env: ManagerBasedRLEnv, + robot_name="robot", + fixed_chop_frame_name="frame_fixed_chop_tip", + free_chop_frame_name="frame_free_chop_tip", +): + robot: RigidObject = env.scene[robot_name] + fixed_chop_frame: FrameTransformer = env.scene[fixed_chop_frame_name] + free_chop_frame: FrameTransformer = env.scene[free_chop_frame_name] + fixed_chop_frame_pos_b, fixed_chop_frame_quat_b = subtract_frame_transforms( + robot.data.root_state_w[:, :3], + robot.data.root_state_w[:, 3:7], + fixed_chop_frame.data.target_pos_w[..., 0, :], + fixed_chop_frame.data.target_quat_w[..., 0, :], + ) + + free_chop_frame_pos_b, _ = subtract_frame_transforms( + robot.data.root_state_w[:, :3], + robot.data.root_state_w[:, 3:7], + free_chop_frame.data.target_pos_w[..., 0, :], + ) + + return torch.cat((fixed_chop_frame_pos_b, free_chop_frame_pos_b, fixed_chop_frame_quat_b), dim=1) + + +def capture_image(env: ManagerBasedRLEnv, camera_key: str): + futures = [] + for i in range(env.num_envs): + # Launch parallel tasks + future = torch.jit.fork(_process_camera_data, env, camera_key, i) + futures.append(future) + + # Wait for all tasks to complete and collect results + results = [torch.jit.wait(f) for f in futures] + + # Concatenate the results along the specified dimension + obs_cam_data = torch.stack(results, dim=0) + return obs_cam_data + + +def _process_camera_data(env, camera_key, idx): + """ + Process camera data for a single environment. + This function is intended to be executed in parallel for each environment. + """ + camera = env.scene[camera_key] + camera.update(dt=env.sim.get_physics_dt()) + cam_data = convert_dict_to_backend(camera.data.output[idx], backend="torch") + cam_data_rgb = cam_data["rgb"][:, :, :3].contiguous() + return cam_data_rgb.view(-1) diff --git a/source/uwlab_assets/uwlab_assets/robots/tycho/mdp/rewards.py b/source/uwlab_assets/uwlab_assets/robots/tycho/mdp/rewards.py new file mode 100644 index 00000000..5969f1fe --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/tycho/mdp/rewards.py @@ -0,0 +1,105 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.assets import Articulation, RigidObject +from isaaclab.managers import SceneEntityCfg +from isaaclab.sensors import FrameTransformer +from isaaclab.utils.math import quat_error_magnitude, quat_mul + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +def punish_hand_tilted( + env: ManagerBasedRLEnv, + fixed_chop_frame_cfg: SceneEntityCfg = SceneEntityCfg("frame_fixed_chop_tip"), + free_chop_frame_cfg: SceneEntityCfg = SceneEntityCfg("frame_free_chop_tip"), +): + fixed_chop_frame: FrameTransformer = env.scene[fixed_chop_frame_cfg.name] + free_chop_frame: FrameTransformer = env.scene[free_chop_frame_cfg.name] + + ee_tip_fixed_height_w = fixed_chop_frame.data.target_pos_w[..., 0, 2] + ee_tip_free_height_w = free_chop_frame.data.target_pos_w[..., 0, 2] + + height_diff = -torch.abs(ee_tip_fixed_height_w - ee_tip_free_height_w) + height_diff_punish = torch.clamp(height_diff * 50, -1) + + return height_diff_punish + + +def punish_touching_ground(env: ManagerBasedRLEnv, robot_str: str = "robot") -> torch.Tensor: + ee_height = env.scene[robot_str].data.body_link_pos_w[..., 8, 2] + + return torch.where((ee_height < 0.015) | (ee_height < 0.015), -1, 0) + + +def punish_bad_elbow_shoulder_posture( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + robot: Articulation = env.scene[robot_cfg.name] + elbow_position = robot.data.joint_pos[:, 2] + shoulder_position = robot.data.joint_pos[:, 1] + # Punish for bad elbow position + elbow_punishment = torch.where(elbow_position < 1.8, 1.8 - elbow_position, torch.tensor(0.0)) + elbow_punishment += torch.where(elbow_position > 2.5, elbow_position - 2.5, torch.tensor(0.0)) + # Punish for bad shoulder position (corrected to use shoulder_position) + shoulder_punishment = torch.where(shoulder_position < 0.8, 0.8 - shoulder_position, torch.tensor(0.0)) + shoulder_punishment += torch.where(shoulder_position > 1.6, shoulder_position - 1.6, torch.tensor(0.0)) + # Combine punishments + total_punishment = elbow_punishment + shoulder_punishment + # Cap the total punishment at a maximum of 1 + total_punishment = torch.clamp(total_punishment, max=1) + + return -total_punishment + + +def from_eeframes_get_tip_mid_pos_w(fixed_chop_tip_frame, free_chop_tip_frame): + ee_tip_fixed_w = fixed_chop_tip_frame.data.target_pos_w[..., 0, :] + ee_tip_free_w = free_chop_tip_frame.data.target_pos_w[..., 0, :] + mid_pos_w = (ee_tip_fixed_w + ee_tip_free_w) / 2.0 + return mid_pos_w + + +def get_objectFrame_eeFrames_distance(object_frame, fixed_chop_tip_frame, free_chop_tip_frame): + # Target object_frame position: (num_envs, 3) + object_frame_pos_w = object_frame.data.target_pos_w[..., 0, :] + # End-effector aiming position: (num_envs, 3) + tip_mid_pos_w = from_eeframes_get_tip_mid_pos_w(fixed_chop_tip_frame, free_chop_tip_frame) + # Distance of the end-effector to the object: (num_envs,) + object_ee_distance = torch.norm(object_frame_pos_w - tip_mid_pos_w, dim=1) + return object_ee_distance + + +def get_object_eeFrame_distance(object, offset, fixed_chop_tip_frame, free_chop_tip_frame): + # End-effector aiming position: (num_envs, 3) + mid_pos_w = from_eeframes_get_tip_mid_pos_w(fixed_chop_tip_frame, free_chop_tip_frame) + # Distance of the end-effector to the object: (num_envs,) + object_ee_distance = torch.norm((object.data.root_pos_w + offset) - mid_pos_w, dim=1) + return object_ee_distance + + +def orientation_command_error_tanh( + env: ManagerBasedRLEnv, std: float, command_name: str, asset_cfg: SceneEntityCfg +) -> torch.Tensor: + """Penalize tracking orientation error using shortest path. + + The function computes the orientation error between the desired orientation (from the command) and the + current orientation of the asset's body (in world frame). The orientation error is computed as the shortest + path between the desired and current orientations. + """ + # extract the asset (to enable type hinting) + asset: RigidObject = env.scene[asset_cfg.name] + 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_link_quat_w[:, asset_cfg.body_ids[0], 3:7] # type: ignore + return 1 - torch.tanh(quat_error_magnitude(curr_quat_w, des_quat_w) / std) diff --git a/source/uwlab_assets/uwlab_assets/robots/tycho/mdp/terminations.py b/source/uwlab_assets/uwlab_assets/robots/tycho/mdp/terminations.py new file mode 100644 index 00000000..2070c5e4 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/tycho/mdp/terminations.py @@ -0,0 +1,43 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Common functions that can be used to activate certain terminations. + +The functions can be passed to the :class:`isaaclab.managers.TerminationTermCfg` object to enable +the termination introduced by the function. +""" + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.assets import Articulation +from isaaclab.managers import SceneEntityCfg + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + +""" +MDP terminations. +""" + + +def terminate_extremely_bad_posture( + env: ManagerBasedRLEnv, probability: float = 0.5, robot_cfg: SceneEntityCfg = SceneEntityCfg("robot") +): + robot: Articulation = env.scene[robot_cfg.name] + + elbow_position = robot.data.joint_pos[:, 2] + shoulder_position = robot.data.joint_pos[:, 1] + + # reset for extremely bad elbow position + elbow_punishment = torch.logical_or(elbow_position < 0.35, elbow_position > 2.9) + + # reset for extremely bad bad shoulder position + shoulder_punishment_mask = torch.logical_or(shoulder_position < 0.1, shoulder_position > 3.0) + bitmask = torch.rand(elbow_punishment.shape, device=env.device) < probability + bad_posture_mask = torch.logical_or(elbow_punishment, shoulder_punishment_mask) + return torch.where(bitmask, bad_posture_mask, False) diff --git a/source/uwlab_assets/uwlab_assets/robots/tycho/teleop.py b/source/uwlab_assets/uwlab_assets/robots/tycho/teleop.py new file mode 100644 index 00000000..4f7221d9 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/tycho/teleop.py @@ -0,0 +1,29 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.managers import SceneEntityCfg +from isaaclab.utils import configclass +from uwlab.devices import KeyboardCfg, TeleopCfg + + +@configclass +class TychoTeleopCfg: + keyboard: TeleopCfg = TeleopCfg( + teleop_devices={ + "device1": TeleopCfg.TeleopDevicesCfg( + attach_body=SceneEntityCfg("robot", body_names="static_chop_tip"), + attach_scope="self", + pose_reference_body=SceneEntityCfg("robot", body_names="base_shoulder"), + reference_axis_remap=("-x", "-y", "z"), + command_type="pose", + debug_vis=True, + teleop_interface_cfg=KeyboardCfg( + pos_sensitivity=0.01, + rot_sensitivity=0.01, + enable_gripper_command=True, + ), + ), + } + ) diff --git a/source/uwlab_assets/uwlab_assets/robots/tycho/tycho.py b/source/uwlab_assets/uwlab_assets/robots/tycho/tycho.py new file mode 100644 index 00000000..179ef2f0 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/tycho/tycho.py @@ -0,0 +1,143 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configuration for the Tycho robot.""" + +from uwlab_assets import UWLAB_CLOUD_ASSETS_DIR + +import isaaclab.sim as sim_utils +from isaaclab.actuators import ImplicitActuatorCfg +from isaaclab.assets.articulation import ArticulationCfg +from isaaclab.sensors import FrameTransformerCfg +from isaaclab.sensors.frame_transformer.frame_transformer_cfg import OffsetCfg + +from isaaclab.markers.config import FRAME_MARKER_CFG # isort: skip + + +## +# Configuration +## + +HEBI_DEFAULT_JOINTPOS = { + "HEBI_base_X8_9": -2.2683857389667805, + "HEBI_shoulder_X8_16": 1.5267610481188283, + "HEBI_elbow_X8_9": 2.115358222505881, + "HEBI_wrist1_X5_1": 0.5894993521468314, + "HEBI_wrist2_X5_1": 0.8740650991816328, + "HEBI_wrist3_X5_1": 0.0014332898815118368, + "HEBI_chopstick_X5_1": -0.36, +} + +HEBI_ARTICULATION = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{UWLAB_CLOUD_ASSETS_DIR}/Robots/HebiRobotic/Tycho/tycho.usd", + activate_contact_sensors=True, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=True, + max_depenetration_velocity=5.0, + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=False, solver_position_iteration_count=32, solver_velocity_iteration_count=1 + ), + ), + init_state=ArticulationCfg.InitialStateCfg(rot=(0.7071068, 0, 0, 0.7071068), joint_pos=HEBI_DEFAULT_JOINTPOS), + soft_joint_pos_limit_factor=1, +) + +HEBI_IMPLICIT_ACTUATOR_CFG = HEBI_ARTICULATION.copy() # type: ignore +HEBI_IMPLICIT_ACTUATOR_CFG.actuators = { + "arm": ImplicitActuatorCfg( + joint_names_expr=["HEBI_.*"], + stiffness={"HEBI_(base|elbow|shoulder).*": 120.0, "HEBI_(wrist|chopstick).*": 40.0}, + damping={"HEBI_(base|elbow|shoulder).*": 20.0, "HEBI_(wrist|chopstick).*": 3.0}, + effort_limit={"HEBI_(base|elbow).*": 23.3, "HEBI_shoulder.*": 44.7632, "HEBI_(wrist|chopstick).*": 2.66}, + velocity_limit=1, + ), +} + + +""" +FRAMES +""" +marker_cfg = FRAME_MARKER_CFG.copy() # type: ignore +marker_cfg.markers["frame"].scale = (0.01, 0.01, 0.01) +marker_cfg.prim_path = "/Visuals/FrameTransformer" + +FRAME_EE = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/base_shoulder", + debug_vis=False, + visualizer_cfg=marker_cfg, + target_frames=[ + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/static_chop_tip", + name="ee", + offset=OffsetCfg( + pos=(0.0, 0.0, 0.0), + rot=(1.0, 0.0, 0.0, 0.0), + ), + ), + ], +) + + +FRAME_FIXED_CHOP_TIP = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/base_shoulder", + debug_vis=False, + visualizer_cfg=marker_cfg, + target_frames=[ + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/wrist3_chopstick", + name="fixed_chop_tip", + offset=OffsetCfg( + pos=(0.13018, 0.07598, 0.06429), + ), + ), + ], +) + +FRAME_FIXED_CHOP_END = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/base_shoulder", + debug_vis=False, + visualizer_cfg=marker_cfg, + target_frames=[ + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/wrist3_chopstick", + name="fixed_chop_end", + offset=OffsetCfg( + pos=(-0.13134, 0.07598, 0.06424), + ), + ), + ], +) + +FRAME_FREE_CHOP_TIP = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/base_shoulder", + debug_vis=False, + visualizer_cfg=marker_cfg, + target_frames=[ + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/end_effector", + name="free_chop_tip", + offset=OffsetCfg( + pos=(0.12001, 0.05445, 0.00229), + ), + ), + ], +) + +FRAME_FREE_CHOP_END = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/base_shoulder", + debug_vis=False, + visualizer_cfg=marker_cfg, + target_frames=[ + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/end_effector", + name="free_chop_end", + offset=OffsetCfg( + pos=(-0.11378, -0.04546, 0.00231), + ), + ), + ], +) diff --git a/source/uwlab_assets/uwlab_assets/robots/ur5/__init__.py b/source/uwlab_assets/uwlab_assets/robots/ur5/__init__.py new file mode 100644 index 00000000..5664a590 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/ur5/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .actions import * +from .ur5 import IMPLICIT_UR5 diff --git a/source/uwlab_assets/uwlab_assets/robots/ur5/actions.py b/source/uwlab_assets/uwlab_assets/robots/ur5/actions.py new file mode 100644 index 00000000..1d42a0e2 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/ur5/actions.py @@ -0,0 +1,103 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from isaaclab.envs.mdp.actions.actions_cfg import ( + BinaryJointPositionActionCfg, + JointPositionActionCfg, + RelativeJointPositionActionCfg, +) +from isaaclab.utils import configclass +from uwlab.controllers.differential_ik_cfg import MultiConstraintDifferentialIKControllerCfg +from uwlab.envs.mdp.actions.actions_cfg import ( + DefaultJointPositionStaticActionCfg, + MultiConstraintsDifferentialInverseKinematicsActionCfg, +) + +""" +UR5 GRIPPER ACTIONS +""" +UR5_JOINT_POSITION: JointPositionActionCfg = JointPositionActionCfg( + asset_name="robot", + joint_names=["shoulder.*", "elbow.*", "wrist.*"], + scale=0.1, + use_default_offset=True, +) + +UR5_RELATIVE_JOINT_POSITION: RelativeJointPositionActionCfg = RelativeJointPositionActionCfg( + asset_name="robot", + joint_names=["shoulder.*", "elbow.*", "wrist.*"], + scale=0.02, + use_zero_offset=True, +) + + +UR5_MC_IKABSOLUTE_ARM = MultiConstraintsDifferentialInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["shoulder.*", "elbow.*", "wrist.*"], + body_name=["robotiq_base_link"], + controller=MultiConstraintDifferentialIKControllerCfg( + command_type="pose", use_relative_mode=False, ik_method="dls" + ), + scale=1, +) + +UR5_MC_IKDELTA_ARM = MultiConstraintsDifferentialInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["joint.*"], + body_name=["robotiq_base_link"], + controller=MultiConstraintDifferentialIKControllerCfg(command_type="pose", use_relative_mode=True, ik_method="dls"), + scale=0.5, +) + +ROBOTIQ_GRIPPER_BINARY_ACTIONS = BinaryJointPositionActionCfg( + asset_name="robot", + joint_names=["finger_joint"], + open_command_expr={"finger_joint": 0.0}, + close_command_expr={"finger_joint": 0.785398}, +) + +ROBOTIQ_COMPLIANT_JOINTS = DefaultJointPositionStaticActionCfg( + asset_name="robot", joint_names=["left_inner_finger_joint", "right_inner_finger_joint"] +) + +ROBOTIQ_MC_IK_ABSOLUTE = MultiConstraintsDifferentialInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["joint.*"], + body_name=["left_inner_finger", "right_inner_finger"], + controller=MultiConstraintDifferentialIKControllerCfg( + command_type="pose", use_relative_mode=False, ik_method="dls" + ), + scale=1, +) + + +@configclass +class Ur5IkAbsoluteAction: + jointpos = UR5_MC_IKABSOLUTE_ARM + gripper = ROBOTIQ_GRIPPER_BINARY_ACTIONS + compliant_joints = ROBOTIQ_COMPLIANT_JOINTS + + +@configclass +class Ur5McIkDeltaAction: + jointpos = UR5_MC_IKDELTA_ARM + gripper = ROBOTIQ_GRIPPER_BINARY_ACTIONS + compliant_joints = ROBOTIQ_COMPLIANT_JOINTS + + +@configclass +class Ur5JointPositionAction: + jointpos = UR5_JOINT_POSITION + gripper = ROBOTIQ_GRIPPER_BINARY_ACTIONS + compliant_joints = ROBOTIQ_COMPLIANT_JOINTS + + +@configclass +class Ur5RelativeJointPositionAction: + jointpos = UR5_RELATIVE_JOINT_POSITION + gripper = ROBOTIQ_GRIPPER_BINARY_ACTIONS + compliant_joints = ROBOTIQ_COMPLIANT_JOINTS diff --git a/source/uwlab_assets/uwlab_assets/robots/ur5/teleop.py b/source/uwlab_assets/uwlab_assets/robots/ur5/teleop.py new file mode 100644 index 00000000..0e150343 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/ur5/teleop.py @@ -0,0 +1,29 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.managers import SceneEntityCfg +from isaaclab.utils import configclass +from uwlab.devices import KeyboardCfg, TeleopCfg + + +@configclass +class Ur5TeleopCfg: + keyboard: TeleopCfg = TeleopCfg( + teleop_devices={ + "device1": TeleopCfg.TeleopDevicesCfg( + attach_body=SceneEntityCfg("robot", body_names="robotiq_base_link"), + attach_scope="self", + pose_reference_body=SceneEntityCfg("robot", body_names="base_link_inertia"), + reference_axis_remap=("-x", "-y", "z"), + command_type="pose", + debug_vis=True, + teleop_interface_cfg=KeyboardCfg( + pos_sensitivity=0.01, + rot_sensitivity=0.04, + enable_gripper_command=True, + ), + ), + } + ) diff --git a/source/uwlab_assets/uwlab_assets/robots/ur5/ur5.py b/source/uwlab_assets/uwlab_assets/robots/ur5/ur5.py new file mode 100644 index 00000000..a42e964c --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/ur5/ur5.py @@ -0,0 +1,74 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configuration for the UR5 robots. + +The following configurations are available: + +* :obj:`UR5_CFG`: Ur5 robot +""" + +from uwlab_assets import UWLAB_CLOUD_ASSETS_DIR + +import isaaclab.sim as sim_utils +from isaaclab.actuators import ImplicitActuatorCfg +from isaaclab.assets.articulation import ArticulationCfg + +UR5_DEFAULT_JOINT_POS = { + "shoulder_pan_joint": 0.0, + "shoulder_lift_joint": -1.5708, + "elbow_joint": 1.5708, + "wrist_1_joint": 4.7112, + "wrist_2_joint": -1.5708, + "wrist_3_joint": -1.5708, + "finger_joint": 0.0, + "right_outer.*": 0.0, + "left_outer.*": 0.0, + "left_inner_finger_knuckle_joint": 0.0, + "right_inner_finger_knuckle_joint": 0.0, + "left_inner_finger_joint": -0.785398, + "right_inner_finger_joint": 0.785398, +} + +UR5_ARTICULATION = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{UWLAB_CLOUD_ASSETS_DIR}/Robots/UniversalRobots/Ur5RobotiqGripper/ur5_robotiq_gripper_backup.usd", + activate_contact_sensors=False, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=True, + max_depenetration_velocity=5.0, + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=False, solver_position_iteration_count=36, solver_velocity_iteration_count=0 + ), + ), + init_state=ArticulationCfg.InitialStateCfg(pos=(0, 0, 0), rot=(0, 0, 0, 1), joint_pos=UR5_DEFAULT_JOINT_POS), + soft_joint_pos_limit_factor=1, +) + +IMPLICIT_UR5 = UR5_ARTICULATION.copy() # type: ignore +IMPLICIT_UR5.actuators = { + "arm": ImplicitActuatorCfg( + joint_names_expr=["shoulder.*", "elbow.*", "wrist.*"], + stiffness=261.8, + damping=26.18, + velocity_limit=3.14, + effort_limit={"shoulder.*": 9000, "elbow.*": 9000, "wrist.*": 1680}, + ), + "gripper": ImplicitActuatorCfg( + joint_names_expr=["finger_joint"], + stiffness=17, + damping=5, + velocity_limit=2.27, + effort_limit=165, + ), + "inner_finger": ImplicitActuatorCfg( + joint_names_expr=[".*_inner_finger_joint"], + stiffness=0.2, + damping=0.02, + velocity_limit=5.3, + effort_limit=0.5, + ), +} diff --git a/source/uwlab_assets/uwlab_assets/robots/xarm/__init__.py b/source/uwlab_assets/uwlab_assets/robots/xarm/__init__.py new file mode 100644 index 00000000..2c1b17f9 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/xarm/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .articulation_drive import * diff --git a/source/uwlab_assets/uwlab_assets/robots/xarm/articulation_drive/__init__.py b/source/uwlab_assets/uwlab_assets/robots/xarm/articulation_drive/__init__.py new file mode 100644 index 00000000..eeb9d2b3 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/xarm/articulation_drive/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .xarm_driver import XarmDriver +from .xarm_driver_cfg import XarmDriverCfg diff --git a/source/uwlab_assets/uwlab_assets/robots/xarm/articulation_drive/installation.md b/source/uwlab_assets/uwlab_assets/robots/xarm/articulation_drive/installation.md new file mode 100644 index 00000000..6d3fc901 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/xarm/articulation_drive/installation.md @@ -0,0 +1,13 @@ +# xArm Hardware Installation Notes + +## Installation Steps + +### 1. Clone the xArm Python SDK repository: +```bash +git clone https://github.com/xArm-Developer/xArm-Python-SDK.git +``` + +### 2. Install the SDK: +```bash +python setup.py install +``` diff --git a/source/uwlab_assets/uwlab_assets/robots/xarm/articulation_drive/xarm_driver.py b/source/uwlab_assets/uwlab_assets/robots/xarm/articulation_drive/xarm_driver.py new file mode 100644 index 00000000..3bdd122e --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/xarm/articulation_drive/xarm_driver.py @@ -0,0 +1,121 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import time +import torch +from typing import TYPE_CHECKING + +from uwlab.assets.articulation.articulation_drive import ArticulationDrive + +if TYPE_CHECKING: + from .xarm_driver_cfg import XarmDriverCfg + + +class XarmDriver(ArticulationDrive): + def __init__(self, cfg: XarmDriverCfg, data_indices: slice = slice(None)): + self.device = torch.device("cpu") + self.cfg = cfg + self.p_gain_scaler = cfg.p_gain_scaler + self.work_space_limit = cfg.work_space_limit + self.data_idx = data_indices + work_space_limit = torch.tensor(cfg.work_space_limit, device=self.device) + self.min_limits = work_space_limit[:, 0] + self.max_limits = work_space_limit[:, 1] + self.is_radian = cfg.is_radian + + self.current_pos = torch.zeros(1, 5, device=self.device) + self.current_vel = torch.zeros(1, 5, device=self.device) + self.current_eff = torch.zeros(1, 5, device=self.device) + self._prepare() + + @property + def ordered_joint_names(self): + return ["joint" + str(i) for i in range(1, 6)] + + def close(self): + self._arm.disconnect() + + def read_dof_states(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Blocking call to get_joint_states, storing the data in local torch Tensors.""" + code = 1 + while code != 0: + code, (pos, vel, eff) = self._arm.get_joint_states(is_radian=True) + if code != 0: + self._log(f"Warning: get_joint_states returned code={code}", is_error=True) + self._prepare() + pos = torch.tensor(pos[:5], device=self.device).view(1, -1) + vel = torch.tensor(vel[:5], device=self.device).view(1, -1) + eff = torch.tensor(eff[:5], device=self.device).view(1, -1) + self.current_pos[:] = pos + self.current_vel[:] = vel + self.current_eff[:] = eff + return pos, vel, eff + + def _get_forward_kinematics(self, pos: torch.Tensor): + code = 1 + while code != 0: + code, pose = self._arm.get_forward_kinematics(pos[0].tolist(), input_is_radian=self.is_radian) + if code != 0: + self._log(f"Warning: get_forward_kinematics returned code={code}", is_error=True) + self._prepare() + ee_pose = torch.tensor(pose, device=self.device) / 1000 # convert to meter + return ee_pose + + def write_dof_targets(self, pos_target: torch.Tensor, vel_target: torch.Tensor, eff_target: torch.Tensor): + # Non-blocking motion + position_error = pos_target - self.current_pos + command = position_error * self.p_gain_scaler + self.current_pos + + ee_pose = self._get_forward_kinematics(command) + within_workspace_limits = ((ee_pose[:3] > self.min_limits) & (ee_pose[:3] < self.max_limits)).all() + + if not within_workspace_limits: + print(f"Arm action {ee_pose[:3]} canceled: arm target end-effector position is out of workspace limits.") + return + code = 1 + while code != 0: + code = self._arm.set_servo_angle_j(command[0].tolist(), is_radian=True, wait=False) + if code != 0: + self._log(f"Warning: set_servo_angle_j returned code={code}", is_error=True) + self._prepare() + + def set_dof_stiffnesses(self, stiffnesses): + pass + + def set_dof_armatures(self, armatures): + pass + + def set_dof_frictions(self, frictions): + pass + + def set_dof_dampings(self, dampings): + pass + + def set_dof_limits(self, limits): + pass + + def _prepare(self): + from xarm.wrapper import XArmAPI + + self._arm = XArmAPI(port=self.cfg.ip, is_radian=self.is_radian) + self._arm.clean_error() + self._arm.clean_warn() + self._arm.motion_enable(enable=True) + self._arm.set_mode(1) + time.sleep(0.50) + self._arm.set_collision_sensitivity(0) + self._arm.set_state(0) + time.sleep(0.50) + + def _log(self, msg, is_error=False): + """ + Simple logging mechanism. + In real code, use 'logging' module or other logging frameworks. + """ + prefix = "[ERROR]" if is_error else "[INFO]" + entry = f"{prefix} {time.strftime('%H:%M:%S')}: {msg}" + print(entry) diff --git a/source/uwlab_assets/uwlab_assets/robots/xarm/articulation_drive/xarm_driver_cfg.py b/source/uwlab_assets/uwlab_assets/robots/xarm/articulation_drive/xarm_driver_cfg.py new file mode 100644 index 00000000..a33c0b81 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/xarm/articulation_drive/xarm_driver_cfg.py @@ -0,0 +1,27 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from dataclasses import MISSING +from typing import Callable + +from isaaclab.utils import configclass +from uwlab.assets.articulation.articulation_drive import ArticulationDriveCfg + +from .xarm_driver import XarmDriver + + +@configclass +class XarmDriverCfg(ArticulationDriveCfg): + class_type: Callable[..., XarmDriver] = XarmDriver + + work_space_limit: list[list[float]] = MISSING # type: ignore + + ip: str = MISSING # type: ignore + + is_radian: bool = True + + p_gain_scaler: float = 0.01 diff --git a/source/uwlab_assets/uwlab_assets/robots/xarm_leap/__init__.py b/source/uwlab_assets/uwlab_assets/robots/xarm_leap/__init__.py new file mode 100644 index 00000000..49f56079 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/xarm_leap/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .actions import * +from .xarm_leap import FRAME_EE, IMPLICIT_XARM_LEAP diff --git a/source/uwlab_assets/uwlab_assets/robots/xarm_leap/actions.py b/source/uwlab_assets/uwlab_assets/robots/xarm_leap/actions.py new file mode 100644 index 00000000..0d66ba12 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/xarm_leap/actions.py @@ -0,0 +1,104 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from uwlab_assets import UWLAB_CLOUD_ASSETS_DIR + +from isaaclab.envs.mdp.actions.actions_cfg import JointPositionActionCfg +from isaaclab.utils import configclass +from uwlab.controllers.differential_ik_cfg import MultiConstraintDifferentialIKControllerCfg +from uwlab.envs.mdp.actions.actions_cfg import ( + MultiConstraintsDifferentialInverseKinematicsActionCfg, + PCAJointPositionActionCfg, +) + +""" +LEAP XARM ACTIONS +""" +XARM_LEAP_JOINT_POSITION: JointPositionActionCfg = JointPositionActionCfg( + asset_name="robot", + joint_names=["joint.*", "j[0-9]+"], + scale=1.0, + use_default_offset=False, +) + + +XARM_LEAP_MC_IKABSOLUTE = MultiConstraintsDifferentialInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["joint.*", "j[0-9]+"], + body_name=["wrist", "pip", "pip_2", "pip_3", "tip", "thumb_tip", "tip_2", "tip_3"], + controller=MultiConstraintDifferentialIKControllerCfg( + command_type="position", use_relative_mode=False, ik_method="dls" + ), + scale=1, +) + +XARM_LEAP_MC_IKABSOLUTE_ARM = MultiConstraintsDifferentialInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["joint.*"], + body_name=["wrist"], + controller=MultiConstraintDifferentialIKControllerCfg( + command_type="pose", use_relative_mode=False, ik_method="dls" + ), + scale=1, +) + + +XARM_LEAP_MC_IKABSOLUTE_FINGER = MultiConstraintsDifferentialInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["j[0-9]+"], + body_name=["pip", "pip_2", "pip_3", "tip", "tip_2", "tip_3", "thumb_tip"], + controller=MultiConstraintDifferentialIKControllerCfg( + command_type="position", use_relative_mode=False, ik_method="dls" + ), + scale=1, +) + + +XARM_LEAP_MC_IKDELTA = MultiConstraintsDifferentialInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["joint.*", "j[0-9]+"], + body_name=["wrist", "pip", "pip_2", "pip_3", "tip", "thumb_tip", "tip_2", "tip_3"], + controller=MultiConstraintDifferentialIKControllerCfg( + command_type="position", use_relative_mode=True, ik_method="dls" + ), + scale=1, +) + + +XARM_LEAP_PCA_JOINT_POSITION = PCAJointPositionActionCfg( + asset_name="robot", + joint_names=["joint.*", "j[0-9]+"], + scale=1.0, + eigenspace_path=f"{UWLAB_CLOUD_ASSETS_DIR}/dataset/misc/hammer_grasping_pca_components.npy", + joint_range=(-3.14, 3.14), +) + + +@configclass +class XarmLeapSeparatedIkAbsoluteAction: + joint_pos = XARM_LEAP_MC_IKABSOLUTE_ARM + finger_pos = XARM_LEAP_MC_IKABSOLUTE_FINGER + + +@configclass +class XarmLeapMcIkAbsoluteAction: + joint_pos = XARM_LEAP_MC_IKABSOLUTE + + +@configclass +class XarmLeapMcIkDeltaAction: + joint_pos = XARM_LEAP_MC_IKDELTA + + +@configclass +class XarmLeapJointPositionAction: + joint_pos = XARM_LEAP_JOINT_POSITION + + +@configclass +class XarmLeapPCAJointPositionAction: + joint_pos = XARM_LEAP_PCA_JOINT_POSITION diff --git a/source/uwlab_assets/uwlab_assets/robots/xarm_leap/teleop.py b/source/uwlab_assets/uwlab_assets/robots/xarm_leap/teleop.py new file mode 100644 index 00000000..0cb9ef43 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/xarm_leap/teleop.py @@ -0,0 +1,89 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.managers import SceneEntityCfg +from isaaclab.utils import configclass +from uwlab.devices import KeyboardCfg, RealsenseT265Cfg, RokokoGlovesCfg, TeleopCfg + + +@configclass +class XarmLeapTeleopCfg: + keyboard_rokoko_glove: TeleopCfg = TeleopCfg( + teleop_devices={ + "device1": TeleopCfg.TeleopDevicesCfg( + attach_body=SceneEntityCfg("robot", body_names="wrist"), + attach_scope="self", + pose_reference_body=SceneEntityCfg("robot", body_names="link_base"), + reference_axis_remap=("x", "y", "z"), + command_type="pose", + debug_vis=True, + teleop_interface_cfg=KeyboardCfg( + pos_sensitivity=0.01, + rot_sensitivity=0.01, + enable_gripper_command=False, + ), + ), + "device2": TeleopCfg.TeleopDevicesCfg( + attach_body=SceneEntityCfg("robot", body_names="wrist"), + attach_scope="descendants", + pose_reference_body=SceneEntityCfg("robot", body_names="link_base"), + reference_axis_remap=("x", "y", "z"), + command_type="position", + debug_vis=True, + teleop_interface_cfg=RokokoGlovesCfg( + UDP_IP="0.0.0.0", + UDP_PORT=14043, + scale=1.55, + thumb_scale=0.9, + right_hand_track=[ + "rightIndexMedial", + "rightMiddleMedial", + "rightRingMedial", + "rightIndexTip", + "rightMiddleTip", + "rightRingTip", + "rightThumbTip", + ], + ), + ), + } + ) + + realsense_rokoko_glove: TeleopCfg = TeleopCfg( + teleop_devices={ + "device1": TeleopCfg.TeleopDevicesCfg( + attach_body=SceneEntityCfg("robot", body_names="wrist"), + attach_scope="self", + pose_reference_body=SceneEntityCfg("robot", body_names="link_base"), + reference_axis_remap=("x", "y", "z"), + command_type="pose", + debug_vis=True, + teleop_interface_cfg=RealsenseT265Cfg(), + ), + "device2": TeleopCfg.TeleopDevicesCfg( + attach_body=SceneEntityCfg("robot", body_names="wrist"), + attach_scope="descendants", + pose_reference_body=SceneEntityCfg("robot", body_names="link_base"), + reference_axis_remap=("x", "y", "z"), + command_type="position", + debug_vis=True, + teleop_interface_cfg=RokokoGlovesCfg( + UDP_IP="0.0.0.0", + UDP_PORT=14043, + scale=1.55, + thumb_scale=0.9, + right_hand_track=[ + "rightIndexMedial", + "rightMiddleMedial", + "rightRingMedial", + "rightIndexTip", + "rightMiddleTip", + "rightRingTip", + "rightThumbTip", + ], + ), + ), + } + ) diff --git a/source/uwlab_assets/uwlab_assets/robots/xarm_leap/xarm_leap.py b/source/uwlab_assets/uwlab_assets/robots/xarm_leap/xarm_leap.py new file mode 100644 index 00000000..199ccbde --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/xarm_leap/xarm_leap.py @@ -0,0 +1,84 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from uwlab_assets import UWLAB_CLOUD_ASSETS_DIR + +import isaaclab.sim as sim_utils +from isaaclab.actuators import ImplicitActuatorCfg +from isaaclab.assets.articulation import ArticulationCfg +from isaaclab.markers.config import FRAME_MARKER_CFG +from isaaclab.sensors import FrameTransformerCfg +from isaaclab.sensors.frame_transformer.frame_transformer_cfg import OffsetCfg + +## +# Configuration +## + +""" +XARM_LEAP +""" +# fmt: off +XARM_LEAP_DEFAULT_JOINT_POS = {".*": 0.0} +# fmt: on + +XARM_LEAP_ARTICULATION = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{UWLAB_CLOUD_ASSETS_DIR}/Robots/UFactory/Xarm5LeapHand/leap_xarm_ikpoints.usd", + activate_contact_sensors=False, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + max_depenetration_velocity=5.0, + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=True, solver_position_iteration_count=1, solver_velocity_iteration_count=0 + ), + ), + init_state=ArticulationCfg.InitialStateCfg(pos=(0, 0, 0), rot=(1, 0, 0, 0), joint_pos=XARM_LEAP_DEFAULT_JOINT_POS), + soft_joint_pos_limit_factor=1, +) + +IMPLICIT_XARM_LEAP = XARM_LEAP_ARTICULATION.copy() # type: ignore +IMPLICIT_XARM_LEAP.actuators = { + "arm1": ImplicitActuatorCfg( + joint_names_expr=["joint.*"], + stiffness={"joint[1-2]": 1000, "joint3": 800, "joint[4-5]": 600}, + damping=100.0, + velocity_limit=3.14, + effort_limit={"joint[1-2]": 50, "joint3": 30, "joint[4-5]": 20}, + ), + "j": ImplicitActuatorCfg( + joint_names_expr=["j[0-9]+"], + stiffness=20.0, + damping=1.0, + armature=0.001, + friction=0.2, + velocity_limit=8.48, + effort_limit=0.95, + ), +} + + +""" +FRAMES +""" +marker_cfg = FRAME_MARKER_CFG.copy() # type: ignore +marker_cfg.markers["frame"].scale = (0.01, 0.01, 0.01) +marker_cfg.prim_path = "/Visuals/FrameTransformer" + +FRAME_EE = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/link_base", + debug_vis=False, + visualizer_cfg=marker_cfg, + target_frames=[ + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/palm_lower", + name="ee", + offset=OffsetCfg( + pos=(-0.028, -0.04, -0.07), + rot=(1.0, 0.0, 0.0, 0.0), + ), + ), + ], +) diff --git a/source/uwlab_assets/uwlab_assets/robots/xarm_uf_gripper/__init__.py b/source/uwlab_assets/uwlab_assets/robots/xarm_uf_gripper/__init__.py new file mode 100644 index 00000000..8aa6f552 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/xarm_uf_gripper/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .actions import * +from .xarm_uf_gripper import FRAME_EE, IMPLICIT_XARM_UF_GRIPPER diff --git a/source/uwlab_assets/uwlab_assets/robots/xarm_uf_gripper/actions.py b/source/uwlab_assets/uwlab_assets/robots/xarm_uf_gripper/actions.py new file mode 100644 index 00000000..e94bf666 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/xarm_uf_gripper/actions.py @@ -0,0 +1,73 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from isaaclab.envs.mdp.actions.actions_cfg import ( + BinaryJointPositionActionCfg, + DifferentialIKControllerCfg, + DifferentialInverseKinematicsActionCfg, + JointEffortActionCfg, + JointPositionActionCfg, +) +from isaaclab.utils import configclass + +""" +XARM GRIPPER ACTIONS +""" +XARM_UF_GRIPPER_JOINT_POSITION: JointPositionActionCfg = JointPositionActionCfg( + asset_name="robot", + joint_names=["joint.*", "drive_joint"], + scale=1.0, + use_default_offset=False, +) + + +XARM_UF_GRIPPER_JOINT_EFFORT: JointEffortActionCfg = JointEffortActionCfg( + asset_name="robot", joint_names=["joint.*", "drive_joint"], scale=0.1 +) + + +XARM_UF_GRIPPER_MC_IKABSOLUTE_ARM = DifferentialInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["joint.*"], + body_name="link_tcp", + controller=DifferentialIKControllerCfg(command_type="pose", use_relative_mode=False, ik_method="dls"), + scale=1, +) + +XARM_UF_GRIPPER_MC_IKDELTA_ARM = DifferentialInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["joint.*"], + body_name="link_tcp", + controller=DifferentialIKControllerCfg(command_type="pose", use_relative_mode=True, ik_method="dls"), + scale=0.5, +) + + +XARM_GRIPPER_BINARY_ACTIONS = BinaryJointPositionActionCfg( + asset_name="robot", + joint_names=["drive_joint"], + open_command_expr={"drive_joint": 0.0}, + close_command_expr={"drive_joint": 1.0}, +) + + +@configclass +class XarmUfGripperIkAbsoluteAction: + joint_pos = XARM_UF_GRIPPER_MC_IKABSOLUTE_ARM + gripper = XARM_GRIPPER_BINARY_ACTIONS + + +@configclass +class XarmUfGripperMcIkDeltaAction: + joint_pos = XARM_UF_GRIPPER_MC_IKDELTA_ARM + gripper = XARM_GRIPPER_BINARY_ACTIONS + + +@configclass +class XarmUfGripperJointPositionAction: + joint_pos = XARM_UF_GRIPPER_JOINT_POSITION + gripper = XARM_GRIPPER_BINARY_ACTIONS diff --git a/source/uwlab_assets/uwlab_assets/robots/xarm_uf_gripper/teleop.py b/source/uwlab_assets/uwlab_assets/robots/xarm_uf_gripper/teleop.py new file mode 100644 index 00000000..b30d788b --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/xarm_uf_gripper/teleop.py @@ -0,0 +1,29 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.managers import SceneEntityCfg +from isaaclab.utils import configclass +from uwlab.devices import KeyboardCfg, TeleopCfg + + +@configclass +class XarmUfGripperTeleopCfg: + keyboard: TeleopCfg = TeleopCfg( + teleop_devices={ + "device1": TeleopCfg.TeleopDevicesCfg( + attach_body=SceneEntityCfg("robot", body_names="link_tcp"), + attach_scope="self", + pose_reference_body=SceneEntityCfg("robot", body_names="link_base"), + reference_axis_remap=("x", "y", "z"), + command_type="pose", + debug_vis=True, + teleop_interface_cfg=KeyboardCfg( + pos_sensitivity=0.01, + rot_sensitivity=0.01, + enable_gripper_command=True, + ), + ), + } + ) diff --git a/source/uwlab_assets/uwlab_assets/robots/xarm_uf_gripper/xarm_uf_gripper.py b/source/uwlab_assets/uwlab_assets/robots/xarm_uf_gripper/xarm_uf_gripper.py new file mode 100644 index 00000000..90502865 --- /dev/null +++ b/source/uwlab_assets/uwlab_assets/robots/xarm_uf_gripper/xarm_uf_gripper.py @@ -0,0 +1,85 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configuration for the Hebi robots. + +The following configurations are available: + +* :obj:`HEBI_CFG`: Hebi robot with chopsticks +""" + +from uwlab_assets import UWLAB_CLOUD_ASSETS_DIR + +import isaaclab.sim as sim_utils +from isaaclab.actuators import ImplicitActuatorCfg +from isaaclab.assets.articulation import ArticulationCfg +from isaaclab.markers.config import FRAME_MARKER_CFG +from isaaclab.sensors import FrameTransformerCfg + +## +# Configuration +## + + +# fmt: off +XARM_UF_GRIPPER_DEFAULT_JOINT_POS = { + "drive_joint": 0.0, "joint1": 0.0, "joint2": 0.0, "joint3": -0.5, "joint4": 0.0, "joint5": 0.0, +} +# fmt: on + +XARM_UF_GRIPPER_ARTICULATION = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{UWLAB_CLOUD_ASSETS_DIR}/Robots/UFactory/Xarm5UfGripper/xarm_gripper.usd", + activate_contact_sensors=False, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + max_depenetration_velocity=5.0, + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=True, solver_position_iteration_count=1, solver_velocity_iteration_count=0 + ), + ), + init_state=ArticulationCfg.InitialStateCfg( + pos=(0, 0, 0), rot=(1, 0, 0, 0), joint_pos=XARM_UF_GRIPPER_DEFAULT_JOINT_POS + ), + soft_joint_pos_limit_factor=1, +) + +IMPLICIT_XARM_UF_GRIPPER = XARM_UF_GRIPPER_ARTICULATION.copy() # type: ignore +IMPLICIT_XARM_UF_GRIPPER.actuators = { + "arm": ImplicitActuatorCfg( + joint_names_expr=["joint.*"], + stiffness={"joint[1-2]": 500, "joint3": 500, "joint[4-5]": 400}, + damping=50.0, + velocity_limit=3.14, + effort_limit={"joint[1-2]": 50, "joint3": 30, "joint[4-5]": 20}, + ), + "gripper": ImplicitActuatorCfg( + joint_names_expr=["drive_joint"], + stiffness=20.0, + damping=1.0, + armature=0.001, + friction=0.2, + velocity_limit=2, + effort_limit=50, + ), +} + + +""" +FRAMES +""" +marker_cfg = FRAME_MARKER_CFG.copy() # type: ignore +marker_cfg.markers["frame"].scale = (0.01, 0.01, 0.01) +marker_cfg.prim_path = "/Visuals/FrameTransformer" + +FRAME_EE = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/link_base", + debug_vis=False, + visualizer_cfg=marker_cfg, + target_frames=[ + FrameTransformerCfg.FrameCfg(prim_path="{ENV_REGEX_NS}/Robot/link_tcp", name="ee"), + ], +) diff --git a/source/uwlab_rl/config/extension.toml b/source/uwlab_rl/config/extension.toml new file mode 100644 index 00000000..9e74cb2e --- /dev/null +++ b/source/uwlab_rl/config/extension.toml @@ -0,0 +1,26 @@ +[package] + +# Note: Semantic Versioning is used: https://semver.org/ +version = "0.1.0" + +# Description +title = "UW Lab RL" +description="Extension containing reinforcement learning related utilities." +readme = "docs/README.md" +repository = "https://github.com/UW-Lab/UWLab" +category = "robotics" +keywords = ["robotics", "rl", "wrappers", "learning"] + +[dependencies] +"isaaclab" = {} +"isaaclab_assets" = {} +"isaaclab_tasks" = {} +"uwlab" = {} +"uwlab_assets" = {} +"uwlab_tasks" = {} + +[core] +reloadable = false + +[[python.module]] +name = "uwlab_rl" diff --git a/source/uwlab_rl/docs/CHANGELOG.rst b/source/uwlab_rl/docs/CHANGELOG.rst new file mode 100644 index 00000000..beed7fc8 --- /dev/null +++ b/source/uwlab_rl/docs/CHANGELOG.rst @@ -0,0 +1,17 @@ +Changelog +--------- + +0.1.0 (2025-03-12) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +Initial version of the extension include the wrapper scripts and extensions for the supported RL libraries. + +Supported RL libraries are: + +* RL Games +* RSL RL +* SKRL +* Stable Baselines3 diff --git a/source/uwlab_rl/pyproject.toml b/source/uwlab_rl/pyproject.toml new file mode 100644 index 00000000..d90ac353 --- /dev/null +++ b/source/uwlab_rl/pyproject.toml @@ -0,0 +1,3 @@ +[build-system] +requires = ["setuptools", "wheel", "toml"] +build-backend = "setuptools.build_meta" diff --git a/source/uwlab_rl/setup.py b/source/uwlab_rl/setup.py new file mode 100644 index 00000000..12c6c119 --- /dev/null +++ b/source/uwlab_rl/setup.py @@ -0,0 +1,57 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Installation script for the 'uwlab_rl' python package.""" + +import itertools +import os +import toml + +from setuptools import setup + +# Obtain the extension data from the extension.toml file +EXTENSION_PATH = os.path.dirname(os.path.realpath(__file__)) +# Read the extension.toml file +EXTENSION_TOML_DATA = toml.load(os.path.join(EXTENSION_PATH, "config", "extension.toml")) + +# Minimum dependencies required prior to installation +INSTALL_REQUIRES = [ + # generic + "wandb>=0.19.6", +] + +PYTORCH_INDEX_URL = ["https://download.pytorch.org/whl/cu118"] + +# Extra dependencies for RL agents +EXTRAS_REQUIRE = {} + +# Cumulation of all extra-requires +EXTRAS_REQUIRE["all"] = list(itertools.chain.from_iterable(EXTRAS_REQUIRE.values())) +# Remove duplicates in the all list to avoid double installations +EXTRAS_REQUIRE["all"] = list(set(EXTRAS_REQUIRE["all"])) + +# Installation operation +setup( + name="uwlab_rl", + author="UW Lab Project Developers", + maintainer="UW Lab Project Developers", + url=EXTENSION_TOML_DATA["package"]["repository"], + version=EXTENSION_TOML_DATA["package"]["version"], + description=EXTENSION_TOML_DATA["package"]["description"], + keywords=EXTENSION_TOML_DATA["package"]["keywords"], + license="BSD-3-Clause", + include_package_data=True, + python_requires=">=3.10", + install_requires=INSTALL_REQUIRES, + dependency_links=PYTORCH_INDEX_URL, + extras_require=EXTRAS_REQUIRE, + packages=["uwlab_rl"], + classifiers=[ + "Natural Language :: English", + "Programming Language :: Python :: 3.10", + "Isaac Sim :: 4.5.0", + ], + zip_safe=False, +) diff --git a/source/uwlab_rl/uwlab_rl/rsl_rl/__init__.py b/source/uwlab_rl/uwlab_rl/rsl_rl/__init__.py new file mode 100644 index 00000000..14d2e333 --- /dev/null +++ b/source/uwlab_rl/uwlab_rl/rsl_rl/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .rl_cfg import BehaviorCloningCfg, OffPolicyAlgorithmCfg, RslRlFancyPpoAlgorithmCfg, SymmetryCfg diff --git a/source/uwlab_rl/uwlab_rl/rsl_rl/ext/__init__.py b/source/uwlab_rl/uwlab_rl/rsl_rl/ext/__init__.py new file mode 100644 index 00000000..4e991c76 --- /dev/null +++ b/source/uwlab_rl/uwlab_rl/rsl_rl/ext/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Main module for the rsl_rl package.""" diff --git a/source/uwlab_rl/uwlab_rl/rsl_rl/ext/algorithms/__init__.py b/source/uwlab_rl/uwlab_rl/rsl_rl/ext/algorithms/__init__.py new file mode 100644 index 00000000..9aef4ca6 --- /dev/null +++ b/source/uwlab_rl/uwlab_rl/rsl_rl/ext/algorithms/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Implementation of different RL agents.""" diff --git a/source/uwlab_rl/uwlab_rl/rsl_rl/ext/algorithms/extensible_ppo.py b/source/uwlab_rl/uwlab_rl/rsl_rl/ext/algorithms/extensible_ppo.py new file mode 100644 index 00000000..ad914ac4 --- /dev/null +++ b/source/uwlab_rl/uwlab_rl/rsl_rl/ext/algorithms/extensible_ppo.py @@ -0,0 +1,640 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.optim as optim +import warnings + +from rsl_rl.modules import ActorCritic +from rsl_rl.modules.rnd import RandomNetworkDistillation +from rsl_rl.utils import string_to_callable + +from ..storage.replay_storage import ReplayStorage +from ..storage.rollout_storage import RolloutStorage + + +class PPO: + """Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347).""" + + actor_critic: ActorCritic + """The actor critic module.""" + + def __init__( + self, + actor_critic, + num_learning_epochs=1, + num_mini_batches=1, + clip_param=0.2, + gamma=0.998, + lam=0.95, + value_loss_coef=1.0, + entropy_coef=0.0, + learning_rate=1e-3, + max_grad_norm=1.0, + use_clipped_value_loss=True, + schedule="fixed", + desired_kl=0.01, + device="cpu", + # BC parameters + behavior_cloning_cfg: dict | None = None, + # RND parameters + rnd_cfg: dict | None = None, + # Symmetry parameters + symmetry_cfg: dict | None = None, + # Offline configuration + offline_algorithm_cfg: dict | None = None, + ): + self.device = device + + self.desired_kl = desired_kl + self.schedule = schedule + self.learning_rate = learning_rate + + # Online configurations + # BC components + if behavior_cloning_cfg is not None: + self.bc = behavior_cloning_cfg + + if self.bc["experts_env_mapping_func"] is not None: + self.experts_env_id_map_fn = self.bc["expert_env_id_map_fn"] + else: + if len(behavior_cloning_cfg["experts_path"]) > 1: + raise ValueError("If you have multiple experts, you need to provide a mapping function.") + self.experts_env_id_map_fn = lambda expert_idx: slice(None) + + if self.bc["experts_observation_func"] is not None: + self.expert_obs_fn = self.bc["expert_obs_fn"] + else: + self.expert_obs_shape = None # same as student observation shape + self.expert_critic_obs_shape = None # same as student critic observation shape + + loader = self.bc["experts_loader"] + if not callable(loader): + loader = eval(loader) + self.experts = [loader(expert_path).to(self.device).eval() for expert_path in self.bc["experts_path"]] + + self.bc_loss_coeff = self.bc["cloning_loss_coeff"] + self.bc_decay = self.bc["loss_decay"] + self.learn_std = self.bc["learn_std"] + + else: + self.bc = None + self.experts = None + self.experts_env_id_map_fn = None + + # RND components + if rnd_cfg is not None: + # Create RND module + self.rnd = RandomNetworkDistillation(device=self.device, **rnd_cfg) + # Create RND optimizer + params = self.rnd.predictor.parameters() + self.rnd_optimizer = optim.Adam(params, lr=rnd_cfg.get("learning_rate", 1e-3)) + else: + self.rnd = None + self.rnd_optimizer = None + + # Symmetry components + if symmetry_cfg is not None: + # Check if symmetry is enabled + use_symmetry = symmetry_cfg["use_data_augmentation"] or symmetry_cfg["use_mirror_loss"] + # Print that we are not using symmetry + if not use_symmetry: + warnings.warn("Symmetry not used for learning. We will use it for logging instead.") + # If function is a string then resolve it to a function + if isinstance(symmetry_cfg["data_augmentation_func"], str): + symmetry_cfg["data_augmentation_func"] = string_to_callable(symmetry_cfg["data_augmentation_func"]) + # Check valid configuration + if symmetry_cfg["use_data_augmentation"] and not callable(symmetry_cfg["data_augmentation_func"]): + raise ValueError( + "Data augmentation enabled but the function is not callable:" + f" {symmetry_cfg['data_augmentation_func']}" + ) + # Store symmetry configuration + self.symmetry = symmetry_cfg + else: + self.symmetry = None + + # PPO components + self.actor_critic = actor_critic + self.actor_critic.to(self.device) + self.storage = None # initialized later + self.optimizer = optim.Adam(self.actor_critic.parameters(), lr=learning_rate) + self.transition = RolloutStorage.Transition() + + # PPO parameters + self.clip_param = clip_param + self.num_learning_epochs = num_learning_epochs + self.num_mini_batches = num_mini_batches + self.value_loss_coef = value_loss_coef + self.entropy_coef = entropy_coef + self.gamma = gamma + self.lam = lam + self.max_grad_norm = max_grad_norm + self.use_clipped_value_loss = use_clipped_value_loss + + # Offline configuration + if offline_algorithm_cfg is not None: + self.offline = True + self.offline_algorithm_cfg = offline_algorithm_cfg + + self.update_counter = 0 + self.update_frequency = self.offline_algorithm_cfg["update_frequencies"] + self.offline_batch_size: int | None = self.offline_algorithm_cfg["batch_size"] + self.offline_num_learning_epochs: int | None = self.offline_algorithm_cfg["num_learning_epochs"] + + # Offline BC + if "behavior_cloning_cfg" in self.offline_algorithm_cfg: + self.offline_bc = self.offline_algorithm_cfg["behavior_cloning_cfg"] + + if self.offline_bc["experts_env_mapping_func"] is not None: + self.offline_experts_env_id_map_fn = self.offline_bc["expert_env_id_map_fn"] + else: + if len(self.offline_bc["experts_path"]) > 1: + raise ValueError("If you have multiple experts, you need to provide a mapping function.") + self.offline_experts_env_id_map_fn = lambda expert_idx: slice(None) + + if self.offline_bc["experts_observation_func"] is not None: + import importlib + + mod, attr_name = self.offline_bc["experts_observation_func"].split(":") + func = getattr(importlib.import_module(mod), attr_name) + self.offline_expert_obs_fn = func + self.offline_expert_obs_shape = self.offline_expert_obs_fn(self.offline_bc["_env"]).shape[1] + else: + self.offline_expert_obs_fn = None + self.offline_expert_obs_shape = None # same as student observation shape + + loader = self.offline_bc["experts_loader"] + if not callable(loader): + loader = eval(loader) + self.offline_experts = [ + loader(expert_path).to(self.device).eval() for expert_path in self.offline_bc["experts_path"] + ] + + self.offline_bc_loss_coeff = self.offline_bc["cloning_loss_coeff"] + self.offline_bc_decay = self.offline_bc["loss_decay"] + self.offline_learn_std = self.offline_bc["learn_std"] + else: + self.offline_bc = False + else: + self.offline = False + self.offline_bc = False # needed the field to exist so can be evaluated with out hasattr(self, offline_bc) + + def init_storage(self, num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, action_shape): + # create memory for RND as well :) + if self.rnd: + rnd_state_shape = [self.rnd.num_states] + else: + rnd_state_shape = None + + expert_mean_action_shape, expert_std_action_shape = None, None + if self.bc: + expert_mean_action_shape = action_shape + expert_std_action_shape = action_shape if self.learn_std else None + if self.offline_bc: + expert_mean_action_shape = action_shape + expert_std_action_shape = action_shape if expert_std_action_shape or self.offline_learn_std else None + + # create rollout storage + self.storage = RolloutStorage( + num_envs, + num_transitions_per_env, + actor_obs_shape, + critic_obs_shape, + action_shape, + rnd_state_shape, + expert_mean_action_shape, + expert_std_action_shape, + self.device, + ) + # create replay storage + if self.offline: + if self.offline_bc: + expert_mean_action_shape = action_shape + expert_std_action_shape = action_shape if self.offline_learn_std else None + else: + expert_mean_action_shape, expert_std_action_shape = None, None + self.replay_storage = ReplayStorage( + num_envs, + num_transitions_per_env * 20, + actor_obs_shape, + critic_obs_shape, + action_shape, + rnd_state_shape, + expert_mean_action_shape, + expert_std_action_shape, + self.device, + ) + + def test_mode(self): + self.actor_critic.test() + + def train_mode(self): + self.actor_critic.train() + + def act(self, obs, critic_obs): + if self.actor_critic.is_recurrent: + self.transition.hidden_states = self.actor_critic.get_hidden_states() + # Compute the actions and values + self.transition.actions = self.actor_critic.act(obs).detach() + self.transition.values = self.actor_critic.evaluate(critic_obs).detach() + self.transition.actions_log_prob = self.actor_critic.get_actions_log_prob(self.transition.actions).detach() + self.transition.action_mean = self.actor_critic.action_mean.detach() + self.transition.action_sigma = self.actor_critic.action_std.detach() + # need to record obs and critic_obs before env.step() + self.transition.observations = obs + self.transition.critic_observations = critic_obs + # BC component + if self.bc: + for i in range(len(self.experts)): + idx_mask = self.experts_env_id_map_fn(i) + self.transition.expert_action_mean = self.experts[i](obs[idx_mask]) + if self.learn_std: + self.transition.expert_action_sigma = self.experts[i].get_actions_log_prob( + self.transition.expert_action_mean + ) + if self.offline_bc: + for i in range(len(self.offline_experts)): + idx_mask = self.offline_experts_env_id_map_fn(i) + expert_obs = obs + if self.offline_expert_obs_fn: + expert_obs = self.offline_expert_obs_fn(self.offline_bc["_env"]) + self.transition.expert_action_mean = self.offline_experts[i](expert_obs[idx_mask]) + if self.offline_learn_std: + self.transition.expert_action_sigma = self.offline_experts[i].get_actions_log_prob( + self.transition.expert_action_mean + ) + + return self.transition.actions + + def process_env_step(self, rewards, dones, infos): + # Record the rewards and dones + # Note: we clone here because later on we bootstrap the rewards based on timeouts + self.transition.rewards = rewards.clone() + self.transition.dones = dones + + # Compute the intrinsic rewards and add to extrinsic rewards + if self.rnd: + # Obtain curiosity gates / observations from infos + rnd_state = infos["observations"]["rnd_state"] + # Compute the intrinsic rewards + # note: rnd_state is the gated_state after normalization if normalization is used + self.intrinsic_rewards, rnd_state = self.rnd.get_intrinsic_reward(rnd_state) + # Add intrinsic rewards to extrinsic rewards + self.transition.rewards += self.intrinsic_rewards + # Record the curiosity gates + self.transition.rnd_state = rnd_state.clone() + + # Bootstrapping on time outs + if "time_outs" in infos: + self.transition.rewards += self.gamma * torch.squeeze( + self.transition.values * infos["time_outs"].unsqueeze(1).to(self.device), 1 + ) + + # Record the transition + self.storage.add_transitions(self.transition) + self.transition.clear() + self.actor_critic.reset(dones) + + def compute_returns(self, last_critic_obs): + # compute value for the last step + last_values = self.actor_critic.evaluate(last_critic_obs).detach() + self.storage.compute_returns(last_values, self.gamma, self.lam) + + def transfer_rollout_to_replay(self, fields: list[str] = ["observations", "privileged_observations"]): + num_steps = self.storage.step # number of valid transitions in rollout + if num_steps == 0: + return + + # Helper to perform vectorized copy into a circular buffer. + def copy_to_replay(replay_field: torch.Tensor, rollout_field: torch.Tensor): + # Detach and clone the slice from rollout storage. + # shape: [num_steps, num_envs, ...] + start_idx = self.replay_storage.step + end_idx = start_idx + num_steps + if end_idx <= self.replay_storage.capacity: + # No wrap-around needed. + replay_field[start_idx:end_idx] = rollout_field[:num_steps].detach().clone() + else: + # Wrap-around: split the batch copy into two parts. + part1 = self.replay_storage.capacity - start_idx + replay_field[start_idx:] = rollout_field[:part1].detach().clone() + replay_field[: end_idx % self.replay_storage.capacity] = rollout_field[part1:].detach().clone() + + # Iterate over the specified fields and transfer them if available. + for field in fields: + copy_to_replay(getattr(self.replay_storage, field), getattr(self.storage, field)) + + # Update circular buffer pointers in replay storage. + self.replay_storage.step = (self.replay_storage.step + num_steps) % self.replay_storage.capacity + self.replay_storage.size = min(self.replay_storage.size + num_steps, self.replay_storage.capacity) + + def update(self): # noqa: C901 + mean_value_loss = 0 + mean_surrogate_loss = 0 + mean_entropy = 0 + # -- BC loss + if self.bc: + mean_bc_loss = 0 + else: + mean_bc_loss = None + # -- RND loss + if self.rnd: + mean_rnd_loss = 0 + else: + mean_rnd_loss = None + # -- Symmetry loss + if self.symmetry: + mean_symmetry_loss = 0 + else: + mean_symmetry_loss = None + + # generator for mini batches + if self.actor_critic.is_recurrent: + generator = self.storage.recurrent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) + else: + generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) + # iterate over batches + for ( + obs_batch, + critic_obs_batch, + actions_batch, + target_values_batch, + advantages_batch, + returns_batch, + old_actions_log_prob_batch, + old_mu_batch, + old_sigma_batch, + hid_states_batch, + masks_batch, + expert_action_mu_batch, + expert_action_sigma_batch, + rnd_state_batch, + ) in generator: + + # number of augmentations per sample + # we start with 1 and increase it if we use symmetry augmentation + num_aug = 1 + # original batch size + original_batch_size = obs_batch.shape[0] + + # Perform symmetric augmentation + if self.symmetry and self.symmetry["use_data_augmentation"]: + # augmentation using symmetry + data_augmentation_func = self.symmetry["data_augmentation_func"] + # returned shape: [batch_size * num_aug, ...] + obs_batch, actions_batch = data_augmentation_func( + obs=obs_batch, actions=actions_batch, env=self.symmetry["_env"], is_critic=False + ) + critic_obs_batch, _ = data_augmentation_func( + obs=critic_obs_batch, actions=None, env=self.symmetry["_env"], is_critic=True + ) + # compute number of augmentations per sample + num_aug = int(obs_batch.shape[0] / original_batch_size) + # repeat the rest of the batch + # -- actor + old_actions_log_prob_batch = old_actions_log_prob_batch.repeat(num_aug, 1) + # -- critic + target_values_batch = target_values_batch.repeat(num_aug, 1) + advantages_batch = advantages_batch.repeat(num_aug, 1) + returns_batch = returns_batch.repeat(num_aug, 1) + + # Recompute actions log prob and entropy for current batch of transitions + # Note: we need to do this because we updated the actor_critic with the new parameters + # -- actor + self.actor_critic.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0]) + actions_log_prob_batch = self.actor_critic.get_actions_log_prob(actions_batch) + # -- critic + value_batch = self.actor_critic.evaluate( + critic_obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1] + ) + # -- entropy + # we only keep the entropy of the first augmentation (the original one) + mu_batch = self.actor_critic.action_mean[:original_batch_size] + sigma_batch = self.actor_critic.action_std[:original_batch_size] + entropy_batch = self.actor_critic.entropy[:original_batch_size] + + # KL + if self.desired_kl is not None and self.schedule == "adaptive": + with torch.inference_mode(): + kl = torch.sum( + torch.log(sigma_batch / old_sigma_batch + 1.0e-5) + + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) + / (2.0 * torch.square(sigma_batch)) + - 0.5, + axis=-1, + ) + kl_mean = torch.mean(kl) + + if kl_mean > self.desired_kl * 2.0: + self.learning_rate = max(1e-5, self.learning_rate / 1.5) + elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0: + self.learning_rate = min(1e-2, self.learning_rate * 1.5) + + for param_group in self.optimizer.param_groups: + param_group["lr"] = self.learning_rate + + # Surrogate loss + ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)) + surrogate = -torch.squeeze(advantages_batch) * ratio + surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp( + ratio, 1.0 - self.clip_param, 1.0 + self.clip_param + ) + surrogate_loss = torch.max(surrogate, surrogate_clipped).mean() + + # Value function loss + if self.use_clipped_value_loss: + value_clipped = target_values_batch + (value_batch - target_values_batch).clamp( + -self.clip_param, self.clip_param + ) + value_losses = (value_batch - returns_batch).pow(2) + value_losses_clipped = (value_clipped - returns_batch).pow(2) + value_loss = torch.max(value_losses, value_losses_clipped).mean() + else: + value_loss = (returns_batch - value_batch).pow(2).mean() + + loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean() + + # Symmetry loss + if self.symmetry: + # obtain the symmetric actions + # if we did augmentation before then we don't need to augment again + if not self.symmetry["use_data_augmentation"]: + data_augmentation_func = self.symmetry["data_augmentation_func"] + obs_batch, _ = data_augmentation_func( + obs=obs_batch, actions=None, env=self.symmetry["_env"], is_critic=False + ) + # compute number of augmentations per sample + num_aug = int(obs_batch.shape[0] / original_batch_size) + + # actions predicted by the actor for symmetrically-augmented observations + mean_actions_batch = self.actor_critic.act_inference(obs_batch.detach().clone()) + + # compute the symmetrically augmented actions + # note: we are assuming the first augmentation is the original one. + # We do not use the action_batch from earlier since that action was sampled from the distribution. + # However, the symmetry loss is computed using the mean of the distribution. + action_mean_orig = mean_actions_batch[:original_batch_size] + _, actions_mean_symm_batch = data_augmentation_func( + obs=None, actions=action_mean_orig, env=self.symmetry["_env"], is_critic=False + ) + + # compute the loss (we skip the first augmentation as it is the original one) + mse_loss = torch.nn.MSELoss() + symmetry_loss = mse_loss( + mean_actions_batch[original_batch_size:], actions_mean_symm_batch.detach()[original_batch_size:] + ) + # add the loss to the total loss + if self.symmetry["use_mirror_loss"]: + loss += self.symmetry["mirror_loss_coeff"] * symmetry_loss + else: + symmetry_loss = symmetry_loss.detach() + # BC loss + if self.bc: + mse_loss = torch.nn.MSELoss() + mean_loss = mse_loss(mu_batch, expert_action_mu_batch) + bc_loss = mean_loss + if self.learn_std: + std_loss = mse_loss(sigma_batch, expert_action_sigma_batch) + bc_loss += std_loss + self.bc_loss_coeff *= self.bc_decay + loss = (1 - self.bc_loss_coeff) * loss + self.bc_loss_coeff * bc_loss + + # Random Network Distillation loss + if self.rnd: + # predict the embedding and the target + predicted_embedding = self.rnd.predictor(rnd_state_batch) + target_embedding = self.rnd.target(rnd_state_batch) + # compute the loss as the mean squared error + mseloss = torch.nn.MSELoss() + rnd_loss = mseloss(predicted_embedding, target_embedding.detach()) + + # Gradient step + # -- For PPO + self.optimizer.zero_grad() + loss.backward() + nn.utils.clip_grad_norm_(self.actor_critic.parameters(), self.max_grad_norm) + self.optimizer.step() + # -- For RND + if self.rnd_optimizer: + self.rnd_optimizer.zero_grad() + rnd_loss.backward() + self.rnd_optimizer.step() + + # Store the losses + mean_value_loss += value_loss.item() + mean_surrogate_loss += surrogate_loss.item() + mean_entropy += entropy_batch.mean().item() + # -- BC loss + if mean_bc_loss is not None: + mean_bc_loss += bc_loss.item() + # -- RND loss + if mean_rnd_loss is not None: + mean_rnd_loss += rnd_loss.item() + # -- Symmetry loss + if mean_symmetry_loss is not None: + mean_symmetry_loss += symmetry_loss.item() + + # -- For PPO + num_updates = self.num_learning_epochs * self.num_mini_batches + mean_value_loss /= num_updates + mean_surrogate_loss /= num_updates + # -- For BC + if mean_bc_loss is not None: + mean_bc_loss /= num_updates + # -- For RND + if mean_rnd_loss is not None: + mean_rnd_loss /= num_updates + # -- For Symmetry + if mean_symmetry_loss is not None: + mean_symmetry_loss /= num_updates + # -- Clear the storage + + if self.offline: + fields_to_transfer = ["observations", "privileged_observations"] + if self.offline_bc: + fields_to_transfer.append("expert_action_mean") + if self.offline_learn_std: + fields_to_transfer.append("expert_action_sigma") + self.transfer_rollout_to_replay(fields_to_transfer) + while self.update_counter >= 0: + capacity_percentage, offline_bc_loss = self.update_offline() + self.update_counter -= 1 / self.update_frequency + self.update_counter += 1 + else: + capacity_percentage, offline_bc_loss = None, None + self.storage.clear() + + return ( + mean_value_loss, + mean_surrogate_loss, + mean_entropy, + mean_bc_loss, + mean_rnd_loss, + mean_symmetry_loss, + capacity_percentage, + offline_bc_loss, + ) + + def update_offline(self): + loss = 0 + storage_data_size = self.replay_storage.num_envs * self.replay_storage.size + + batch_size = ( + self.offline_batch_size + if self.offline_batch_size + else self.storage.num_envs * self.storage.num_transitions_per_env // self.num_mini_batches + ) + num_mini_batches = storage_data_size // batch_size + num_learning_epoch = ( + self.offline_num_learning_epochs if self.offline_num_learning_epochs else self.num_learning_epochs + ) + + if self.actor_critic.is_recurrent: + generator = self.replay_storage.recurrent_mini_batch_generator(num_mini_batches, num_learning_epoch) + else: + generator = self.replay_storage.mini_batch_generator(num_mini_batches, num_learning_epoch) + + # -- For BC + if self.offline_bc: + mse_loss_fn = nn.MSELoss() + mean_bc_loss = 0.0 + else: + mean_bc_loss = None + + for obs_batch, _, _, _, _, _, _, _, _, _, _, expert_action_mu_batch, expert_action_sigma_batch, _ in generator: + # -- For BC + if self.offline_bc: + predicted_actions = self.actor_critic.act_inference(obs_batch) + bc_loss = mse_loss_fn(predicted_actions, expert_action_mu_batch) + if self.offline_learn_std: + predicted_std = self.actor_critic.action_std[: obs_batch.shape[0]] + bc_loss += mse_loss_fn(predicted_std, expert_action_sigma_batch) + self.offline_bc_loss_coeff *= self.offline_bc_decay + loss = (1 - self.offline_bc_loss_coeff) * loss + self.offline_bc_loss_coeff * bc_loss + + self.optimizer.zero_grad() + bc_loss.backward() + nn.utils.clip_grad_norm_(self.actor_critic.parameters(), self.max_grad_norm) + self.optimizer.step() + + # --BC loss + if mean_bc_loss is not None: + mean_bc_loss += bc_loss.item() + + num_updates = num_mini_batches * num_learning_epoch + # For BC + if mean_bc_loss is not None: + mean_bc_loss /= num_updates + capacity_percentage = self.replay_storage.size / self.replay_storage.capacity + return capacity_percentage, mean_bc_loss diff --git a/source/uwlab_rl/uwlab_rl/rsl_rl/ext/algorithms/original_ppo.py b/source/uwlab_rl/uwlab_rl/rsl_rl/ext/algorithms/original_ppo.py new file mode 100644 index 00000000..8b3ec743 --- /dev/null +++ b/source/uwlab_rl/uwlab_rl/rsl_rl/ext/algorithms/original_ppo.py @@ -0,0 +1,378 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.optim as optim +import warnings + +from rsl_rl.modules import ActorCritic +from rsl_rl.modules.rnd import RandomNetworkDistillation +from rsl_rl.storage import RolloutStorage +from rsl_rl.utils import string_to_callable + + +class PPO: + """Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347).""" + + actor_critic: ActorCritic + """The actor critic module.""" + + def __init__( + self, + actor_critic, + num_learning_epochs=1, + num_mini_batches=1, + clip_param=0.2, + gamma=0.998, + lam=0.95, + value_loss_coef=1.0, + entropy_coef=0.0, + learning_rate=1e-3, + max_grad_norm=1.0, + use_clipped_value_loss=True, + schedule="fixed", + desired_kl=0.01, + device="cpu", + # RND parameters + rnd_cfg: dict | None = None, + # Symmetry parameters + symmetry_cfg: dict | None = None, + ): + self.device = device + + self.desired_kl = desired_kl + self.schedule = schedule + self.learning_rate = learning_rate + + # RND components + if rnd_cfg is not None: + # Create RND module + self.rnd = RandomNetworkDistillation(device=self.device, **rnd_cfg) + # Create RND optimizer + params = self.rnd.predictor.parameters() + self.rnd_optimizer = optim.Adam(params, lr=rnd_cfg.get("learning_rate", 1e-3)) + else: + self.rnd = None + self.rnd_optimizer = None + + # Symmetry components + if symmetry_cfg is not None: + # Check if symmetry is enabled + use_symmetry = symmetry_cfg["use_data_augmentation"] or symmetry_cfg["use_mirror_loss"] + # Print that we are not using symmetry + if not use_symmetry: + warnings.warn("Symmetry not used for learning. We will use it for logging instead.") + # If function is a string then resolve it to a function + if isinstance(symmetry_cfg["data_augmentation_func"], str): + symmetry_cfg["data_augmentation_func"] = string_to_callable(symmetry_cfg["data_augmentation_func"]) + # Check valid configuration + if symmetry_cfg["use_data_augmentation"] and not callable(symmetry_cfg["data_augmentation_func"]): + raise ValueError( + "Data augmentation enabled but the function is not callable:" + f" {symmetry_cfg['data_augmentation_func']}" + ) + # Store symmetry configuration + self.symmetry = symmetry_cfg + else: + self.symmetry = None + + # PPO components + self.actor_critic = actor_critic + self.actor_critic.to(self.device) + self.storage = None # initialized later + self.optimizer = optim.Adam(self.actor_critic.parameters(), lr=learning_rate) + self.transition = RolloutStorage.Transition() + + # PPO parameters + self.clip_param = clip_param + self.num_learning_epochs = num_learning_epochs + self.num_mini_batches = num_mini_batches + self.value_loss_coef = value_loss_coef + self.entropy_coef = entropy_coef + self.gamma = gamma + self.lam = lam + self.max_grad_norm = max_grad_norm + self.use_clipped_value_loss = use_clipped_value_loss + + def init_storage(self, num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, action_shape): + # create memory for RND as well :) + if self.rnd: + rnd_state_shape = [self.rnd.num_states] + else: + rnd_state_shape = None + # create rollout storage + self.storage = RolloutStorage( + num_envs, + num_transitions_per_env, + actor_obs_shape, + critic_obs_shape, + action_shape, + rnd_state_shape, + self.device, + ) + + def test_mode(self): + self.actor_critic.test() + + def train_mode(self): + self.actor_critic.train() + + def act(self, obs, critic_obs): + if self.actor_critic.is_recurrent: + self.transition.hidden_states = self.actor_critic.get_hidden_states() + # Compute the actions and values + self.transition.actions = self.actor_critic.act(obs).detach() + self.transition.values = self.actor_critic.evaluate(critic_obs).detach() + self.transition.actions_log_prob = self.actor_critic.get_actions_log_prob(self.transition.actions).detach() + self.transition.action_mean = self.actor_critic.action_mean.detach() + self.transition.action_sigma = self.actor_critic.action_std.detach() + # need to record obs and critic_obs before env.step() + self.transition.observations = obs + self.transition.critic_observations = critic_obs + return self.transition.actions + + def process_env_step(self, rewards, dones, infos): + # Record the rewards and dones + # Note: we clone here because later on we bootstrap the rewards based on timeouts + self.transition.rewards = rewards.clone() + self.transition.dones = dones + + # Compute the intrinsic rewards and add to extrinsic rewards + if self.rnd: + # Obtain curiosity gates / observations from infos + rnd_state = infos["observations"]["rnd_state"] + # Compute the intrinsic rewards + # note: rnd_state is the gated_state after normalization if normalization is used + self.intrinsic_rewards, rnd_state = self.rnd.get_intrinsic_reward(rnd_state) + # Add intrinsic rewards to extrinsic rewards + self.transition.rewards += self.intrinsic_rewards + # Record the curiosity gates + self.transition.rnd_state = rnd_state.clone() + + # Bootstrapping on time outs + if "time_outs" in infos: + self.transition.rewards += self.gamma * torch.squeeze( + self.transition.values * infos["time_outs"].unsqueeze(1).to(self.device), 1 + ) + + # Record the transition + self.storage.add_transitions(self.transition) + self.transition.clear() + self.actor_critic.reset(dones) + + def compute_returns(self, last_critic_obs): + # compute value for the last step + last_values = self.actor_critic.evaluate(last_critic_obs).detach() + self.storage.compute_returns(last_values, self.gamma, self.lam) + + def update(self): # noqa: C901 + mean_value_loss = 0 + mean_surrogate_loss = 0 + mean_entropy = 0 + # -- RND loss + if self.rnd: + mean_rnd_loss = 0 + else: + mean_rnd_loss = None + # -- Symmetry loss + if self.symmetry: + mean_symmetry_loss = 0 + else: + mean_symmetry_loss = None + + # generator for mini batches + if self.actor_critic.is_recurrent: + generator = self.storage.recurrent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) + else: + generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) + + # iterate over batches + for ( + obs_batch, + critic_obs_batch, + actions_batch, + target_values_batch, + advantages_batch, + returns_batch, + old_actions_log_prob_batch, + old_mu_batch, + old_sigma_batch, + hid_states_batch, + masks_batch, + rnd_state_batch, + ) in generator: + + # number of augmentations per sample + # we start with 1 and increase it if we use symmetry augmentation + num_aug = 1 + # original batch size + original_batch_size = obs_batch.shape[0] + + # Perform symmetric augmentation + if self.symmetry and self.symmetry["use_data_augmentation"]: + # augmentation using symmetry + data_augmentation_func = self.symmetry["data_augmentation_func"] + # returned shape: [batch_size * num_aug, ...] + obs_batch, actions_batch = data_augmentation_func( + obs=obs_batch, actions=actions_batch, env=self.symmetry["_env"], is_critic=False + ) + critic_obs_batch, _ = data_augmentation_func( + obs=critic_obs_batch, actions=None, env=self.symmetry["_env"], is_critic=True + ) + # compute number of augmentations per sample + num_aug = int(obs_batch.shape[0] / original_batch_size) + # repeat the rest of the batch + # -- actor + old_actions_log_prob_batch = old_actions_log_prob_batch.repeat(num_aug, 1) + # -- critic + target_values_batch = target_values_batch.repeat(num_aug, 1) + advantages_batch = advantages_batch.repeat(num_aug, 1) + returns_batch = returns_batch.repeat(num_aug, 1) + + # Recompute actions log prob and entropy for current batch of transitions + # Note: we need to do this because we updated the actor_critic with the new parameters + # -- actor + self.actor_critic.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0]) + actions_log_prob_batch = self.actor_critic.get_actions_log_prob(actions_batch) + # -- critic + value_batch = self.actor_critic.evaluate( + critic_obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1] + ) + # -- entropy + # we only keep the entropy of the first augmentation (the original one) + mu_batch = self.actor_critic.action_mean[:original_batch_size] + sigma_batch = self.actor_critic.action_std[:original_batch_size] + entropy_batch = self.actor_critic.entropy[:original_batch_size] + + # KL + if self.desired_kl is not None and self.schedule == "adaptive": + with torch.inference_mode(): + kl = torch.sum( + torch.log(sigma_batch / old_sigma_batch + 1.0e-5) + + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) + / (2.0 * torch.square(sigma_batch)) + - 0.5, + axis=-1, + ) + kl_mean = torch.mean(kl) + + if kl_mean > self.desired_kl * 2.0: + self.learning_rate = max(1e-5, self.learning_rate / 1.5) + elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0: + self.learning_rate = min(1e-2, self.learning_rate * 1.5) + + for param_group in self.optimizer.param_groups: + param_group["lr"] = self.learning_rate + + # Surrogate loss + ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)) + surrogate = -torch.squeeze(advantages_batch) * ratio + surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp( + ratio, 1.0 - self.clip_param, 1.0 + self.clip_param + ) + surrogate_loss = torch.max(surrogate, surrogate_clipped).mean() + + # Value function loss + if self.use_clipped_value_loss: + value_clipped = target_values_batch + (value_batch - target_values_batch).clamp( + -self.clip_param, self.clip_param + ) + value_losses = (value_batch - returns_batch).pow(2) + value_losses_clipped = (value_clipped - returns_batch).pow(2) + value_loss = torch.max(value_losses, value_losses_clipped).mean() + else: + value_loss = (returns_batch - value_batch).pow(2).mean() + + loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean() + + # Symmetry loss + if self.symmetry: + # obtain the symmetric actions + # if we did augmentation before then we don't need to augment again + if not self.symmetry["use_data_augmentation"]: + data_augmentation_func = self.symmetry["data_augmentation_func"] + obs_batch, _ = data_augmentation_func( + obs=obs_batch, actions=None, env=self.symmetry["_env"], is_critic=False + ) + # compute number of augmentations per sample + num_aug = int(obs_batch.shape[0] / original_batch_size) + + # actions predicted by the actor for symmetrically-augmented observations + mean_actions_batch = self.actor_critic.act_inference(obs_batch.detach().clone()) + + # compute the symmetrically augmented actions + # note: we are assuming the first augmentation is the original one. + # We do not use the action_batch from earlier since that action was sampled from the distribution. + # However, the symmetry loss is computed using the mean of the distribution. + action_mean_orig = mean_actions_batch[:original_batch_size] + _, actions_mean_symm_batch = data_augmentation_func( + obs=None, actions=action_mean_orig, env=self.symmetry["_env"], is_critic=False + ) + + # compute the loss (we skip the first augmentation as it is the original one) + mse_loss = torch.nn.MSELoss() + symmetry_loss = mse_loss( + mean_actions_batch[original_batch_size:], actions_mean_symm_batch.detach()[original_batch_size:] + ) + # add the loss to the total loss + if self.symmetry["use_mirror_loss"]: + loss += self.symmetry["mirror_loss_coeff"] * symmetry_loss + else: + symmetry_loss = symmetry_loss.detach() + + # Random Network Distillation loss + if self.rnd: + # predict the embedding and the target + predicted_embedding = self.rnd.predictor(rnd_state_batch) + target_embedding = self.rnd.target(rnd_state_batch) + # compute the loss as the mean squared error + mseloss = torch.nn.MSELoss() + rnd_loss = mseloss(predicted_embedding, target_embedding.detach()) + + # Gradient step + # -- For PPO + self.optimizer.zero_grad() + loss.backward() + nn.utils.clip_grad_norm_(self.actor_critic.parameters(), self.max_grad_norm) + self.optimizer.step() + # -- For RND + if self.rnd_optimizer: + self.rnd_optimizer.zero_grad() + rnd_loss.backward() + self.rnd_optimizer.step() + + # Store the losses + mean_value_loss += value_loss.item() + mean_surrogate_loss += surrogate_loss.item() + mean_entropy += entropy_batch.mean().item() + # -- RND loss + if mean_rnd_loss is not None: + mean_rnd_loss += rnd_loss.item() + # -- Symmetry loss + if mean_symmetry_loss is not None: + mean_symmetry_loss += symmetry_loss.item() + + # -- For PPO + num_updates = self.num_learning_epochs * self.num_mini_batches + mean_value_loss /= num_updates + mean_surrogate_loss /= num_updates + # -- For RND + if mean_rnd_loss is not None: + mean_rnd_loss /= num_updates + # -- For Symmetry + if mean_symmetry_loss is not None: + mean_symmetry_loss /= num_updates + # -- Clear the storage + self.storage.clear() + + return mean_value_loss, mean_surrogate_loss, mean_entropy, mean_rnd_loss, mean_symmetry_loss diff --git a/source/uwlab_rl/uwlab_rl/rsl_rl/ext/runners/on_policy_runner.py b/source/uwlab_rl/uwlab_rl/rsl_rl/ext/runners/on_policy_runner.py new file mode 100644 index 00000000..b0335d9c --- /dev/null +++ b/source/uwlab_rl/uwlab_rl/rsl_rl/ext/runners/on_policy_runner.py @@ -0,0 +1,472 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +import statistics +import time +import torch +from collections import deque + +import rsl_rl +from rsl_rl.env import VecEnv +from rsl_rl.modules import ActorCritic, ActorCriticRecurrent, EmpiricalNormalization +from rsl_rl.utils import store_code_state + +from ..algorithms.extensible_ppo import PPO + + +class OnPolicyRunner: + """On-policy runner for training and evaluation.""" + + def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): + self.cfg = train_cfg + self.alg_cfg = train_cfg["algorithm"] + self.policy_cfg = train_cfg["policy"] + self.device = device + self.env = env + + # resolve dimensions of observations + obs, extras = self.env.get_observations() + num_obs = obs.shape[1] + if "critic" in extras["observations"]: + num_critic_obs = extras["observations"]["critic"].shape[1] + else: + num_critic_obs = num_obs + actor_critic_class = eval(self.policy_cfg.pop("class_name")) # ActorCritic + actor_critic: ActorCritic | ActorCriticRecurrent = actor_critic_class( + num_obs, num_critic_obs, self.env.num_actions, **self.policy_cfg + ).to(self.device) + + # resolve dimension of rnd gated state + if "rnd_cfg" in self.alg_cfg: + # check if rnd gated state is present + rnd_state = extras["observations"].get("rnd_state") + if rnd_state is None: + raise ValueError("Observations for they key 'rnd_state' not found in infos['observations'].") + # get dimension of rnd gated state + num_rnd_state = rnd_state.shape[1] + # add rnd gated state to config + self.alg_cfg["rnd_cfg"]["num_state"] = num_rnd_state + # scale down the rnd weight with timestep (similar to how rewards are scaled down in legged_gym envs) + self.alg_cfg["rnd_cfg"]["weight"] *= env.dt + + # if using symmetry then pass the environment config object + if "symmetry_cfg" in self.alg_cfg: + # this is used by the symmetry function for handling different observation terms + self.alg_cfg["symmetry_cfg"]["_env"] = env + + if "offline_algorithm_cfg" in self.alg_cfg: + if "behavior_cloning_cfg" in self.alg_cfg["offline_algorithm_cfg"]: + # this is used by the symmetry function for handling different observation terms + self.alg_cfg["offline_algorithm_cfg"]["behavior_cloning_cfg"]["_env"] = env + + # init algorithm + alg_class = eval(self.alg_cfg.pop("class_name")) # PPO + self.alg: PPO = alg_class(actor_critic, device=self.device, **self.alg_cfg) + + # store training configuration + self.num_steps_per_env = self.cfg["num_steps_per_env"] + self.save_interval = self.cfg["save_interval"] + self.empirical_normalization = self.cfg["empirical_normalization"] + if self.empirical_normalization: + self.obs_normalizer = EmpiricalNormalization(shape=[num_obs], until=1.0e8).to(self.device) + self.critic_obs_normalizer = EmpiricalNormalization(shape=[num_critic_obs], until=1.0e8).to(self.device) + else: + self.obs_normalizer = torch.nn.Identity().to(self.device) # no normalization + self.critic_obs_normalizer = torch.nn.Identity().to(self.device) # no normalization + # init storage and model + self.alg.init_storage( + self.env.num_envs, + self.num_steps_per_env, + [num_obs], + [num_critic_obs], + [self.env.num_actions], + ) + + # Log + self.log_dir = log_dir + self.writer = None + self.tot_timesteps = 0 + self.tot_time = 0 + self.current_learning_iteration = 0 + self.git_status_repos = [rsl_rl.__file__] + + def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): + # initialize writer + if self.log_dir is not None and self.writer is None: + # Launch either Tensorboard or Neptune & Tensorboard summary writer(s), default: Tensorboard. + self.logger_type = self.cfg.get("logger", "tensorboard") + self.logger_type = self.logger_type.lower() + + if self.logger_type == "neptune": + from rsl_rl.utils.neptune_utils import NeptuneSummaryWriter + + self.writer = NeptuneSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) + elif self.logger_type == "wandb": + from rsl_rl.utils.wandb_utils import WandbSummaryWriter + + self.writer = WandbSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) + elif self.logger_type == "tensorboard": + from torch.utils.tensorboard import SummaryWriter + + self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10) + else: + raise ValueError("Logger type not found. Please choose 'neptune', 'wandb' or 'tensorboard'.") + + # randomize initial episode lengths (for exploration) + if init_at_random_ep_len: + self.env.episode_length_buf = torch.randint_like( + self.env.episode_length_buf, high=int(self.env.max_episode_length) + ) + + # start learning + obs, extras = self.env.get_observations() + critic_obs = extras["observations"].get("critic", obs) + obs, critic_obs = obs.to(self.device), critic_obs.to(self.device) + self.train_mode() # switch to train mode (for dropout for example) + + # Book keeping + ep_infos = [] + rewbuffer = deque(maxlen=100) + lenbuffer = deque(maxlen=100) + cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + # create buffers for logging extrinsic and intrinsic rewards + if self.alg.rnd: + erewbuffer = deque(maxlen=100) + irewbuffer = deque(maxlen=100) + cur_ereward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_ireward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + + start_iter = self.current_learning_iteration + tot_iter = start_iter + num_learning_iterations + for it in range(start_iter, tot_iter): + start = time.time() + # Rollout + with torch.inference_mode(): + for _ in range(self.num_steps_per_env): + # Sample actions from policy + actions = self.alg.act(obs, critic_obs) + # Step environment + obs, rewards, dones, infos = self.env.step(actions.to(self.env.device)) + + # Move to the agent device + obs, rewards, dones = obs.to(self.device), rewards.to(self.device), dones.to(self.device) + + # Normalize observations + obs = self.obs_normalizer(obs) + # Extract critic observations and normalize + if "critic" in infos["observations"]: + critic_obs = self.critic_obs_normalizer(infos["observations"]["critic"].to(self.device)) + else: + critic_obs = obs + + # Intrinsic rewards (extracted here only for logging)! + intrinsic_rewards = self.alg.intrinsic_rewards if self.alg.rnd else None + + # Process env step and store in buffer + self.alg.process_env_step(rewards, dones, infos) + + if self.log_dir is not None: + # Book keeping + if "episode" in infos: + ep_infos.append(infos["episode"]) + elif "log" in infos: + ep_infos.append(infos["log"]) + # Update rewards + if self.alg.rnd: + cur_ereward_sum += rewards + cur_ireward_sum += intrinsic_rewards # type: ignore + cur_reward_sum += rewards + intrinsic_rewards + else: + cur_reward_sum += rewards + # Update episode length + cur_episode_length += 1 + # Clear data for completed episodes + # -- common + new_ids = (dones > 0).nonzero(as_tuple=False) + rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) + lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) + cur_reward_sum[new_ids] = 0 + cur_episode_length[new_ids] = 0 + # -- intrinsic and extrinsic rewards + if self.alg.rnd: + erewbuffer.extend(cur_ereward_sum[new_ids][:, 0].cpu().numpy().tolist()) + irewbuffer.extend(cur_ireward_sum[new_ids][:, 0].cpu().numpy().tolist()) + cur_ereward_sum[new_ids] = 0 + cur_ireward_sum[new_ids] = 0 + + stop = time.time() + collection_time = stop - start + + # Learning step + start = stop + self.alg.compute_returns(critic_obs) + + # Update policy + # Note: we keep arguments here since locals() loads them + ( + mean_value_loss, + mean_surrogate_loss, + mean_entropy, + mean_bc_loss, + mean_rnd_loss, + mean_symmetry_loss, + capacity_percentage, + offline_bc_loss, + ) = self.alg.update() + stop = time.time() + learn_time = stop - start + self.current_learning_iteration = it + + # Logging info and save checkpoint + if self.log_dir is not None: + # Log information + self.log(locals()) + # Save model + if it % self.save_interval == 0: + self.save(os.path.join(self.log_dir, f"model_{it}.pt")) + + # Clear episode infos + ep_infos.clear() + + # Save code state + if it == start_iter: + # obtain all the diff files + git_file_paths = store_code_state(self.log_dir, self.git_status_repos) + # if possible store them to wandb + if self.logger_type in ["wandb", "neptune"] and git_file_paths: + for path in git_file_paths: + self.writer.save_file(path) + + # Save the final model after training + if self.log_dir is not None: + self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) + + def log(self, locs: dict, width: int = 80, pad: int = 35): + self.tot_timesteps += self.num_steps_per_env * self.env.num_envs + self.tot_time += locs["collection_time"] + locs["learn_time"] + iteration_time = locs["collection_time"] + locs["learn_time"] + + # -- Episode info + ep_string = "" + if locs["ep_infos"]: + for key in locs["ep_infos"][0]: + infotensor = torch.tensor([], device=self.device) + for ep_info in locs["ep_infos"]: + # handle scalar and zero dimensional tensor infos + if key not in ep_info: + continue + if not isinstance(ep_info[key], torch.Tensor): + ep_info[key] = torch.Tensor([ep_info[key]]) + if len(ep_info[key].shape) == 0: + ep_info[key] = ep_info[key].unsqueeze(0) + infotensor = torch.cat((infotensor, ep_info[key].to(self.device))) + value = torch.mean(infotensor) + # log to logger and terminal + if "/" in key: + self.writer.add_scalar(key, value, locs["it"]) + ep_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" + else: + self.writer.add_scalar("Episode/" + key, value, locs["it"]) + ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n""" + mean_std = self.alg.actor_critic.std.mean() + fps = int(self.num_steps_per_env * self.env.num_envs / (locs["collection_time"] + locs["learn_time"])) + + # -- Losses + self.writer.add_scalar("Loss/value_function", locs["mean_value_loss"], locs["it"]) + self.writer.add_scalar("Loss/surrogate", locs["mean_surrogate_loss"], locs["it"]) + self.writer.add_scalar("Loss/entropy", locs["mean_entropy"], locs["it"]) + self.writer.add_scalar("Loss/learning_rate", self.alg.learning_rate, locs["it"]) + if self.alg.bc: + self.writer.add_scalar("Loss/bc", locs["mean_bc_loss"], locs["it"]) + if self.alg.rnd: + self.writer.add_scalar("Loss/rnd", locs["mean_rnd_loss"], locs["it"]) + if self.alg.symmetry: + self.writer.add_scalar("Loss/symmetry", locs["mean_symmetry_loss"], locs["it"]) + if self.alg.offline: + if self.alg.offline_bc: + self.writer.add_scalar("Loss/offline_bc", locs["offline_bc_loss"], locs["it"]) + + # -- Policy + self.writer.add_scalar("Policy/mean_noise_std", mean_std.item(), locs["it"]) + + # -- Performance + self.writer.add_scalar("Perf/total_fps", fps, locs["it"]) + self.writer.add_scalar("Perf/collection time", locs["collection_time"], locs["it"]) + self.writer.add_scalar("Perf/learning_time", locs["learn_time"], locs["it"]) + + # -- Training + if len(locs["rewbuffer"]) > 0: + # separate logging for intrinsic and extrinsic rewards + if self.alg.rnd: + self.writer.add_scalar("Rnd/mean_extrinsic_reward", statistics.mean(locs["erewbuffer"]), locs["it"]) + self.writer.add_scalar("Rnd/mean_intrinsic_reward", statistics.mean(locs["irewbuffer"]), locs["it"]) + self.writer.add_scalar("Rnd/weight", self.alg.rnd.weight, locs["it"]) + # everything else + self.writer.add_scalar("Train/mean_reward", statistics.mean(locs["rewbuffer"]), locs["it"]) + self.writer.add_scalar("Train/mean_episode_length", statistics.mean(locs["lenbuffer"]), locs["it"]) + if self.logger_type != "wandb": # wandb does not support non-integer x-axis logging + self.writer.add_scalar("Train/mean_reward/time", statistics.mean(locs["rewbuffer"]), self.tot_time) + self.writer.add_scalar( + "Train/mean_episode_length/time", statistics.mean(locs["lenbuffer"]), self.tot_time + ) + + str = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m " + + if len(locs["rewbuffer"]) > 0: + log_string = ( + f"""{'#' * width}\n""" + f"""{str.center(width, ' ')}\n\n""" + f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ + 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" + f"""{'Value function loss:':>{pad}} {locs['mean_value_loss']:.4f}\n""" + f"""{'Surrogate loss:':>{pad}} {locs['mean_surrogate_loss']:.4f}\n""" + ) + + # -- For BC + if self.alg.bc: + log_string += f"""{'Behavior cloning loss:':>{pad}} {locs['mean_bc_loss']:.4f}\n""" + + # -- For symmetry + if self.alg.symmetry: + log_string += f"""{'Symmetry loss:':>{pad}} {locs['mean_symmetry_loss']:.4f}\n""" + + log_string += f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" + + # -- For RND + if self.alg.rnd: + log_string += ( + f"""{'Mean extrinsic reward:':>{pad}} {statistics.mean(locs['erewbuffer']):.2f}\n""" + f"""{'Mean intrinsic reward:':>{pad}} {statistics.mean(locs['irewbuffer']):.2f}\n""" + ) + + # -- For Offline + if self.alg.offline: + log_string += f"""{'Replay Capacity:':>{pad}} {locs['capacity_percentage']:.4f}\n""" + if self.alg.offline_bc: + log_string += f"""{'Offline BC loss:':>{pad}} {locs['offline_bc_loss']:.4f}\n""" + + log_string += f"""{'Mean total reward:':>{pad}} {statistics.mean(locs['rewbuffer']):.2f}\n""" + log_string += f"""{'Mean episode length:':>{pad}} {statistics.mean(locs['lenbuffer']):.2f}\n""" + # f"""{'Mean reward/step:':>{pad}} {locs['mean_reward']:.2f}\n""" + # f"""{'Mean episode length/episode:':>{pad}} {locs['mean_trajectory_length']:.2f}\n""") + else: + log_string = ( + f"""{'#' * width}\n""" + f"""{str.center(width, ' ')}\n\n""" + f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ + 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" + f"""{'Value function loss:':>{pad}} {locs['mean_value_loss']:.4f}\n""" + f"""{'Surrogate loss:':>{pad}} {locs['mean_surrogate_loss']:.4f}\n""" + ) + # -- For symmetry + if self.alg.symmetry: + log_string += f"""{'Symmetry loss:':>{pad}} {locs['mean_symmetry_loss']:.4f}\n""" + + log_string += f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" + + # f"""{'Mean reward/step:':>{pad}} {locs['mean_reward']:.2f}\n""" + # f"""{'Mean episode length/episode:':>{pad}} {locs['mean_trajectory_length']:.2f}\n""") + + log_string += ep_string + log_string += ( + f"""{'-' * width}\n""" + f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n""" + f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n""" + f"""{'Total time:':>{pad}} {self.tot_time:.2f}s\n""" + f"""{'ETA:':>{pad}} {self.tot_time / (locs['it'] + 1) * ( + locs['num_learning_iterations'] - locs['it']):.1f}s\n""" + ) + print(log_string) + + def save(self, path: str, infos=None): + # -- Save PPO model + saved_dict = { + "model_state_dict": self.alg.actor_critic.state_dict(), + "optimizer_state_dict": self.alg.optimizer.state_dict(), + "iter": self.current_learning_iteration, + "infos": infos, + } + # -- Save RND model if used + if self.alg.rnd: + saved_dict["rnd_state_dict"] = self.alg.rnd.state_dict() + saved_dict["rnd_optimizer_state_dict"] = self.alg.rnd_optimizer.state_dict() + # -- Save observation normalizer if used + if self.empirical_normalization: + saved_dict["obs_norm_state_dict"] = self.obs_normalizer.state_dict() + saved_dict["critic_obs_norm_state_dict"] = self.critic_obs_normalizer.state_dict() + torch.save(saved_dict, path) + + # Upload model to external logging service + if self.logger_type in ["neptune", "wandb"]: + self.writer.save_model(path, self.current_learning_iteration) + + def load(self, path: str, load_optimizer: bool = True): + loaded_dict = torch.load(path, weights_only=False) + # -- Load PPO model + self.alg.actor_critic.load_state_dict(loaded_dict["model_state_dict"]) + # -- Load RND model if used + if self.alg.rnd: + self.alg.rnd.load_state_dict(loaded_dict["rnd_state_dict"]) + # -- Load observation normalizer if used + if self.empirical_normalization: + self.obs_normalizer.load_state_dict(loaded_dict["obs_norm_state_dict"]) + self.critic_obs_normalizer.load_state_dict(loaded_dict["critic_obs_norm_state_dict"]) + # -- Load optimizer if used + if load_optimizer: + # -- PPO + self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"]) + # -- RND optimizer if used + if self.alg.rnd: + self.alg.rnd_optimizer.load_state_dict(loaded_dict["rnd_optimizer_state_dict"]) + # -- Load current learning iteration + self.current_learning_iteration = loaded_dict["iter"] + return loaded_dict["infos"] + + def get_inference_policy(self, device=None): + self.eval_mode() # switch to evaluation mode (dropout for example) + if device is not None: + self.alg.actor_critic.to(device) + policy = self.alg.actor_critic.act_inference + if self.cfg["empirical_normalization"]: + if device is not None: + self.obs_normalizer.to(device) + policy = lambda x: self.alg.actor_critic.act_inference(self.obs_normalizer(x)) # noqa: E731 + return policy + + def train_mode(self): + # -- PPO + self.alg.actor_critic.train() + # -- RND + if self.alg.rnd: + self.alg.rnd.train() + # -- Normalization + if self.empirical_normalization: + self.obs_normalizer.train() + self.critic_obs_normalizer.train() + + def eval_mode(self): + # -- PPO + self.alg.actor_critic.eval() + # -- RND + if self.alg.rnd: + self.alg.rnd.eval() + # -- Normalization + if self.empirical_normalization: + self.obs_normalizer.eval() + self.critic_obs_normalizer.eval() + + def add_git_repo_to_log(self, repo_file_path): + self.git_status_repos.append(repo_file_path) diff --git a/source/uwlab_rl/uwlab_rl/rsl_rl/ext/runners/original_on_policy_runner.py b/source/uwlab_rl/uwlab_rl/rsl_rl/ext/runners/original_on_policy_runner.py new file mode 100644 index 00000000..9092662b --- /dev/null +++ b/source/uwlab_rl/uwlab_rl/rsl_rl/ext/runners/original_on_policy_runner.py @@ -0,0 +1,442 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +import statistics +import time +import torch +from collections import deque + +import rsl_rl +from rsl_rl.algorithms import PPO +from rsl_rl.env import VecEnv +from rsl_rl.modules import ActorCritic, ActorCriticRecurrent, EmpiricalNormalization +from rsl_rl.utils import store_code_state + + +class OnPolicyRunner: + """On-policy runner for training and evaluation.""" + + def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): + self.cfg = train_cfg + self.alg_cfg = train_cfg["algorithm"] + self.policy_cfg = train_cfg["policy"] + self.device = device + self.env = env + + # resolve dimensions of observations + obs, extras = self.env.get_observations() + num_obs = obs.shape[1] + if "critic" in extras["observations"]: + num_critic_obs = extras["observations"]["critic"].shape[1] + else: + num_critic_obs = num_obs + actor_critic_class = eval(self.policy_cfg.pop("class_name")) # ActorCritic + actor_critic: ActorCritic | ActorCriticRecurrent = actor_critic_class( + num_obs, num_critic_obs, self.env.num_actions, **self.policy_cfg + ).to(self.device) + + # resolve dimension of rnd gated state + if "rnd_cfg" in self.alg_cfg: + # check if rnd gated state is present + rnd_state = extras["observations"].get("rnd_state") + if rnd_state is None: + raise ValueError("Observations for they key 'rnd_state' not found in infos['observations'].") + # get dimension of rnd gated state + num_rnd_state = rnd_state.shape[1] + # add rnd gated state to config + self.alg_cfg["rnd_cfg"]["num_state"] = num_rnd_state + # scale down the rnd weight with timestep (similar to how rewards are scaled down in legged_gym envs) + self.alg_cfg["rnd_cfg"]["weight"] *= env.dt + + # if using symmetry then pass the environment config object + if "symmetry_cfg" in self.alg_cfg: + # this is used by the symmetry function for handling different observation terms + self.alg_cfg["symmetry_cfg"]["_env"] = env + + # init algorithm + alg_class = eval(self.alg_cfg.pop("class_name")) # PPO + self.alg: PPO = alg_class(actor_critic, device=self.device, **self.alg_cfg) + + # store training configuration + self.num_steps_per_env = self.cfg["num_steps_per_env"] + self.save_interval = self.cfg["save_interval"] + self.empirical_normalization = self.cfg["empirical_normalization"] + if self.empirical_normalization: + self.obs_normalizer = EmpiricalNormalization(shape=[num_obs], until=1.0e8).to(self.device) + self.critic_obs_normalizer = EmpiricalNormalization(shape=[num_critic_obs], until=1.0e8).to(self.device) + else: + self.obs_normalizer = torch.nn.Identity().to(self.device) # no normalization + self.critic_obs_normalizer = torch.nn.Identity().to(self.device) # no normalization + # init storage and model + self.alg.init_storage( + self.env.num_envs, + self.num_steps_per_env, + [num_obs], + [num_critic_obs], + [self.env.num_actions], + ) + + # Log + self.log_dir = log_dir + self.writer = None + self.tot_timesteps = 0 + self.tot_time = 0 + self.current_learning_iteration = 0 + self.git_status_repos = [rsl_rl.__file__] + + def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): + # initialize writer + if self.log_dir is not None and self.writer is None: + # Launch either Tensorboard or Neptune & Tensorboard summary writer(s), default: Tensorboard. + self.logger_type = self.cfg.get("logger", "tensorboard") + self.logger_type = self.logger_type.lower() + + if self.logger_type == "neptune": + from rsl_rl.utils.neptune_utils import NeptuneSummaryWriter + + self.writer = NeptuneSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) + elif self.logger_type == "wandb": + from rsl_rl.utils.wandb_utils import WandbSummaryWriter + + self.writer = WandbSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) + elif self.logger_type == "tensorboard": + from torch.utils.tensorboard import SummaryWriter + + self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10) + else: + raise ValueError("Logger type not found. Please choose 'neptune', 'wandb' or 'tensorboard'.") + + # randomize initial episode lengths (for exploration) + if init_at_random_ep_len: + self.env.episode_length_buf = torch.randint_like( + self.env.episode_length_buf, high=int(self.env.max_episode_length) + ) + + # start learning + obs, extras = self.env.get_observations() + critic_obs = extras["observations"].get("critic", obs) + obs, critic_obs = obs.to(self.device), critic_obs.to(self.device) + self.train_mode() # switch to train mode (for dropout for example) + + # Book keeping + ep_infos = [] + rewbuffer = deque(maxlen=100) + lenbuffer = deque(maxlen=100) + cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + # create buffers for logging extrinsic and intrinsic rewards + if self.alg.rnd: + erewbuffer = deque(maxlen=100) + irewbuffer = deque(maxlen=100) + cur_ereward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_ireward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + + start_iter = self.current_learning_iteration + tot_iter = start_iter + num_learning_iterations + for it in range(start_iter, tot_iter): + start = time.time() + # Rollout + with torch.inference_mode(): + for _ in range(self.num_steps_per_env): + # Sample actions from policy + actions = self.alg.act(obs, critic_obs) + # Step environment + obs, rewards, dones, infos = self.env.step(actions.to(self.env.device)) + + # Move to the agent device + obs, rewards, dones = obs.to(self.device), rewards.to(self.device), dones.to(self.device) + + # Normalize observations + obs = self.obs_normalizer(obs) + # Extract critic observations and normalize + if "critic" in infos["observations"]: + critic_obs = self.critic_obs_normalizer(infos["observations"]["critic"].to(self.device)) + else: + critic_obs = obs + + # Intrinsic rewards (extracted here only for logging)! + intrinsic_rewards = self.alg.intrinsic_rewards if self.alg.rnd else None + + # Process env step and store in buffer + self.alg.process_env_step(rewards, dones, infos) + + if self.log_dir is not None: + # Book keeping + if "episode" in infos: + ep_infos.append(infos["episode"]) + elif "log" in infos: + ep_infos.append(infos["log"]) + # Update rewards + if self.alg.rnd: + cur_ereward_sum += rewards + cur_ireward_sum += intrinsic_rewards # type: ignore + cur_reward_sum += rewards + intrinsic_rewards + else: + cur_reward_sum += rewards + # Update episode length + cur_episode_length += 1 + # Clear data for completed episodes + # -- common + new_ids = (dones > 0).nonzero(as_tuple=False) + rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) + lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) + cur_reward_sum[new_ids] = 0 + cur_episode_length[new_ids] = 0 + # -- intrinsic and extrinsic rewards + if self.alg.rnd: + erewbuffer.extend(cur_ereward_sum[new_ids][:, 0].cpu().numpy().tolist()) + irewbuffer.extend(cur_ireward_sum[new_ids][:, 0].cpu().numpy().tolist()) + cur_ereward_sum[new_ids] = 0 + cur_ireward_sum[new_ids] = 0 + + stop = time.time() + collection_time = stop - start + + # Learning step + start = stop + self.alg.compute_returns(critic_obs) + + # Update policy + # Note: we keep arguments here since locals() loads them + mean_value_loss, mean_surrogate_loss, mean_entropy, mean_rnd_loss, mean_symmetry_loss = self.alg.update() + stop = time.time() + learn_time = stop - start + self.current_learning_iteration = it + + # Logging info and save checkpoint + if self.log_dir is not None: + # Log information + self.log(locals()) + # Save model + if it % self.save_interval == 0: + self.save(os.path.join(self.log_dir, f"model_{it}.pt")) + + # Clear episode infos + ep_infos.clear() + + # Save code state + if it == start_iter: + # obtain all the diff files + git_file_paths = store_code_state(self.log_dir, self.git_status_repos) + # if possible store them to wandb + if self.logger_type in ["wandb", "neptune"] and git_file_paths: + for path in git_file_paths: + self.writer.save_file(path) + + # Save the final model after training + if self.log_dir is not None: + self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) + + def log(self, locs: dict, width: int = 80, pad: int = 35): + self.tot_timesteps += self.num_steps_per_env * self.env.num_envs + self.tot_time += locs["collection_time"] + locs["learn_time"] + iteration_time = locs["collection_time"] + locs["learn_time"] + + # -- Episode info + ep_string = "" + if locs["ep_infos"]: + for key in locs["ep_infos"][0]: + infotensor = torch.tensor([], device=self.device) + for ep_info in locs["ep_infos"]: + # handle scalar and zero dimensional tensor infos + if key not in ep_info: + continue + if not isinstance(ep_info[key], torch.Tensor): + ep_info[key] = torch.Tensor([ep_info[key]]) + if len(ep_info[key].shape) == 0: + ep_info[key] = ep_info[key].unsqueeze(0) + infotensor = torch.cat((infotensor, ep_info[key].to(self.device))) + value = torch.mean(infotensor) + # log to logger and terminal + if "/" in key: + self.writer.add_scalar(key, value, locs["it"]) + ep_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" + else: + self.writer.add_scalar("Episode/" + key, value, locs["it"]) + ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n""" + mean_std = self.alg.actor_critic.std.mean() + fps = int(self.num_steps_per_env * self.env.num_envs / (locs["collection_time"] + locs["learn_time"])) + + # -- Losses + self.writer.add_scalar("Loss/value_function", locs["mean_value_loss"], locs["it"]) + self.writer.add_scalar("Loss/surrogate", locs["mean_surrogate_loss"], locs["it"]) + self.writer.add_scalar("Loss/entropy", locs["mean_entropy"], locs["it"]) + self.writer.add_scalar("Loss/learning_rate", self.alg.learning_rate, locs["it"]) + if self.alg.rnd: + self.writer.add_scalar("Loss/rnd", locs["mean_rnd_loss"], locs["it"]) + if self.alg.symmetry: + self.writer.add_scalar("Loss/symmetry", locs["mean_symmetry_loss"], locs["it"]) + + # -- Policy + self.writer.add_scalar("Policy/mean_noise_std", mean_std.item(), locs["it"]) + + # -- Performance + self.writer.add_scalar("Perf/total_fps", fps, locs["it"]) + self.writer.add_scalar("Perf/collection time", locs["collection_time"], locs["it"]) + self.writer.add_scalar("Perf/learning_time", locs["learn_time"], locs["it"]) + + # -- Training + if len(locs["rewbuffer"]) > 0: + # separate logging for intrinsic and extrinsic rewards + if self.alg.rnd: + self.writer.add_scalar("Rnd/mean_extrinsic_reward", statistics.mean(locs["erewbuffer"]), locs["it"]) + self.writer.add_scalar("Rnd/mean_intrinsic_reward", statistics.mean(locs["irewbuffer"]), locs["it"]) + self.writer.add_scalar("Rnd/weight", self.alg.rnd.weight, locs["it"]) + # everything else + self.writer.add_scalar("Train/mean_reward", statistics.mean(locs["rewbuffer"]), locs["it"]) + self.writer.add_scalar("Train/mean_episode_length", statistics.mean(locs["lenbuffer"]), locs["it"]) + if self.logger_type != "wandb": # wandb does not support non-integer x-axis logging + self.writer.add_scalar("Train/mean_reward/time", statistics.mean(locs["rewbuffer"]), self.tot_time) + self.writer.add_scalar( + "Train/mean_episode_length/time", statistics.mean(locs["lenbuffer"]), self.tot_time + ) + + str = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m " + + if len(locs["rewbuffer"]) > 0: + log_string = ( + f"""{'#' * width}\n""" + f"""{str.center(width, ' ')}\n\n""" + f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ + 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" + f"""{'Value function loss:':>{pad}} {locs['mean_value_loss']:.4f}\n""" + f"""{'Surrogate loss:':>{pad}} {locs['mean_surrogate_loss']:.4f}\n""" + ) + + # -- For symmetry + if self.alg.symmetry: + log_string += f"""{'Symmetry loss:':>{pad}} {locs['mean_symmetry_loss']:.4f}\n""" + + log_string += f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" + + # -- For RND + if self.alg.rnd: + log_string += ( + f"""{'Mean extrinsic reward:':>{pad}} {statistics.mean(locs['erewbuffer']):.2f}\n""" + f"""{'Mean intrinsic reward:':>{pad}} {statistics.mean(locs['irewbuffer']):.2f}\n""" + ) + + log_string += f"""{'Mean total reward:':>{pad}} {statistics.mean(locs['rewbuffer']):.2f}\n""" + log_string += f"""{'Mean episode length:':>{pad}} {statistics.mean(locs['lenbuffer']):.2f}\n""" + # f"""{'Mean reward/step:':>{pad}} {locs['mean_reward']:.2f}\n""" + # f"""{'Mean episode length/episode:':>{pad}} {locs['mean_trajectory_length']:.2f}\n""") + else: + log_string = ( + f"""{'#' * width}\n""" + f"""{str.center(width, ' ')}\n\n""" + f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ + 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" + f"""{'Value function loss:':>{pad}} {locs['mean_value_loss']:.4f}\n""" + f"""{'Surrogate loss:':>{pad}} {locs['mean_surrogate_loss']:.4f}\n""" + ) + # -- For symmetry + if self.alg.symmetry: + log_string += f"""{'Symmetry loss:':>{pad}} {locs['mean_symmetry_loss']:.4f}\n""" + + log_string += f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" + + # f"""{'Mean reward/step:':>{pad}} {locs['mean_reward']:.2f}\n""" + # f"""{'Mean episode length/episode:':>{pad}} {locs['mean_trajectory_length']:.2f}\n""") + + log_string += ep_string + log_string += ( + f"""{'-' * width}\n""" + f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n""" + f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n""" + f"""{'Total time:':>{pad}} {self.tot_time:.2f}s\n""" + f"""{'ETA:':>{pad}} {self.tot_time / (locs['it'] + 1) * ( + locs['num_learning_iterations'] - locs['it']):.1f}s\n""" + ) + print(log_string) + + def save(self, path: str, infos=None): + # -- Save PPO model + saved_dict = { + "model_state_dict": self.alg.actor_critic.state_dict(), + "optimizer_state_dict": self.alg.optimizer.state_dict(), + "iter": self.current_learning_iteration, + "infos": infos, + } + # -- Save RND model if used + if self.alg.rnd: + saved_dict["rnd_state_dict"] = self.alg.rnd.state_dict() + saved_dict["rnd_optimizer_state_dict"] = self.alg.rnd_optimizer.state_dict() + # -- Save observation normalizer if used + if self.empirical_normalization: + saved_dict["obs_norm_state_dict"] = self.obs_normalizer.state_dict() + saved_dict["critic_obs_norm_state_dict"] = self.critic_obs_normalizer.state_dict() + torch.save(saved_dict, path) + + # Upload model to external logging service + if self.logger_type in ["neptune", "wandb"]: + self.writer.save_model(path, self.current_learning_iteration) + + def load(self, path: str, load_optimizer: bool = True): + loaded_dict = torch.load(path, weights_only=False) + # -- Load PPO model + self.alg.actor_critic.load_state_dict(loaded_dict["model_state_dict"]) + # -- Load RND model if used + if self.alg.rnd: + self.alg.rnd.load_state_dict(loaded_dict["rnd_state_dict"]) + # -- Load observation normalizer if used + if self.empirical_normalization: + self.obs_normalizer.load_state_dict(loaded_dict["obs_norm_state_dict"]) + self.critic_obs_normalizer.load_state_dict(loaded_dict["critic_obs_norm_state_dict"]) + # -- Load optimizer if used + if load_optimizer: + # -- PPO + self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"]) + # -- RND optimizer if used + if self.alg.rnd: + self.alg.rnd_optimizer.load_state_dict(loaded_dict["rnd_optimizer_state_dict"]) + # -- Load current learning iteration + self.current_learning_iteration = loaded_dict["iter"] + return loaded_dict["infos"] + + def get_inference_policy(self, device=None): + self.eval_mode() # switch to evaluation mode (dropout for example) + if device is not None: + self.alg.actor_critic.to(device) + policy = self.alg.actor_critic.act_inference + if self.cfg["empirical_normalization"]: + if device is not None: + self.obs_normalizer.to(device) + policy = lambda x: self.alg.actor_critic.act_inference(self.obs_normalizer(x)) # noqa: E731 + return policy + + def train_mode(self): + # -- PPO + self.alg.actor_critic.train() + # -- RND + if self.alg.rnd: + self.alg.rnd.train() + # -- Normalization + if self.empirical_normalization: + self.obs_normalizer.train() + self.critic_obs_normalizer.train() + + def eval_mode(self): + # -- PPO + self.alg.actor_critic.eval() + # -- RND + if self.alg.rnd: + self.alg.rnd.eval() + # -- Normalization + if self.empirical_normalization: + self.obs_normalizer.eval() + self.critic_obs_normalizer.eval() + + def add_git_repo_to_log(self, repo_file_path): + self.git_status_repos.append(repo_file_path) diff --git a/source/uwlab_rl/uwlab_rl/rsl_rl/ext/storage/replay_storage.py b/source/uwlab_rl/uwlab_rl/rsl_rl/ext/storage/replay_storage.py new file mode 100644 index 00000000..46969591 --- /dev/null +++ b/source/uwlab_rl/uwlab_rl/rsl_rl/ext/storage/replay_storage.py @@ -0,0 +1,357 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch + +from rsl_rl.utils import split_and_pad_trajectories + + +class ReplayStorage: + class Transition: + def __init__(self): + self.observations = None + self.critic_observations = None + self.actions = None + self.rewards = None + self.dones = None + + # For Policy Gradient-like algorithms + self.values = None + self.actions_log_prob = None + self.action_mean = None + self.action_sigma = None + + # For RNN-based policies + self.hidden_states = None + + # For RND + self.rnd_state = None + + # For BC (behavior cloning) + self.expert_action_mean = None + self.expert_action_sigma = None + + def clear(self): + self.__init__() + + def __init__( + self, + num_envs, + capacity: int, + obs_shape=None, + privileged_obs_shape=None, + actions_shape=None, + rnd_state_shape=None, + expert_actions_shape=None, + expert_actions_sigma_shape=None, + device="cpu", + ): + # store inputs + self.device = device + self.capacity = capacity + self.num_envs = num_envs + self.obs_shape = obs_shape + self.privileged_obs_shape = privileged_obs_shape + self.rnd_state_shape = rnd_state_shape + self.actions_shape = actions_shape + self.expert_action_mean_shape = expert_actions_shape + self.expert_action_sigma_shape = expert_actions_sigma_shape + + # Core + if obs_shape is not None: + self.observations = torch.zeros(capacity, num_envs, *obs_shape, device=device) + if privileged_obs_shape is not None: + self.privileged_observations = torch.zeros(capacity, num_envs, *privileged_obs_shape, device=device) + else: + self.privileged_observations = None + + if actions_shape is not None: + self.rewards = torch.zeros(capacity, num_envs, 1, device=self.device) + self.actions = torch.zeros(capacity, num_envs, *actions_shape, device=self.device) + self.dones = torch.zeros(capacity, num_envs, 1, device=self.device).byte() + + # For Policy-Gradient-like algorithms + self.actions_log_prob = torch.zeros(capacity, num_envs, 1, device=self.device) + self.values = torch.zeros(capacity, num_envs, 1, device=self.device) + self.returns = torch.zeros(capacity, num_envs, 1, device=self.device) + self.advantages = torch.zeros(capacity, num_envs, 1, device=self.device) + self.mu = torch.zeros(capacity, num_envs, *actions_shape, device=self.device) + self.sigma = torch.zeros(capacity, num_envs, *actions_shape, device=self.device) + + # For BC + if expert_actions_shape is not None: + self.expert_action_mean = torch.zeros(capacity, num_envs, *expert_actions_shape, device=self.device) + if expert_actions_sigma_shape is not None: + self.expert_action_sigma = torch.zeros(capacity, num_envs, *expert_actions_sigma_shape, device=self.device) + + # For RND + if rnd_state_shape is not None: + self.rnd_state = torch.zeros(capacity, num_envs, *rnd_state_shape, device=self.device) + + # For RNN networks + self.saved_hidden_states_a = None + self.saved_hidden_states_c = None + + # Circular buffer pointers + self.step = 0 + self.size = 0 + self.is_full = False + + def add_transitions(self, transition: Transition): + # check if the transition is valid + if self.size >= self.capacity: + self.is_full = True + # Core + if self.obs_shape is not None: + self.observations[self.step].copy_(transition.observations) + if self.privileged_observations is not None: + self.privileged_observations[self.step].copy_(transition.critic_observations) + if self.actions_shape is not None: + self.actions[self.step].copy_(transition.actions) + self.rewards[self.step].copy_(transition.rewards.view(-1, 1)) + self.dones[self.step].copy_(transition.dones.view(-1, 1)) + + # For Gradient-like algorithms + self.values[self.step].copy_(transition.values) + self.actions_log_prob[self.step].copy_(transition.actions_log_prob.view(-1, 1)) + self.mu[self.step].copy_(transition.action_mean) + self.sigma[self.step].copy_(transition.action_sigma) + + # For BC + if transition.expert_action_mean is not None and self.expert_action_mean is not None: + self.expert_action_mean[self.step].copy_(transition.expert_action_mean) + if transition.expert_action_sigma is not None and self.expert_action_sigma is not None: + self.expert_action_sigma[self.step].copy_(transition.expert_action_sigma) + + # For RND + if self.rnd_state_shape is not None: + self.rnd_state[self.step].copy_(transition.rnd_state) + + # For RNN networks + self._save_hidden_states(transition.hidden_states) + + # Update circular buffer pointers + self.step = (self.step + 1) % self.capacity + self.size = min(self.size + 1, self.capacity) + + def _save_hidden_states(self, hidden_states): + if hidden_states is None or hidden_states == (None, None): + return + # make a tuple out of GRU hidden state sto match the LSTM format + hid_a = hidden_states[0] if isinstance(hidden_states[0], tuple) else (hidden_states[0],) + hid_c = hidden_states[1] if isinstance(hidden_states[1], tuple) else (hidden_states[1],) + + # initialize if needed + if self.saved_hidden_states_a is None: + self.saved_hidden_states_a = [ + torch.zeros(self.observations.shape[0], *hid_a[i].shape, device=self.device) for i in range(len(hid_a)) + ] + self.saved_hidden_states_c = [ + torch.zeros(self.observations.shape[0], *hid_c[i].shape, device=self.device) for i in range(len(hid_c)) + ] + # copy the states + for i in range(len(hid_a)): + self.saved_hidden_states_a[i][self.step].copy_(hid_a[i]) + self.saved_hidden_states_c[i][self.step].copy_(hid_c[i]) + + def clear(self): + self.step = 0 + self.size = 0 + self.is_full = False + + def compute_returns(self, last_values, gamma, lam): + advantage = 0 + for step in reversed(range(self.size)): + # if we are at the last step, bootstrap the return value + if step == self.size - 1: + next_values = last_values + else: + next_values = self.values[step + 1] + # 1 if we are not in a terminal state, 0 otherwise + next_is_not_terminal = 1.0 - self.dones[step].float() + # TD error: r_t + gamma * V(s_{t+1}) - V(s_t) + delta = self.rewards[step] + next_is_not_terminal * gamma * next_values - self.values[step] + # Advantage: A(s_t, a_t) = delta_t + gamma * lambda * A(s_{t+1}, a_{t+1}) + advantage = delta + next_is_not_terminal * gamma * lam * advantage + # Return: R_t = A(s_t, a_t) + V(s_t) + self.returns[step] = advantage + self.values[step] + + # Compute and normalize the advantages + self.advantages = self.returns[: self.size] - self.values[: self.size] + self.advantages[: self.size] = (self.advantages - self.advantages.mean()) / (self.advantages.std() + 1e-8) + + def get_statistics(self): + done = self.dones[: self.size] + done[-1] = 1 + flat_dones = done.permute(1, 0, 2).reshape(-1, 1) + done_indices = torch.cat( + (flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero(as_tuple=False)[:, 0]) + ) + trajectory_lengths = done_indices[1:] - done_indices[:-1] + return trajectory_lengths.float().mean(), self.rewards[: self.size].mean() + + def mini_batch_generator(self, num_mini_batches, num_epochs=8): + batch_size = self.num_envs * self.size + mini_batch_size = batch_size // num_mini_batches + indices = torch.randperm(num_mini_batches * mini_batch_size, requires_grad=False, device=self.device) + + # Core + if self.observations is not None: + observations = self.observations[: self.size].flatten(0, 1) + if self.privileged_observations is not None: + critic_observations = self.privileged_observations[: self.size].flatten(0, 1) + else: + critic_observations = observations + + if self.actions is not None: + actions = self.actions[: self.size].flatten(0, 1) + values = self.values[: self.size].flatten(0, 1) + returns = self.returns[: self.size].flatten(0, 1) + + # For PPO + old_actions_log_prob = self.actions_log_prob[: self.size].flatten(0, 1) + advantages = self.advantages[: self.size].flatten(0, 1) + old_mu = self.mu[: self.size].flatten(0, 1) + old_sigma = self.sigma[: self.size].flatten(0, 1) + # For BC + if self.expert_action_mean_shape is not None: + expert_action_mu = self.expert_action_mean[: self.size].flatten(0, 1) + if self.expert_action_sigma_shape is not None: + expert_action_sigma = self.expert_action_sigma[: self.size].flatten(0, 1) + + # For RND + if self.rnd_state_shape is not None: + rnd_state = self.rnd_state[: self.size].flatten(0, 1) + + for epoch in range(num_epochs): + for i in range(num_mini_batches): + # Select the indices for the mini-batch + start = i * mini_batch_size + end = (i + 1) * mini_batch_size + batch_idx = indices[start:end] + + # Create the mini-batch + # -- Core + obs_batch, critic_observations_batch, actions_batch = None, None, None + if self.observations is not None: + obs_batch = observations[batch_idx] + critic_observations_batch = critic_observations[batch_idx] + if self.actions is not None: + actions_batch = actions[batch_idx] + + # -- For PPO + target_values_batch, returns_batch, advantages_batch = None, None, None + old_actions_log_prob_batch, old_mu_batch, old_sigma_batch = None, None, None + if self.actions is not None: + target_values_batch = values[batch_idx] + returns_batch = returns[batch_idx] + old_actions_log_prob_batch = old_actions_log_prob[batch_idx] + advantages_batch = advantages[batch_idx] + old_mu_batch = old_mu[batch_idx] + old_sigma_batch = old_sigma[batch_idx] + # -- For BC + expert_action_mu_batch, expert_action_sigma_batch = None, None + if self.expert_action_mean_shape is not None: + expert_action_mu_batch = expert_action_mu[batch_idx] + if self.expert_action_sigma_shape is not None: + expert_action_sigma_batch = expert_action_sigma[batch_idx] + + # -- For RND + rnd_state_batch = None + if self.rnd_state_shape is not None: + rnd_state_batch = rnd_state[batch_idx] + + # Yield the mini-batch + yield obs_batch, critic_observations_batch, actions_batch, target_values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( + None, + None, + ), None, expert_action_mu_batch, expert_action_sigma_batch, rnd_state_batch + + # for RNNs only + def recurrent_mini_batch_generator(self, num_mini_batches, num_epochs=8): + padded_obs_trajectories, trajectory_masks = split_and_pad_trajectories(self.observations, self.dones) + if self.privileged_observations is not None: + padded_critic_obs_trajectories, _ = split_and_pad_trajectories(self.privileged_observations, self.dones) + else: + padded_critic_obs_trajectories = padded_obs_trajectories + + if self.rnd_state_shape is not None: + padded_rnd_state_trajectories, _ = split_and_pad_trajectories(self.rnd_state, self.dones) + else: + padded_rnd_state_trajectories = None + + mini_batch_size = self.num_envs // num_mini_batches + for ep in range(num_epochs): + first_traj = 0 + for i in range(num_mini_batches): + start = i * mini_batch_size + stop = (i + 1) * mini_batch_size + + dones = self.dones.squeeze(-1) + last_was_done = torch.zeros_like(dones, dtype=torch.bool) + last_was_done[1:] = dones[:-1] + last_was_done[0] = True + trajectories_batch_size = torch.sum(last_was_done[:, start:stop]) + last_traj = first_traj + trajectories_batch_size + + masks_batch = trajectory_masks[:, first_traj:last_traj] + obs_batch = padded_obs_trajectories[:, first_traj:last_traj] + critic_obs_batch = padded_critic_obs_trajectories[:, first_traj:last_traj] + + # For BC + if self.expert_action_mean_shape is not None: + expert_action_mu_batch = self.expert_action_mean[:, start:stop] + else: + expert_action_mu_batch = None + if self.expert_action_sigma_shape is not None: + expert_action_sigma_batch = self.expert_action_sigma[:, start:stop] + else: + expert_action_sigma_batch = None + + if padded_rnd_state_trajectories is not None: + rnd_state_batch = padded_rnd_state_trajectories[:, first_traj:last_traj] + else: + rnd_state_batch = None + actions_batch = self.actions[:, start:stop] + old_mu_batch = self.mu[:, start:stop] + old_sigma_batch = self.sigma[:, start:stop] + returns_batch = self.returns[:, start:stop] + advantages_batch = self.advantages[:, start:stop] + values_batch = self.values[:, start:stop] + old_actions_log_prob_batch = self.actions_log_prob[:, start:stop] + + # reshape to [num_envs, time, num layers, hidden dim] (original shape: [time, num_layers, num_envs, hidden_dim]) + # then take only time steps after dones (flattens num envs and time dimensions), + # take a batch of trajectories and finally reshape back to [num_layers, batch, hidden_dim] + last_was_done = last_was_done.permute(1, 0) + hid_a_batch = [ + saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] + .transpose(1, 0) + .contiguous() + for saved_hidden_states in self.saved_hidden_states_a + ] + hid_c_batch = [ + saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] + .transpose(1, 0) + .contiguous() + for saved_hidden_states in self.saved_hidden_states_c + ] + # remove the tuple for GRU + hid_a_batch = hid_a_batch[0] if len(hid_a_batch) == 1 else hid_a_batch + hid_c_batch = hid_c_batch[0] if len(hid_c_batch) == 1 else hid_c_batch + + yield obs_batch, critic_obs_batch, actions_batch, values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( + hid_a_batch, + hid_c_batch, + ), masks_batch, expert_action_mu_batch, expert_action_sigma_batch, rnd_state_batch + + first_traj = last_traj diff --git a/source/uwlab_rl/uwlab_rl/rsl_rl/ext/storage/rollout_storage.py b/source/uwlab_rl/uwlab_rl/rsl_rl/ext/storage/rollout_storage.py new file mode 100644 index 00000000..eca28463 --- /dev/null +++ b/source/uwlab_rl/uwlab_rl/rsl_rl/ext/storage/rollout_storage.py @@ -0,0 +1,344 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch + +from rsl_rl.utils import split_and_pad_trajectories + + +class RolloutStorage: + class Transition: + def __init__(self): + self.observations = None + self.critic_observations = None + self.actions = None + self.rewards = None + self.dones = None + self.values = None + self.actions_log_prob = None + self.action_mean = None + self.action_sigma = None + self.hidden_states = None + self.rnd_state = None + self.expert_action_mean = None + self.expert_action_sigma = None + + def clear(self): + self.__init__() + + def __init__( + self, + num_envs, + num_transitions_per_env, + obs_shape, + privileged_obs_shape, + actions_shape, + rnd_state_shape=None, + expert_action_mean_shape=None, + expert_action_sigma_shape=None, + device="cpu", + ): + # store inputs + self.device = device + self.num_transitions_per_env = num_transitions_per_env + self.num_envs = num_envs + self.obs_shape = obs_shape + self.privileged_obs_shape = privileged_obs_shape + self.rnd_state_shape = rnd_state_shape + self.actions_shape = actions_shape + + # For bc + self.expert_action_mean_shape = expert_action_mean_shape + self.expert_action_sigma_shape = expert_action_sigma_shape + + # Core + self.observations = torch.zeros(num_transitions_per_env, num_envs, *obs_shape, device=self.device) + if privileged_obs_shape is not None: + self.privileged_observations = torch.zeros( + num_transitions_per_env, num_envs, *privileged_obs_shape, device=self.device + ) + else: + self.privileged_observations = None + self.rewards = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) + self.dones = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device).byte() + + # For PPO + self.actions_log_prob = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.values = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.returns = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.advantages = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.mu = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) + self.sigma = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) + + # For BC + if expert_action_mean_shape is not None: + self.expert_action_mean = torch.zeros( + num_transitions_per_env, num_envs, *expert_action_mean_shape, device=self.device + ) + if expert_action_sigma_shape is not None: + self.expert_action_sigma = torch.zeros( + num_transitions_per_env, num_envs, *expert_action_sigma_shape, device=self.device + ) + + # For RND + if rnd_state_shape is not None: + self.rnd_state = torch.zeros(num_transitions_per_env, num_envs, *rnd_state_shape, device=self.device) + + # For RNN networks + self.saved_hidden_states_a = None + self.saved_hidden_states_c = None + # counter for the number of transitions stored + self.step = 0 + + def add_transitions(self, transition: Transition): + # check if the transition is valid + if self.step >= self.num_transitions_per_env: + raise OverflowError("Rollout buffer overflow! You should call clear() before adding new transitions.") + + # Core + self.observations[self.step].copy_(transition.observations) + if self.privileged_observations is not None: + self.privileged_observations[self.step].copy_(transition.critic_observations) + self.actions[self.step].copy_(transition.actions) + self.rewards[self.step].copy_(transition.rewards.view(-1, 1)) + self.dones[self.step].copy_(transition.dones.view(-1, 1)) + + # For PPO + self.values[self.step].copy_(transition.values) + self.actions_log_prob[self.step].copy_(transition.actions_log_prob.view(-1, 1)) + self.mu[self.step].copy_(transition.action_mean) + self.sigma[self.step].copy_(transition.action_sigma) + + # For BC + if transition.expert_action_mean is not None: + self.expert_action_mean[self.step].copy_(transition.expert_action_mean) + if transition.expert_action_sigma is not None: + self.expert_action_sigma[self.step].copy_(transition.expert_action_sigma) + + # For RND + if self.rnd_state_shape is not None: + self.rnd_state[self.step].copy_(transition.rnd_state) + + # For RNN networks + self._save_hidden_states(transition.hidden_states) + + # increment the counter + self.step += 1 + + def _save_hidden_states(self, hidden_states): + if hidden_states is None or hidden_states == (None, None): + return + # make a tuple out of GRU hidden state sto match the LSTM format + hid_a = hidden_states[0] if isinstance(hidden_states[0], tuple) else (hidden_states[0],) + hid_c = hidden_states[1] if isinstance(hidden_states[1], tuple) else (hidden_states[1],) + + # initialize if needed + if self.saved_hidden_states_a is None: + self.saved_hidden_states_a = [ + torch.zeros(self.observations.shape[0], *hid_a[i].shape, device=self.device) for i in range(len(hid_a)) + ] + self.saved_hidden_states_c = [ + torch.zeros(self.observations.shape[0], *hid_c[i].shape, device=self.device) for i in range(len(hid_c)) + ] + # copy the states + for i in range(len(hid_a)): + self.saved_hidden_states_a[i][self.step].copy_(hid_a[i]) + self.saved_hidden_states_c[i][self.step].copy_(hid_c[i]) + + def clear(self): + self.step = 0 + + def compute_returns(self, last_values, gamma, lam): + advantage = 0 + for step in reversed(range(self.num_transitions_per_env)): + # if we are at the last step, bootstrap the return value + if step == self.num_transitions_per_env - 1: + next_values = last_values + else: + next_values = self.values[step + 1] + # 1 if we are not in a terminal state, 0 otherwise + next_is_not_terminal = 1.0 - self.dones[step].float() + # TD error: r_t + gamma * V(s_{t+1}) - V(s_t) + delta = self.rewards[step] + next_is_not_terminal * gamma * next_values - self.values[step] + # Advantage: A(s_t, a_t) = delta_t + gamma * lambda * A(s_{t+1}, a_{t+1}) + advantage = delta + next_is_not_terminal * gamma * lam * advantage + # Return: R_t = A(s_t, a_t) + V(s_t) + self.returns[step] = advantage + self.values[step] + + # Compute and normalize the advantages + self.advantages = self.returns - self.values + self.advantages = (self.advantages - self.advantages.mean()) / (self.advantages.std() + 1e-8) + + def get_statistics(self): + done = self.dones + done[-1] = 1 + flat_dones = done.permute(1, 0, 2).reshape(-1, 1) + done_indices = torch.cat( + (flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero(as_tuple=False)[:, 0]) + ) + trajectory_lengths = done_indices[1:] - done_indices[:-1] + return trajectory_lengths.float().mean(), self.rewards.mean() + + def mini_batch_generator(self, num_mini_batches, num_epochs=8): + batch_size = self.num_envs * self.num_transitions_per_env + mini_batch_size = batch_size // num_mini_batches + indices = torch.randperm(num_mini_batches * mini_batch_size, requires_grad=False, device=self.device) + + # Core + observations = self.observations.flatten(0, 1) + if self.privileged_observations is not None: + critic_observations = self.privileged_observations.flatten(0, 1) + else: + critic_observations = observations + + actions = self.actions.flatten(0, 1) + values = self.values.flatten(0, 1) + returns = self.returns.flatten(0, 1) + + # For PPO + old_actions_log_prob = self.actions_log_prob.flatten(0, 1) + advantages = self.advantages.flatten(0, 1) + old_mu = self.mu.flatten(0, 1) + old_sigma = self.sigma.flatten(0, 1) + # For BC + if self.expert_action_mean_shape is not None: + expert_action_mu = self.expert_action_mean.flatten(0, 1) + if self.expert_action_sigma_shape is not None: + expert_action_sigma = self.expert_action_sigma.flatten(0, 1) + + # For RND + if self.rnd_state_shape is not None: + rnd_state = self.rnd_state.flatten(0, 1) + + for epoch in range(num_epochs): + for i in range(num_mini_batches): + # Select the indices for the mini-batch + start = i * mini_batch_size + end = (i + 1) * mini_batch_size + batch_idx = indices[start:end] + + # Create the mini-batch + # -- Core + obs_batch = observations[batch_idx] + critic_observations_batch = critic_observations[batch_idx] + actions_batch = actions[batch_idx] + + # -- For PPO + target_values_batch = values[batch_idx] + returns_batch = returns[batch_idx] + old_actions_log_prob_batch = old_actions_log_prob[batch_idx] + advantages_batch = advantages[batch_idx] + old_mu_batch = old_mu[batch_idx] + old_sigma_batch = old_sigma[batch_idx] + # -- For BC + if self.expert_action_mean_shape is not None: + expert_action_mu_batch = expert_action_mu[batch_idx] + else: + expert_action_mu_batch = None + if self.expert_action_sigma_shape is not None: + expert_action_sigma_batch = expert_action_sigma[batch_idx] + else: + expert_action_sigma_batch = None + + # -- For RND + if self.rnd_state_shape is not None: + rnd_state_batch = rnd_state[batch_idx] + else: + rnd_state_batch = None + + # Yield the mini-batch + yield obs_batch, critic_observations_batch, actions_batch, target_values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( + None, + None, + ), None, expert_action_mu_batch, expert_action_sigma_batch, rnd_state_batch + + # for RNNs only + def recurrent_mini_batch_generator(self, num_mini_batches, num_epochs=8): + padded_obs_trajectories, trajectory_masks = split_and_pad_trajectories(self.observations, self.dones) + if self.privileged_observations is not None: + padded_critic_obs_trajectories, _ = split_and_pad_trajectories(self.privileged_observations, self.dones) + else: + padded_critic_obs_trajectories = padded_obs_trajectories + + if self.rnd_state_shape is not None: + padded_rnd_state_trajectories, _ = split_and_pad_trajectories(self.rnd_state, self.dones) + else: + padded_rnd_state_trajectories = None + + mini_batch_size = self.num_envs // num_mini_batches + for ep in range(num_epochs): + first_traj = 0 + for i in range(num_mini_batches): + start = i * mini_batch_size + stop = (i + 1) * mini_batch_size + + dones = self.dones.squeeze(-1) + last_was_done = torch.zeros_like(dones, dtype=torch.bool) + last_was_done[1:] = dones[:-1] + last_was_done[0] = True + trajectories_batch_size = torch.sum(last_was_done[:, start:stop]) + last_traj = first_traj + trajectories_batch_size + + masks_batch = trajectory_masks[:, first_traj:last_traj] + obs_batch = padded_obs_trajectories[:, first_traj:last_traj] + critic_obs_batch = padded_critic_obs_trajectories[:, first_traj:last_traj] + + # For BC + if self.expert_action_mean_shape is not None: + expert_action_mu_batch = self.expert_action_mean[:, start:stop] + else: + expert_action_mu_batch = None + if self.expert_action_sigma_shape is not None: + expert_action_sigma_batch = self.expert_action_sigma[:, start:stop] + else: + expert_action_sigma_batch = None + + if padded_rnd_state_trajectories is not None: + rnd_state_batch = padded_rnd_state_trajectories[:, first_traj:last_traj] + else: + rnd_state_batch = None + + actions_batch = self.actions[:, start:stop] + old_mu_batch = self.mu[:, start:stop] + old_sigma_batch = self.sigma[:, start:stop] + returns_batch = self.returns[:, start:stop] + advantages_batch = self.advantages[:, start:stop] + values_batch = self.values[:, start:stop] + old_actions_log_prob_batch = self.actions_log_prob[:, start:stop] + + # reshape to [num_envs, time, num layers, hidden dim] (original shape: [time, num_layers, num_envs, hidden_dim]) + # then take only time steps after dones (flattens num envs and time dimensions), + # take a batch of trajectories and finally reshape back to [num_layers, batch, hidden_dim] + last_was_done = last_was_done.permute(1, 0) + hid_a_batch = [ + saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] + .transpose(1, 0) + .contiguous() + for saved_hidden_states in self.saved_hidden_states_a + ] + hid_c_batch = [ + saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] + .transpose(1, 0) + .contiguous() + for saved_hidden_states in self.saved_hidden_states_c + ] + # remove the tuple for GRU + hid_a_batch = hid_a_batch[0] if len(hid_a_batch) == 1 else hid_a_batch + hid_c_batch = hid_c_batch[0] if len(hid_c_batch) == 1 else hid_c_batch + + yield obs_batch, critic_obs_batch, actions_batch, values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( + hid_a_batch, + hid_c_batch, + ), masks_batch, expert_action_mu_batch, expert_action_sigma_batch, rnd_state_batch + + first_traj = last_traj diff --git a/source/uwlab_rl/uwlab_rl/rsl_rl/rl_cfg.py b/source/uwlab_rl/uwlab_rl/rsl_rl/rl_cfg.py new file mode 100644 index 00000000..018cd20c --- /dev/null +++ b/source/uwlab_rl/uwlab_rl/rsl_rl/rl_cfg.py @@ -0,0 +1,77 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import MISSING + +from isaaclab.utils import configclass + +from isaaclab_rl.rsl_rl import RslRlPpoAlgorithmCfg # noqa: F401 + + +@configclass +class SymmetryCfg: + use_data_augmentation: bool = False + + use_mirror_loss: bool = False + + data_augmentation_func: callable = None + + +@configclass +class BehaviorCloningCfg: + experts_path: list[str] = MISSING # type: ignore + """The path to the expert data.""" + + experts_loader: callable = "torch.jit.load" + """The function to construct the expert. Default is None, for which is loaded in the same way student is loaded.""" + + experts_env_mapping_func: callable = None + """The function to map the expert to env_ids. Default is None, for which is mapped to all env_ids""" + + experts_observation_group_cfg: str | None = None + """The observation group of the expert which may be different from student""" + + experts_observation_func: callable = None + """The function that returns expert observation data, default is None, same as student observation.""" + + learn_std: bool = False + """Whether to learn the standard deviation of the expert policy.""" + + cloning_loss_coeff: float = MISSING # type: ignore + """The coefficient for the cloning loss.""" + + loss_decay: float = 1.0 + """The decay for the cloning loss coefficient. default to 1, no decay.""" + + +@configclass +class OffPolicyAlgorithmCfg: + """Configuration for the off-policy algorithm.""" + + update_frequencies: float = 1 + """The frequency to update relative to online update.""" + + batch_size: int | None = None + """The batch size for the offline algorithm update, default to None, same of online size.""" + + num_learning_epochs: int | None = None + """The number of learning epochs for the offline algorithm update.""" + + behavior_cloning_cfg: BehaviorCloningCfg | None = None + """The configuration for the offline behavior cloning(dagger).""" + + +@configclass +class RslRlFancyPpoAlgorithmCfg(RslRlPpoAlgorithmCfg): + """Configuration for the PPO algorithm.""" + + symmetry_cfg: SymmetryCfg | None = None + """The configuration for the symmetry.""" + + behavior_cloning_cfg: BehaviorCloningCfg | None = None + """The configuration for the online behavior cloning.""" + + offline_algorithm_cfg: OffPolicyAlgorithmCfg | None = None + """The configuration for the offline algorithms.""" diff --git a/source/uwlab_rl/uwlab_rl/skrl/__init__.py b/source/uwlab_rl/uwlab_rl/skrl/__init__.py new file mode 100644 index 00000000..937c60e1 --- /dev/null +++ b/source/uwlab_rl/uwlab_rl/skrl/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .ppo_cfg import * diff --git a/source/uwlab_rl/uwlab_rl/skrl/extensions/__init__.py b/source/uwlab_rl/uwlab_rl/skrl/extensions/__init__.py new file mode 100644 index 00000000..b2e10416 --- /dev/null +++ b/source/uwlab_rl/uwlab_rl/skrl/extensions/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .ext_cfg import ( + ContextInitializerCfg, + SupplementaryLossesCfg, + SupplementarySampleTermsCfg, + SupplementaryTrainingCfg, +) +from .loss_ext import * +from .patches import patch_agent_with_supplementary_loss +from .sample_ext import * diff --git a/source/uwlab_rl/uwlab_rl/skrl/extensions/ext_cfg.py b/source/uwlab_rl/uwlab_rl/skrl/extensions/ext_cfg.py new file mode 100644 index 00000000..3e9a9398 --- /dev/null +++ b/source/uwlab_rl/uwlab_rl/skrl/extensions/ext_cfg.py @@ -0,0 +1,60 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from dataclasses import MISSING +from typing import Any, Callable + +from isaaclab.utils import configclass + + +@configclass +class ContextInitializerCfg: + """context initializer for the supplementary training.""" + + context_identifier: str = MISSING + + context_initializer: Callable[..., dict[str, Any]] = MISSING + + context_params: dict[str, Any] = MISSING + + +@configclass +class SupplementaryLossesCfg: + """additional loss term for the training.""" + + loss_term: str = MISSING + + loss_f_creator: Callable[..., Callable[[dict[str, torch.Tensor]], float]] = MISSING + + loss_params: dict[str, Any] = MISSING + + +@configclass +class SupplementarySampleTermsCfg: + """additional sample term for the training.""" + + sample_term: str = MISSING + + sample_size_f_creator: Callable[..., Callable[[], int]] = MISSING + + sample_size_params: dict[str, Any] = MISSING + + sample_f_creator: Callable[..., Callable[[dict[str, torch.Tensor]], torch.Tensor]] = MISSING + + sample_params: dict[str, Any] = MISSING + + +@configclass +class SupplementaryTrainingCfg: + context_manager: Callable = MISSING + + context_initializers: list[ContextInitializerCfg] = MISSING + + supplementary_losses: list[SupplementaryLossesCfg] = MISSING + + supplementary_sample_terms: list[SupplementarySampleTermsCfg] = MISSING diff --git a/source/uwlab_rl/uwlab_rl/skrl/extensions/loss_ext.py b/source/uwlab_rl/uwlab_rl/skrl/extensions/loss_ext.py new file mode 100644 index 00000000..f8e14b82 --- /dev/null +++ b/source/uwlab_rl/uwlab_rl/skrl/extensions/loss_ext.py @@ -0,0 +1,60 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING, Any, Callable + +if TYPE_CHECKING: + from skrl.agents.torch.base import Agent + + from uwlab.envs import DataManagerBasedRLEnv +""" +loss extension functions return a loss value given a batch of samples. +they can be used to add additional loss terms to the agent's loss function. +""" + + +def expert_distillation_loss_f( + env: DataManagerBasedRLEnv, + agent: Agent, + context: dict[str, Any], + criterion: torch.nn.Module, + student_action_key: str = "actions", + expert_action_key: str = "expert_actions", +) -> Callable[[dict[str, torch.Tensor]], float]: + """ + Creates a distillation loss function to align student actions with expert actions. + + Parameters: + ---------- + criterion : torch.nn.Module + Loss function (e.g., `torch.nn.MSELoss`) to measure the difference between expert and student actions. + student_action_key : str, optional + Key to access student actions in the batch dictionary, default is "actions". + expert_action_key : str, optional + Key to access expert actions in the batch dictionary, default is "expert_actions". + + Returns: + ------- + Callable[[dict], torch.Tensor] + A function `calc_distillation_loss` that computes the distillation loss from a batch dictionary + containing student and expert actions. + + Usage: + ------ + ``` + distillation_loss_fn = expert_distillation_loss_f(criterion) + loss = distillation_loss_fn(batch) + ``` + """ + + def calc_distillation_loss(batch): + expert_action, stu_actions = batch[expert_action_key], batch[student_action_key] + distillation_loss = criterion()(stu_actions, expert_action) + return distillation_loss.item() + + return calc_distillation_loss diff --git a/source/uwlab_rl/uwlab_rl/skrl/extensions/patches.py b/source/uwlab_rl/uwlab_rl/skrl/extensions/patches.py new file mode 100644 index 00000000..a6e60c92 --- /dev/null +++ b/source/uwlab_rl/uwlab_rl/skrl/extensions/patches.py @@ -0,0 +1,101 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch +import types +from contextlib import contextmanager +from typing import Any, Dict, List + +from skrl.agents.torch.base import Agent + +from uwlab.envs import DataManagerBasedRLEnv + +from .ext_cfg import ContextInitializerCfg, SupplementaryTrainingCfg + + +class AgentPatcher: + def __init__( + self, locs: Dict[str, Any], env: DataManagerBasedRLEnv, agent: Agent, suppl_train_cfg: SupplementaryTrainingCfg + ): + self.env = env + self.agent = agent + self.context = self.context_init(locs, self.env, self.agent, suppl_train_cfg.context_initializers) + self.dict_mem_func, self.dict_sample_func, self.dict_loss_func = self.init_contextual_func(suppl_train_cfg) + + self.original__update = agent._update + self.original_memory_add_sample = agent.memory.add_samples + self.patch_agent() + + def context_init( + self, + locs: Dict[str, Any], + env: DataManagerBasedRLEnv, + agent: Agent, + context_initializers: List[ContextInitializerCfg], + ): + context = {} + for context_initializer in context_initializers: + context_identifier = context_initializer.context_identifier + context_initializer_f = context_initializer.context_initializer + context_params = context_initializer.context_params + context[context_identifier] = context_initializer_f(env, agent, locs, **context_params) + return context + + def init_contextual_func(self, suppl_train_cfg: SupplementaryTrainingCfg): + dict_mem_func = {} + dict_sample_func = {} + dict_loss_func = {} + + for sample_term in suppl_train_cfg.supplementary_sample_terms: + sample_f = sample_term.sample_f_creator(self.env, self.agent, self.context, **sample_term.sample_params) + sample_size_f = sample_term.sample_size_f_creator( + self.env, self.agent, self.context, **sample_term.sample_size_params + ) + dict_sample_func[sample_term.sample_term] = sample_f + dict_mem_func[sample_term.sample_term] = sample_size_f + + for loss_terms in suppl_train_cfg.supplementary_losses: + loss_f = loss_terms.loss_f_creator(self.env, self.agent, self.context, **loss_terms.loss_params) + dict_loss_func[loss_terms.loss_term] = loss_f + + return dict_mem_func, dict_sample_func, dict_loss_func + + def patch_agent(self): + for term, size_f in self.dict_mem_func.items(): + size = size_f() + self.agent.memory.create_tensor(name=term, size=size, dtype=torch.float32) + self.agent._tensors_names += [memory_term for memory_term, _ in self.dict_mem_func.items()] + + def patched_add_samples(self, *args, **kwargs): + for mem_term, sample_f in self.dict_sample_func.items(): + sample = sample_f(kwargs) + kwargs[mem_term] = sample + return self.original_memory_add_sample(**kwargs) + + def patched__update(self, *args, **kwargs): + """ + Patched _update method that injects 'suppl_loss' if not already provided. + """ + return self.original__update(*(args[1:]), suppl_loss=self.dict_loss_func) + + def apply_patch(self): + # Replace the agent's _update method with the patched one + self.agent._update = types.MethodType(self.patched__update, self.agent) + self.agent.memory.add_samples = types.MethodType(self.patched_add_samples, self.agent.memory) + + def remove_patch(self): + # Restore the original _update method + self.agent._update = self.original__update + self.agent.memory.add_samples = self.original_memory_add_sample + + +@contextmanager +def patch_agent_with_supplementary_loss(locs, env, agent, suppl_loss): + patcher = AgentPatcher(locs, env, agent, suppl_loss) + patcher.apply_patch() + try: + yield + finally: + patcher.remove_patch() diff --git a/source/uwlab_rl/uwlab_rl/skrl/extensions/sample_ext.py b/source/uwlab_rl/uwlab_rl/skrl/extensions/sample_ext.py new file mode 100644 index 00000000..1fa14149 --- /dev/null +++ b/source/uwlab_rl/uwlab_rl/skrl/extensions/sample_ext.py @@ -0,0 +1,97 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING, Callable + +if TYPE_CHECKING: + from skrl.agents.torch.base import Agent + + from uwlab.envs import DataManagerBasedRLEnv + +""" +Sample extension returns a function that takes a batch of samples +and return the additional samples wish to be added to the memory. +""" + + +def experts_act_f( + env: DataManagerBasedRLEnv, + agent: Agent, + context: dict, + map_encoding_to_expert_key: str, +): + """ + Creates an action function for agents using terrain-specific expert policies based on encoding-based mapping. + + Parameters: + ---------- + env : DataManagerBasedRLEnv + The IsaacLab Manager-based RL environment + + agent : Agent + The SKRL Agent + + context : dict + A dictionary containing various configuration and context-specific information, including expert policies. + + map_encoding_to_expert_key : str + The key used to access `expert_encoding_policies_dict` from `context`. + + Returns: + ------- + Callable[[dict], torch.Tensor] + An action function `act` that takes a batch of states and returns terrain-based actions specific to each agent. + + Detailed Functionality: + ----------------------- + - Retrieves expert policies using `map_encoding_to_expert_key` from `context`. + - Maps environment terrain encodings to expert policies, aligning each agent with a policy based on their + current terrain. + - The `act` function generates actions by processing a batch of states and returning terrain-specific actions + for each agent according to `agent_terrain_id` from the environment. + + Example Usage: + -------------- + ``` + action_function = experts_act_f(env, agent, context, "terrain_expert_policies") + actions = action_function({'states': torch.rand(env.num_envs, observation_dim)}) + ``` + + Notes: + ------ + - The function utilizes terrain encodings from the environment to match agents with appropriate expert policies + from `expert_encoding_policies_dict`. + - The `act` function recalculates agent-appropriate actions based on the `agent_terrain_id` at each step, + allowing dynamic policy application if agent terrain changes over time. + - The `act` function returns a clone of the generated action tensor to ensure isolation from in-place modifications. + """ + + expert_encoding_policies_dict: dict[tuple, Callable[[torch.Tensor], torch.Tensor]] = context[ + map_encoding_to_expert_key + ] + + # Order of .keys() and .values() is guaranteed to be the same in Python 3.7+ + expert_encodings = list(expert_encoding_policies_dict.keys()) + expert_policies = list(expert_encoding_policies_dict.values()) + expert_encodings = torch.tensor(expert_encodings, device=env.device) + terrain_encodings_tensor = env.extensions["terrain_encoding_cache"] + order = [torch.where((expert_encodings == t).all(dim=1))[0][0] for t in terrain_encodings_tensor] + ordered_expert_policies = [expert_policies[i] for i in order] + + act_dim = env.action_space.shape[0] # type: ignore + actions = torch.zeros((env.num_envs, act_dim)).to(env.device) + + def act(batch): + observations = batch["states"] + # query agent_terrain_id in the loop to consider the case where agent changes terrain + agent_terrain_id = env.scene.terrain.agent_terrain_id # type: ignore + for i in range(len(ordered_expert_policies)): + actions[agent_terrain_id == i] = ordered_expert_policies[i](observations[agent_terrain_id == i]) + return actions.clone() + + return act diff --git a/source/uwlab_rl/uwlab_rl/skrl/ppo_cfg.py b/source/uwlab_rl/uwlab_rl/skrl/ppo_cfg.py new file mode 100644 index 00000000..6054f09a --- /dev/null +++ b/source/uwlab_rl/uwlab_rl/skrl/ppo_cfg.py @@ -0,0 +1,127 @@ +# base_config.py +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import MISSING, field +from typing import Dict, List, Literal, Optional, Union + +from isaaclab.utils import configclass + +from .extensions import SupplementaryTrainingCfg + + +@configclass +class ModelPolicyCfg: + """Base configuration for the policy model.""" + + policy_class: Literal["GaussianMixin", "MultivariateGaussianMixin"] = field(default="GaussianMixin") + clip_actions: bool = MISSING + clip_log_std: bool = MISSING + initial_log_std: float = MISSING + min_log_std: float = MISSING + max_log_std: float = MISSING + input_shape: str = MISSING + hiddens: List[int] = MISSING + hidden_activation: List[str] = MISSING + output_shape: str = MISSING + output_activation: str = MISSING + output_scale: float = MISSING + + +@configclass +class ModelValueCfg: + """Base configuration for the value model.""" + + value_class: str = MISSING + clip_actions: bool = MISSING + input_shape: str = MISSING + hiddens: List[int] = MISSING + hidden_activation: List[str] = MISSING + output_shape: str = MISSING + output_activation: str = MISSING + output_scale: float = MISSING + + +@configclass +class ModelsCfg: + """Base configuration for the model instantiators.""" + + separate: bool = MISSING + policy: ModelPolicyCfg = MISSING + value: ModelValueCfg = MISSING + + +@configclass +class ExperimentCfg: + """Base configuration for experiment logging and checkpoints.""" + + directory: str = MISSING + experiment_name: str = MISSING + write_interval: int = MISSING + checkpoint_interval: int = MISSING + wandb: bool = MISSING + wandb_kwargs: Dict[str, str] = field(default_factory=dict) + + +@configclass +class PPOAgentCfg: + """Base configuration for the PPO agent.""" + + agent_class: str = MISSING + rollouts: int = MISSING + learning_epochs: int = MISSING + mini_batches: int = MISSING + discount_factor: float = MISSING + lambda_: float = MISSING + learning_rate: float = MISSING + learning_rate_scheduler: str = MISSING + learning_rate_scheduler_kwargs: Dict[str, Union[float, str]] = field(default_factory=dict) + state_preprocessor: str = MISSING + state_preprocessor_kwargs: Optional[Dict] = None + value_preprocessor: str = MISSING + value_preprocessor_kwargs: Optional[Dict] = None + random_timesteps: int = MISSING + learning_starts: int = MISSING + grad_norm_clip: float = MISSING + ratio_clip: float = MISSING + value_clip: float = MISSING + clip_predicted_values: bool = MISSING + entropy_loss_scale: float = MISSING + value_loss_scale: float = MISSING + kl_threshold: float = MISSING + rewards_shaper_scale: float = MISSING + experiment: ExperimentCfg = MISSING + + +@configclass +class TrainerCfg: + """Base configuration for the sequential trainer.""" + + timesteps: int = MISSING + environment_info: str = MISSING + close_environment_at_exit: bool = True + + +@configclass +class ExtensionCfg: + supplementary_training_cfg: SupplementaryTrainingCfg = MISSING + + +@configclass +class SKRLConfig: + """Base configuration for the RL agent and training setup.""" + + seed: int = MISSING + models: ModelsCfg = MISSING + agent: PPOAgentCfg = MISSING + trainer: TrainerCfg = MISSING + extension: Optional[ExtensionCfg] = None + + def to_skrl_dict(self): + dict = self.to_dict() + dict["agent"]["class"] = dict["agent"].pop("agent_class", None) + dict["models"]["policy"]["class"] = dict["models"]["policy"].pop("policy_class", None) + dict["models"]["value"]["class"] = dict["models"]["value"].pop("value_class", None) + return dict diff --git a/source/uwlab_tasks/config/extension.toml b/source/uwlab_tasks/config/extension.toml new file mode 100644 index 00000000..60e26cd0 --- /dev/null +++ b/source/uwlab_tasks/config/extension.toml @@ -0,0 +1,25 @@ +[package] + +# Semantic Versioning is used: https://semver.org/ +version = "0.13.0" + +# Description +title = "UW Lab Tasks" +readme = "docs/README.md" +description="Extension for Isaac Lab" +repository = "https://github.com/UW-Lab/UWLab" +category = "robotics" +keywords = ["robotics", "rl", "il", "learning"] + +[dependencies] +"isaaclab" = {} +"isaaclab_assets" = {} +"uwlab" = {} +"uwlab_assets" = {} +# NOTE: Add additional dependencies here + +[core] +reloadable = false + +[[python.module]] +name = "uwlab_tasks" diff --git a/source/uwlab_tasks/docs/CHANGELOG.rst b/source/uwlab_tasks/docs/CHANGELOG.rst new file mode 100644 index 00000000..4db507a4 --- /dev/null +++ b/source/uwlab_tasks/docs/CHANGELOG.rst @@ -0,0 +1,467 @@ +Changelog +--------- + +0.13.1 (2025-01-13) + +Fixed +^^^^^ + +* Fixed LiftHammer environment to use isaac lab native multi-asset spawning configuration +* Fixed LiftHammer environment to updated tiled camera configuration + + +0.13.0 (2024-11-10) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* introducing storage manager as a module in lab_task under :class:`uwlab_apps.utils.storage_manager` + + +0.12.1 (2024-10-27) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* added skrl ppo config for single cake decoration environment + +0.12.0 (2024-10-27) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Introducing stateful configuratble skrl training workflow pipeline +* tested to be compatible with current evolution branch + +0.11.0 (2024-10-20) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* merged franka workshop environment, frank multi cake environment into uwlab, Thanks Yufeng! + +0.10.0 (2024-10-20) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* merged skrl workflow pipeline into uwlab + +0.9.6 (2024-09-06) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* The agent experiment name for rsl_rl gym registration was not correct was "rough" now is "terrain_gen" + changes at :func:`uwlab_tasks.tasks.locomotion.fetching.config.a1.__init__.py` + +0.9.5 (2024-09-02) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ +* in favoring uwlab having reset_from_demostration and record_state_configuration functions +* remove functions :func:`uwlab_tasks.tasks.manipulation.cake_decoration.mdp.reset_from_demostration` +* remove functions :func:`uwlab_tasks.tasks.manipulation.cake_decoration.mdp.record_state_configuration` +* remove functions :func:`uwlab_tasks.tasks.manipulation.clockHand.mdp.reset_from_demostration` +* remove functions :func:`uwlab_tasks.tasks.manipulation.clockHand.mdp.record_state_configuration` + + +0.9.4 (2024-08-28) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ +* adding functions at ``uwlab_tasks.uwlab_apps.utils.cfg_utils.py`` and enable easy rotation + modification to environments + + +0.9.3 (2024-08-24) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ +* separated out lift objects environment from lift hammer environment at tasks.manipulation.lift_objects + +0.9.2 (2024-08-19) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ +* fixed problem where the order of tycho gripper joint action idex and body joint pos are reversed + :class:`uwlab_tasks.tasks.manipulation.cake_decoration.config.hebi.tycho_joint_pos.IkdeltaAction` + and :class:`uwlab_tasks.tasks.manipulation.cake_decoration.config.hebi.tycho_joint_pos.IkabsoluteAction` + +0.9.1 (2024-08-06) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^^^ +* Added necessary mdps for :folder:`uwlab_tasks.tasks.locomotion` tasks + +0.9.0 (2024-08-06) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ +* rename unitree_a1, unitree_go1, unitree_go2 to a1, a2, a3 under + :file:`uwlab_tasks.tasks.locomotion` + + +0.8.3 (2024-08-06) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ +* added terrain_gen environment as separate task in + :file:`uwlab_tasks.tasks.locomotion.fetching.fetching_terrain_gen_env` + +Changed +^^^^^^^ +* renamed ``uwlab_tasks.tasks.locomotion.fetching.rough_env_cfg`` to + ``fetching_env_cfg`` to show its difference from locomotion Velocity tasks + + +0.8.2 (2024-08-06) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ +* added coefficient as input argument in + :func:`uwlab_tasks.tasks.locomotion.fetching.mdp.rewards.track_interpolated_lin_vel_xy_exp` + + +0.8.1 (2024-08-06) +~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ +* ui_extension is deleted to prevent the buggy import +* :file:`uwlab_tasks.uwlab_tasks.__init__.py` does not import ui_extension + + +0.8.0 (2024-07-29) +~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ +* :file:`uwlab_tasks.uwlab_tasks.__init__.py` did not import tasks folder + now it is imported + + +0.8.0 (2024-07-29) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ +* updated dependency and meta information to isaac sim 4.1.0 + + + +0.7.0 (2024-07-29) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ +* added Unitree Go1 Go2 and spot for Fetching task at + :folder:`uwlab_tasks.tasks.locomotion.fetching` + + +0.6.1 (2024-07-29) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ +* bug fix in logging name unitree a1 agent, flat config should log flat instead of rough at + at :class:`uwlab_tasks.tasks.locomotion.fetching.config.unitree_a1.agents.rsl_rl_cfg.UnitreeA1FlatPPORunnerCfg` + + +0.6.0 (2024-07-28) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ +* restructured fetching task to new architecture and added Unitree A1 + for fetching task + + +0.5.2 (2024-07-28) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ +* merge all gym registering tasks to one whole name unseparated by "-" + what used to be 'UW-Lift-Objects-XarmLeap-IkDel-v0' now becomes + 'UW-LiftObjects-XarmLeap-IkDel-v0' + +0.5.1 (2024-07-28) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ +* support IkDelta action for environment LiftObjectsXarmLeap at + :folder:`uwlab_tasks.tasks.manipulation.lift_objects` + + +0.5.0 (2024-07-28) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ +* adopting new environment structure for task track_goal + + +0.4.3 (2024-07-28) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ +* fix several minor bugs that introduced when migrating for new environment structure for tasks lift_objects + + +0.4.2 (2024-07-28) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ +* added fetching task specific reward at :func:`uwlab_tasks.locomotion.fetching.mdp.track_interpolated_lin_vel_xy_exp` + and :func:`uwlab_tasks.locomotion.fetching.mdp.track_interpolated_ang_vel_z_exp` + + +0.4.1 (2024-07-27) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* update track_goal tasks under folder :folder:`uwlab_tasks.tasks.manipulation.track_goal` + + +0.4.0 (2024-07-27) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* renaming :folder:`uwlab_tasks.tasks.manipulation.lift_cube` as + :folder:`uwlab_tasks.tasks.manipulation.lift_objects` +* separates lift_cube and lift_multiobjects as two different environments + +* adopting new environment structure for task lift_objects + + +0.3.0 (2024-07-27) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* renaming :folder:`uwlab_tasks.tasks.manipulation.craneberryLavaChocoCake` as + :folder:`uwlab_tasks.tasks.manipulation.cake_decoration` + +* adopting new environment structure for task cake_decoration + + +0.2.3 (2024-07-27) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* sketched Fetching as a separate locomotion task, instead of being a part of + :folder:`uwlab_tasks.tasks.locomotion.velocity` + + +0.2.2 (2024-07-27) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* dropped dependency of :folder:`uwlab_tasks.cfg` in favor of extension ``uwlab_assets`` + + + +0.2.1 (2024-07-27) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* added UW as author and maintainer to :file:`uwlab_tasks.setup.py` + +0.2.0 (2024-07-14) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* added support for register gym environment with MultiConstraintDifferentialIKController for leap_hand_xarm at + :file:`uwlab_tasks.tasks.maniputation.lift_cube.config.leap_hand_xarm.__init__` + + +0.2.0 (2024-07-14) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* added leap hand xarm reward :func:`uwlab_tasks.cfgs.robots.leap_hand_xarm.mdp.rewards.reward_fingers_object_distance` +* tuned liftCube environment reward function for LeapHandXarm environments + reward_fingers_object_distance scale was 1.5, now 5 + reward_object_ee_distance scale was 1, now 3 + reward_fingers_object_distance tanh return std was 0.1 now 0.2 + +0.1.9 (2024-07-13) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* added leap hand xarm reward :func:`uwlab_tasks.cfgs.robots.leap_hand_xarm.mdp.rewards.reward_cross_finger_similarity` +* added leap hand xarm reward :func:`uwlab_tasks.cfgs.robots.leap_hand_xarm.mdp.rewards.reward_intra_finger_similarity` +* added leap hand xarm event :func:`uwlab_tasks.cfgs.robots.leap_hand_xarm.mdp.events.reset_joints_by_offset` which accepts + additional joint ids +* changed cube lift environment cube size to be a bit larger +* added mass randomization cfg in cube lift environment :field:`uwlab_tasks.tasks.manipulation.lift_cube.` + + +0.1.8 (2024-07-12) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* added leap hand xarm robot cfg and dynamic at :file:`uwlab_tasks.cfgs.robots.leap_hand.robot_cfg.py` and + :file:`uwlab_tasks.cfgs.robots.leap_hand_xarm.robot_dynamics.py` +* added environment :file:`uwlab_tasks.tasks.manipulation.lift_cube.track_goal.config.leap_hand_xarm.LeapHandXarm_JointPos_GoalTracking_Env.py` +* added environment :file:`uwlab_tasks.tasks.manipulation.lift_cube.lift_cube.config.leap_hand_xarm.LeapHandXarm_JointPos_LiftCube_Env.py` + + +0.1.7 (2024-07-08) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Hebi Gravity Enabled now becomes default +* orbid_mdp changed to lab_mdp in :file:`uwlab_tasks.cfgs.robots.leap_hand.robot_dynamics.py` +* Removed Leap hand standard ik absolute and ik delta in :file:`uwlab_tasks.cfgs.robots.leap_hand.robot_dynamics.py` +* Reflect support of RokokoGloveKeyboard in :func:`workflows.teleoperation.teleop_se3_agent_absolute.main` + + +Added +^^^^^ +* Added experiments run script :file:`workflows.experiments.idealpd_experiments.py` +* Added experiments :file:`uwlab_tasks.tasks.manipulation.track_goal.config.hebi.idealpd_scale_experiments.py` + + +0.1.6 (2024-07-07) +~~~~~~~~~~~~~~~~~~ + +memo: +^^^^^ + +* Termination term should be carefully considered along with the punishment reward functions. + When there are too many negative reward in the beginning, agent would prefer to die sooner by + exploiting the termination condition, and this would lead to the agent not learning the task. + +* tips: + When designing the reward function, try be incentive than punishment. + +Changed +^^^^^^^ + +* Changed :class:`uwlab_tasks.cfgs.robots.hebi.robot_dynamics.RobotTerminationsCfg` to include DoneTerm: robot_extremely_bad_posture +* Changed :function:`uwlab_tasks.cfgs.robots.hebi.mdp.terminations.terminate_extremely_bad_posture` to be probabilistic +* Changed :field:`uwlab_tasks.tasks.manipulation.track_goal.config.hebi.Hebi_JointPos_GoalTracking_Env.RewardsCfg.end_effector_position_tracking` + and :field:`uwlab_tasks.tasks.manipulation.track_goal.config.hebi.Hebi_JointPos_GoalTracking_Env.RewardsCfg.end_effector_orientation_tracking` + to be incentive reward instead of punishment reward. +* Renamed orbit_mdp to lab_mdp in :file:`uwlab_tasks.tasks.manipulation.track_goal.config.Hebi_JointPos_GoalTracking_Env` + +Added +^^^^^ + +* Added hebi reward term :func:`uwlab_tasks.cfgs.robots.hebi.mdp.rewards.orientation_command_error_tanh` +* Added experiments run script :file:`workflows.experiments.strategy4_scale_experiments.py` +* Added experiments :file:`uwlab_tasks.tasks.manipulation.track_goal.config.hebi.strategy4_scale_experiments.py` + +0.1.5 (2024-07-06) +~~~~~~~~~~~~~~~~~~ + + +Added +^^^^^ + +* Added experiments run script :file:`workflows.experiments.actuator_experiments.py` +* Added experiments run script :file:`workflows.experiments.agent_update_frequency_experiments.py` +* Added experiments run script :file:`workflows.experiments.decimation_experiments.py` +* Added experiments run script :file:`workflows.experiments.strategy3_scale_experiments.py` +* Added experiments :file:`uwlab_tasks.tasks.manipulation.track_goal.config.hebi.agent_update_rate_experiments.py` +* Added experiments :file:`uwlab_tasks.tasks.manipulation.track_goal.config.hebi.decimation_experiments.py` +* Added experiments :file:`uwlab_tasks.tasks.manipulation.track_goal.config.hebi.strategy3_scale_experiments.py` +* Modified :file:`uwlab_tasks.tasks.manipulation.track_goal.config.hebi.agents.rsl_rl_agent_cfg`, and + :file:`uwlab_tasks.tasks.manipulation.track_goal.config.hebi.__init__` with logging name consistent to experiments + + +0.1.4 (2024-07-05) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* :const:`uwlab_tasks.cfgs.robots.hebi.robot_cfg.HEBI_STRATEGY3_CFG` + :const:`uwlab_tasks.cfgs.robots.hebi.robot_cfg.HEBI_STRATEGY4_CFG` + changed from manually editing scaling factor to cfg specifying scaling factor. +* :const:`uwlab_tasks.cfgs.robots.hebi.robot_cfg.robot_dynamic` +* :func:`workflows.teleoperation.teleop_se3_agent_absolute.main` added visualization for full gloves data + +0.1.3 (2024-06-29) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* updated :func:`workflows.teleoperation.teleop_se3_agent_absolute.main` gloves device to match updated + requirement needed for rokoko gloves. New version can define port usage, output parts + + + + +0.1.2 (2024-06-28) +~~~~~~~~~~~~~~~~~~ + + +Changed +^^^^^^^ + +* Restructured lab to accommodate new extension lab environments +* renamed the repository from lab.tycho to lab.envs +* removed :func:`workflows.teleoperation.teleop_se3_agent_absolute_leap.main` as it has been integrated + into :func:`workflows.teleoperation.teleop_se3_agent_absolute.main` + + +0.1.1 (2024-06-27) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* teleoperation absolute ik control for leap hand at :func:`workflows.teleoperation.teleop_se3_agent_absolute_leap.main` + + +0.1.0 (2024-06-11) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Performed tycho migration. Done with Tasks: cake, liftcube, clock, meat, Goal Tracking +* Need to check: meat seems to have a bit of issue +* Plan to do: Learn a mujoco motor model, test out dreamerv3, refactorization continue diff --git a/source/uwlab_tasks/pyproject.toml b/source/uwlab_tasks/pyproject.toml new file mode 100644 index 00000000..d90ac353 --- /dev/null +++ b/source/uwlab_tasks/pyproject.toml @@ -0,0 +1,3 @@ +[build-system] +requires = ["setuptools", "wheel", "toml"] +build-backend = "setuptools.build_meta" diff --git a/source/uwlab_tasks/setup.py b/source/uwlab_tasks/setup.py new file mode 100644 index 00000000..fc668c2e --- /dev/null +++ b/source/uwlab_tasks/setup.py @@ -0,0 +1,41 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Installation script for the 'uwlab_tasks' python package.""" + +import os +import toml + +from setuptools import setup + +# Obtain the extension data from the extension.toml file +EXTENSION_PATH = os.path.dirname(os.path.realpath(__file__)) +# Read the extension.toml file +EXTENSION_TOML_DATA = toml.load(os.path.join(EXTENSION_PATH, "config", "extension.toml")) + +# Minimum dependencies required prior to installation +INSTALL_REQUIRES = [] + +# Installation operation +setup( + name="uwlab_tasks", + author="UW and Isaac Lab Project Developers", + maintainer="UW and Isaac Lab Project Developers", + url=EXTENSION_TOML_DATA["package"]["repository"], + version=EXTENSION_TOML_DATA["package"]["version"], + description=EXTENSION_TOML_DATA["package"]["description"], + keywords=EXTENSION_TOML_DATA["package"]["keywords"], + license="BSD-3-Clause", + include_package_data=True, + python_requires=">=3.10", + install_requires=INSTALL_REQUIRES, + packages=["uwlab_tasks"], + classifiers=[ + "Natural Language :: English", + "Programming Language :: Python :: 3.10", + "Isaac Sim :: 4.5.0", + ], + zip_safe=False, +) diff --git a/source/uwlab_tasks/uwlab_tasks/__init__.py b/source/uwlab_tasks/uwlab_tasks/__init__.py new file mode 100644 index 00000000..6e24c56a --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/__init__.py @@ -0,0 +1,30 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Package containing task implementations for various robotic environments.""" + +import os +import toml + +# Conveniences to other module directories via relative paths +UWLAB_TASKS_EXT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "../")) +"""Path to the extension source directory.""" + +UWLAB_TASKS_METADATA = toml.load(os.path.join(UWLAB_TASKS_EXT_DIR, "config", "extension.toml")) +"""Extension metadata dictionary parsed from the extension.toml file.""" + +# Configure the module-level variables +__version__ = UWLAB_TASKS_METADATA["package"]["version"] + +## +# Register Gym environments. +## + +from isaaclab_tasks.utils import import_packages + +# The blacklist is used to prevent importing configs from sub-packages +_BLACKLIST_PKGS = ["utils"] +# Import all configs in this package +import_packages(__name__, _BLACKLIST_PKGS) diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/__init__.py b/source/uwlab_tasks/uwlab_tasks/manager_based/__init__.py new file mode 100644 index 00000000..b25784ab --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2024-2025, The UW Lab Project Developers. +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +""" +Config-based workflow environments. +""" + +import gymnasium as gym diff --git a/uwlab.sh b/uwlab.sh new file mode 100755 index 00000000..67d6ca2a --- /dev/null +++ b/uwlab.sh @@ -0,0 +1,82 @@ +#!/usr/bin/env bash +# Exit immediately if a command fails +set -e + +# Define a function to print help +print_help() { + echo "Usage: $(basename "$0") [option]" + echo -e "\t-h, --help Display this help message." + echo -e "\t-i, --install [LIB] Install the uwlab packages. Optionally specify a LIB." + echo -e "\t-f, --format Format the uwlab code." + echo -e "\t-t, --test Run tests for uwlab." +} + +# If no arguments are provided, show the help and exit. +if [ $# -eq 0 ]; then + print_help + exit 0 +fi + +# Base path to the packages (adjust as needed) +BASE_PATH="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )/source" + +# Process the command line arguments +while [[ $# -gt 0 ]]; do + key="$1" + case $key in + -h|--help) + print_help + exit 0 + ;; + -i|--install) + echo "[INFO] Installing uwlab packages..." + echo "Installing uwlab..." + pip install -e "${BASE_PATH}/uwlab" + echo "Installing uwlab_assets..." + pip install -e "${BASE_PATH}/uwlab_assets" + echo "Installing uwlab_tasks..." + pip install -e "${BASE_PATH}/uwlab_tasks" + pip install -e "${BASE_PATH}/uwlab_rl" + echo "[INFO] All packages have been installed in editable mode." + ;; + -f|--format) + echo "[INFO] Formatting uwlab code..." + # Reset the PYTHONPATH if using a conda environment to avoid conflicts with pre-commit + if [ -n "${CONDA_DEFAULT_ENV}" ]; then + cache_pythonpath=${PYTHONPATH} + export PYTHONPATH="" + fi + + # Ensure pre-commit is installed + if ! command -v pre-commit &>/dev/null; then + echo "[INFO] Installing pre-commit..." + pip install pre-commit + fi + + echo "[INFO] Formatting the repository..." + # Determine the repository root directory (assumes this script is at the repo root) + UWLAB_PATH="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )" + cd ${UWLAB_PATH} + pre-commit run --all-files + cd - > /dev/null + + # Restore the PYTHONPATH if it was modified + if [ -n "${CONDA_DEFAULT_ENV}" ]; then + export PYTHONPATH=${cache_pythonpath} + fi + shift + break + ;; + -t|--test) + echo "[INFO] Running tests..." + # Insert your test command here (for example, pytest or another test runner) + # pytest + ;; + *) + echo "[ERROR] Unknown option: $key" + print_help + exit 1 + ;; + esac + shift +done