├── .flake8 ├── .gitattributes ├── .gitignore ├── .gitmodules ├── .pre-commit-config.yaml ├── CONTRIBUTING.md ├── Dockerfile ├── LICENCE ├── README.md ├── docs ├── robots_in_action.gif ├── stable_punch.gif └── stable_wave.gif ├── install_deps.sh ├── neural_wbc ├── core │ ├── README.md │ ├── neural_wbc │ │ └── core │ │ │ ├── __init__.py │ │ │ ├── body_state.py │ │ │ ├── environment_wrapper.py │ │ │ ├── evaluator.py │ │ │ ├── mask.py │ │ │ ├── math_utils.py │ │ │ ├── modes.py │ │ │ ├── observations │ │ │ ├── __init__.py │ │ │ ├── student_history.py │ │ │ ├── student_observations.py │ │ │ └── teacher_observations.py │ │ │ ├── reference_motion.py │ │ │ ├── robot_wrapper.py │ │ │ ├── termination.py │ │ │ └── util.py │ ├── setup.py │ └── tests │ │ ├── __init__.py │ │ ├── test_body_state.py │ │ ├── test_evaluator.py │ │ ├── test_mask.py │ │ ├── test_reference_motion_manager.py │ │ ├── test_student_history.py │ │ ├── test_student_observations.py │ │ ├── test_teacher_observations.py │ │ └── test_termination.py ├── data │ ├── README.md │ ├── data │ │ ├── motion_lib │ │ │ └── h1.xml │ │ └── mujoco │ │ │ └── models │ │ │ ├── assets │ │ │ ├── left_ankle_link.stl │ │ │ ├── left_elbow_link.stl │ │ │ ├── left_hip_pitch_link.stl │ │ │ ├── left_hip_roll_link.stl │ │ │ ├── left_hip_yaw_link.stl │ │ │ ├── left_knee_link.stl │ │ │ ├── left_shoulder_pitch_link.stl │ │ │ ├── left_shoulder_roll_link.stl │ │ │ ├── left_shoulder_yaw_link.stl │ │ │ ├── logo_link.stl │ │ │ ├── pelvis.stl │ │ │ ├── right_ankle_link.stl │ │ │ ├── right_elbow_link.stl │ │ │ ├── right_hip_pitch_link.stl │ │ │ ├── right_hip_roll_link.stl │ │ │ ├── right_hip_yaw_link.stl │ │ │ ├── right_knee_link.stl │ │ │ ├── right_shoulder_pitch_link.stl │ │ │ ├── right_shoulder_roll_link.stl │ │ │ ├── right_shoulder_yaw_link.stl │ │ │ └── torso_link.stl │ │ │ ├── h1.xml │ │ │ └── scene.xml │ ├── neural_wbc │ │ └── data │ │ │ └── __init__.py │ ├── setup.py │ └── tests │ │ ├── __init__.py │ │ └── test_get_file_path.py ├── hw_wrappers │ ├── README.md │ ├── docs │ │ ├── push_recovery.gif │ │ └── unitree_h1_setup.png │ ├── hw_wrappers │ │ ├── __init__.py │ │ ├── h1_sdk_wrapper.py │ │ └── unitree_h1.py │ ├── setup.py │ └── tests │ │ ├── __init__.py │ │ └── test_unitree_h1.py ├── inference_env │ ├── __init__.py │ ├── inference_env │ │ ├── README.md │ │ ├── __init__.py │ │ ├── deployment_player.py │ │ ├── neural_wbc_env.py │ │ ├── neural_wbc_env_cfg.py │ │ ├── neural_wbc_env_cfg_h1.py │ │ ├── neural_wbc_env_cfg_real_h1.py │ │ └── utils.py │ ├── scripts │ │ ├── eval.py │ │ ├── mujoco_viewer_player.py │ │ └── s2r_player.py │ ├── setup.py │ └── tests │ │ ├── __init__.py │ │ ├── test_evaluation_pipeline.py │ │ └── test_neural_wbc_env.py ├── isaac_lab_wrapper │ ├── neural_wbc │ │ └── isaac_lab_wrapper │ │ │ ├── __init__.py │ │ │ ├── body_state.py │ │ │ ├── control.py │ │ │ ├── events │ │ │ ├── __init__.py │ │ │ ├── event_cfg.py │ │ │ └── events.py │ │ │ ├── neural_wbc_env.py │ │ │ ├── neural_wbc_env_cfg.py │ │ │ ├── neural_wbc_env_cfg_h1.py │ │ │ ├── observations.py │ │ │ ├── rewards │ │ │ ├── __init__.py │ │ │ ├── reward_cfg.py │ │ │ └── rewards.py │ │ │ ├── terrain.py │ │ │ ├── utils.py │ │ │ └── visualization.py │ ├── setup.py │ └── tests │ │ ├── __init__.py │ │ ├── test_main.py │ │ ├── test_neural_wbc_env.py │ │ ├── test_neural_wbc_env_cfg.py │ │ └── test_observations.py ├── mujoco_wrapper │ ├── README.md │ ├── mujoco_wrapper │ │ ├── __init__.py │ │ ├── control.py │ │ ├── mujoco_robot.py │ │ ├── mujoco_simulator.py │ │ ├── mujoco_utils.py │ │ ├── utils.py │ │ └── visualization.py │ ├── setup.py │ └── tests │ │ ├── __init__.py │ │ ├── test_mujoco_simulator.py │ │ └── test_mujoco_utils.py └── student_policy │ ├── README.md │ ├── neural_wbc │ ├── student_policy │ │ ├── __init__.py │ │ ├── policy.py │ │ ├── storage.py │ │ ├── student_policy_trainer.py │ │ ├── student_policy_trainer_cfg.py │ │ └── teacher_policy.py │ └── tests │ │ ├── __init__.py │ │ ├── test_storage.py │ │ └── test_student_policy_trainer.py │ └── setup.py ├── punch.yaml ├── pyproject.toml ├── requirements.txt ├── requirements_deploy.txt ├── retarget_h1.sh ├── run_e2e_tests.sh ├── run_unit_tests.sh ├── scripts └── rsl_rl │ ├── config_overwrites │ └── sample.overwrite.yaml │ ├── eval.py │ ├── play.py │ ├── players.py │ ├── teacher_policy_cfg.py │ ├── tests │ ├── __init__.py │ ├── student_evaluation_pipeline_test.py │ ├── student_training_pipeline_test.py │ ├── teacher_evaluation_pipeline_test.py │ ├── teacher_training_pipeline_test.py │ ├── test_teacher_policy_cfg.py │ └── test_utils.py │ ├── train_student_policy.py │ ├── train_teacher_policy.py │ ├── utils.py │ └── vecenv_wrapper.py └── third_party ├── grad_fit_h1.patch ├── mujoco_viewer ├── LICENCE ├── __init__.py ├── mujoco_viewer │ ├── __init__.py │ ├── callbacks.py │ └── mujoco_viewer.py └── setup.py ├── phc_torch_utils.patch ├── requirements.patch └── rsl_rl ├── .gitignore ├── LICENSE ├── README.md ├── licenses └── dependencies │ ├── numpy_license.txt │ └── torch_license.txt ├── rsl_rl ├── __init__.py ├── algorithms │ ├── __init__.py │ └── ppo.py ├── env │ ├── __init__.py │ └── vec_env.py ├── modules │ ├── __init__.py │ └── actor_critic.py ├── runners │ ├── __init__.py │ └── on_policy_runner.py ├── storage │ ├── __init__.py │ └── rollout_storage.py └── utils │ ├── __init__.py │ └── utils.py └── setup.py /.flake8: -------------------------------------------------------------------------------- 1 | # copied from https://github.com/isaac-sim/IsaacLab/blob/main/.flake8 2 | 3 | [flake8] 4 | show-source=True 5 | statistics=True 6 | per-file-ignores=*/__init__.py:F401 7 | # E402: Module level import not at top of file 8 | # E501: Line too long 9 | # W503: Line break before binary operator 10 | # E203: Whitespace before ':' -> conflicts with black 11 | # D401: First line should be in imperative mood 12 | # R504: Unnecessary variable assignment before return statement. 13 | # R505: Unnecessary elif after return statement 14 | # SIM102: Use a single if-statement instead of nested if-statements 15 | # SIM117: Merge with statements for context managers that have same scope. 16 | ignore=E402,E501,W503,E203,D401,R504,R505,SIM102,SIM117 17 | max-line-length = 120 18 | max-complexity = 30 19 | exclude=_*,.vscode,.git,docs/** 20 | # docstrings 21 | docstring-convention=google 22 | # annotations 23 | suppress-none-returning=True 24 | allow-star-arg-any=True 25 | -------------------------------------------------------------------------------- /.gitattributes: -------------------------------------------------------------------------------- 1 | # copied and modified from https://github.com/NVIDIA/warp/blob/main/.gitattributes 2 | * text=auto 3 | *.sh text eol=LF 4 | # copied from https://github.com/isaac-sim/IsaacLab/blob/main/.gitattributes 5 | *.usd filter=lfs diff=lfs merge=lfs -text 6 | *.dae filter=lfs diff=lfs merge=lfs -text 7 | *.mtl filter=lfs diff=lfs merge=lfs -text 8 | *.obj filter=lfs diff=lfs merge=lfs -text 9 | *.gif filter=lfs diff=lfs merge=lfs -text 10 | *.png filter=lfs diff=lfs merge=lfs -text 11 | *.jpg filter=lfs diff=lfs merge=lfs -text 12 | *.psd filter=lfs diff=lfs merge=lfs -text 13 | *.mp4 filter=lfs diff=lfs merge=lfs -text 14 | *.usda filter=lfs diff=lfs merge=lfs -text 15 | *.hdr filter=lfs diff=lfs merge=lfs -text 16 | *.pt filter=lfs diff=lfs merge=lfs -text 17 | *.jit filter=lfs diff=lfs merge=lfs -text 18 | neural_wbc/data/data/policy/h1:student/model_40000.pt filter=lfs diff=lfs merge=lfs -text 19 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | **/__pycache__/ 2 | *.egg-info 3 | logs/ 4 | /data/ 5 | third_party/human2humanoid 6 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "third_party/poselib"] 2 | path = third_party/poselib 3 | url = ../../gr00t-control/poselib.git 4 | [submodule "third_party/human2humanoid"] 5 | path = third_party/human2humanoid 6 | url = https://github.com/LeCAR-Lab/human2humanoid 7 | ignore = dirty 8 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | exclude: (third_party/.*) 2 | repos: 3 | - repo: https://github.com/python/black 4 | rev: 23.10.1 5 | hooks: 6 | - id: black 7 | args: ["--line-length", "120", "--preview"] 8 | - repo: https://github.com/pycqa/flake8 9 | rev: 6.1.0 10 | hooks: 11 | - id: flake8 12 | additional_dependencies: [flake8-simplify, flake8-return] 13 | - repo: https://github.com/pre-commit/pre-commit-hooks 14 | rev: v4.5.0 15 | hooks: 16 | - id: trailing-whitespace 17 | - id: check-symlinks 18 | - id: destroyed-symlinks 19 | - id: check-yaml 20 | - id: check-merge-conflict 21 | - id: check-case-conflict 22 | - id: check-executables-have-shebangs 23 | - id: check-toml 24 | - id: end-of-file-fixer 25 | - id: check-shebang-scripts-are-executable 26 | - id: detect-private-key 27 | - id: debug-statements 28 | - repo: https://github.com/pycqa/isort 29 | rev: 5.12.0 30 | hooks: 31 | - id: isort 32 | name: isort (python) 33 | args: ["--profile", "black", "--filter-files"] 34 | - repo: https://github.com/asottile/pyupgrade 35 | rev: v3.15.0 36 | hooks: 37 | - id: pyupgrade 38 | args: ["--py37-plus"] 39 | - repo: https://github.com/codespell-project/codespell 40 | rev: v2.2.6 41 | hooks: 42 | - id: codespell 43 | additional_dependencies: 44 | - tomli 45 | - repo: https://github.com/pre-commit/pygrep-hooks 46 | rev: v1.10.0 47 | hooks: 48 | - id: rst-backticks 49 | - id: rst-directive-colons 50 | - id: rst-inline-touching-normal 51 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # How to Contribute 2 | 3 | We'd love to accept your patches and contributions. 4 | 5 | ## Code Reviews 6 | 7 | All submissions, including submissions by project members, require review. We use GitHub pull requests for this purpose. Consult 8 | [GitHub Help](https://help.github.com/articles/about-pull-requests/) for more information on using pull requests. 9 | 10 | ## Signing Your Work 11 | 12 | * We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license. 13 | 14 | * Any contribution which contains commits that are not Signed-Off will not be accepted. 15 | 16 | * To sign off on a commit you simply use the `--signoff` (or `-s`) option when committing your changes: 17 | ```bash 18 | $ git commit -s -m "Add cool feature." 19 | ``` 20 | This will append the following to your commit message: 21 | ``` 22 | Signed-off-by: Your Name 23 | ``` 24 | 25 | * Full text of the DCO: 26 | 27 | ``` 28 | Developer Certificate of Origin 29 | Version 1.1 30 | 31 | Copyright (C) 2004, 2006 The Linux Foundation and its contributors. 32 | 1 Letterman Drive 33 | Suite D4700 34 | San Francisco, CA, 94129 35 | 36 | Everyone is permitted to copy and distribute verbatim copies of this license document, but changing it is not allowed. 37 | ``` 38 | 39 | ``` 40 | Developer's Certificate of Origin 1.1 41 | 42 | By making a contribution to this project, I certify that: 43 | 44 | (a) The contribution was created in whole or in part by me and I have the right to submit it under the open source license indicated in the file; or 45 | 46 | (b) The contribution is based upon previous work that, to the best of my knowledge, is covered under an appropriate open source license and I have the right under that license to submit that work with modifications, whether created in whole or in part by me, under the same open source license (unless I am permitted to submit under a different license), as indicated in the file; or 47 | 48 | (c) The contribution was provided directly to me by some other person who certified (a), (b) or (c) and I have not modified it. 49 | 50 | (d) I understand and agree that this project and the contribution are public and that a record of the contribution (including all personal information I submit with it, including my sign-off) is maintained indefinitely and may be redistributed consistent with this project or the open source license(s) involved. 51 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | # syntax=docker.io/docker/dockerfile:1.7-labs 2 | FROM nvcr.io/nvidia/isaac-lab:2.0.0 3 | 4 | # Hide conflicting Vulkan files, if needed. 5 | RUN if [ -e "/usr/share/vulkan" ] && [ -e "/etc/vulkan" ]; then \ 6 | mv /usr/share/vulkan /usr/share/vulkan_hidden; \ 7 | fi 8 | 9 | # Install vs code server and file browser to interact with the running workflow. 10 | RUN curl -fsSL https://code-server.dev/install.sh | bash 11 | RUN curl -fsSL https://raw.githubusercontent.com/filebrowser/get/master/get.sh | bash 12 | 13 | WORKDIR /workspace/neural_wbc 14 | 15 | COPY \ 16 | --exclude=Dockerfile \ 17 | --exclude=data/ \ 18 | --exclude=logs/ \ 19 | --exclude=release/ \ 20 | --exclude=run.py \ 21 | --exclude=workflows/ \ 22 | . /workspace/neural_wbc 23 | 24 | RUN ./install_deps.sh 25 | 26 | # Bake in a resume policy. This could be a teacher or student policy. We set this first since we 27 | # expect it to change less frequently than the policy below. 28 | ARG RESUME_PATH=neural_wbc/data/data/policy/h1:student/ 29 | COPY ${RESUME_PATH} /workspace/neural_wbc/policy/resume 30 | ENV RESUME_PATH=/workspace/neural_wbc/policy/resume 31 | ARG RESUME_CHECKPOINT=model_4000.pt 32 | ENV RESUME_CHECKPOINT=${RESUME_CHECKPOINT} 33 | 34 | # Bake in the main policy. This could be a teacher or student policy. Depending on the use case this 35 | # is used as the teacher policy or the evaluation policy. 36 | ARG POLICY_PATH=neural_wbc/data/data/policy/h1:teacher/ 37 | COPY ${POLICY_PATH} /workspace/neural_wbc/policy/ 38 | ENV POLICY_PATH=/workspace/neural_wbc/policy/ 39 | ARG POLICY_CHECKPOINT=model_76000.pt 40 | ENV POLICY_CHECKPOINT=${POLICY_CHECKPOINT} 41 | -------------------------------------------------------------------------------- /docs/robots_in_action.gif: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:e3cd11985adaae50a8b8be10b0d6e376ac55c7045e86e436f44dc2c1a37ce137 3 | size 14723710 4 | -------------------------------------------------------------------------------- /docs/stable_punch.gif: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:41518ac929aa1a70d5b1e13b6729cb31274d87423b2bed5e80f7994d14082628 3 | size 16462671 4 | -------------------------------------------------------------------------------- /docs/stable_wave.gif: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:0003a174698056e11d7911967ecd696b5337a4b644412eaa67929301d009d50b 3 | size 6818271 4 | -------------------------------------------------------------------------------- /install_deps.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # 3 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 4 | # SPDX-License-Identifier: Apache-2.0 5 | # 6 | # Licensed under the Apache License, Version 2.0 (the "License"); 7 | # you may not use this file except in compliance with the License. 8 | # You may obtain a copy of the License at 9 | # 10 | # http://www.apache.org/licenses/LICENSE-2.0 11 | # 12 | # Unless required by applicable law or agreed to in writing, software 13 | # distributed under the License is distributed on an "AS IS" BASIS, 14 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 | # See the License for the specific language governing permissions and 16 | # limitations under the License. 17 | set -e 18 | 19 | # We allow this to fail since in docker we only have the files without the git 20 | # info. 21 | git submodule update --init --recursive || true 22 | 23 | echo "Resetting changes in third_party/human2humanoid..." 24 | pushd third_party/human2humanoid 25 | git reset --hard || true 26 | popd 27 | 28 | # Apply patch to files 29 | if patch --dry-run --silent -f third_party/human2humanoid/phc/phc/utils/torch_utils.py < third_party/phc_torch_utils.patch; then 30 | echo "Dry run succeeded. Applying the patch..." 31 | patch third_party/human2humanoid/phc/phc/utils/torch_utils.py < third_party/phc_torch_utils.patch 32 | echo "Patch applied successfully." 33 | else 34 | echo "Dry run failed. Patch was not applied. Check if the patch file contains errors." 35 | exit 1 36 | fi 37 | 38 | # Install libraries. 39 | ${ISAACLAB_PATH:?}/isaaclab.sh -p -m ensurepip 40 | ${ISAACLAB_PATH:?}/isaaclab.sh -p -m pip install --upgrade pip 41 | ${ISAACLAB_PATH:?}/isaaclab.sh -p -m pip install wheel 42 | ${ISAACLAB_PATH:?}/isaaclab.sh -p -m pip install -e . 43 | ${ISAACLAB_PATH:?}/isaaclab.sh -p -m pip install -r requirements.txt 44 | 45 | 46 | # Check that IsaacLab is using the correct version. 47 | if [ -d ${ISAACLAB_PATH}/.git ]; then 48 | expected_isaac_lab_tag="v2.0.0" 49 | if ! git -C ${ISAACLAB_PATH} tag -l "${expected_isaac_lab_tag}" | grep -q "${expected_isaac_lab_tag}"; then 50 | echo "Error: IsaacLab does not have this tag." 51 | echo "Expected tag: ${expected_isaac_lab_tag}" 52 | exit 1 53 | fi 54 | fi 55 | -------------------------------------------------------------------------------- /neural_wbc/core/README.md: -------------------------------------------------------------------------------- 1 | # Neural WBC - Core 2 | 3 | The `neural_wbc.core` package provides core functionalities for implementing neural whole-body 4 | control for humanoid robots. It includes 5 | 6 | * **Data Structures** 7 | - `BodyState` for managing robot body states. 8 | - `ReferenceMotionState` for managing motion states from a dataset. 9 | * **Algorithm Functions**: Key functions for observations and termination conditions. 10 | * **Motion Dataset Reader**: `ReferenceMotionManager` for reading and querying motion datasets. 11 | * **Environment Wrapper**: Ensures compatibility with different simulators. 12 | * **Evaluator**: Facilitates easy evaluation of trained policies. 13 | 14 | ## Installation 15 | To use the `neural_wbc.core` package, install it with 16 | 17 | ```bash 18 | ${ISAACLAB_PATH:?}/isaaclab.sh -p -m pip install -e neural_wbc/core 19 | ``` 20 | 21 | ## Usage 22 | Refer to the Isaac Lab implementation of the OmniH2O environment, as well as the training and 23 | testing workflows, for examples of how to use the module. 24 | 25 | ## Evaluation and Metrics 26 | The `Evaluator` class is used to collect data and evaluate the motion tracking performance of a 27 | reinforcement learning (RL) policy in a simulated environment. Example usage: 28 | 29 | ```python 30 | from neural_wbc.core import Evaluator 31 | from your_module import YourEnvironmentWrapper 32 | 33 | # Step 1: Initialization 34 | env_wrapper = YourEnvironmentWrapper() # Assume this is correctly set up 35 | metrics_path = "path/to/save/metrics" # Optional 36 | evaluator = Evaluator(env_wrapper=env_wrapper, metrics_path=metrics_path) 37 | 38 | # Simulate environment loop 39 | while not evaluator.is_evaluation_complete(): 40 | # Get observations, rewards, dones, and extra info from the environment 41 | actions = policy(observations) 42 | observations, rewards, dones, extras = env_wrapper.step(actions) 43 | 44 | # Step 2: Collect data from each step 45 | reset_env = evaluator.collect(dones=dones, info=extras) 46 | 47 | # Step 3: Reset environment if needed 48 | if reset_env: 49 | evaluator.forward_motion_samples() 50 | observations, _ = env_wrapper.reset() 51 | 52 | # Step 4: Concluding the evaluation 53 | evaluator.conclude() 54 | ``` 55 | 56 | The evaluator provides the following metrics: 57 | 58 | * **Success Rate [%]**: The percentage of motion tracking episodes that are successfully completed. An 59 | episode is considered successful if it follows the reference motion from start to finish without 60 | losing balance and avoiding collisions on specific body parts. 61 | * **mpjpe_g [mm]**: The global mean per-joint position error, which measures the policy’s ability to 62 | imitate the reference motion globally. 63 | * **mpjpe_l [mm]**: The root-relative mean per-joint position error, which measures the policy’s ability 64 | to imitate the reference motion locally. 65 | * **mpjpe_pa [mm]**: The procrustes aligned mean per-joint position error, which aligns the links with 66 | the ground truth before calculating the errors. 67 | * **accel_dist [mm/frame^2]**: The average joint acceleration error. 68 | * **vel_dist [mm/frame]**: The average joint velocity error. 69 | * **root_r_error [radians]**: The average torso roll error. 70 | * **root_p_error [radians]**: The average torso pitch error. 71 | * **root_y_error [radians]**: The average torso yaw error. 72 | * **root_vel_error [m/frame]**: The average torso velocity error. 73 | * **root_height_error [m]**: The average torso height error. 74 | 75 | 76 | ## Unit Tests 77 | The tests are located in the `tests` directory and can be run with the `unittest` module. 78 | 79 | ### Running the Tests 80 | To run the unit tests, you can use the `unittest` module. 81 | 82 | ```bash 83 | cd neural_wbc/core 84 | ${ISAACLAB_PATH:?}/isaaclab.sh -p -m unittest 85 | ``` 86 | -------------------------------------------------------------------------------- /neural_wbc/core/neural_wbc/core/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | __path__ = __import__("pkgutil").extend_path(__path__, __name__) 17 | 18 | import sys 19 | import types 20 | 21 | 22 | def create_dummy_warp_module(): 23 | """Creates a dummy 'warp' module with necessary attributes and submodules. 24 | 25 | This function is created because the 'warp' module can only be properly imported 26 | when Isaac Sim is also launched, which is not always required when running unit 27 | tests or simulations with other simulators. By creating a dummy module, we can 28 | avoid import errors and ensure that the dependent modules can be tested without 29 | requiring Isaac Sim to be running. 30 | """ 31 | # Step 1: Create the main dummy module 32 | dummy_module = types.ModuleType("warp") 33 | 34 | # Step 2: Create the submodule 'torch' 35 | torch_submodule = types.ModuleType("torch") 36 | 37 | # Step 3: Define the array class 38 | class Array: 39 | def __init__(self, value): 40 | self.value = value 41 | 42 | def __repr__(self): 43 | return f"Array(value={self.value})" 44 | 45 | # Define torch.to_torch and torch.from_torch functions 46 | def to_torch(value): 47 | return f"Converted {value} to torch" 48 | 49 | def from_torch(value): 50 | return f"Converted {value} from torch" 51 | 52 | # Step 4: Add the class to the main dummy module 53 | dummy_module.array = Array 54 | 55 | # Add the functions to the torch submodule 56 | torch_submodule.to_torch = to_torch 57 | torch_submodule.from_torch = from_torch 58 | 59 | # Add the torch submodule to the main dummy module 60 | dummy_module.torch = torch_submodule 61 | 62 | # Step 5: Insert the main dummy module into sys.modules 63 | sys.modules["warp"] = dummy_module 64 | 65 | 66 | from .body_state import BodyState # noqa 67 | from .environment_wrapper import EnvironmentWrapper # noqa 68 | from .evaluator import Evaluator # noqa 69 | from .modes import NeuralWBCModes # noqa 70 | from .reference_motion import ReferenceMotionManager, ReferenceMotionManagerCfg, ReferenceMotionState # noqa 71 | from .termination import check_termination_conditions # noqa 72 | -------------------------------------------------------------------------------- /neural_wbc/core/neural_wbc/core/environment_wrapper.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | import torch 16 | 17 | from .modes import NeuralWBCModes 18 | from .reference_motion import ReferenceMotionManager 19 | 20 | 21 | class EnvironmentWrapper: 22 | """Class that provides an interface to perform training and evaluation in different simulators.""" 23 | 24 | num_envs: int # Number of environments running in the simulator 25 | device: torch.device 26 | reference_motion_manager: ReferenceMotionManager 27 | 28 | def __init__(self, mode: NeuralWBCModes): 29 | self._mode = mode 30 | 31 | def step(self, actions: torch.Tensor) -> tuple[dict, torch.Tensor, torch.Tensor, dict]: 32 | """Performs one step of simulation. 33 | Returns: 34 | * Observations as a dict or an object (TBD). 35 | * Rewards of shape (num_envs,) 36 | * "Dones": a boolean tensor representing termination of an episode in each environment. 37 | * Extra information captured, as a dict. 38 | """ 39 | raise NotImplementedError 40 | 41 | def reset(self, env_ids: list | torch.Tensor): 42 | """Resets environment specified by env_ids.""" 43 | raise NotImplementedError 44 | 45 | def get_observations(self) -> torch.Tensor: 46 | """Gets policy observations for each environment based on the mode.""" 47 | if self._mode.is_distill_mode(): 48 | return self.get_student_observations() 49 | return self.get_teacher_observations() 50 | 51 | def get_teacher_observations(self) -> torch.Tensor: 52 | """Gets teacher policy observations for each environment.""" 53 | raise NotImplementedError 54 | 55 | def get_student_observations(self) -> torch.Tensor: 56 | """Gets student policy observations for each environment.""" 57 | raise NotImplementedError 58 | -------------------------------------------------------------------------------- /neural_wbc/core/neural_wbc/core/modes.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | import enum 16 | 17 | 18 | class NeuralWBCModes(enum.Enum): 19 | TRAIN = 0 20 | TEST = 1 21 | DISTILL = 2 22 | DISTILL_TEST = 3 23 | 24 | def is_distill_mode(self): 25 | return self in {self.DISTILL, self.DISTILL_TEST} 26 | 27 | def is_training_mode(self): 28 | return self in {self.DISTILL, self.TRAIN} 29 | 30 | def is_test_mode(self): 31 | return self in {self.TEST, self.DISTILL_TEST} 32 | 33 | def is_distill_test_mode(self): 34 | return self == self.DISTILL_TEST 35 | -------------------------------------------------------------------------------- /neural_wbc/core/neural_wbc/core/observations/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | from .student_history import StudentHistory 18 | from .student_observations import compute_student_observations 19 | from .teacher_observations import compute_teacher_observations 20 | -------------------------------------------------------------------------------- /neural_wbc/core/neural_wbc/core/observations/student_history.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | import torch 16 | 17 | 18 | class StudentHistory: 19 | def __init__(self, num_envs: int, device: torch.device, entry_length: int, max_entries: int): 20 | """ 21 | Args: 22 | env (EnvironmentWrapper): An instance of the environment wrapper. 23 | entry_length (int): The length of a single entry. 24 | max_entries (int): The maximum number of entries to keep in the history. 25 | """ 26 | self._entries = torch.zeros( 27 | num_envs, 28 | entry_length * max_entries, 29 | device=device, 30 | ) 31 | self._entry_length = entry_length 32 | self._entry_shape = torch.Size([num_envs, entry_length]) 33 | 34 | @property 35 | def entries(self): 36 | return self._entries.clone() 37 | 38 | def update(self, obs_dict: dict[str, torch.Tensor]): 39 | """Updates the history with a new entry. 40 | 41 | Args: 42 | obs_dict (dict[str, torch.Tensor]): A dictionary containing the latest student observations. 43 | Expected keys are "distilled_robot_state" and "distilled_last_action". 44 | 45 | Raises: 46 | AssertionError: If the new entry does not match the expected shape. 47 | KeyError: If the observation does not contain expected keys. 48 | """ 49 | new_entry = torch.cat( 50 | [obs_dict["distilled_robot_state"], obs_dict["distilled_last_action"]], 51 | dim=-1, 52 | ) 53 | assert new_entry.shape == self._entry_shape, "New entry has an unexpected shape." 54 | self._entries[:, self._entry_length :] = self._entries[:, : -self._entry_length].clone() 55 | self._entries[:, : self._entry_length] = new_entry.clone() 56 | 57 | def reset(self, env_ids: torch.Tensor | None): 58 | """Resets the history for specified environments. 59 | 60 | Args: 61 | env_ids (torch.Tensor | None): A tensor containing the IDs of the environments to reset. 62 | If None, all environments are reset. 63 | 64 | Example: 65 | history.reset(None) # Resets history for all environments 66 | history.reset(torch.tensor([0, 2])) # Resets history for environments with IDs 0 and 2 67 | """ 68 | if env_ids is None: 69 | self._entries[:] = 0.0 70 | self._entries[env_ids, ...] = 0.0 71 | -------------------------------------------------------------------------------- /neural_wbc/core/neural_wbc/core/util.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | # Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | # SPDX-License-Identifier: Apache-2.0 17 | import re 18 | 19 | 20 | def assert_equal(lhs: any, rhs: any, name: str): 21 | """Assert that 2 values are equal and provide a useful error if not. 22 | 23 | Args: 24 | lhs: First value to compare 25 | rhs: Second value to compare 26 | name: Description of what is being compared, used in error messages 27 | """ 28 | # Handle dictionary comparisons 29 | if isinstance(lhs, dict) and isinstance(rhs, dict): 30 | _assert_dicts_equal(lhs, rhs, name) 31 | # Handle numeric comparisons 32 | elif isinstance(lhs, (int, float)) and isinstance(rhs, (int, float)): 33 | _assert_numbers_equal(lhs, rhs, name) 34 | # Handle all other types 35 | else: 36 | _assert_values_equal(lhs, rhs, name) 37 | 38 | 39 | def _assert_dicts_equal(lhs: dict, rhs: dict, name: str): 40 | """Compare two dictionaries and raise assertion error with details if not equal.""" 41 | lhs_keys = set(lhs.keys()) 42 | rhs_keys = set(rhs.keys()) 43 | 44 | # Check for missing keys 45 | only_in_lhs = lhs_keys - rhs_keys 46 | only_in_rhs = rhs_keys - lhs_keys 47 | 48 | # Check for value differences 49 | diff_values = _get_differing_values(lhs, rhs, lhs_keys & rhs_keys) 50 | 51 | # Build error message if there are any differences 52 | error_parts = [] 53 | if only_in_lhs: 54 | error_parts.append(f"Keys only in first dict: {only_in_lhs}") 55 | if only_in_rhs: 56 | error_parts.append(f"Keys only in second dict: {only_in_rhs}") 57 | if diff_values: 58 | error_parts.append(f"Keys with different values: {diff_values}") 59 | 60 | if error_parts: 61 | raise AssertionError(f"{name}: Dictionaries are not equal:\n" + "\n".join(error_parts)) 62 | 63 | 64 | def _get_differing_values(lhs: dict, rhs: dict, common_keys: set) -> dict: 65 | """Compare values for common keys between two dicts, return dict of differences.""" 66 | diff_values = {} 67 | for key in common_keys: 68 | if isinstance(lhs[key], (int, float)) and isinstance(rhs[key], (int, float)): 69 | if abs(lhs[key] - rhs[key]) >= 1e-6: 70 | diff_values[key] = (lhs[key], rhs[key]) 71 | elif lhs[key] != rhs[key]: 72 | diff_values[key] = (lhs[key], rhs[key]) 73 | return diff_values 74 | 75 | 76 | def _assert_numbers_equal(lhs: float, rhs: float, name: str): 77 | """Assert that two numbers are equal within a small tolerance.""" 78 | if abs(lhs - rhs) >= 1e-6: 79 | raise AssertionError(f"{name}: Values are not equal within tolerance: {lhs} != {rhs}") 80 | 81 | 82 | def _assert_values_equal(lhs: any, rhs: any, name: str): 83 | """Assert that two non-numeric values are exactly equal.""" 84 | if lhs != rhs: 85 | raise AssertionError(f"{name}: Values are not equal: {lhs} != {rhs}") 86 | 87 | 88 | def get_matching_indices(patterns: list[str], values: list[str], allow_empty: bool = False) -> list[int]: 89 | """Get indices of all elements in values that match any of the regex patterns.""" 90 | all_indices = set() 91 | for pattern in patterns: 92 | regex = re.compile(pattern) 93 | indices = [i for i, v in enumerate(values) if regex.match(v)] 94 | if len(indices) == 0 and not allow_empty: 95 | raise ValueError(f"No matching indices found for pattern {pattern} in {values}") 96 | all_indices.update(indices) 97 | return list(all_indices) 98 | -------------------------------------------------------------------------------- /neural_wbc/core/setup.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | from setuptools import find_namespace_packages, setup 18 | 19 | setup( 20 | name="neural_wbc_core", 21 | version="0.1.0", 22 | packages=find_namespace_packages(), 23 | description="Core library for neural whole body control", 24 | python_requires=">=3.10", 25 | install_requires=[ 26 | "torch", 27 | "tqdm", 28 | "numpy", 29 | ], 30 | ) 31 | -------------------------------------------------------------------------------- /neural_wbc/core/tests/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | -------------------------------------------------------------------------------- /neural_wbc/core/tests/test_body_state.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import torch 18 | import unittest 19 | 20 | from neural_wbc.core.body_state import BodyState 21 | 22 | 23 | def random_quaternions(n=1): 24 | u1, u2, u3 = torch.rand(n), torch.rand(n), torch.rand(n) 25 | q = torch.zeros((n, 4)) 26 | q[:, 0] = torch.sqrt(1 - u1) * torch.sin(2 * torch.pi * u2) 27 | q[:, 1] = torch.sqrt(1 - u1) * torch.cos(2 * torch.pi * u2) 28 | q[:, 2] = torch.sqrt(u1) * torch.sin(2 * torch.pi * u3) 29 | q[:, 3] = torch.sqrt(u1) * torch.cos(2 * torch.pi * u3) 30 | return q 31 | 32 | 33 | class TestBodyState(unittest.TestCase): 34 | def setUp(self): 35 | self.num_envs = 2 36 | self.num_bodies = 3 37 | self.num_joints = 2 38 | self.num_extended_bodies = 1 39 | self.root_id = 0 40 | 41 | # Create random tensors for initial states 42 | self.body_pos = torch.rand(self.num_envs, self.num_bodies, 3) 43 | self.body_rot = random_quaternions(self.num_envs * self.num_bodies).view(self.num_envs, self.num_bodies, 4) 44 | self.body_lin_vel = torch.rand(self.num_envs, self.num_bodies, 3) 45 | self.body_ang_vel = torch.rand(self.num_envs, self.num_bodies, 3) 46 | self.joint_pos = torch.rand(self.num_envs, self.num_joints) 47 | self.joint_vel = torch.rand(self.num_envs, self.num_joints) 48 | 49 | # Create an instance of BodyState 50 | self.body_state = BodyState( 51 | body_pos=self.body_pos, 52 | body_rot=self.body_rot, 53 | body_lin_vel=self.body_lin_vel, 54 | body_ang_vel=self.body_ang_vel, 55 | joint_pos=self.joint_pos, 56 | joint_vel=self.joint_vel, 57 | root_id=self.root_id, 58 | ) 59 | 60 | def test_initialization(self): 61 | self.assertTrue(torch.equal(self.body_state.body_pos, self.body_pos)) 62 | self.assertTrue(torch.equal(self.body_state.body_rot, self.body_rot)) 63 | self.assertTrue(torch.equal(self.body_state.body_lin_vel, self.body_lin_vel)) 64 | self.assertTrue(torch.equal(self.body_state.body_ang_vel, self.body_ang_vel)) 65 | self.assertTrue(torch.equal(self.body_state.joint_pos, self.joint_pos)) 66 | self.assertTrue(torch.equal(self.body_state.joint_vel, self.joint_vel)) 67 | 68 | # Without extension, body_*_extend is the same as body_* 69 | self.assertTrue(torch.equal(self.body_state.body_pos_extend, self.body_pos)) 70 | self.assertTrue(torch.equal(self.body_state.body_rot_extend, self.body_rot)) 71 | self.assertTrue(torch.equal(self.body_state.body_lin_vel_extend, self.body_lin_vel)) 72 | self.assertTrue(torch.equal(self.body_state.body_ang_vel_extend, self.body_ang_vel)) 73 | 74 | def test_extend_body_states(self): 75 | extend_body_pos = torch.rand(self.num_envs, self.num_extended_bodies, 3) 76 | extend_body_parent_ids = [0] 77 | 78 | self.body_state.extend_body_states(extend_body_pos, extend_body_parent_ids) 79 | 80 | self.assertEqual( 81 | self.body_state.body_pos_extend.shape, (self.num_envs, self.num_bodies + self.num_extended_bodies, 3) 82 | ) 83 | self.assertEqual( 84 | self.body_state.body_rot_extend.shape, (self.num_envs, self.num_bodies + self.num_extended_bodies, 4) 85 | ) 86 | self.assertEqual( 87 | self.body_state.body_lin_vel_extend.shape, (self.num_envs, self.num_bodies + self.num_extended_bodies, 3) 88 | ) 89 | self.assertEqual( 90 | self.body_state.body_ang_vel_extend.shape, (self.num_envs, self.num_bodies + self.num_extended_bodies, 3) 91 | ) 92 | 93 | 94 | if __name__ == "__main__": 95 | unittest.main() 96 | -------------------------------------------------------------------------------- /neural_wbc/core/tests/test_reference_motion_manager.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import torch 18 | import unittest 19 | 20 | from neural_wbc.core.reference_motion import ReferenceMotionManager, ReferenceMotionManagerCfg, ReferenceMotionState 21 | from neural_wbc.data import get_data_path 22 | 23 | 24 | class TestReferenceMotionManager(unittest.TestCase): 25 | num_envs = 10 26 | device = torch.device("cpu") 27 | 28 | def _create_reference_motion_manager(self): 29 | cfg = ReferenceMotionManagerCfg() 30 | cfg.motion_path = get_data_path("motions/stable_punch.pkl") 31 | cfg.skeleton_path = get_data_path("motion_lib/h1.xml") 32 | 33 | return ReferenceMotionManager( 34 | cfg=cfg, 35 | device=self.device, 36 | num_envs=self.num_envs, 37 | random_sample=False, 38 | extend_head=False, 39 | dt=0.005, 40 | # TODO: add RayCaster 41 | ) 42 | 43 | def test_load_reference_motion(self): 44 | mgr = self._create_reference_motion_manager() 45 | self.assertEqual(mgr.motion_lib.motion_ids.numel(), self.num_envs) 46 | 47 | def test_get_state_from_motion_lib_cache(self): 48 | mgr = self._create_reference_motion_manager() 49 | env_ids = torch.arange(self.num_envs).long() 50 | mgr.reset_motion_start_times(env_ids=env_ids, sample=True) 51 | episode_length_buf = torch.Tensor([0] * self.num_envs, device=self.device).long() 52 | ref_motion_state = mgr.get_state_from_motion_lib_cache(episode_length_buf=episode_length_buf) 53 | for attr_name in dir(ref_motion_state): 54 | if attr_name.startswith("__"): 55 | continue 56 | value = getattr(ref_motion_state, attr_name) 57 | self.assertEqual(value.shape[0], self.num_envs) 58 | 59 | def test_reference_motion_state(self): 60 | mgr = self._create_reference_motion_manager() 61 | env_ids = torch.arange(self.num_envs).long() 62 | mgr.reset_motion_start_times(env_ids=env_ids, sample=False) 63 | episode_length_buf = torch.Tensor([0] * self.num_envs, device=self.device).long() 64 | ref_motion_state: ReferenceMotionState = mgr.get_state_from_motion_lib_cache( 65 | episode_length_buf=episode_length_buf 66 | ) 67 | self.assertAlmostEqual(ref_motion_state.body_rot[0, 0, 0], 1) 68 | self.assertAlmostEqual(ref_motion_state.body_rot[0, 0, 1], 0) 69 | self.assertAlmostEqual(ref_motion_state.body_rot[0, 0, 2], 0) 70 | self.assertAlmostEqual(ref_motion_state.body_rot[0, 0, 3], 0) 71 | 72 | def test_reset_motion_start_times_after_load_motion(self): 73 | mgr = self._create_reference_motion_manager() 74 | env_ids = torch.arange(self.num_envs).long() 75 | mgr.reset_motion_start_times(env_ids=env_ids, sample=True) 76 | mgr.load_motions(random_sample=False, start_idx=0) 77 | self.assertTrue(torch.all(mgr.motion_start_times == 0)) 78 | 79 | def test_reference_motion_state_reset(self): 80 | mgr = self._create_reference_motion_manager() 81 | episode_length_buf = torch.Tensor([0] * self.num_envs, device=self.device).long() 82 | # Read the first frame 83 | ref_motion_state1: ReferenceMotionState = mgr.get_state_from_motion_lib_cache( 84 | episode_length_buf=episode_length_buf 85 | ) 86 | # Load a new motion and read the first frame. 87 | mgr.load_motions(random_sample=False, start_idx=1) 88 | ref_motion_state2: ReferenceMotionState = mgr.get_state_from_motion_lib_cache( 89 | episode_length_buf=episode_length_buf 90 | ) 91 | # Confirm that the motion is updated properly 92 | self.assertFalse(torch.equal(ref_motion_state1.joint_pos, ref_motion_state2.joint_pos)) 93 | -------------------------------------------------------------------------------- /neural_wbc/core/tests/test_student_history.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import torch 18 | import unittest 19 | 20 | from neural_wbc.core.observations import StudentHistory 21 | 22 | 23 | class TestStudentHistory(unittest.TestCase): 24 | def setUp(self): 25 | self.num_envs = 5 26 | self.entry_length = 6 27 | self.max_entries = 5 28 | self.device = torch.device("cpu") 29 | 30 | self.history = StudentHistory(self.num_envs, self.device, self.entry_length, self.max_entries) 31 | 32 | def test_initialization(self): 33 | self.assertEqual(self.history.entries.shape, (self.num_envs, self.entry_length * self.max_entries)) 34 | self.assertEqual(self.history._entry_length, self.entry_length) 35 | self.assertEqual(self.history._entry_shape, torch.Size([self.num_envs, self.entry_length])) 36 | 37 | def test_update_correct_shape(self): 38 | obs_dict = { 39 | "distilled_robot_state": torch.rand(self.num_envs, 4), 40 | "distilled_last_action": torch.rand(self.num_envs, 2), 41 | } 42 | self.history.update(obs_dict) 43 | self.assertTrue( 44 | torch.equal( 45 | self.history.entries[:, : self.entry_length], 46 | torch.cat([obs_dict["distilled_robot_state"], obs_dict["distilled_last_action"]], dim=-1), 47 | ) 48 | ) 49 | 50 | def test_update_incorrect_shape(self): 51 | obs_dict = { 52 | "distilled_robot_state": torch.rand(self.num_envs, 3), 53 | "distilled_last_action": torch.rand(self.num_envs, 2), 54 | } 55 | with self.assertRaises(AssertionError): 56 | self.history.update(obs_dict) 57 | 58 | def test_update_missing_key(self): 59 | obs_dict = { 60 | "distilled_robot_state": torch.rand(self.num_envs, 4), 61 | # "distilled_last_action" key is missing 62 | } 63 | with self.assertRaises(KeyError): 64 | self.history.update(obs_dict) 65 | 66 | def test_history_rollover(self): 67 | for _ in range(self.max_entries + 1): 68 | obs_dict = { 69 | "distilled_robot_state": torch.rand(self.num_envs, 4), 70 | "distilled_last_action": torch.rand(self.num_envs, 2), 71 | } 72 | self.history.update(obs_dict) 73 | 74 | self.assertEqual(self.history.entries.shape, (self.num_envs, self.entry_length * self.max_entries)) 75 | 76 | def test_reset_all(self): 77 | self.history._entries = torch.rand_like(self.history.entries) 78 | self.history.reset(None) 79 | self.assertTrue(torch.equal(self.history.entries, torch.zeros_like(self.history.entries))) 80 | 81 | def test_reset_specific_envs(self): 82 | self.history._entries = torch.rand_like(self.history.entries) 83 | env_ids = torch.tensor([1, 3]) 84 | self.history.reset(env_ids) 85 | for env_id in env_ids: 86 | self.assertTrue(torch.equal(self.history.entries[env_id], torch.zeros_like(self.history.entries[env_id]))) 87 | for env_id in set(range(self.num_envs)) - set(env_ids.tolist()): 88 | self.assertFalse(torch.equal(self.history.entries[env_id], torch.zeros_like(self.history.entries[env_id]))) 89 | 90 | 91 | if __name__ == "__main__": 92 | unittest.main() 93 | -------------------------------------------------------------------------------- /neural_wbc/data/README.md: -------------------------------------------------------------------------------- 1 | # Neural WBC - Data 2 | 3 | The `neural_wbc.data` package provides utility functions to retrieve data files. It is designed to 4 | help manage file paths in a structured and organized manner. 5 | 6 | ## Installation 7 | To use the `neural_wbc.data` package, install it with 8 | 9 | ```bash 10 | ${ISAACLAB_PATH:?}/isaaclab.sh -p -m pip install -e neural_wbc/data 11 | ``` 12 | 13 | ## Usage 14 | You can use the `get_data_path` function to retrieve the absolute path of a data file located in the 15 | `neural_wbc/data/data` directory. The function raises a `FileNotFoundError` if the specified file 16 | does not exist. 17 | 18 | ### Example 19 | ```python 20 | from neural_wbc.data import get_data_path 21 | 22 | try: 23 | data_file_path = get_data_path('mujoco/models/h1.xml') 24 | print(f"Data file path: {data_file_path}") 25 | except FileNotFoundError as e: 26 | print(e) 27 | ``` 28 | 29 | ## Data Files 30 | The `data` directory contains the following subdirectories: 31 | * `motions`: Contains motion datasets. 32 | * `motion_lib`: Contains skeleton files used for loading motion datasets. 33 | * `mujoco`: Contains files specific to the MuJoCo simulator, such as models and assets. 34 | * `policy`: Contains policy checkpoints. 35 | 36 | 37 | ## Unit Tests 38 | The tests are located in the `tests` directory and can be run with the `unittest` module. 39 | 40 | ```bash 41 | cd neural_wbc/data 42 | ${ISAACLAB_PATH:?}/isaaclab.sh -p -m unittest 43 | ``` 44 | -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/left_ankle_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/left_ankle_link.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/left_elbow_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/left_elbow_link.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/left_hip_pitch_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/left_hip_pitch_link.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/left_hip_roll_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/left_hip_roll_link.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/left_hip_yaw_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/left_hip_yaw_link.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/left_knee_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/left_knee_link.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/left_shoulder_pitch_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/left_shoulder_pitch_link.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/left_shoulder_roll_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/left_shoulder_roll_link.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/left_shoulder_yaw_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/left_shoulder_yaw_link.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/logo_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/logo_link.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/pelvis.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/pelvis.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/right_ankle_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/right_ankle_link.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/right_elbow_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/right_elbow_link.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/right_hip_pitch_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/right_hip_pitch_link.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/right_hip_roll_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/right_hip_roll_link.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/right_hip_yaw_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/right_hip_yaw_link.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/right_knee_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/right_knee_link.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/right_shoulder_pitch_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/right_shoulder_pitch_link.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/right_shoulder_roll_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/right_shoulder_roll_link.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/right_shoulder_yaw_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/right_shoulder_yaw_link.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/assets/torso_link.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/data/data/mujoco/models/assets/torso_link.stl -------------------------------------------------------------------------------- /neural_wbc/data/data/mujoco/models/scene.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 34 | 35 | 36 | -------------------------------------------------------------------------------- /neural_wbc/data/neural_wbc/data/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import os 18 | 19 | 20 | def get_data_path(rel_path: str) -> str: 21 | """ 22 | Get the absolute path to a data file located in the 'neural_wbc/data/data' directory. 23 | 24 | Args: 25 | rel_path (str): The relative path to the data file from the data directory. 26 | 27 | Returns: 28 | str: The absolute path to the data file. 29 | 30 | Raises: 31 | FileNotFoundError: If the specified file does not exist. 32 | """ 33 | file_dir = os.path.dirname(os.path.abspath(__file__)) 34 | data_path = os.path.join(file_dir, "../../data", rel_path) 35 | if not os.path.exists(data_path): 36 | raise FileNotFoundError(f"The file at '{rel_path}' does not exist in neural_wbc_data.") 37 | return data_path 38 | 39 | 40 | def get_default_teacher_policy() -> tuple[str, str]: 41 | """Get the default teacher policy path and checkpoint.""" 42 | return get_data_path("policy/h1:teacher"), "model_110000.pt" 43 | 44 | 45 | def get_default_omnih2o_student_policy() -> tuple[str, str]: 46 | """Get the default omnih2o student policy path and checkpoint.""" 47 | return get_data_path("policy/h1:student/omnih2o"), "model_7500.pt" 48 | 49 | 50 | def get_default_hover_student_policy() -> tuple[str, str]: 51 | """Get the default hover student policy path and checkpoint.""" 52 | return get_data_path("policy/h1:student/hover"), "model_51000.pt" 53 | -------------------------------------------------------------------------------- /neural_wbc/data/setup.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | from setuptools import find_namespace_packages, setup 18 | 19 | setup( 20 | name="neural_wbc_data", 21 | version="0.1.0", 22 | packages=find_namespace_packages(), 23 | description="Data and assets used for neural whole body control.", 24 | python_requires=">=3.10", 25 | ) 26 | -------------------------------------------------------------------------------- /neural_wbc/data/tests/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | -------------------------------------------------------------------------------- /neural_wbc/data/tests/test_get_file_path.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import os 18 | import pathlib 19 | import unittest 20 | 21 | from neural_wbc.data import get_data_path 22 | 23 | 24 | class TestGetDataPath(unittest.TestCase): 25 | def setUp(self): 26 | self.data_dir = pathlib.Path(get_data_path("")).resolve() 27 | self.existing_files = [str(f) for f in self.data_dir.glob("**/*") if f.is_file()] 28 | self.non_existing_file = "non_existing_file.txt" 29 | 30 | def test_existing_files(self): 31 | self.assertTrue(str(self.data_dir).endswith("neural_wbc/data/data")) 32 | 33 | for rel_path in self.existing_files: 34 | with self.subTest(rel_path=rel_path): 35 | data_path = get_data_path(rel_path) 36 | self.assertTrue(os.path.isfile(data_path), f"File should exist: {data_path}") 37 | 38 | def test_non_existing_file(self): 39 | with self.assertRaises(FileNotFoundError): 40 | get_data_path(self.non_existing_file) 41 | 42 | 43 | if __name__ == "__main__": 44 | unittest.main() 45 | -------------------------------------------------------------------------------- /neural_wbc/hw_wrappers/docs/push_recovery.gif: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:f61b33eeab39ad9b97b96d1140dc34fd9f36653710f721c6a3468a01251b221d 3 | size 14250580 4 | -------------------------------------------------------------------------------- /neural_wbc/hw_wrappers/docs/unitree_h1_setup.png: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:467354633d7902146cf1d0dff5f1d005f31e099d10923dc33056f5671b8cff9a 3 | size 175647 4 | -------------------------------------------------------------------------------- /neural_wbc/hw_wrappers/hw_wrappers/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | # Needed to have the robot registry 'activated' 17 | from hw_wrappers.unitree_h1 import UnitreeH1 # noqa 18 | -------------------------------------------------------------------------------- /neural_wbc/hw_wrappers/setup.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | from setuptools import find_packages, setup 18 | 19 | setup( 20 | name="hw_wrappers", 21 | version="0.1.0", 22 | packages=find_packages(), 23 | description="HW wrappers for different robots used in Sim-to-Real validation", 24 | python_requires=">=3.10", 25 | install_requires=[ 26 | "mujoco", 27 | "mujoco-python-viewer", 28 | "numpy", 29 | "torch", 30 | "tyro", 31 | "neural_wbc_core", 32 | "neural_wbc_mujoco_wrapper", 33 | ], 34 | ) 35 | -------------------------------------------------------------------------------- /neural_wbc/hw_wrappers/tests/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | -------------------------------------------------------------------------------- /neural_wbc/inference_env/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/inference_env/__init__.py -------------------------------------------------------------------------------- /neural_wbc/inference_env/inference_env/README.md: -------------------------------------------------------------------------------- 1 | # Neural WBC - Inference environment 2 | 3 | High level inference environment used for evaluation and deployment. 4 | 5 | ## Installation 6 | 7 | Install it with 8 | 9 | ```bash 10 | ${ISAACLAB_PATH:?}/isaaclab.sh -p -m pip install -e neural_wbc/inference_env 11 | ``` 12 | 13 | ## Usage 14 | 15 | - First ensure that all dependencies are installed using the common script 16 | [./install_deps.sh](../../install_deps.sh) 17 | 18 | - Run the simple viewer with: 19 | 20 | ```sh 21 | ${ISAACLAB_PATH}/isaaclab.sh -p neural_wbc/inference_env/scripts/mujoco_viewer_player.py 22 | ``` 23 | 24 | The viewer is paused on start by default, press `SPACE` key to start simulation or `RIGHT` arrow key 25 | to step forward once. By default this script will load the UniTree H1 robot model scene from 26 | [data/mujoco/models](../data/data/mujoco/models/scene.xml). 27 | 28 | Additional scripts are provides for running comprehensive evaluation in [scripts](./scripts/) 29 | directory 30 | 31 | ## Unit Tests 32 | 33 | The tests are located in the `tests` directory and can be run with the `unittest` module. 34 | 35 | ```bash 36 | cd neural_wbc/inference_env 37 | ${ISAACLAB_PATH:?}/isaaclab.sh -p -m unittest 38 | ``` 39 | -------------------------------------------------------------------------------- /neural_wbc/inference_env/inference_env/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/inference_env/inference_env/__init__.py -------------------------------------------------------------------------------- /neural_wbc/inference_env/inference_env/neural_wbc_env_cfg.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | from dataclasses import dataclass 18 | from typing import Literal, Protocol 19 | 20 | from neural_wbc.core.modes import NeuralWBCModes 21 | from neural_wbc.core.reference_motion import ReferenceMotionManagerCfg 22 | 23 | 24 | @dataclass 25 | class NeuralWBCEnvCfg(Protocol): 26 | model_xml_path: str 27 | enable_viewer: bool = False 28 | mode = NeuralWBCModes.DISTILL_TEST 29 | 30 | decimation = 4 31 | dt = 0.005 32 | max_episode_length_s = 3600 33 | action_scale = 0.25 34 | ctrl_delay_step_range = [0, 0] 35 | default_rfi_lim = 0.0 36 | 37 | # Distillation parameters: 38 | single_history_dim: int = 63 39 | observation_history_length: int = 25 40 | 41 | reference_motion_cfg = ReferenceMotionManagerCfg() 42 | 43 | # Termination conditions 44 | gravity_x_threshold = 0.7 45 | gravity_y_threshold = 0.7 46 | max_ref_motion_dist = 0.5 47 | 48 | # control type: the action type from the policy 49 | # "Pos": target joint pos, "Torque": joint torques 50 | # 'None': by passes control and returns same value 51 | control_type: Literal["Pos", "Torque", "None"] = "Pos" 52 | -------------------------------------------------------------------------------- /neural_wbc/inference_env/inference_env/utils.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | import argparse 17 | 18 | 19 | def get_player_args(description: str): 20 | parser = argparse.ArgumentParser(description=description) 21 | parser.add_argument("--num_envs", type=int, default=1, help="Number of environments to simulate.") 22 | parser.add_argument("--headless", action="store_true", help="Whether to show mujoco visualizer.") 23 | parser.add_argument("--reference_motion_path", type=str, default=None, help="Path to the reference motion dataset.") 24 | parser.add_argument("--student_path", type=str, default=None, help="The path for the student policy.") 25 | parser.add_argument("--student_checkpoint", type=str, default=None, help="The exact checkpoint.") 26 | parser.add_argument("--robot", type=str, default="mujoco_robot", help="The name of registered robot class.") 27 | parser.add_argument("--max_iterations", type=int, default=100, help="Number steps to run the demo example.") 28 | parser.add_argument( 29 | "--env_config_overwrite", 30 | type=str, 31 | default=None, 32 | help="Path to yaml file with overwriting configuration values.", 33 | ) 34 | return parser 35 | -------------------------------------------------------------------------------- /neural_wbc/inference_env/scripts/eval.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import pprint 18 | import yaml 19 | 20 | from inference_env.deployment_player import DeploymentPlayer 21 | from inference_env.neural_wbc_env_cfg_h1 import NeuralWBCEnvCfgH1 22 | from inference_env.utils import get_player_args 23 | 24 | from neural_wbc.core.evaluator import Evaluator 25 | from neural_wbc.data import get_data_path 26 | 27 | # add argparse arguments 28 | parser = get_player_args(description="Evaluates motion tracking policy and collects metrics in MuJoCo.") 29 | parser.add_argument("--metrics_path", type=str, default=None, help="Path to store metrics in.") 30 | args_cli = parser.parse_args() 31 | 32 | 33 | def main(): 34 | custom_config = None 35 | if args_cli.env_config_overwrite is not None: 36 | with open(args_cli.env_config_overwrite) as fh: 37 | custom_config = yaml.safe_load(fh) 38 | print("[INFO]: Using custom configuration:") 39 | pprint.pprint(custom_config) 40 | 41 | # Initialize the player with the updated properties 42 | player = DeploymentPlayer( 43 | args_cli=args_cli, 44 | env_cfg=NeuralWBCEnvCfgH1(model_xml_path=get_data_path("mujoco/models/scene.xml")), 45 | custom_config=custom_config, 46 | ) 47 | evaluator = Evaluator(env_wrapper=player.env, metrics_path=args_cli.metrics_path) 48 | 49 | should_stop = False 50 | while not should_stop: 51 | _, obs, dones, extras = player.play_once() 52 | 53 | reset_env = evaluator.collect(dones=dones, info=extras) 54 | if reset_env and not evaluator.is_evaluation_complete(): 55 | evaluator.forward_motion_samples() 56 | _ = player.reset() 57 | 58 | should_stop = evaluator.is_evaluation_complete() 59 | 60 | evaluator.conclude() 61 | 62 | 63 | if __name__ == "__main__": 64 | main() 65 | -------------------------------------------------------------------------------- /neural_wbc/inference_env/scripts/mujoco_viewer_player.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import pprint 18 | import time 19 | import torch 20 | import yaml 21 | 22 | from inference_env.deployment_player import DeploymentPlayer 23 | from inference_env.neural_wbc_env_cfg_h1 import NeuralWBCEnvCfgH1 24 | from inference_env.utils import get_player_args 25 | 26 | from neural_wbc.data import get_data_path 27 | 28 | # add argparse arguments 29 | parser = get_player_args(description="Evaluates motion tracking policy and collects metrics in MuJoCo.") 30 | args_cli = parser.parse_args() 31 | 32 | CTRL_FREQ = 200.0 33 | 34 | 35 | def prep_joint_pos_cmd(env, joint_pos_cmd: dict[str, float] | None = None) -> torch.Tensor: 36 | full_joint_pos_cmd = env.robot.joint_positions 37 | if joint_pos_cmd is not None: 38 | with torch.inference_mode(): 39 | for joint_name, joint_cmd in joint_pos_cmd.items(): 40 | joint_id = env.robot.get_joint_ids(joint_names=[joint_name])[joint_name] 41 | full_joint_pos_cmd[:, joint_id] = joint_cmd 42 | return full_joint_pos_cmd 43 | 44 | 45 | def main(): 46 | custom_config = None 47 | if args_cli.env_config_overwrite is not None: 48 | with open(args_cli.env_config_overwrite) as fh: 49 | custom_config = yaml.safe_load(fh) 50 | print("[INFO]: Using custom configuration:") 51 | pprint.pprint(custom_config) 52 | 53 | # Delta joint position command of the robot. These will be applied based on the PD positional controller 54 | # defined in control.py. 55 | joint_pos_cmd = { 56 | "right_elbow": 2.0, 57 | "left_elbow": -2.0, 58 | "torso": 2.0, 59 | "right_knee": 2.0, 60 | "left_knee": 2.0, 61 | "right_ankle": 1.0, 62 | "left_ankle": 1.0, 63 | } 64 | 65 | player = DeploymentPlayer( 66 | args_cli=args_cli, 67 | env_cfg=NeuralWBCEnvCfgH1(model_xml_path=get_data_path("mujoco/models/scene.xml")), 68 | custom_config=custom_config, 69 | demo_mode=True, 70 | ) 71 | for i in range(args_cli.max_iterations): 72 | _, obs, dones, extras = player.play_once(prep_joint_pos_cmd(player.env, joint_pos_cmd)) 73 | time.sleep(1.0 / CTRL_FREQ) 74 | 75 | 76 | if __name__ == "__main__": 77 | # Run the main execution 78 | main() 79 | -------------------------------------------------------------------------------- /neural_wbc/inference_env/scripts/s2r_player.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import pprint 18 | import time 19 | import torch 20 | import yaml 21 | 22 | from hw_wrappers.unitree_h1 import UnitreeH1 # noqa 23 | from inference_env.deployment_player import DeploymentPlayer 24 | from inference_env.neural_wbc_env_cfg_real_h1 import NeuralWBCEnvCfgRealH1 25 | from inference_env.utils import get_player_args 26 | 27 | from neural_wbc.data import get_data_path 28 | 29 | parser = get_player_args(description="Basic deployment player for running HOVER policy on real robots.") 30 | args_cli = parser.parse_args() 31 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 32 | 33 | 34 | def main(): 35 | custom_config = None 36 | if args_cli.env_config_overwrite is not None: 37 | with open(args_cli.env_config_overwrite) as fh: 38 | custom_config = yaml.safe_load(fh) 39 | print("[INFO]: Using custom configuration:") 40 | pprint.pprint(custom_config) 41 | 42 | env_cfg = NeuralWBCEnvCfgRealH1(model_xml_path=get_data_path("mujoco/models/scene.xml")) 43 | 44 | player = DeploymentPlayer( 45 | args_cli=args_cli, 46 | env_cfg=env_cfg, 47 | custom_config=custom_config, 48 | ) 49 | 50 | inference_time = env_cfg.decimation * env_cfg.dt 51 | print("Deploying policy on real robot.") 52 | start_time = time.time() 53 | elapsed_time = 0.0 54 | for i in range(args_cli.max_iterations): 55 | print( 56 | "\rActual loop frequency: {:.2f} Hz | Update time: {:.2f}s".format( 57 | 1 / (time.time() - start_time), elapsed_time 58 | ), 59 | end="", 60 | flush=True, 61 | ) 62 | start_time = time.time() 63 | player.play_once() 64 | elapsed_time = time.time() - start_time 65 | remaining_time = inference_time - elapsed_time 66 | if remaining_time > 0.0: 67 | time.sleep(remaining_time) 68 | 69 | 70 | if __name__ == "__main__": 71 | main() 72 | -------------------------------------------------------------------------------- /neural_wbc/inference_env/setup.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | from setuptools import find_packages, setup 18 | 19 | setup( 20 | name="neural_wbc_inference_env", 21 | version="0.1.0", 22 | packages=find_packages(), 23 | description="Environment used for inference e.g. at deployment and evaluation", 24 | python_requires=">=3.10", 25 | install_requires=[ 26 | "mujoco", 27 | "mujoco-python-viewer", 28 | "neural_wbc_mujoco_wrapper", 29 | "neural_wbc_student_policy", 30 | "numpy", 31 | "torch", 32 | "tyro", 33 | "neural_wbc_core", 34 | ], 35 | ) 36 | -------------------------------------------------------------------------------- /neural_wbc/inference_env/tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/HOVER/c14bb7e56a691da7f1f906997d966567ef138108/neural_wbc/inference_env/tests/__init__.py -------------------------------------------------------------------------------- /neural_wbc/inference_env/tests/test_evaluation_pipeline.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import json 18 | import os 19 | import pprint 20 | import subprocess 21 | import sys 22 | import tempfile 23 | 24 | import neural_wbc.data 25 | from neural_wbc.core import util 26 | 27 | 28 | def get_json_files(directory): 29 | json_files = [] 30 | for filename in os.listdir(directory): 31 | if filename.endswith(".json"): 32 | json_files.append(os.path.join(directory, filename)) 33 | return json_files 34 | 35 | 36 | def compare_generated_metrics(generated_metrics): 37 | # Make sure we always have 100% success rate with the stable_punch dataset and default model 38 | util.assert_equal(generated_metrics["success_rate"], 1.0, "Success rate") 39 | # Make sure all and success metrics are the same 40 | util.assert_equal(generated_metrics["all"], generated_metrics["success"], "All metrics vs Success metrics") 41 | # Make sure we played the correct number of reference motions 42 | util.assert_equal(generated_metrics["num_motions"], 29, "Number of motions") 43 | 44 | 45 | if __name__ == "__main__": 46 | metrics_path = tempfile.TemporaryDirectory() 47 | policy_path, checkpoint = neural_wbc.data.get_default_omnih2o_student_policy() 48 | eval_pipeline_command = [ 49 | "${ISAACLAB_PATH}/isaaclab.sh -p neural_wbc/inference_env/scripts/eval.py " 50 | "--headless " 51 | f"--student_path={policy_path} " 52 | f"--student_checkpoint={checkpoint} " 53 | f"--metrics_path={metrics_path.name} " 54 | ] 55 | return_value = subprocess.call(eval_pipeline_command, shell=True) 56 | json_files = get_json_files(metrics_path.name) 57 | assert len(json_files) == 1 58 | with open(json_files[0], encoding="utf-8") as fh: 59 | generated_metrics = json.load(fh) 60 | print("Generated metrics:") 61 | pprint.pprint(generated_metrics) 62 | compare_generated_metrics(generated_metrics) 63 | sys.exit(return_value) 64 | -------------------------------------------------------------------------------- /neural_wbc/inference_env/tests/test_neural_wbc_env.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import torch 18 | import unittest 19 | 20 | from inference_env.neural_wbc_env import NeuralWBCEnv 21 | from inference_env.neural_wbc_env_cfg_h1 import NeuralWBCEnvCfgH1 22 | 23 | from neural_wbc.data import get_data_path 24 | 25 | 26 | class TestNeuralWBCEnv(unittest.TestCase): 27 | def setUp(self): 28 | # create environment configuration 29 | self.env_cfg = NeuralWBCEnvCfgH1(model_xml_path=get_data_path("mujoco/models/h1.xml")) 30 | 31 | # setup RL environment 32 | self.device = torch.device("cpu") 33 | self.env = NeuralWBCEnv(cfg=self.env_cfg, device=self.device) 34 | 35 | self.num_instances = self.env.robot.num_instances 36 | self.num_controls = self.env.robot.num_controls 37 | 38 | def test_reset(self): 39 | # Reset without specifying reset positions 40 | joint_pos_before_reset = self.env.robot.joint_positions.clone() 41 | self.env.reset() 42 | joint_pos_after_reset = self.env.robot.joint_positions.clone() 43 | self.assertTrue(torch.equal(joint_pos_before_reset, joint_pos_after_reset)) 44 | 45 | # Reset WBCMujoco with specified reset positions 46 | joint_pos_before_reset = self.env.robot.joint_positions.clone() 47 | reset_pos = ( 48 | torch.from_numpy(self.env.robot.internal_sim.data.qpos.copy()) 49 | .to(dtype=torch.float, device=self.device) 50 | .expand(1, -1) 51 | ) 52 | reset_pos[:, : self.num_controls] = torch.rand(self.num_instances, self.num_controls, device=self.device) 53 | self.env.robot.reset(qpos=reset_pos) 54 | joint_pos_after_reset = self.env.robot.joint_positions.clone() 55 | self.assertFalse(torch.equal(joint_pos_before_reset, joint_pos_after_reset)) 56 | 57 | def test_step(self): 58 | joint_pos_before_step = self.env.robot.joint_positions.clone() 59 | actions = torch.zeros(self.num_instances, self.num_controls, device=self.device) 60 | self.env.step(actions=actions) 61 | joint_pos_after_step = self.env.robot.joint_positions.clone() 62 | self.assertFalse(torch.equal(joint_pos_before_step, joint_pos_after_step)) 63 | 64 | def test_action_delay(self): 65 | # Set the action delay. 66 | self.env_cfg.ctrl_delay_step_range = [1, 1] 67 | self.env = NeuralWBCEnv(cfg=self.env_cfg, device=self.device) 68 | actions = torch.rand(self.num_instances, self.num_controls, device=self.device) 69 | self.env.step(actions=actions) 70 | # First action should be zero. 71 | self.assertTrue(torch.equal(self.env.actions, torch.zeros_like(actions))) 72 | 73 | # Now should be equal. 74 | self.env.step(actions=actions) 75 | self.assertTrue(torch.equal(self.env.actions, actions)) 76 | 77 | def test_effort_limit(self): 78 | # Disable the action delay to test action limit. 79 | self.env_cfg.ctrl_delay_step_range = [0, 0] 80 | self.env = NeuralWBCEnv(cfg=self.env_cfg, device=self.device) 81 | # Create a very large value action and make sure it is properly clamped 82 | actions = 1000 * torch.ones(self.num_instances, self.num_controls, device=self.device) 83 | self.env.step(actions=actions) 84 | 85 | effort_limits = torch.zeros_like(actions) 86 | for joint_name, effort_limit in self.env_cfg.effort_limit.items(): 87 | joint_id = self.env.robot.get_joint_ids([joint_name])[joint_name] 88 | effort_limits[:, joint_id] = effort_limit 89 | 90 | self.assertTrue(torch.equal(self.env._processed_action, effort_limits)) 91 | 92 | 93 | if __name__ == "__main__": 94 | unittest.main() 95 | -------------------------------------------------------------------------------- /neural_wbc/isaac_lab_wrapper/neural_wbc/isaac_lab_wrapper/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import gymnasium as gym 18 | 19 | from . import neural_wbc_env_cfg_h1 20 | 21 | ## 22 | # Register Gym environments. 23 | ## 24 | 25 | gym.register( 26 | id="NeuralWBC-H1-v0", 27 | entry_point="neural_wbc.isaac_lab_wrapper.neural_wbc_env:NeuralWBCEnv", 28 | disable_env_checker=True, 29 | kwargs={ 30 | "env_cfg_entry_point": neural_wbc_env_cfg_h1.NeuralWBCEnvCfgH1, 31 | }, 32 | ) 33 | -------------------------------------------------------------------------------- /neural_wbc/isaac_lab_wrapper/neural_wbc/isaac_lab_wrapper/body_state.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import torch 18 | 19 | from neural_wbc.core.body_state import BodyState 20 | 21 | from isaaclab.assets import ArticulationData 22 | 23 | 24 | def build_body_state( 25 | data: ArticulationData, 26 | root_id: int, 27 | body_ids: list[int] | None = None, 28 | joint_ids: list[int] | None = None, 29 | extend_body_pos: torch.Tensor | None = None, 30 | extend_body_parent_ids: list[int] | None = None, 31 | ) -> BodyState: 32 | """Creates a body state from Isaac Lab articulation data. 33 | 34 | Args: 35 | data (ArticulationData): Articulation data containing robot's body and joint states. 36 | body_ids (list[int] | None, optional): The desired order of bodies. If not, the order in body states is preserved. 37 | Defaults to None. 38 | joint_ids (list[int] | None, optional): The desired order of joint. If not, the order in joint states is preserved. 39 | Defaults to None. 40 | extend_body_parent_ids (list[int] | None, optional): ID of the bodies to extend. Defaults to None. 41 | extend_body_pos (torch.Tensor | None, optional): Position of the extended bodies from their parent bodies. Defaults to None. 42 | 43 | Returns: 44 | BodyState: The constructed BodyState object containing the reordered and extended body and joint states. 45 | """ 46 | if body_ids is None: 47 | num_bodies = data.body_pos_w.shape[1] 48 | body_ids = list(range(0, num_bodies)) 49 | 50 | if joint_ids is None: 51 | num_joints = data.joint_pos.shape[1] 52 | joint_ids = list(range(0, num_joints)) 53 | 54 | body_state = BodyState( 55 | body_pos=data.body_pos_w[:, body_ids, :], 56 | body_rot=data.body_quat_w[:, body_ids, :], 57 | body_lin_vel=data.body_lin_vel_w[:, body_ids, :], 58 | body_ang_vel=data.body_ang_vel_w[:, body_ids, :], 59 | joint_pos=data.joint_pos[:, joint_ids], 60 | joint_vel=data.joint_vel[:, joint_ids], 61 | root_id=root_id, 62 | ) 63 | 64 | if (extend_body_pos is not None) and (extend_body_parent_ids is not None): 65 | body_state.extend_body_states(extend_body_pos=extend_body_pos, extend_body_parent_ids=extend_body_parent_ids) 66 | 67 | return body_state 68 | -------------------------------------------------------------------------------- /neural_wbc/isaac_lab_wrapper/neural_wbc/isaac_lab_wrapper/control.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | """Common functions that can be used to enable different controller. 18 | 19 | Controller will compute the torque that will be applied to the joint. The implemented controller are PD controllers that track 20 | desire pos/vel and torques. By using different controller, the trained RL policy will be considered outputting different 21 | action. For example, when using Pos Controller, the trained policy is considered outputting desire pos. 22 | 23 | """ 24 | 25 | from __future__ import annotations 26 | 27 | import torch 28 | from typing import TYPE_CHECKING, Literal 29 | 30 | from isaaclab.assets import Articulation 31 | from isaaclab.managers import SceneEntityCfg 32 | 33 | if TYPE_CHECKING: 34 | from isaaclab.envs import DirectRLEnv 35 | 36 | 37 | def position_pd_control( 38 | env: DirectRLEnv, actions_scaled: torch.Tensor, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), joint_ids=None 39 | ): 40 | asset: Articulation = env.scene.articulations[asset_cfg.name] 41 | default_joint_pos = asset.data.default_joint_pos 42 | joint_pos = asset.data.joint_pos 43 | joint_vel = asset.data.joint_vel 44 | 45 | if joint_ids: 46 | default_joint_pos = default_joint_pos[:, joint_ids] 47 | joint_pos = joint_pos[:, joint_ids] 48 | joint_vel = joint_vel[:, joint_ids] 49 | 50 | torques = ( 51 | env.kp_scale * env._p_gains * (actions_scaled + default_joint_pos - joint_pos) 52 | - env.kd_scale * env._d_gains * joint_vel 53 | ) 54 | return torques 55 | 56 | 57 | def velocity_pd_control( 58 | env: DirectRLEnv, actions_scaled: torch.Tensor, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), joint_ids=None 59 | ): 60 | asset: Articulation = env.scene.articulations[asset_cfg.name] 61 | joint_vel = asset.data.joint_vel 62 | joint_acc = asset.data.joint_acc 63 | 64 | if joint_ids: 65 | joint_vel = joint_vel[:, joint_ids] 66 | joint_acc = joint_acc[:, joint_ids] 67 | 68 | torques = env.kp_scale * env._p_gains * (actions_scaled - joint_vel) - env.kd_scale * env._d_gains * joint_acc 69 | 70 | return torques 71 | 72 | 73 | def torque_control( 74 | env: DirectRLEnv, actions_scaled: torch.Tensor, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), joint_ids=None 75 | ): 76 | torques = actions_scaled 77 | 78 | return torques 79 | 80 | 81 | def resolve_control_fn( 82 | control_type: Literal["Pos", "Vel", "Torque"] = "Pos", 83 | ): 84 | control_fn = position_pd_control 85 | 86 | if control_type == "Pos": 87 | print("[INFO]: Setting up pos control") 88 | control_fn = position_pd_control 89 | elif control_type == "Vel": 90 | print("[INFO]: Setting up vel control") 91 | control_fn = velocity_pd_control 92 | elif control_type == "Torque": 93 | print("[INFO]: Setting up torque control") 94 | control_fn = torque_control 95 | else: 96 | raise ValueError(f"Unrecognized control type {control_type}") 97 | 98 | return control_fn 99 | -------------------------------------------------------------------------------- /neural_wbc/isaac_lab_wrapper/neural_wbc/isaac_lab_wrapper/events/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | from .event_cfg import * # noqa 18 | -------------------------------------------------------------------------------- /neural_wbc/isaac_lab_wrapper/neural_wbc/isaac_lab_wrapper/observations.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | from __future__ import annotations 18 | 19 | import torch 20 | from typing import TYPE_CHECKING 21 | 22 | from isaaclab.assets import Articulation 23 | from isaaclab.managers import SceneEntityCfg 24 | 25 | if TYPE_CHECKING: 26 | from .neural_wbc_env import NeuralWBCEnv 27 | 28 | from neural_wbc.core import ReferenceMotionState 29 | from neural_wbc.core.body_state import BodyState 30 | from neural_wbc.core.observations import compute_student_observations, compute_teacher_observations 31 | 32 | 33 | def compute_observations( 34 | env: NeuralWBCEnv, 35 | body_state: BodyState, 36 | ref_motion_state: ReferenceMotionState, 37 | asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), 38 | ): 39 | asset: Articulation = env.scene.articulations[asset_cfg.name] 40 | 41 | obs_dict = {} 42 | 43 | # First collect teacher policy observations. 44 | teacher_obs, teacher_obs_dict = compute_teacher_observations( 45 | body_state=body_state, 46 | ref_motion_state=ref_motion_state, 47 | tracked_body_ids=env._tracked_body_ids, 48 | ref_episodic_offset=env.ref_episodic_offset, 49 | last_actions=env.actions, 50 | ) 51 | obs_dict.update(teacher_obs_dict) 52 | obs_dict["teacher_policy"] = teacher_obs 53 | 54 | # Then the privileged observations. 55 | privileged_obs, privileged_obs_dict = compute_privileged_observations(env=env, asset=asset) 56 | obs_dict.update(privileged_obs_dict) 57 | obs_dict["critic"] = torch.cat([teacher_obs, privileged_obs], dim=1) 58 | 59 | # If we are in a distill mode, add student observations. 60 | if env.cfg.mode.is_distill_mode(): 61 | base_id = env._body_names.index(env.base_name) 62 | student_obs, student_obs_dict = compute_student_observations( 63 | base_id=base_id, 64 | body_state=body_state, 65 | ref_motion_state=ref_motion_state, 66 | projected_gravity=asset.data.projected_gravity_b, 67 | last_actions=env.actions, 68 | history=env.history.entries, 69 | mask=env.mask, 70 | ref_episodic_offset=env.ref_episodic_offset, 71 | ) 72 | 73 | obs_dict.update(student_obs_dict) 74 | obs_dict["student_policy"] = student_obs 75 | 76 | return obs_dict 77 | 78 | 79 | def compute_privileged_observations(env: NeuralWBCEnv, asset: Articulation): 80 | contact_forces = env.contact_sensor.data.net_forces_w[:, env.feet_ids, :] 81 | 82 | privileged_obs_dict = { 83 | "base_com_bias": env.base_com_bias.to(env.device), 84 | "ground_friction_values": asset.data.joint_friction[:, env.feet_ids], 85 | "body_mass_scale": env.body_mass_scale, 86 | "kp_scale": env.kp_scale, 87 | "kd_scale": env.kd_scale, 88 | "rfi_lim_scale": env.rfi_lim / env.cfg.default_rfi_lim, 89 | "contact_forces": contact_forces.reshape(contact_forces.shape[0], -1), 90 | "recovery_counters": torch.clamp_max(env.recovery_counters.unsqueeze(1), 1), 91 | } 92 | privileged_obs = torch.cat( 93 | [tensor for tensor in privileged_obs_dict.values()], 94 | dim=-1, 95 | ) 96 | return privileged_obs, privileged_obs_dict 97 | -------------------------------------------------------------------------------- /neural_wbc/isaac_lab_wrapper/neural_wbc/isaac_lab_wrapper/rewards/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | from .reward_cfg import NeuralWBCRewardCfg # noqa: F401, F403 18 | from .rewards import NeuralWBCRewards # noqa: F401, F403 19 | -------------------------------------------------------------------------------- /neural_wbc/isaac_lab_wrapper/neural_wbc/isaac_lab_wrapper/rewards/reward_cfg.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | from __future__ import annotations 18 | 19 | from isaaclab.utils import configclass 20 | 21 | 22 | @configclass 23 | class NeuralWBCRewardCfg: 24 | # Reward and penalty scales 25 | scales = { 26 | "reward_track_joint_positions": 32.0, 27 | "reward_track_joint_velocities": 16.0, 28 | "reward_track_body_velocities": 8.0, 29 | "reward_track_body_angular_velocities": 8.0, 30 | "reward_track_body_position_extended": 30.0, 31 | "reward_track_body_position_vr_key_points": 50.0, 32 | "penalize_torques": -0.0001, 33 | "penalize_by_torque_limits": -2, 34 | "penalize_joint_accelerations": -0.000011, 35 | "penalize_joint_velocities": -0.004, 36 | "penalize_lower_body_action_changes": -3.0, 37 | "penalize_upper_body_action_changes": -0.625, 38 | "penalize_by_joint_pos_limits": -125.0, 39 | "penalize_by_joint_velocity_limits": -50.0, 40 | "penalize_early_termination": -250.0, 41 | "penalize_feet_contact_forces": -0.75, 42 | "penalize_stumble": -1000.0, 43 | "penalize_slippage": -37.5, 44 | "penalize_feet_orientation": -62.5, 45 | "penalize_feet_air_time": 1000.0, 46 | "penalize_both_feet_in_air": -200.0, 47 | "penalize_orientation": -200.0, 48 | "penalize_max_feet_height_before_contact": -2500.0, 49 | } 50 | 51 | # Sigmas for exponential terms 52 | body_pos_lower_body_sigma = 0.5 53 | body_pos_upper_body_sigma = 0.03 54 | body_pos_vr_key_points_sigma = 0.03 55 | body_rot_sigma = 0.1 56 | body_vel_sigma = 10 57 | body_ang_vel_sigma = 10 58 | joint_pos_sigma = 0.5 59 | joint_vel_sigma = 1.0 60 | 61 | # Weights for weighted sums 62 | body_pos_lower_body_weight = 0.5 63 | body_pos_upper_body_weight = 1.0 64 | 65 | # Limits 66 | torque_limits_scale = 0.85 67 | # The order here follows the order in cfg.joint_names 68 | torque_limits = [ 69 | 200.0, 70 | 200.0, 71 | 200.0, 72 | 300.0, 73 | 40.0, 74 | 200.0, 75 | 200.0, 76 | 200.0, 77 | 300.0, 78 | 40.0, 79 | 200.0, 80 | 40.0, 81 | 40.0, 82 | 18.0, 83 | 18.0, 84 | 40.0, 85 | 40.0, 86 | 18.0, 87 | 18.0, 88 | ] 89 | # Joint pos limits, in the form of (lower_limit, upper_limit) 90 | joint_pos_limits = [ 91 | (-0.43, 0.43), 92 | (-0.43, 0.43), 93 | (-1.57, 1.57), 94 | (-0.26, 2.05), 95 | (-0.87, 0.52), 96 | (-0.43, 0.43), 97 | (-0.43, 0.43), 98 | (-1.57, 1.57), 99 | (-0.26, 2.05), 100 | (-0.87, 0.52), 101 | (-2.35, 2.35), 102 | (-2.87, 2.87), 103 | (-0.34, 3.11), 104 | (-1.3, 4.45), 105 | (-1.25, 2.61), 106 | (-2.87, 2.87), 107 | (-3.11, 0.34), 108 | (-4.45, 1.3), 109 | (-1.25, 2.61), 110 | ] 111 | joint_vel_limits_scale = 0.85 112 | joint_vel_limits = [ 113 | 23.0, 114 | 23.0, 115 | 23.0, 116 | 14.0, 117 | 9.0, 118 | 23.0, 119 | 23.0, 120 | 23.0, 121 | 14.0, 122 | 9.0, 123 | 23.0, 124 | 9.0, 125 | 9.0, 126 | 20.0, 127 | 20.0, 128 | 9.0, 129 | 9.0, 130 | 20.0, 131 | 20.0, 132 | ] 133 | max_contact_force = 500.0 134 | max_feet_height_limit_before_contact = 0.25 135 | -------------------------------------------------------------------------------- /neural_wbc/isaac_lab_wrapper/neural_wbc/isaac_lab_wrapper/utils.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import torch 18 | 19 | 20 | def convert_tensors_and_slices_to_serializable(d): 21 | """Recursively convert torch tensors to lists and handle slice objects in a nested dictionary, including lists and tuples.""" 22 | if isinstance(d, dict): 23 | return {k: convert_tensors_and_slices_to_serializable(v) for k, v in d.items()} 24 | elif isinstance(d, torch.Tensor): 25 | return {"__type__": "tensor", "data": d.tolist()} 26 | elif isinstance(d, list): 27 | return [convert_tensors_and_slices_to_serializable(item) for item in d] 28 | elif isinstance(d, tuple): 29 | return {"__type__": "tuple", "data": tuple(convert_tensors_and_slices_to_serializable(item) for item in d)} 30 | elif isinstance(d, slice): 31 | return {"__type__": "slice", "start": d.start, "stop": d.stop, "step": d.step} 32 | else: 33 | return d 34 | 35 | 36 | def convert_serializable_to_tensors_and_slices(d): 37 | """Recursively convert lists back to torch tensors and dictionaries back to slice objects.""" 38 | if isinstance(d, dict): 39 | if "__type__" in d: 40 | if d["__type__"] == "tensor": 41 | return torch.tensor(d["data"]) 42 | elif d["__type__"] == "slice": 43 | return slice(d["start"], d["stop"], d["step"]) 44 | elif d["__type__"] == "tuple": 45 | return tuple(convert_serializable_to_tensors_and_slices(item) for item in d["data"]) 46 | else: 47 | return {k: convert_serializable_to_tensors_and_slices(v) for k, v in d.items()} 48 | elif isinstance(d, list): 49 | return [convert_serializable_to_tensors_and_slices(item) for item in d] 50 | else: 51 | return d 52 | -------------------------------------------------------------------------------- /neural_wbc/isaac_lab_wrapper/neural_wbc/isaac_lab_wrapper/visualization.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | import torch 17 | 18 | from neural_wbc.core import ReferenceMotionState 19 | 20 | import isaaclab.sim as sim_utils 21 | from isaaclab.markers import VisualizationMarkers 22 | from isaaclab.markers.config import DEFORMABLE_TARGET_MARKER_CFG 23 | 24 | 25 | class RefMotionVisualizer: 26 | """Visualizer of the reference motions.""" 27 | 28 | def __init__(self): 29 | self._initialized = False 30 | self._active_ref_motion_markers = None 31 | self._inactive_ref_motion_markers = None 32 | 33 | def _initialize_ref_motion_markers(self): 34 | print("Initialize markers for ref motion joints.") 35 | # Visualizer for the active reference body positions. 36 | active_marker_cfg = DEFORMABLE_TARGET_MARKER_CFG.copy() 37 | active_marker_cfg.markers["target"].radius = 0.05 38 | active_marker_cfg.markers["target"].visual_material = sim_utils.PreviewSurfaceCfg( 39 | diffuse_color=(0.0, 0.0, 1.0) 40 | ) # blue 41 | active_marker_cfg.prim_path = "/Visuals/Command/active_ref_motion" 42 | self._active_ref_motion_markers = VisualizationMarkers(active_marker_cfg) 43 | self._active_ref_motion_markers.set_visibility(True) 44 | 45 | # Visualizere for the inactive reference body positions. 46 | inactive_marker_cfg = DEFORMABLE_TARGET_MARKER_CFG.copy() 47 | inactive_marker_cfg.markers["target"].radius = 0.03 48 | inactive_marker_cfg.markers["target"].visual_material = sim_utils.PreviewSurfaceCfg( 49 | diffuse_color=(1.0, 0.0, 0.0) 50 | ) # red 51 | inactive_marker_cfg.prim_path = "/Visuals/Command/inactive_ref_motion" 52 | self._inactive_ref_motion_markers = VisualizationMarkers(inactive_marker_cfg) 53 | self._inactive_ref_motion_markers.set_visibility(True) 54 | 55 | self._initialized = True 56 | 57 | def visualize(self, ref_motion: ReferenceMotionState, mask: torch.Tensor | None = None): 58 | if not self._initialized: 59 | self._initialize_ref_motion_markers() 60 | 61 | # Split the active and inactive body references. 62 | if mask is None: 63 | active_body_pos = ref_motion.body_pos_extend.view(-1, 3) 64 | device = active_body_pos.device 65 | inactive_body_pos = torch.zeros(0, 3, device=device) 66 | else: 67 | num_bodies = ref_motion.body_pos_extend.shape[1] 68 | mask = mask[:, :num_bodies] 69 | # We reshape to still have the correct shape even if the mask is set to all false. 70 | # Else we will try to visualize non-existing bodies below. 71 | mask_flat = mask.flatten() 72 | active_body_pos = ref_motion.body_pos_extend.view(-1, 3)[mask_flat, :] 73 | inactive_body_pos = ref_motion.body_pos_extend.view(-1, 3)[~mask_flat, :] 74 | 75 | # Update the position of the visualization markers. 76 | if active_body_pos.shape[0] != 0: 77 | # Need to set visibility to True to ensure the markers are visible. 78 | self._active_ref_motion_markers.set_visibility(True) 79 | self._active_ref_motion_markers.visualize(active_body_pos) 80 | else: 81 | self._active_ref_motion_markers.set_visibility(False) 82 | if inactive_body_pos.shape[0] != 0: 83 | # Need to set visibility to True to ensure the markers are visible. 84 | self._inactive_ref_motion_markers.set_visibility(True) 85 | self._inactive_ref_motion_markers.visualize(inactive_body_pos) 86 | else: 87 | self._inactive_ref_motion_markers.set_visibility(False) 88 | -------------------------------------------------------------------------------- /neural_wbc/isaac_lab_wrapper/setup.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | from setuptools import find_namespace_packages, setup 17 | 18 | setup( 19 | name="neural_wbc_isaac_lab_wrapper", 20 | version="0.1.0", 21 | packages=find_namespace_packages(), 22 | description="IsaacLab wrapper for neural whole body control", 23 | python_requires=">=3.10", 24 | install_requires=[ 25 | "psutil", 26 | ], 27 | ) 28 | -------------------------------------------------------------------------------- /neural_wbc/isaac_lab_wrapper/tests/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | -------------------------------------------------------------------------------- /neural_wbc/isaac_lab_wrapper/tests/test_main.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import unittest 18 | 19 | from isaaclab.app import AppLauncher 20 | 21 | # launch the simulator 22 | app_launcher = AppLauncher(headless=True) 23 | simulation_app = app_launcher.app 24 | APP_IS_READY = True 25 | 26 | if __name__ == "__main__": 27 | unittest.main(exit=False) 28 | simulation_app.close() 29 | -------------------------------------------------------------------------------- /neural_wbc/isaac_lab_wrapper/tests/test_neural_wbc_env.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import torch 18 | import unittest 19 | 20 | from .test_main import APP_IS_READY 21 | 22 | if APP_IS_READY: 23 | from neural_wbc.core.modes import NeuralWBCModes 24 | from neural_wbc.isaac_lab_wrapper.neural_wbc_env import NeuralWBCEnv 25 | from neural_wbc.isaac_lab_wrapper.neural_wbc_env_cfg_h1 import NeuralWBCEnvCfgH1 26 | 27 | 28 | class TestNeuralWBCEnv(unittest.TestCase): 29 | def _check_observations(self, env: NeuralWBCEnv, obs: dict[str, torch.Tensor]): 30 | for key, value in obs.items(): 31 | # Make sure that all observations do not contain NaN values. 32 | self.assertFalse(torch.any(torch.isnan(value)), msg=f"Observation {key} has NaN values {value}.") 33 | # Ensure that the policy and critic observations have shapes matching the configuration. 34 | self.assertEqual(obs["teacher_policy"].shape, env.observation_space.shape) 35 | self.assertEqual(obs["critic"].shape, env.state_space.shape) 36 | 37 | def test_step(self): 38 | # create environment configuration 39 | env_cfg = NeuralWBCEnvCfgH1() 40 | env_cfg.scene.num_envs = 100 41 | env_cfg.scene.env_spacing = 20 42 | env_cfg.mode = NeuralWBCModes.DISTILL # Set modes to distill so that all observations are tested. 43 | 44 | # number of episodes to collect 45 | num_demos = 0 46 | 47 | # setup RL environment 48 | env = NeuralWBCEnv(cfg=env_cfg) 49 | env.reset() 50 | 51 | # sample actions 52 | joint_efforts = torch.zeros((env.num_envs, env.num_actions)) 53 | 54 | # simulate physics 55 | while num_demos < 10 * env_cfg.scene.num_envs: 56 | with torch.inference_mode(): 57 | 58 | # store signals 59 | # -- obs 60 | obs = env._get_observations() 61 | self._check_observations(env=env, obs=obs) 62 | 63 | # step the environment 64 | obs, _, terminated, truncated, _ = env.step(joint_efforts) 65 | num_demos += len(terminated) + len(truncated) 66 | self._check_observations(env=env, obs=obs) 67 | -------------------------------------------------------------------------------- /neural_wbc/mujoco_wrapper/README.md: -------------------------------------------------------------------------------- 1 | # Neural WBC - Mujoco Wrapper 2 | 3 | High level wrapper of Mujoco simulator for use in validating Whole Body Control (WBC) policies 4 | trained in IsaacLab. The wrapper provides convenient interfaces for accessing data and manipulating 5 | the underlying models in Mujoco, as well as an environment that allows comparison with IsaacLab. 6 | 7 | ```mermaid 8 | graph LR 9 | 10 | subgraph MJW[Mujoco Wrapper] 11 | style MJW fill:#76B900,stroke:#333,stroke-width:2px 12 | A[Mujoco] 13 | style A fill:#CDCDCD,stroke:#333,stroke-width:2px 14 | end 15 | 16 | MJW -->|Data & Model| C[ Mujoco Visualizer] 17 | MJW -->|Data| D[Evaluator] 18 | style D fill:#76B900,stroke:#333,stroke-width:2px 19 | style C fill:#CDCDCD,stroke:#333,stroke-width:2px 20 | 21 | E[IsaacLab] 22 | E[IsaacLab] -->|WBC Policy| MJW 23 | style E fill:#76B900,stroke:#333,stroke-width:2px 24 | 25 | RFM[Reference Motion Manager] 26 | style RFM fill:#76B900,stroke:#333,stroke-width:2px 27 | 28 | RFM -->| Reference Motions| MJW 29 | ``` 30 | 31 | ## Installation 32 | 33 | To use the `neural_wbc.mujoco_wrapper` package, install it with 34 | 35 | ```bash 36 | ${ISAACLAB_PATH:?}/isaaclab.sh -p -m pip install -e neural_wbc/mujoco_wrapper 37 | ``` 38 | 39 | Also install the inference environment with 40 | 41 | ```bash 42 | ${ISAACLAB_PATH:?}/isaaclab.sh -p -m pip install -e neural_wbc/inference_env 43 | ``` 44 | 45 | ## Usage 46 | 47 | - First ensure that all dependencies are installed using the common script 48 | [./install_deps.sh](../../install_deps.sh) 49 | 50 | - Run the simple viewer with: 51 | 52 | ```sh 53 | ${ISAACLAB_PATH}/isaaclab.sh -p neural_wbc/inference_env/scripts/mujoco_viewer_player.py 54 | ``` 55 | 56 | The viewer is paused on start by default, press `SPACE` key to start simulation or `RIGHT` arrow key 57 | to step forward once. By default this script will load the UniTree H1 robot model scene from 58 | [data/mujoco/models](../data/data/mujoco/models/scene.xml). Additional scripts are provided for running comprehensive evaluation in [scripts](../inference_env/scripts/) 59 | directory. 60 | 61 | **NOTE:** The visualization of models is done using a modified version of [mujoco-python-viewer](https://github.com/rohanpsingh/mujoco-python-viewer) with changes to allow 62 | persistent markers for reference motion states. These modifications are bundled in [third_party/mujoco_viewer](../../third_party/mujoco_viewer/) for self containment. 63 | 64 | **NOTE:** During internal testing, a few users reported segmentation fault errors when visualization was enabled, although others did not encounter this issue. It is likely due to a dependency mismatch (such as `opengl`, `mujoco`, `nvidia-driver`, etc.), possibly compounded by complications with Conda. We haven't been able to pinpoint the exact cause at the time of release. For users facing this issue, here are a few mitigation options: 65 | 1. On top of the `third_party/mujoco_viewer/mujoco_viewer.py`, before importing mujoco, add the following lines: 66 | ```python 67 | import os 68 | os.environ["MUJOCO_GL"] = "egl" 69 | ``` 70 | This forces MuJoCo to use EGL instead of the default GLX for OpenGL context creation, bypassing display server dependencies. 71 | 2. In case option 1 failed, in the `third_party/mujoco_viewer/mujoco_viewer.py:line 460`, you can comment out the following lines: 72 | ```python 73 | # marker items 74 | for _ , marker in self._markers_dict.items(): 75 | self._add_marker_to_scene(marker) 76 | ``` 77 | This will disable the marker rendering but the robot will still be visualized. 78 | 79 | ## Deployment 80 | 81 | Once sim-to-sim validation is done, the policy can be deployed in simulation or on real robot. 82 | This involves loading the relevant checkpoint and generating actions based on the observations 83 | from the target robot. For real robots, where some quantities are not easily observable like 84 | body orientations, use of a forward dynamics model is necessary. 85 | 86 | ## Unit Tests 87 | 88 | The tests are located in the `tests` directory and can be run with the `unittest` module. 89 | 90 | ```bash 91 | cd neural_wbc/mujoco_wrapper 92 | ${ISAACLAB_PATH:?}/isaaclab.sh -p -m unittest 93 | ``` 94 | -------------------------------------------------------------------------------- /neural_wbc/mujoco_wrapper/mujoco_wrapper/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | # Needed to have the robot registry 'activated' 17 | from mujoco_wrapper.mujoco_robot import MujocoRobot # noqa 18 | -------------------------------------------------------------------------------- /neural_wbc/mujoco_wrapper/mujoco_wrapper/control.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | from __future__ import annotations 18 | 19 | import torch 20 | from typing import TYPE_CHECKING, Literal 21 | 22 | if TYPE_CHECKING: 23 | from mujoco_wrapper.neural_wbc_env import NeuralWBCEnv 24 | 25 | 26 | def position_pd_control(env: NeuralWBCEnv, pos_actions: torch.Tensor, joint_ids=None): 27 | """Calculates the PD control torque based on the network output position actions""" 28 | robot = env.robot 29 | joint_pos = robot.joint_positions 30 | joint_vel = robot.joint_velocities 31 | 32 | if joint_ids: 33 | joint_pos = joint_pos[:, joint_ids] 34 | joint_vel = joint_vel[:, joint_ids] 35 | 36 | torques = env.p_gains * (pos_actions - joint_pos) - env.d_gains * joint_vel 37 | return torques 38 | 39 | 40 | def torque_control(env: NeuralWBCEnv, torque_actions: torch.Tensor, joint_ids=None): 41 | """Calculates desired torques based on the output of the network""" 42 | torques = torque_actions 43 | 44 | return torques 45 | 46 | 47 | def null_control(env: NeuralWBCEnv, actions_scaled: torch.Tensor, joint_ids=None): 48 | return actions_scaled 49 | 50 | 51 | def resolve_control_fn( 52 | control_type: Literal["Pos", "Vel", "Torque", "None"] = "Pos", 53 | ): 54 | control_fn = position_pd_control 55 | 56 | if control_type == "Pos": 57 | print("[INFO]: Setting up pos control") 58 | control_fn = position_pd_control 59 | elif control_type == "Torque": 60 | print("[INFO]: Setting up torque control") 61 | control_fn = torque_control 62 | elif control_type == "None": 63 | print("[INFO]: Setting up no control") 64 | control_fn = null_control 65 | else: 66 | raise ValueError(f"Unrecognized control type {control_type}") 67 | 68 | return control_fn 69 | -------------------------------------------------------------------------------- /neural_wbc/mujoco_wrapper/mujoco_wrapper/mujoco_utils.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | from __future__ import annotations 18 | 19 | import mujoco as mj 20 | 21 | OBJECT_MAP = { 22 | "joint": mj.mjtObj.mjOBJ_JOINT, 23 | "body": mj.mjtObj.mjOBJ_BODY, 24 | "actuator": mj.mjtObj.mjOBJ_ACTUATOR, 25 | } 26 | 27 | 28 | def get_entity_name(model: mj.MjModel, entity_type: str, entity_id: int) -> str: 29 | """Gets name of an entity based on ID 30 | 31 | Args: 32 | model (mj.MjModel): model 33 | entity_type (str): entity type 34 | entity_id (int): entity id 35 | 36 | Returns: 37 | str: entity name 38 | """ 39 | if entity_type == "body": 40 | return model.body(entity_id).name 41 | return mj.mj_id2name(model, OBJECT_MAP[entity_type], entity_id) 42 | 43 | 44 | def get_entity_id(model: mj.MjModel, entity_type: str, entity_name: str) -> int: 45 | """Gets name of an entity based on ID 46 | 47 | Args: 48 | model (mj.MjModel): model 49 | entity_type (str): entity type 50 | entity_name (str): entity name 51 | 52 | Returns: 53 | int: entity id 54 | 55 | Notes: 56 | If the entity does not exist, returns -1 57 | """ 58 | return mj.mj_name2id(model, OBJECT_MAP[entity_type], entity_name) 59 | 60 | 61 | def has_free_joint(model: mj.MjModel) -> bool: 62 | """Check if the model has a free joint 63 | 64 | Returns: 65 | bool: True if the model has a free joint 66 | """ 67 | return model.nv != model.nq 68 | -------------------------------------------------------------------------------- /neural_wbc/mujoco_wrapper/mujoco_wrapper/utils.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | import torch 17 | 18 | 19 | def to_numpy(x): 20 | """ 21 | Check if input is a PyTorch tensor and convert to numpy array if true. 22 | Otherwise return the input unchanged. 23 | 24 | Args: 25 | x: Input to check and potentially convert 26 | 27 | Returns: 28 | numpy array if input was torch tensor, otherwise original input 29 | """ 30 | if isinstance(x, torch.Tensor): 31 | return x.detach().cpu().numpy() 32 | return x 33 | 34 | 35 | def squeeze_if_tensor(x, dim: int = 0): 36 | """ 37 | Check if input is a PyTorch tensor and squeeze along the given dim if true. 38 | 39 | Args: 40 | x: Input to check and potentially convert 41 | dim: Dimension to squeeze 42 | 43 | Returns: 44 | numpy array if input was torch tensor, otherwise original input 45 | """ 46 | if isinstance(x, torch.Tensor): 47 | return x.squeeze(dim=dim) 48 | return x 49 | -------------------------------------------------------------------------------- /neural_wbc/mujoco_wrapper/mujoco_wrapper/visualization.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import numpy as np 18 | 19 | import mujoco as mj 20 | import mujoco_viewer.mujoco_viewer as mjv 21 | 22 | from neural_wbc.core.reference_motion import ReferenceMotionState 23 | 24 | 25 | class MujocoVisualizer: 26 | """A basic Mujoco visualizer""" 27 | 28 | def __init__( 29 | self, 30 | model: mj.MjModel, 31 | data: mj.MjData, 32 | width: int = 1400, 33 | height: int = 1200, 34 | start_paused: bool = True, 35 | ): 36 | self._model = model 37 | self._data = data 38 | self._viewer = mjv.MujocoViewer( 39 | self._model, 40 | self._data, 41 | width=width, 42 | height=height, 43 | hide_menus=True, 44 | ) 45 | 46 | self._viewer._paused = start_paused # pylint: disable=W0212 47 | 48 | def update(self): 49 | """Update the Mujoco viewer.""" 50 | if self._viewer.is_alive: 51 | self._viewer.render() 52 | 53 | def draw_reference_state(self, state: ReferenceMotionState): 54 | """Visualize the reference state in Mujoco.""" 55 | 56 | body_pos_np = np.squeeze(state.body_pos.detach().cpu().numpy()) 57 | body_pos_extend_np = np.squeeze(state.body_pos_extend.detach().cpu().numpy()) 58 | body_pos = np.vstack([body_pos_np, body_pos_extend_np]) 59 | 60 | for i in range(body_pos.shape[0]): 61 | self._viewer.add_marker( 62 | pos=body_pos[i], 63 | size=0.05, 64 | rgba=(1, 0, 0, 1), 65 | type=mj.mjtGeom.mjGEOM_SPHERE, 66 | label="", 67 | id=i, 68 | ) 69 | 70 | def close(self): 71 | self._viewer.close() 72 | -------------------------------------------------------------------------------- /neural_wbc/mujoco_wrapper/setup.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | from setuptools import find_packages, setup 18 | 19 | setup( 20 | name="neural_wbc_mujoco_wrapper", 21 | version="0.1.0", 22 | packages=find_packages(), 23 | description="Mujoco wrapper for sim-to-sim validation", 24 | python_requires=">=3.10", 25 | install_requires=[ 26 | "mujoco", 27 | "mujoco-python-viewer", 28 | "numpy", 29 | "torch", 30 | "tyro", 31 | "neural_wbc_core", 32 | "neural_wbc_data", 33 | ], 34 | ) 35 | -------------------------------------------------------------------------------- /neural_wbc/mujoco_wrapper/tests/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | -------------------------------------------------------------------------------- /neural_wbc/student_policy/README.md: -------------------------------------------------------------------------------- 1 | # Neural WBC - Student Policy 2 | 3 | The `neural_wbc.policy` package provides tooling to train a student policy for neural whole body 4 | controllers. 5 | 6 | ## Installation 7 | To use the `neural_wbc.student_policy` package, install it with 8 | 9 | ```bash 10 | ${ISAACLAB_PATH:?}/isaaclab.sh -p -m pip install -e neural_wbc/student_policy 11 | ``` 12 | 13 | ## Usage 14 | Refer to the Isaac Lab implementation of the neural WBC environment, as well as the training and 15 | testing workflows, for examples of how to use the module. 16 | 17 | 18 | ## Unit Tests 19 | The tests are located in the `tests` directory and can be run with the `unittest` module. 20 | 21 | ```bash 22 | cd neural_wbc/student_policy 23 | ${ISAACLAB_PATH:?}/isaaclab.sh -p -m unittest 24 | ``` 25 | -------------------------------------------------------------------------------- /neural_wbc/student_policy/neural_wbc/student_policy/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | from .storage import Storage # noqa 18 | from .student_policy_trainer import StudentPolicyTrainer # noqa 19 | from .student_policy_trainer_cfg import StudentPolicyTrainerCfg # noqa 20 | from .teacher_policy import TeacherPolicy # noqa 21 | -------------------------------------------------------------------------------- /neural_wbc/student_policy/neural_wbc/student_policy/policy.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import torch 18 | import torch.nn as nn 19 | from torch.distributions import Normal 20 | 21 | 22 | class StudentPolicy(nn.Module): 23 | def __init__(self, num_obs, num_actions, policy_hidden_dims=[256, 256, 256], activation="elu", noise_std=0.001): 24 | super().__init__() 25 | 26 | activation = get_activation(activation) 27 | 28 | mlp_input_dim_a = num_obs 29 | 30 | # Policy 31 | policy_layers = [] 32 | policy_layers.append(nn.Linear(mlp_input_dim_a, policy_hidden_dims[0])) 33 | policy_layers.append(activation) 34 | for layer in range(len(policy_hidden_dims)): 35 | if layer == len(policy_hidden_dims) - 1: 36 | policy_layers.append(nn.Linear(policy_hidden_dims[layer], num_actions)) 37 | else: 38 | policy_layers.append(nn.Linear(policy_hidden_dims[layer], policy_hidden_dims[layer + 1])) 39 | policy_layers.append(activation) 40 | self.policy = nn.Sequential(*policy_layers) 41 | 42 | print(f"Student Policy MLP: {self.policy}") 43 | 44 | # Action noise 45 | self.std = nn.Parameter(noise_std * torch.ones(num_actions)) 46 | self.std.requires_grad = False 47 | self.distribution = None 48 | # disable args validation for speedup 49 | Normal.set_default_validate_args = False 50 | 51 | def reset(self, dones=None): 52 | pass 53 | 54 | def forward(self): 55 | raise NotImplementedError 56 | 57 | @property 58 | def action_mean(self): 59 | return self.distribution.mean 60 | 61 | @property 62 | def action_std(self): 63 | return self.distribution.stddev 64 | 65 | @property 66 | def entropy(self): 67 | return self.distribution.entropy().sum(dim=-1) 68 | 69 | def update_distribution(self, observations): 70 | mean = self.policy(observations) 71 | self.distribution = Normal(mean, mean * 0.0 + self.std) 72 | 73 | def act(self, observations, **kwargs): 74 | self.update_distribution(observations) 75 | return self.distribution.sample() 76 | 77 | def get_actions_log_prob(self, actions): 78 | return self.distribution.log_prob(actions).sum(dim=-1) 79 | 80 | def act_inference(self, observations): 81 | actions_mean = self.policy(observations) 82 | return actions_mean 83 | 84 | def load(self, path, device): 85 | self.to(device) 86 | loaded_dict = torch.load(path, map_location=device) 87 | self.load_state_dict(loaded_dict["model_state_dict"]) 88 | 89 | 90 | def get_activation(act_name): 91 | activation_functions = { 92 | "elu": nn.ELU, 93 | "selu": nn.SELU, 94 | "relu": nn.ReLU, 95 | "crelu": nn.ReLU, 96 | "lrelu": nn.LeakyReLU, 97 | "tanh": nn.Tanh, 98 | "sigmoid": nn.Sigmoid, 99 | } 100 | 101 | return activation_functions.get(act_name, lambda: print(f"invalid activation function {act_name}!"))() 102 | -------------------------------------------------------------------------------- /neural_wbc/student_policy/neural_wbc/student_policy/storage.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import torch 18 | from dataclasses import dataclass 19 | 20 | 21 | @dataclass 22 | class Slice: 23 | policy_observations: torch.Tensor 24 | student_observations: torch.Tensor 25 | ground_truth_actions: torch.Tensor 26 | applied_actions: torch.Tensor 27 | 28 | 29 | @dataclass 30 | class Storage: 31 | step: int 32 | policy_observations: torch.Tensor 33 | student_observations: torch.Tensor 34 | ground_truth_actions: torch.Tensor 35 | applied_actions: torch.Tensor 36 | 37 | def __init__( 38 | self, 39 | max_steps: int, 40 | num_envs: int, 41 | device: torch.device, 42 | policy_obs_shape, 43 | student_obs_shape, 44 | actions_shape, 45 | ): 46 | self._max_steps = max_steps 47 | self._num_envs = num_envs 48 | self._device = device 49 | self._policy_obs_shape = policy_obs_shape 50 | self._student_obs_shape = student_obs_shape 51 | self._actions_shape = actions_shape 52 | 53 | self.reset() 54 | 55 | def is_full(self): 56 | return self.step >= self._max_steps 57 | 58 | def add(self, slice: Slice): 59 | if self.is_full(): 60 | raise AssertionError("Rollout buffer overflow") 61 | self.policy_observations[self.step].copy_(slice.policy_observations) 62 | self.student_observations[self.step].copy_(slice.student_observations) 63 | self.ground_truth_actions[self.step].copy_(slice.ground_truth_actions) 64 | self.applied_actions[self.step].copy_(slice.applied_actions) 65 | self.step += 1 66 | 67 | def reset(self): 68 | self.step = 0 69 | self.policy_observations = torch.zeros( 70 | self._max_steps, self._num_envs, *self._policy_obs_shape, device=self._device 71 | ) 72 | self.student_observations = torch.zeros( 73 | self._max_steps, self._num_envs, *self._student_obs_shape, device=self._device 74 | ) 75 | self.ground_truth_actions = torch.zeros( 76 | self._max_steps, self._num_envs, *self._actions_shape, device=self._device 77 | ) 78 | self.applied_actions = torch.zeros(self._max_steps, self._num_envs, *self._actions_shape, device=self._device) 79 | 80 | def mini_batch_generator(self, num_mini_batches, num_epochs=8): 81 | batch_size = self._num_envs * self._max_steps 82 | mini_batch_size = batch_size // num_mini_batches 83 | indices = torch.randperm(num_mini_batches * mini_batch_size, requires_grad=False, device=self._device) 84 | 85 | policy_observations = self.policy_observations.flatten(0, 1) 86 | student_observations = self.student_observations.flatten(0, 1) 87 | ground_truth_actions = self.ground_truth_actions.flatten(0, 1) 88 | applied_actions = self.applied_actions.flatten(0, 1) 89 | 90 | for epoch in range(num_epochs): 91 | for i in range(num_mini_batches): 92 | start = i * mini_batch_size 93 | end = (i + 1) * mini_batch_size 94 | batch_idx = indices[start:end] 95 | 96 | policy_observations_batch = policy_observations[batch_idx] 97 | student_observations_batch = student_observations[batch_idx] 98 | ground_truth_actions_batch = ground_truth_actions[batch_idx] 99 | applied_actions_batch = applied_actions[batch_idx] 100 | 101 | yield policy_observations_batch, student_observations_batch, ground_truth_actions_batch, applied_actions_batch 102 | -------------------------------------------------------------------------------- /neural_wbc/student_policy/neural_wbc/student_policy/teacher_policy.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | from typing import Optional 18 | 19 | 20 | class TeacherPolicy: 21 | # Path where the teacher policy is loaded from 22 | path: Optional[str] = None 23 | 24 | def __init__(self): 25 | pass 26 | 27 | def load(self, path, **kwargs): 28 | # load the policy 29 | raise NotImplementedError(f"Please implement the 'load' method for {self.__class__.__name__}.") 30 | 31 | def act_rollout(self, observations, **kwargs): 32 | # resolve the action that used for data collection, like action with random 33 | raise NotImplementedError(f"Please implement the 'act_rollout' method for {self.__class__.__name__}.") 34 | 35 | def act(self, observations, **kwargs): 36 | # resolve the action that used for labeling 37 | raise NotImplementedError(f"Please implement the 'act' method for {self.__class__.__name__}.") 38 | -------------------------------------------------------------------------------- /neural_wbc/student_policy/neural_wbc/tests/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | -------------------------------------------------------------------------------- /neural_wbc/student_policy/neural_wbc/tests/test_storage.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import torch 18 | import unittest 19 | 20 | from neural_wbc.student_policy.storage import Slice, Storage 21 | 22 | 23 | class TestStorage(unittest.TestCase): 24 | """ 25 | Unit tests for the Storage class. 26 | 27 | This module contains unit tests to validate the functionality of the Storage class. 28 | All test values, including numbers and shapes, are arbitrary and do not correspond 29 | to any specific type of robots or real-world scenarios. 30 | """ 31 | 32 | def test_add(self): 33 | max_steps = 10 34 | num_envs = 3 35 | device = torch.device("cpu") 36 | policy_obs_shape = torch.Size([19, 3]) 37 | student_obs_shape = torch.Size([9, 3]) 38 | actions_shape = torch.Size([10]) 39 | 40 | storage = Storage( 41 | max_steps=max_steps, 42 | num_envs=num_envs, 43 | device=device, 44 | policy_obs_shape=policy_obs_shape, 45 | student_obs_shape=student_obs_shape, 46 | actions_shape=actions_shape, 47 | ) 48 | self.assertEqual(storage.step, 0) 49 | 50 | policy_observations = torch.randn(max_steps, num_envs, *policy_obs_shape) 51 | student_observations = torch.randn(max_steps, num_envs, *student_obs_shape) 52 | ground_truth_actions = torch.randn(max_steps, num_envs, *actions_shape) 53 | student_actions = torch.randn(max_steps, num_envs, *actions_shape) 54 | 55 | for step in range(max_steps): 56 | slice = Slice( 57 | policy_observations=policy_observations[step, ...], 58 | student_observations=student_observations[step, ...], 59 | ground_truth_actions=ground_truth_actions[step, ...], 60 | applied_actions=student_actions[step, ...], 61 | ) 62 | storage.add(slice) 63 | self.assertEqual(storage.step, step + 1) 64 | 65 | self.assertTrue(torch.all(torch.eq(storage.policy_observations, policy_observations))) 66 | self.assertTrue(torch.all(torch.eq(storage.student_observations, student_observations))) 67 | self.assertTrue(torch.all(torch.eq(storage.ground_truth_actions, ground_truth_actions))) 68 | self.assertTrue(torch.all(torch.eq(storage.applied_actions, student_actions))) 69 | 70 | self.assertRaises(AssertionError, storage.add, slice) 71 | 72 | storage.reset() 73 | self.assertEqual(storage.step, 0) 74 | 75 | 76 | if __name__ == "__main__": 77 | unittest.main() 78 | -------------------------------------------------------------------------------- /neural_wbc/student_policy/setup.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | from setuptools import find_namespace_packages, setup 18 | 19 | setup( 20 | name="neural_wbc_student_policy", 21 | version="0.0.0", 22 | packages=find_namespace_packages(), 23 | namespace_packages=["neural_wbc"], 24 | description="Neural WBC student policy trainer", 25 | python_requires=">=3.10", 26 | ) 27 | -------------------------------------------------------------------------------- /punch.yaml: -------------------------------------------------------------------------------- 1 | - "ACCAD/Male2MartialArtsPunches_c3d/E1 - Jab left_poses.npz" 2 | - "ACCAD/Male2MartialArtsPunches_c3d/E10 - body hook right_poses.npz" 3 | - "ACCAD/Male2MartialArtsPunches_c3d/E11 - backfist left_poses.npz" 4 | - "ACCAD/Male2MartialArtsPunches_c3d/E11 - body hook right t2_poses.npz" 5 | - "ACCAD/Male2MartialArtsPunches_c3d/E12 - backfist right_poses.npz" 6 | - "ACCAD/Male2MartialArtsPunches_c3d/E14 - body cross left t2_poses.npz" 7 | - "ACCAD/Male2MartialArtsPunches_c3d/E14 - body cross right_poses.npz" 8 | - "ACCAD/Male2MartialArtsPunches_c3d/E16 - body jab left_poses.npz" 9 | - "ACCAD/Male2MartialArtsPunches_c3d/E2 - Jab right_poses.npz" 10 | - "ACCAD/Male2MartialArtsPunches_c3d/E3 - cross left_poses.npz" 11 | - "ACCAD/Male2MartialArtsPunches_c3d/E4 - cross right_poses.npz" 12 | - "ACCAD/Male2MartialArtsPunches_c3d/E5 - hook left_poses.npz" 13 | - "ACCAD/Male2MartialArtsPunches_c3d/E6 - hook right_poses.npz" 14 | - "ACCAD/Male2MartialArtsPunches_c3d/E7 - uppercut left_poses.npz" 15 | - "ACCAD/Male2MartialArtsPunches_c3d/E8 - uppercut right_poses.npz" 16 | - "ACCAD/Male2MartialArtsPunches_c3d/E9 - body hook left_poses.npz" 17 | - "KIT/291/punching03_poses.npz" 18 | - "KIT/3/punch_left04_poses.npz" 19 | - "KIT/3/punch_right01_poses.npz" 20 | - "KIT/442/Punch01_poses.npz" 21 | - "KIT/572/punch_*_poses.npz" 22 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["setuptools", "toml"] 3 | build-backend = "setuptools.build_meta" 4 | 5 | [project] 6 | name = "neural_wbc" 7 | version = "0.1.0" 8 | 9 | [tool.isort] 10 | atomic = true 11 | profile = "black" 12 | line_length = 120 13 | py_version = 310 14 | skip_glob = ["docs/*", "logs/*"] 15 | group_by_package = true 16 | 17 | # These should be under tool.isort, not directly in the file 18 | known_firstparty = ["neural_wbc"] # Changed to list format 19 | known_local_folder = ["config"] # Changed to list format 20 | 21 | sections = [ 22 | "FUTURE", 23 | "STDLIB", 24 | "THIRDPARTY", 25 | "FIRSTPARTY", 26 | "ISAACLABPARTY", 27 | "LOCALFOLDER", 28 | ] 29 | extra_standard_library = [ 30 | "numpy", 31 | "h5py", 32 | "open3d", 33 | "torch", 34 | "tensordict", 35 | "bpy", 36 | "matplotlib", 37 | "gymnasium", 38 | "gym", 39 | "scipy", 40 | "hid", 41 | "yaml", 42 | "prettytable", 43 | "toml", 44 | "trimesh", 45 | "tqdm", 46 | ] 47 | known_thirdparty = [ 48 | "omni.isaac.core", 49 | "omni.replicator.isaac", 50 | "omni.replicator.core", 51 | "pxr", 52 | "omni.kit.*", 53 | "warp", 54 | "carb", 55 | ] 56 | known_isaaclabparty = [ 57 | "isaaclab", 58 | "isaaclab_tasks", 59 | "isaaclab_assets" 60 | ] 61 | 62 | [tool.setuptools] 63 | package-dir = {"rsl_rl" = "third_party/rsl_rl/rsl_rl"} 64 | packages = ["rsl_rl"] 65 | 66 | [tool.pyright] 67 | exclude = [ 68 | "**/__pycache__", 69 | "**/docs", 70 | "**/logs", 71 | ".git", 72 | ".vscode", 73 | ] 74 | 75 | typeCheckingMode = "basic" 76 | pythonVersion = "3.10" 77 | pythonPlatform = "Linux" 78 | enableTypeIgnoreComments = true 79 | 80 | reportMissingImports = "none" 81 | reportMissingModuleSource = "none" 82 | reportGeneralTypeIssues = "none" 83 | reportOptionalMemberAccess = "warning" 84 | reportPrivateUsage = "warning" 85 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | joblib>=1.2.0 2 | wheel # Without this chumpy (which is a dep of smplsim) fails to build. 3 | git+https://github.com/ZhengyiLuo/SMPLSim.git@dd65a86 4 | easydict 5 | warp-lang 6 | dataclass-wizard 7 | 8 | -e neural_wbc/core 9 | -e neural_wbc/data 10 | -e neural_wbc/isaac_lab_wrapper 11 | -e neural_wbc/mujoco_wrapper 12 | -e neural_wbc/inference_env 13 | -e neural_wbc/student_policy 14 | -e third_party/human2humanoid/phc 15 | -e third_party/mujoco_viewer 16 | -e third_party/rsl_rl 17 | -------------------------------------------------------------------------------- /requirements_deploy.txt: -------------------------------------------------------------------------------- 1 | -e git+https://github.com/unitreerobotics/unitree_sdk2_python.git@3af1610#egg=unitree_sdk2py 2 | 3 | -e neural_wbc/hw_wrappers 4 | -------------------------------------------------------------------------------- /run_e2e_tests.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | # SPDX-License-Identifier: Apache-2.0 4 | # 5 | # Licensed under the Apache License, Version 2.0 (the "License"); 6 | # you may not use this file except in compliance with the License. 7 | # You may obtain a copy of the License at 8 | # 9 | # http://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, software 12 | # distributed under the License is distributed on an "AS IS" BASIS, 13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | # See the License for the specific language governing permissions and 15 | # limitations under the License. 16 | 17 | # Function to run a command and exit on failure 18 | run_command() { 19 | local cmd="$1" 20 | echo "Running: $cmd" 21 | $cmd 22 | if [ $? -ne 0 ]; then 23 | echo "Command failed: $cmd" 24 | exit 1 25 | fi 26 | } 27 | 28 | # List of test scripts to run 29 | test_scripts=( 30 | "neural_wbc/inference_env/tests/test_evaluation_pipeline.py" 31 | "scripts/rsl_rl/tests/student_evaluation_pipeline_test.py" 32 | "scripts/rsl_rl/tests/student_training_pipeline_test.py" 33 | "scripts/rsl_rl/tests/teacher_evaluation_pipeline_test.py" 34 | "scripts/rsl_rl/tests/teacher_training_pipeline_test.py" 35 | ) 36 | 37 | # Execute each test script 38 | for script in "${test_scripts[@]}"; do 39 | cmd="${ISAACLAB_PATH}/isaaclab.sh -p ${script}" 40 | run_command "$cmd" 41 | done 42 | 43 | echo "All commands executed successfully." 44 | -------------------------------------------------------------------------------- /run_unit_tests.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # 3 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 4 | # SPDX-License-Identifier: Apache-2.0 5 | # 6 | # Licensed under the Apache License, Version 2.0 (the "License"); 7 | # you may not use this file except in compliance with the License. 8 | # You may obtain a copy of the License at 9 | # 10 | # http://www.apache.org/licenses/LICENSE-2.0 11 | # 12 | # Unless required by applicable law or agreed to in writing, software 13 | # distributed under the License is distributed on an "AS IS" BASIS, 14 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 | # See the License for the specific language governing permissions and 16 | # limitations under the License. 17 | set -e 18 | 19 | run_unittests_in_folder() ( 20 | # We use a subshell so we can cd into the folder. 21 | local folder=${1:?} 22 | cd ${folder} 23 | echo "Running unit tests in '${folder}'." 24 | ${ISAACLAB_PATH:?}/isaaclab.sh -p -m unittest 25 | ) 26 | 27 | run_unittests_in_file() { 28 | local file=${1:?} 29 | echo "Running unit tests in '${file}'." 30 | ${ISAACLAB_PATH}/isaaclab.sh -p -m unittest "${file}" 31 | } 32 | 33 | script_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )" 34 | cd ${script_dir} 35 | 36 | # Directories containing unit tests. 37 | test_folders=( 38 | "." 39 | "neural_wbc/core" 40 | "neural_wbc/data" 41 | "neural_wbc/mujoco_wrapper" 42 | "neural_wbc/student_policy" 43 | "neural_wbc/inference_env" 44 | ) 45 | 46 | # Run unit tests in each directory. 47 | for folder in "${test_folders[@]}"; do 48 | run_unittests_in_folder "$folder" 49 | done 50 | 51 | # Additional files containing unit tests. 52 | test_files=( 53 | "scripts/rsl_rl/tests/test_teacher_policy_cfg.py" 54 | ) 55 | 56 | # Run unit tests in each file. 57 | for file in "${test_files[@]}"; do 58 | run_unittests_in_file "$file" 59 | done 60 | -------------------------------------------------------------------------------- /scripts/rsl_rl/config_overwrites/sample.overwrite.yaml: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | sim.dt: 0.017 18 | decimation: 4 19 | add_policy_obs_noise: False 20 | default_rfi_limit: 0.1 21 | ctrl_delay_step_range: [0, 3] 22 | -------------------------------------------------------------------------------- /scripts/rsl_rl/eval.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import pprint 18 | import yaml 19 | 20 | from teacher_policy_cfg import TeacherPolicyCfg 21 | 22 | from isaaclab.app import AppLauncher 23 | 24 | # local imports 25 | from utils import get_player_args # isort: skip 26 | 27 | # add argparse arguments 28 | parser = get_player_args(description="Evaluates motion tracking policy and collects metrics.") 29 | parser.add_argument("--metrics_path", type=str, default=None, help="Path to store metrics in.") 30 | parser.add_argument( 31 | "--env_config_overwrite", type=str, default=None, help="Path to yaml file with overwriting configuration values." 32 | ) 33 | 34 | # append RSL-RL cli arguments 35 | TeacherPolicyCfg.add_args_to_parser(parser=parser) 36 | # append AppLauncher cli args 37 | AppLauncher.add_app_launcher_args(parser) 38 | args_cli = parser.parse_args() 39 | 40 | # launch omniverse app 41 | app_launcher = AppLauncher(args_cli) 42 | simulation_app = app_launcher.app 43 | 44 | from players import EvaluationPlayer 45 | 46 | 47 | def main(): 48 | custom_config = None 49 | if args_cli.env_config_overwrite is not None: 50 | with open(args_cli.env_config_overwrite) as fh: 51 | custom_config = yaml.safe_load(fh) 52 | print("[INFO]: Using custom configuration:") 53 | pprint.pprint(custom_config) 54 | 55 | # Initialize the player with the updated properties 56 | player = EvaluationPlayer(args_cli=args_cli, metrics_path=args_cli.metrics_path, custom_config=custom_config) 57 | player.play(simulation_app=simulation_app) 58 | 59 | 60 | if __name__ == "__main__": 61 | # Run the main execution 62 | main() 63 | # Close sim app 64 | simulation_app.close() 65 | -------------------------------------------------------------------------------- /scripts/rsl_rl/play.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | from teacher_policy_cfg import TeacherPolicyCfg 18 | 19 | from isaaclab.app import AppLauncher 20 | 21 | # local imports 22 | from utils import get_player_args # isort: skip 23 | 24 | 25 | # add argparse arguments 26 | parser = get_player_args(description="Plays motion tracking policy in Isaac Lab.") 27 | parser.add_argument("--randomize", action="store_true", help="Whether to randomize reference motion while playing.") 28 | 29 | # append RSL-RL cli arguments 30 | TeacherPolicyCfg.add_args_to_parser(parser) 31 | # append AppLauncher cli args 32 | AppLauncher.add_app_launcher_args(parser) 33 | args_cli = parser.parse_args() 34 | 35 | # launch omniverse app 36 | app_launcher = AppLauncher(args_cli) 37 | simulation_app = app_launcher.app 38 | 39 | from players import DemoPlayer 40 | 41 | 42 | def main(): 43 | player = DemoPlayer(args_cli=args_cli, randomize=args_cli.randomize) 44 | player.play(simulation_app=simulation_app) 45 | 46 | 47 | if __name__ == "__main__": 48 | # run the main execution 49 | main() 50 | # close sim app 51 | simulation_app.close() 52 | -------------------------------------------------------------------------------- /scripts/rsl_rl/tests/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | -------------------------------------------------------------------------------- /scripts/rsl_rl/tests/student_evaluation_pipeline_test.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import json 18 | import os 19 | import pprint 20 | import subprocess 21 | import sys 22 | import tempfile 23 | 24 | import neural_wbc.data 25 | from neural_wbc.core import util 26 | 27 | 28 | def get_json_files(directory): 29 | json_files = [] 30 | for filename in os.listdir(directory): 31 | if filename.endswith(".json"): 32 | json_files.append(os.path.join(directory, filename)) 33 | return json_files 34 | 35 | 36 | def compare_generated_metrics(generated_metrics): 37 | # Make sure we always have 100% success rate with the stable_punch dataset and default model 38 | util.assert_equal(generated_metrics["success_rate"], 1.0, "Success rate") 39 | # Make sure all and success metrics are the same 40 | util.assert_equal(generated_metrics["all"], generated_metrics["success"], "All metrics vs Success metrics") 41 | # Make sure we played the correct number of reference motions 42 | util.assert_equal(generated_metrics["num_motions"], 29, "Number of motions") 43 | 44 | 45 | if __name__ == "__main__": 46 | # First find the evaluation script 47 | metrics_path = tempfile.TemporaryDirectory() 48 | policy_path, checkpoint = neural_wbc.data.get_default_omnih2o_student_policy() 49 | eval_pipeline_command = [ 50 | "${ISAACLAB_PATH}/isaaclab.sh -p scripts/rsl_rl/eval.py " 51 | "--headless " 52 | "--num_envs 10 " 53 | "--student_player " 54 | f"--metrics_path={metrics_path.name} " 55 | f"--student_path={policy_path} " 56 | f"--student_checkpoint={checkpoint} " 57 | ] 58 | return_value = subprocess.call(eval_pipeline_command, shell=True) 59 | json_files = get_json_files(metrics_path.name) 60 | assert len(json_files) == 1 61 | with open(json_files[0]) as fh: 62 | generated_metrics = json.load(fh) 63 | print("Generated metrics:") 64 | pprint.pprint(generated_metrics) 65 | compare_generated_metrics(generated_metrics) 66 | sys.exit(return_value) 67 | -------------------------------------------------------------------------------- /scripts/rsl_rl/tests/student_training_pipeline_test.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import os 18 | import subprocess 19 | import tempfile 20 | 21 | import neural_wbc.data 22 | from neural_wbc.data import get_data_path 23 | 24 | 25 | def main(): 26 | motion_dataset_path = get_data_path("motions/stable_punch.pkl") 27 | policy_path, checkpoint = neural_wbc.data.get_default_teacher_policy() 28 | # Create a temporary directory 29 | with tempfile.TemporaryDirectory() as temp_dir: 30 | command = ( 31 | "${ISAACLAB_PATH}/isaaclab.sh -p scripts/rsl_rl/train_student_policy.py " 32 | "--headless " 33 | "--num_envs=10 " 34 | "--student_policy.max_iteration=3 " 35 | "--student_policy.save_iteration=2 " 36 | f"--reference_motion_path={motion_dataset_path} " 37 | f"--student_policy.root_path={temp_dir} " 38 | f"--teacher_policy.resume_path {policy_path} " 39 | f"--teacher_policy.checkpoint {checkpoint} " 40 | ) 41 | 42 | # Run the command 43 | subprocess.run(command, shell=True, check=True) 44 | 45 | # Find the result directory 46 | result_directories = [d for d in os.listdir(temp_dir) if os.path.isdir(os.path.join(temp_dir, d))] 47 | assert len(result_directories) == 1, "There should be exactly one result directory." 48 | result_dir_path = os.path.join(temp_dir, result_directories[0]) 49 | 50 | # Verify the files exist in the result directory under the temporary directory 51 | expected_files = ["model_2.pt", "final_model.pt", "config.json"] 52 | for file in expected_files: 53 | file_path = os.path.join(result_dir_path, file) 54 | assert os.path.exists(file_path), f"{file} does NOT exist in the result directory." 55 | 56 | print("All expected files are present in the result directory.") 57 | 58 | 59 | if __name__ == "__main__": 60 | main() 61 | -------------------------------------------------------------------------------- /scripts/rsl_rl/tests/teacher_evaluation_pipeline_test.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import json 18 | import os 19 | import pprint 20 | import subprocess 21 | import sys 22 | import tempfile 23 | 24 | import neural_wbc.data 25 | from neural_wbc.core import util 26 | 27 | 28 | def get_json_files(directory): 29 | json_files = [] 30 | for filename in os.listdir(directory): 31 | if filename.endswith(".json"): 32 | json_files.append(os.path.join(directory, filename)) 33 | return json_files 34 | 35 | 36 | def compare_generated_metrics(generated_metrics): 37 | # Make sure we always have 100% success rate with the stable_punch dataset and default model 38 | util.assert_equal(generated_metrics["success_rate"], 1.0, "Success rate") 39 | # Make sure all and success metrics are the same 40 | util.assert_equal(generated_metrics["all"], generated_metrics["success"], "All metrics vs Success metrics") 41 | # Make sure we played the correct number of reference motions 42 | util.assert_equal(generated_metrics["num_motions"], 29, "Number of motions") 43 | 44 | 45 | if __name__ == "__main__": 46 | # First find the evaluation script 47 | metrics_path = tempfile.TemporaryDirectory() 48 | policy_path, checkpoint = neural_wbc.data.get_default_teacher_policy() 49 | eval_pipeline_command = [ 50 | "${ISAACLAB_PATH}/isaaclab.sh -p scripts/rsl_rl/eval.py " 51 | "--headless " 52 | "--num_envs 10 " 53 | f"--metrics_path={metrics_path.name} " 54 | f"--teacher_policy.resume_path={policy_path} " 55 | f"--teacher_policy.checkpoint={checkpoint} " 56 | ] 57 | return_value = subprocess.call(eval_pipeline_command, shell=True) 58 | json_files = get_json_files(metrics_path.name) 59 | assert len(json_files) == 1 60 | with open(json_files[0]) as fh: 61 | generated_metrics = json.load(fh) 62 | print("Generated metrics:") 63 | pprint.pprint(generated_metrics) 64 | compare_generated_metrics(generated_metrics) 65 | sys.exit(return_value) 66 | -------------------------------------------------------------------------------- /scripts/rsl_rl/tests/teacher_training_pipeline_test.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | import glob 18 | import os 19 | import subprocess 20 | import tempfile 21 | 22 | import neural_wbc.data 23 | from neural_wbc.data import get_data_path 24 | 25 | 26 | def main(): 27 | motion_dataset_path = get_data_path("motions/stable_punch.pkl") 28 | policy_path, checkpoint = neural_wbc.data.get_default_teacher_policy() 29 | # Create a temporary directory 30 | with tempfile.TemporaryDirectory() as temp_dir: 31 | command = ( 32 | "${ISAACLAB_PATH}/isaaclab.sh -p scripts/rsl_rl/train_teacher_policy.py " 33 | "--headless " 34 | "--num_envs=10 " 35 | "--teacher_policy.max_iterations=10 " 36 | "--teacher_policy.resume " 37 | f"--reference_motion_path={motion_dataset_path} " 38 | f"--teacher_policy.path={temp_dir} " 39 | f"--teacher_policy.resume_path={policy_path} " 40 | f"--teacher_policy.checkpoint={checkpoint} " 41 | ) 42 | 43 | # Run the command 44 | subprocess.run(command, shell=True, check=True) 45 | 46 | # Find the results directory 47 | results_path_pattern = os.path.join(temp_dir, "*") 48 | results_directories = glob.glob(results_path_pattern) 49 | assert len(results_directories) == 1, "There should be exactly one results directory." 50 | results_dir_path = results_directories[0] 51 | 52 | # Verify the files exist in the results directory under the temporary directory 53 | expected_files = ["model_0.pt", "model_10.pt", "config.json"] 54 | for file in expected_files: 55 | file_path = os.path.join(results_dir_path, file) 56 | assert os.path.exists(file_path), f"{file} does NOT exist in the results directory." 57 | 58 | print("All expected files are present in the results directory.") 59 | 60 | 61 | if __name__ == "__main__": 62 | main() 63 | -------------------------------------------------------------------------------- /scripts/rsl_rl/tests/test_utils.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | import torch 16 | import unittest 17 | from unittest.mock import MagicMock, mock_open, patch 18 | 19 | from ..teacher_policy_cfg import AlgorithmCfg, PolicyCfg, RunnerCfg, TeacherPolicyCfg 20 | from ..utils import get_customized_rsl_rl, get_ppo_runner_and_checkpoint_path 21 | 22 | get_customized_rsl_rl() 23 | from rsl_rl.runners import OnPolicyRunner 24 | 25 | 26 | class TestLoadFromCheckpoint(unittest.TestCase): 27 | # Mocking open and load functions to avoid specifying files 28 | @patch("builtins.open", new_callable=mock_open, read_data='{"policy": {}}') 29 | @patch("json.load") 30 | @patch("torch.optim.Adam") 31 | def test_load_from_checkpoint_success(self, mock_adam, mock_json_load, mock_open): 32 | # Mock arguments 33 | teacher_policy_cfg = TeacherPolicyCfg( 34 | seed=1, 35 | policy=PolicyCfg(), 36 | algorithm=AlgorithmCfg(gamma=1.9), 37 | runner=RunnerCfg(resume_path="/path/to/resume", checkpoint="checkpoint_file", num_steps_per_env=38), 38 | ) 39 | wrapped_env = MagicMock() 40 | wrapped_env.reset = MagicMock(return_value=[None, None]) 41 | device = torch.device("cpu") 42 | log_dir = "/path/to/log" 43 | 44 | # Mock return values 45 | mock_json_load.return_value = {"policy": {"init_noise_std": 10.0, "activation": "sigmoid"}} 46 | 47 | # Call the function 48 | ppo_runner, checkpoint_path = get_ppo_runner_and_checkpoint_path( 49 | teacher_policy_cfg, wrapped_env, device, log_dir 50 | ) 51 | 52 | # Assertions 53 | self.assertEqual(checkpoint_path, "/path/to/resume/checkpoint_file") 54 | self.assertIsInstance(ppo_runner, OnPolicyRunner) 55 | mock_open.assert_called_once_with("/path/to/resume/config.json", encoding="utf-8") 56 | 57 | # Check overwrite values 58 | self.assertEqual(teacher_policy_cfg.policy.init_noise_std, 10.0) 59 | self.assertEqual(teacher_policy_cfg.policy.activation, "sigmoid") 60 | 61 | # Compare teacher policy config and OnPolicyPlayer's config 62 | teacher_policy_dict = teacher_policy_cfg.to_dict() 63 | self.assertDictEqual(teacher_policy_dict["policy"], ppo_runner.policy_cfg) 64 | self.assertDictEqual(teacher_policy_dict["algorithm"], ppo_runner.alg_cfg) 65 | self.assertDictEqual(teacher_policy_dict["runner"], ppo_runner.cfg) 66 | 67 | def test_load_from_checkpoint_no_resume_path(self): 68 | teacher_policy_cfg = TeacherPolicyCfg( 69 | seed=1, 70 | policy=PolicyCfg(), 71 | algorithm=AlgorithmCfg(), 72 | runner=RunnerCfg(resume_path="", checkpoint="checkpoint_file"), 73 | ) 74 | wrapped_env = MagicMock() 75 | device = torch.device("cpu") 76 | log_dir = "/path/to/log" 77 | 78 | with self.assertRaises(ValueError) as context: 79 | get_ppo_runner_and_checkpoint_path(teacher_policy_cfg, wrapped_env, device, log_dir) 80 | self.assertEqual(str(context.exception), "teacher_policy.resume_path is not specified") 81 | 82 | def test_load_from_checkpoint_no_checkpoint(self): 83 | teacher_policy_cfg = TeacherPolicyCfg( 84 | seed=1, 85 | policy=PolicyCfg(), 86 | algorithm=AlgorithmCfg(), 87 | runner=RunnerCfg(resume_path="/path/to/resume", checkpoint=""), 88 | ) 89 | wrapped_env = MagicMock() 90 | device = torch.device("cpu") 91 | log_dir = "/path/to/log" 92 | 93 | with self.assertRaises(ValueError) as context: 94 | get_ppo_runner_and_checkpoint_path(teacher_policy_cfg, wrapped_env, device, log_dir) 95 | self.assertEqual(str(context.exception), "teacher_policy.checkpoint is not specified") 96 | -------------------------------------------------------------------------------- /scripts/rsl_rl/utils.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | 17 | from __future__ import annotations 18 | 19 | import argparse 20 | import os 21 | from typing import TYPE_CHECKING 22 | 23 | if TYPE_CHECKING: 24 | import torch 25 | 26 | from teacher_policy_cfg import TeacherPolicyCfg 27 | 28 | from neural_wbc.core import EnvironmentWrapper 29 | 30 | 31 | def get_customized_rsl_rl(): 32 | """Helper function to ensure the correct version of rsl_rl is imported. 33 | 34 | This function does the following: 35 | 1. Gets the installed rsl_rl package location and adds it to sys.path 36 | 2. Removes any existing rsl_rl and submodules from sys.modules to force reimporting 37 | """ 38 | import sys 39 | 40 | import pkg_resources 41 | 42 | dist = pkg_resources.require("rsl_rl")[0] 43 | sys.path.insert(0, dist.location) 44 | 45 | # Remove 'rsl_rl' from sys.modules if it was already imported 46 | modules_to_remove = [key for key in sys.modules if key.startswith("rsl_rl")] 47 | for module in modules_to_remove: 48 | print(f"Removing {module} from sys.modules") 49 | del sys.modules[module] 50 | 51 | 52 | def get_player_args(description: str) -> argparse.ArgumentParser: 53 | parser = argparse.ArgumentParser(description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter) 54 | parser.add_argument("--num_envs", type=int, default=2, help="Number of environments to simulate.") 55 | parser.add_argument("--env_spacing", type=int, default=5, help="Distance between environments in simulator.") 56 | parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment") 57 | parser.add_argument("--reference_motion_path", type=str, default=None, help="Path to the reference motion dataset.") 58 | parser.add_argument("--robot", type=str, choices=["h1", "gr1"], default="h1", help="Robot used in environment") 59 | parser.add_argument( 60 | "--student_player", action="store_true", help="Whether the evaluated policy is a student policy." 61 | ) 62 | parser.add_argument("--student_path", type=str, default=None, help="The path for the student policy.") 63 | parser.add_argument("--student_checkpoint", type=str, default=None, help="The exact checkpoint.") 64 | return parser 65 | 66 | 67 | get_customized_rsl_rl() 68 | from rsl_rl.runners import OnPolicyRunner 69 | 70 | 71 | def get_ppo_runner_and_checkpoint_path( 72 | teacher_policy_cfg: TeacherPolicyCfg, 73 | wrapped_env: EnvironmentWrapper, 74 | device: torch.device, 75 | log_dir: str | None = None, 76 | ) -> tuple[OnPolicyRunner, str]: 77 | resume_path = teacher_policy_cfg.runner.resume_path 78 | if not resume_path: 79 | raise ValueError("teacher_policy.resume_path is not specified") 80 | checkpoint = teacher_policy_cfg.runner.checkpoint 81 | if not checkpoint: 82 | raise ValueError("teacher_policy.checkpoint is not specified") 83 | # specify directory for logging experiments 84 | print(f"[INFO] Loading experiment from directory: {teacher_policy_cfg.runner.path}") 85 | checkpoint_path = os.path.join(resume_path, teacher_policy_cfg.runner.checkpoint) 86 | print(f"[INFO]: Loading model checkpoint from: {checkpoint_path}") 87 | 88 | # Try overwrite policy configuration with the content from {log_root_path}/config.json. 89 | teacher_policy_cfg.overwrite_policy_cfg_from_file(os.path.join(resume_path, "config.json")) 90 | 91 | # load previously trained model 92 | ppo_runner = OnPolicyRunner(wrapped_env, teacher_policy_cfg.to_dict(), log_dir=log_dir, device=device) 93 | 94 | return ppo_runner, checkpoint_path 95 | -------------------------------------------------------------------------------- /third_party/grad_fit_h1.patch: -------------------------------------------------------------------------------- 1 | 116a117,118 2 | > if not amass_data: 3 | > continue 4 | 169,176c171,172 5 | < } 6 | < 7 | < print(f"dumping {data_key} for testing, remove the line if you want to process all data") 8 | < import ipdb; ipdb.set_trace() 9 | < joblib.dump(data_dump, "data/h1/test.pkl") 10 | < 11 | < 12 | < import ipdb; ipdb.set_trace() 13 | --- 14 | > } 15 | > 16 | -------------------------------------------------------------------------------- /third_party/mujoco_viewer/LICENCE: -------------------------------------------------------------------------------- 1 | BSD 2-Clause License 2 | 3 | Copyright (c) 2022, Rohan P. Singh 4 | All rights reserved. 5 | 6 | Redistribution and use in source and binary forms, with or without 7 | modification, are permitted provided that the following conditions are met: 8 | 9 | 1. Redistributions of source code must retain the above copyright notice, this 10 | list of conditions and the following disclaimer. 11 | 12 | 2. Redistributions in binary form must reproduce the above copyright notice, 13 | this list of conditions and the following disclaimer in the documentation 14 | and/or other materials provided with the distribution. 15 | 16 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 17 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 18 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 19 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 20 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 21 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 22 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 23 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 24 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -------------------------------------------------------------------------------- /third_party/mujoco_viewer/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | # Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | # SPDX-License-Identifier: Apache-2.0 17 | -------------------------------------------------------------------------------- /third_party/mujoco_viewer/mujoco_viewer/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES 2 | # Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | # SPDX-License-Identifier: Apache-2.0 17 | -------------------------------------------------------------------------------- /third_party/mujoco_viewer/setup.py: -------------------------------------------------------------------------------- 1 | 2 | import os 3 | from setuptools import find_packages, setup 4 | 5 | VERSION = '0.1.4' 6 | 7 | INSTALL_REQUIRES = ( 8 | ['mujoco >= 2.1.5', 9 | 'glfw >= 2.5.0', 10 | 'imageio', 11 | 'pyyaml'] 12 | ) 13 | 14 | setup( 15 | name='mujoco-python-viewer', 16 | version=VERSION, 17 | author='Rohan P. Singh', 18 | author_email='rohan565singh@gmail.com', 19 | url='https://github.com/rohanpsingh/mujoco-python-viewer', 20 | description='Interactive renderer for MuJoCo Python', 21 | long_description='Interactive renderer for MuJoCo Python', 22 | install_requires=INSTALL_REQUIRES, 23 | packages=find_packages(), 24 | python_requires='>=3.6', 25 | classifiers=[ 26 | 'Programming Language :: Python :: 3', 27 | 'Programming Language :: Python :: 3.6', 28 | 'Programming Language :: Python :: 3.7', 29 | 'Programming Language :: Python :: 3 :: Only', 30 | ], 31 | zip_safe=False, 32 | ) 33 | -------------------------------------------------------------------------------- /third_party/phc_torch_utils.patch: -------------------------------------------------------------------------------- 1 | 31d30 2 | < from isaacgym.torch_utils import * 3 | 38c37,102 4 | < def project_to_norm(x, norm=5, z_type = "sphere"): 5 | --- 6 | > # The following functions are copied from Isaac Lab to avoid dependency on it. 7 | > # ---- Start of functions from Isaac Lab ----- 8 | > @torch.jit.script 9 | > def quat_from_euler_xyz(roll: torch.Tensor, pitch: torch.Tensor, yaw: torch.Tensor) -> torch.Tensor: 10 | > """Convert rotations given as Euler angles in radians to Quaternions. 11 | > 12 | > Note: 13 | > The euler angles are assumed in XYZ convention. 14 | > 15 | > Args: 16 | > roll: Rotation around x-axis (in radians). Shape is (N,). 17 | > pitch: Rotation around y-axis (in radians). Shape is (N,). 18 | > yaw: Rotation around z-axis (in radians). Shape is (N,). 19 | > 20 | > Returns: 21 | > The quaternion in (w, x, y, z). Shape is (N, 4). 22 | > """ 23 | > cy = torch.cos(yaw * 0.5) 24 | > sy = torch.sin(yaw * 0.5) 25 | > cr = torch.cos(roll * 0.5) 26 | > sr = torch.sin(roll * 0.5) 27 | > cp = torch.cos(pitch * 0.5) 28 | > sp = torch.sin(pitch * 0.5) 29 | > # compute quaternion 30 | > qw = cy * cr * cp + sy * sr * sp 31 | > qx = cy * sr * cp - sy * cr * sp 32 | > qy = cy * cr * sp + sy * sr * cp 33 | > qz = sy * cr * cp - cy * sr * sp 34 | > 35 | > return torch.stack([qw, qx, qy, qz], dim=-1) 36 | > 37 | > 38 | > @torch.jit.script 39 | > def normalize(x: torch.Tensor, eps: float = 1e-9) -> torch.Tensor: 40 | > """Normalizes a given input tensor to unit length. 41 | > 42 | > Args: 43 | > x: Input tensor of shape (N, dims). 44 | > eps: A small value to avoid division by zero. Defaults to 1e-9. 45 | > 46 | > Returns: 47 | > Normalized tensor of shape (N, dims). 48 | > """ 49 | > return x / x.norm(p=2, dim=-1).clamp(min=eps, max=None).unsqueeze(-1) 50 | > 51 | > 52 | > @torch.jit.script 53 | > def quat_from_angle_axis(angle: torch.Tensor, axis: torch.Tensor) -> torch.Tensor: 54 | > """Convert rotations given as angle-axis to quaternions. 55 | > 56 | > Args: 57 | > angle: The angle turned anti-clockwise in radians around the vector's direction. Shape is (N,). 58 | > axis: The axis of rotation. Shape is (N, 3). 59 | > 60 | > Returns: 61 | > The quaternion in (w, x, y, z). Shape is (N, 4). 62 | > """ 63 | > theta = (angle / 2).unsqueeze(-1) 64 | > xyz = normalize(axis) * theta.sin() 65 | > w = theta.cos() 66 | > return normalize(torch.cat([w, xyz], dim=-1)) 67 | > 68 | > # ---- End of functions from Isaac Lab ----- 69 | > 70 | > 71 | > def project_to_norm(x, norm=5, z_type="sphere"): 72 | 44a109 73 | > 74 | 52,54c117 75 | < c = q_vec * \ 76 | < torch.bmm(q_vec.view(shape[0], 1, 3), v.view( 77 | < shape[0], 3, 1)).squeeze(-1) * 2.0 78 | --- 79 | > c = q_vec * torch.bmm(q_vec.view(shape[0], 1, 3), v.view(shape[0], 3, 1)).squeeze(-1) * 2.0 80 | 56a120,125 81 | > 82 | > @torch.jit.script 83 | > def normalize_angle(x): 84 | > return torch.atan2(torch.sin(x), torch.cos(x)) 85 | > 86 | > 87 | 62c131 88 | < # ZL: could have issues. 89 | --- 90 | > # ZL: could have issues. 91 | 144c213,214 92 | < q = quat_from_euler_xyz(roll, pitch, yaw) 93 | --- 94 | > q = quat_from_euler_xyz(roll, pitch, yaw) # IsaacLab return wxyz 95 | > q = q.roll(-1, dims=-1) 96 | 172c242,243 97 | < q = quat_from_angle_axis(angle, axis) 98 | --- 99 | > q = quat_from_angle_axis(angle, axis) # IsaacLab return wxyz 100 | > q = q.roll(-1, dims=-1) 101 | 226c297,298 102 | < heading_q = quat_from_angle_axis(heading, axis) 103 | --- 104 | > heading_q = quat_from_angle_axis(heading, axis) # IsaacLab return wxyz 105 | > heading_q = heading_q.roll(-1, dims=-1) 106 | 240c312,313 107 | < heading_q = quat_from_angle_axis(-heading, axis) 108 | --- 109 | > heading_q = quat_from_angle_axis(-heading, axis) # IsaacLab return wxyz 110 | > heading_q = heading_q.roll(-1, dims=-1) 111 | 242a316 112 | > 113 | 244c318 114 | < if act_name == 'relu': 115 | --- 116 | > if act_name == "relu": 117 | 246c320 118 | < elif act_name == 'tanh': 119 | --- 120 | > elif act_name == "tanh": 121 | 248c322 122 | < elif act_name == 'sigmoid': 123 | --- 124 | > elif act_name == "sigmoid": 125 | 261c335 126 | < return nn.Identity 127 | \ No newline at end of file 128 | --- 129 | > return nn.Identity 130 | -------------------------------------------------------------------------------- /third_party/requirements.patch: -------------------------------------------------------------------------------- 1 | 40d39 2 | < -e legged_gym 3 | 45d43 4 | < pytorch3d 5 | -------------------------------------------------------------------------------- /third_party/rsl_rl/.gitignore: -------------------------------------------------------------------------------- 1 | # IDEs 2 | .idea 3 | 4 | # builds 5 | *.egg-info 6 | 7 | # cache 8 | __pycache__ 9 | .pytest_cache 10 | 11 | # vs code 12 | .vscode -------------------------------------------------------------------------------- /third_party/rsl_rl/LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2021, ETH Zurich, Nikita Rudin 2 | Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES 3 | All rights reserved. 4 | 5 | Redistribution and use in source and binary forms, with or without modification, 6 | are permitted provided that the following conditions are met: 7 | 8 | 1. Redistributions of source code must retain the above copyright notice, 9 | this list of conditions and the following disclaimer. 10 | 11 | 2. Redistributions in binary form must reproduce the above copyright notice, 12 | this list of conditions and the following disclaimer in the documentation 13 | and/or other materials provided with the distribution. 14 | 15 | 3. Neither the name of the copyright holder nor the names of its contributors 16 | may be used to endorse or promote products derived from this software without 17 | specific prior written permission. 18 | 19 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 20 | ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 21 | WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 22 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR 23 | ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES 24 | (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 25 | LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND 26 | ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 27 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 28 | SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 29 | 30 | See licenses/dependencies for license information of dependencies of this package. -------------------------------------------------------------------------------- /third_party/rsl_rl/README.md: -------------------------------------------------------------------------------- 1 | # RSL RL 2 | Fast and simple implementation of RL algorithms, designed to run fully on GPU. 3 | This code is an evolution of `rl-pytorch` provided with NVIDIA's Isaac GYM. 4 | 5 | Only PPO is implemented for now. More algorithms will be added later. 6 | Contributions are welcome. 7 | 8 | ## Setup 9 | 10 | ``` 11 | git clone https://github.com/leggedrobotics/rsl_rl 12 | cd rsl_rl 13 | pip install -e . 14 | ``` 15 | 16 | ### Useful Links ### 17 | Example use case: https://github.com/leggedrobotics/legged_gym 18 | Project website: https://leggedrobotics.github.io/legged_gym/ 19 | Paper: https://arxiv.org/abs/2109.11978 20 | 21 | **Maintainer**: Nikita Rudin 22 | **Affiliation**: Robotic Systems Lab, ETH Zurich & NVIDIA 23 | **Contact**: rudinn@ethz.ch -------------------------------------------------------------------------------- /third_party/rsl_rl/licenses/dependencies/numpy_license.txt: -------------------------------------------------------------------------------- 1 | Copyright (c) 2005-2021, NumPy Developers. 2 | All rights reserved. 3 | 4 | Redistribution and use in source and binary forms, with or without 5 | modification, are permitted provided that the following conditions are 6 | met: 7 | 8 | * Redistributions of source code must retain the above copyright 9 | notice, this list of conditions and the following disclaimer. 10 | 11 | * Redistributions in binary form must reproduce the above 12 | copyright notice, this list of conditions and the following 13 | disclaimer in the documentation and/or other materials provided 14 | with the distribution. 15 | 16 | * Neither the name of the NumPy Developers nor the names of any 17 | contributors may be used to endorse or promote products derived 18 | from this software without specific prior written permission. 19 | 20 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 21 | "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 22 | LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 23 | A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT 24 | OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 25 | SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT 26 | LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 27 | DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 28 | THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 29 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 30 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -------------------------------------------------------------------------------- /third_party/rsl_rl/licenses/dependencies/torch_license.txt: -------------------------------------------------------------------------------- 1 | From PyTorch: 2 | 3 | Copyright (c) 2016- Facebook, Inc (Adam Paszke) 4 | Copyright (c) 2014- Facebook, Inc (Soumith Chintala) 5 | Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert) 6 | Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu) 7 | Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu) 8 | Copyright (c) 2011-2013 NYU (Clement Farabet) 9 | Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston) 10 | Copyright (c) 2006 Idiap Research Institute (Samy Bengio) 11 | Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz) 12 | 13 | From Caffe2: 14 | 15 | Copyright (c) 2016-present, Facebook Inc. All rights reserved. 16 | 17 | All contributions by Facebook: 18 | Copyright (c) 2016 Facebook Inc. 19 | 20 | All contributions by Google: 21 | Copyright (c) 2015 Google Inc. 22 | All rights reserved. 23 | 24 | All contributions by Yangqing Jia: 25 | Copyright (c) 2015 Yangqing Jia 26 | All rights reserved. 27 | 28 | All contributions by Kakao Brain: 29 | Copyright 2019-2020 Kakao Brain 30 | 31 | All contributions from Caffe: 32 | Copyright(c) 2013, 2014, 2015, the respective contributors 33 | All rights reserved. 34 | 35 | All other contributions: 36 | Copyright(c) 2015, 2016 the respective contributors 37 | All rights reserved. 38 | 39 | Caffe2 uses a copyright model similar to Caffe: each contributor holds 40 | copyright over their contributions to Caffe2. The project versioning records 41 | all such contribution and copyright details. If a contributor wants to further 42 | mark their specific copyright on a particular contribution, they should 43 | indicate their copyright solely in the commit message of the change when it is 44 | committed. 45 | 46 | All rights reserved. 47 | 48 | Redistribution and use in source and binary forms, with or without 49 | modification, are permitted provided that the following conditions are met: 50 | 51 | 1. Redistributions of source code must retain the above copyright 52 | notice, this list of conditions and the following disclaimer. 53 | 54 | 2. Redistributions in binary form must reproduce the above copyright 55 | notice, this list of conditions and the following disclaimer in the 56 | documentation and/or other materials provided with the distribution. 57 | 58 | 3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America 59 | and IDIAP Research Institute nor the names of its contributors may be 60 | used to endorse or promote products derived from this software without 61 | specific prior written permission. 62 | 63 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 64 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 65 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 66 | ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE 67 | LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 68 | CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 69 | SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 70 | INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 71 | CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 72 | ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 73 | POSSIBILITY OF SUCH DAMAGE. -------------------------------------------------------------------------------- /third_party/rsl_rl/rsl_rl/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: BSD-3-Clause 3 | # 4 | # Redistribution and use in source and binary forms, with or without 5 | # modification, are permitted provided that the following conditions are met: 6 | # 7 | # 1. Redistributions of source code must retain the above copyright notice, this 8 | # list of conditions and the following disclaimer. 9 | # 10 | # 2. Redistributions in binary form must reproduce the above copyright notice, 11 | # this list of conditions and the following disclaimer in the documentation 12 | # and/or other materials provided with the distribution. 13 | # 14 | # 3. Neither the name of the copyright holder nor the names of its 15 | # contributors may be used to endorse or promote products derived from 16 | # this software without specific prior written permission. 17 | # 18 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | # 29 | # Copyright (c) 2021 ETH Zurich, Nikita Rudin -------------------------------------------------------------------------------- /third_party/rsl_rl/rsl_rl/algorithms/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: BSD-3-Clause 3 | # 4 | # Redistribution and use in source and binary forms, with or without 5 | # modification, are permitted provided that the following conditions are met: 6 | # 7 | # 1. Redistributions of source code must retain the above copyright notice, this 8 | # list of conditions and the following disclaimer. 9 | # 10 | # 2. Redistributions in binary form must reproduce the above copyright notice, 11 | # this list of conditions and the following disclaimer in the documentation 12 | # and/or other materials provided with the distribution. 13 | # 14 | # 3. Neither the name of the copyright holder nor the names of its 15 | # contributors may be used to endorse or promote products derived from 16 | # this software without specific prior written permission. 17 | # 18 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | # 29 | # Copyright (c) 2021 ETH Zurich, Nikita Rudin 30 | 31 | from .ppo import PPO 32 | -------------------------------------------------------------------------------- /third_party/rsl_rl/rsl_rl/env/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: BSD-3-Clause 3 | # 4 | # Redistribution and use in source and binary forms, with or without 5 | # modification, are permitted provided that the following conditions are met: 6 | # 7 | # 1. Redistributions of source code must retain the above copyright notice, this 8 | # list of conditions and the following disclaimer. 9 | # 10 | # 2. Redistributions in binary form must reproduce the above copyright notice, 11 | # this list of conditions and the following disclaimer in the documentation 12 | # and/or other materials provided with the distribution. 13 | # 14 | # 3. Neither the name of the copyright holder nor the names of its 15 | # contributors may be used to endorse or promote products derived from 16 | # this software without specific prior written permission. 17 | # 18 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | # 29 | # Copyright (c) 2021 ETH Zurich, Nikita Rudin 30 | 31 | from .vec_env import VecEnv -------------------------------------------------------------------------------- /third_party/rsl_rl/rsl_rl/env/vec_env.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: BSD-3-Clause 3 | # 4 | # Redistribution and use in source and binary forms, with or without 5 | # modification, are permitted provided that the following conditions are met: 6 | # 7 | # 1. Redistributions of source code must retain the above copyright notice, this 8 | # list of conditions and the following disclaimer. 9 | # 10 | # 2. Redistributions in binary form must reproduce the above copyright notice, 11 | # this list of conditions and the following disclaimer in the documentation 12 | # and/or other materials provided with the distribution. 13 | # 14 | # 3. Neither the name of the copyright holder nor the names of its 15 | # contributors may be used to endorse or promote products derived from 16 | # this software without specific prior written permission. 17 | # 18 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | # 29 | # Copyright (c) 2021 ETH Zurich, Nikita Rudin 30 | 31 | from abc import ABC, abstractmethod 32 | import torch 33 | from typing import Tuple, Union 34 | 35 | # minimal interface of the environment 36 | class VecEnv(ABC): 37 | num_envs: int 38 | num_obs: int 39 | num_privileged_obs: int 40 | num_actions: int 41 | max_episode_length: int 42 | privileged_obs_buf: torch.Tensor 43 | obs_buf: torch.Tensor 44 | rew_buf: torch.Tensor 45 | reset_buf: torch.Tensor 46 | episode_length_buf: torch.Tensor # current episode duration 47 | extras: dict 48 | device: torch.device 49 | @abstractmethod 50 | def step(self, actions: torch.Tensor) -> Tuple[torch.Tensor, Union[torch.Tensor, None], torch.Tensor, torch.Tensor, dict]: 51 | pass 52 | @abstractmethod 53 | def reset(self, env_ids: Union[list, torch.Tensor]): 54 | pass 55 | @abstractmethod 56 | def get_observations(self) -> torch.Tensor: 57 | pass 58 | @abstractmethod 59 | def get_privileged_observations(self) -> Union[torch.Tensor, None]: 60 | pass -------------------------------------------------------------------------------- /third_party/rsl_rl/rsl_rl/modules/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: BSD-3-Clause 3 | # 4 | # Redistribution and use in source and binary forms, with or without 5 | # modification, are permitted provided that the following conditions are met: 6 | # 7 | # 1. Redistributions of source code must retain the above copyright notice, this 8 | # list of conditions and the following disclaimer. 9 | # 10 | # 2. Redistributions in binary form must reproduce the above copyright notice, 11 | # this list of conditions and the following disclaimer in the documentation 12 | # and/or other materials provided with the distribution. 13 | # 14 | # 3. Neither the name of the copyright holder nor the names of its 15 | # contributors may be used to endorse or promote products derived from 16 | # this software without specific prior written permission. 17 | # 18 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | # 29 | # Copyright (c) 2021 ETH Zurich, Nikita Rudin 30 | 31 | from .actor_critic import ActorCritic 32 | -------------------------------------------------------------------------------- /third_party/rsl_rl/rsl_rl/runners/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: BSD-3-Clause 3 | # 4 | # Redistribution and use in source and binary forms, with or without 5 | # modification, are permitted provided that the following conditions are met: 6 | # 7 | # 1. Redistributions of source code must retain the above copyright notice, this 8 | # list of conditions and the following disclaimer. 9 | # 10 | # 2. Redistributions in binary form must reproduce the above copyright notice, 11 | # this list of conditions and the following disclaimer in the documentation 12 | # and/or other materials provided with the distribution. 13 | # 14 | # 3. Neither the name of the copyright holder nor the names of its 15 | # contributors may be used to endorse or promote products derived from 16 | # this software without specific prior written permission. 17 | # 18 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | # 29 | # Copyright (c) 2021 ETH Zurich, Nikita Rudin 30 | 31 | from .on_policy_runner import OnPolicyRunner 32 | -------------------------------------------------------------------------------- /third_party/rsl_rl/rsl_rl/storage/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 ETH Zurich, NVIDIA CORPORATION 2 | # SPDX-License-Identifier: BSD-3-Clause 3 | 4 | from .rollout_storage import RolloutStorage 5 | -------------------------------------------------------------------------------- /third_party/rsl_rl/rsl_rl/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: BSD-3-Clause 3 | # 4 | # Redistribution and use in source and binary forms, with or without 5 | # modification, are permitted provided that the following conditions are met: 6 | # 7 | # 1. Redistributions of source code must retain the above copyright notice, this 8 | # list of conditions and the following disclaimer. 9 | # 10 | # 2. Redistributions in binary form must reproduce the above copyright notice, 11 | # this list of conditions and the following disclaimer in the documentation 12 | # and/or other materials provided with the distribution. 13 | # 14 | # 3. Neither the name of the copyright holder nor the names of its 15 | # contributors may be used to endorse or promote products derived from 16 | # this software without specific prior written permission. 17 | # 18 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | # 29 | # Copyright (c) 2021 ETH Zurich, Nikita Rudin 30 | 31 | from .utils import split_and_pad_trajectories, unpad_trajectories -------------------------------------------------------------------------------- /third_party/rsl_rl/rsl_rl/utils/utils.py: -------------------------------------------------------------------------------- 1 | # SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # SPDX-License-Identifier: BSD-3-Clause 3 | # 4 | # Redistribution and use in source and binary forms, with or without 5 | # modification, are permitted provided that the following conditions are met: 6 | # 7 | # 1. Redistributions of source code must retain the above copyright notice, this 8 | # list of conditions and the following disclaimer. 9 | # 10 | # 2. Redistributions in binary form must reproduce the above copyright notice, 11 | # this list of conditions and the following disclaimer in the documentation 12 | # and/or other materials provided with the distribution. 13 | # 14 | # 3. Neither the name of the copyright holder nor the names of its 15 | # contributors may be used to endorse or promote products derived from 16 | # this software without specific prior written permission. 17 | # 18 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | # 29 | # Copyright (c) 2021 ETH Zurich, Nikita Rudin 30 | 31 | import torch 32 | 33 | def split_and_pad_trajectories(tensor, dones): 34 | """ Splits trajectories at done indices. Then concatenates them and padds with zeros up to the length og the longest trajectory. 35 | Returns masks corresponding to valid parts of the trajectories 36 | Example: 37 | Input: [ [a1, a2, a3, a4 | a5, a6], 38 | [b1, b2 | b3, b4, b5 | b6] 39 | ] 40 | 41 | Output:[ [a1, a2, a3, a4], | [ [True, True, True, True], 42 | [a5, a6, 0, 0], | [True, True, False, False], 43 | [b1, b2, 0, 0], | [True, True, False, False], 44 | [b3, b4, b5, 0], | [True, True, True, False], 45 | [b6, 0, 0, 0] | [True, False, False, False], 46 | ] | ] 47 | 48 | Assumes that the inputy has the following dimension order: [time, number of envs, aditional dimensions] 49 | """ 50 | dones = dones.clone() 51 | dones[-1] = 1 52 | # Permute the buffers to have order (num_envs, num_transitions_per_env, ...), for correct reshaping 53 | flat_dones = dones.transpose(1, 0).reshape(-1, 1) 54 | 55 | # Get length of trajectory by counting the number of successive not done elements 56 | done_indices = torch.cat((flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero()[:, 0])) 57 | trajectory_lengths = done_indices[1:] - done_indices[:-1] 58 | trajectory_lengths_list = trajectory_lengths.tolist() 59 | # Extract the individual trajectories 60 | trajectories = torch.split(tensor.transpose(1, 0).flatten(0, 1),trajectory_lengths_list) 61 | padded_trajectories = torch.nn.utils.rnn.pad_sequence(trajectories) 62 | 63 | 64 | trajectory_masks = trajectory_lengths > torch.arange(0, tensor.shape[0], device=tensor.device).unsqueeze(1) 65 | return padded_trajectories, trajectory_masks 66 | 67 | def unpad_trajectories(trajectories, masks): 68 | """ Does the inverse operation of split_and_pad_trajectories() 69 | """ 70 | # Need to transpose before and after the masking to have proper reshaping 71 | return trajectories.transpose(1, 0)[masks.transpose(1, 0)].view(-1, trajectories.shape[0], trajectories.shape[-1]).transpose(1, 0) -------------------------------------------------------------------------------- /third_party/rsl_rl/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | 3 | setup( 4 | name="rsl_rl", 5 | version="1.0.2", 6 | author="Nikita Rudin", 7 | author_email="rudinn@ethz.ch", 8 | license="BSD-3-Clause", 9 | packages=find_packages(), 10 | description="Fast and simple RL algorithms implemented in pytorch", 11 | python_requires=">=3.6", 12 | install_requires=["torch>=1.4.0", "torchvision>=0.5.0", "numpy>=1.16.4"], 13 | ) 14 | --------------------------------------------------------------------------------