Sim2Real Deployment: Functionality to deploy the policy to Unitree H1

From pull request #6

Author: steple

Diff for: README.md

@@ -304,7 +304,23 @@ To apply these custom settings, pass the path to your YAML file using the `--env
 option when running the script. If the YAML file contains keys that do not exist in the default
 configuration, those keys will be ignored.
 
-# Sim-to-Sim Validation
+# Validation
+The trained policy in Isaac Lab can be validated in two ways - Sim-to-Sim and Sim-to-Real.
+
+<div align="center">
+<table>
+  <tr>
+    <td><img src="docs/stable_punch.gif" width="400"/></td>
+    <td><img src="docs/stable_wave.gif" width="385"/></td>
+  </tr>
+  <tr>
+    <td align="center"><font size="1"><em>Stable Punch - Mujoco (left) & Real Robot (right)</em></font></td>
+    <td align="center"><font size="1"><em>Stable Wave - Mujoco (left) & Real Robot (right)</em></font></td>
+  </tr>
+</table>
+</div>
+
+## Sim-to-Sim Validation
 
 We also provide a [Mujoco environment](./neural_wbc/mujoco_wrapper/) for conducting sim-to-sim validation of the trained policy.
 To run the evaluation of Sim2Sim,
@@ -317,7 +333,26 @@ ${ISAACLAB_PATH:?}/isaaclab.sh -p neural_wbc/inference_env/scripts/eval.py \
     --student_checkpoint model_<iteration_number>.pt
 ```
 
-Please be aware that the `mujoco_wrapper` only supports one environment at a time. For a reference, it will take up to `5h` to evaluate `8k` reference motions. The inference_env is designed for maximum versatility. It's also suitable for Sim2Real deployment when the user provides the real robot interface.
+Please be aware that the `mujoco_wrapper` only supports one environment at a time. For a reference, it will take up to `5h` to evaluate `8k` reference motions. The inference_env is designed for maximum versatility.
+
+## Sim-to-Real Deployment
+
+For sim-to-real validation, we provide a [Hardware environment](./neural_wbc/hw_wrappers/) for Unitree H1.
+To install the required dependencies and environment setup, please refer to the [README of sim2real deployment](./neural_wbc/hw_wrappers/README.md).
+
+To deploy the trained policy on the [Unitree H1 robot](https://unitree.com/h1),
+
+```bash
+${ISAACLAB_PATH:?}/isaaclab.sh -p neural_wbc/inference_env/scripts/s2r_player.py \
+    --student_path neural_wbc/data/data/policy/h1:student/ \
+    --student_checkpoint model_<iteration_number>.pt \
+    --reference_motion_path neural_wbc/data/data/motions/<motion_name>.pkl \
+    --robot unitree_h1 \
+    --max_iterations 5000 \
+    --num_envs 1
+```
+
+> **_NOTE:_** The sim-to-real deployment wrapper currently only supports the Unitree H1 robot. It can be extended to other robots by implementing the corresponding hardware wrapper interface.
 
 # Developing
 
@@ -410,6 +445,7 @@ We would like to acknowledge the following projects where parts of the codes in
 - [Mujoco Python Viewer](https://github.com/rohanpsingh/mujoco-python-viewer)
 - [RSL RL](https://github.com/leggedrobotics/rsl_rl)
 - [human2humanoid](https://github.com/LeCAR-Lab/human2humanoid)
+- [Unitree Python SDK](https://github.com/unitreerobotics/unitree_sdk2_python)
 
 [omnih2o_paper]: https://arxiv.org/abs/2406.08858
 [hover_paper]: https://arxiv.org/abs/2410.21229

Diff for: docs/stable_punch.gif

@@ -33,13 +33,15 @@ def compute_student_observations(
     history: torch.Tensor,
     mask: torch.Tensor,
     ref_episodic_offset: torch.Tensor | None = None,
+    local_base_ang_velocity: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
     """Computes observations for a student policy."""
     obs_dict = {
         "distilled_robot_state": compute_distilled_robot_state_observation(
             body_state=body_state,
             base_id=base_id,
             projected_gravity=projected_gravity,
+            local_base_ang_velocity=local_base_ang_velocity,
         ),
         "distilled_imitation": compute_distilled_imitation_observations(
             ref_motion_state=ref_motion_state,
@@ -193,14 +195,18 @@ def compute_distilled_robot_state_observation(
     body_state: BodyState,
     base_id: int,
     projected_gravity: torch.Tensor,
+    local_base_ang_velocity: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Root body state in the robot root frame."""
     # for normalization
     joint_pos = body_state.joint_pos.clone()
     joint_vel = body_state.joint_vel.clone()
 
-    base_rot = body_state.body_rot_extend[:, base_id, :].clone()
-    base_ang_vel = body_state.body_ang_vel_extend[:, base_id, :].clone()
-    local_base_ang_vel = math_utils.quat_rotate_inverse(base_rot, base_ang_vel)
+    local_base_ang_vel = local_base_ang_velocity
+
+    if local_base_ang_velocity is None:
+        local_base_ang_vel = math_utils.quat_rotate_inverse(
+            body_state.body_rot_extend[:, base_id, :], body_state.body_ang_vel_extend[:, base_id, :]
+        )
 
     return torch.cat([joint_pos, joint_vel, local_base_ang_vel, projected_gravity], dim=-1)

Diff for: docs/stable_wave.gif

@@ -0,0 +1,98 @@
+# Neural WBC - Hardware Wrapper for Unitree H1
+
+This module is a high-level wrapper for the [Unitree H1](https://unitree.com/h1) robot hardware. It is designed to provide convenient interfaces to deploy the trained policy on the Unitree H1 robot. The main purpose of this wrapper is to validate Whole Body Control (WBC) policies trained in IsaacLab on the real robot.
+> **Note:**  The sim-to-real deployment wrapper currently only supports the Unitree H1 robot. It can be extended to other robots by implementing the corresponding hardware wrapper interface.
+
+## Setup
+Our deployment setup uses an Ubuntu PC with GPU for running policy inference, connected to the H1 robot via Ethernet. For network configuration, please refer to the [Unitree H1 Quick Start Guide](https://support.unitree.com/home/en/H1_developer/start).
+<div align="center">
+<div style="background-color: #f5f5f5; padding: 20px; border-radius: 8px; display: inline-block;">
+<img src="docs/unitree_h1_setup.png" width="400"/>
+<br>
+<font size="1"><em>Unitree H1 setup</em></font>
+</div>
+</div>
+
+> **Note:** This guide assumes familiarity with basic Unitree H1 operation. Please refer to the official [Unitree H1 Operation Guide](https://support.unitree.com/home/en/H1_developer/quick_start) for fundamental robot operations, as these are outside the scope of this guide.
+
+## Installation
+
+To install the required dependencies:
+
+```bash
+${ISAACLAB_PATH:?}/isaaclab.sh -p -m pip install -r requirements_deploy.txt
+```
+For any issues with Unitree SDK installation, see the [Unitree Python SDK](https://github.com/unitreerobotics/unitree_sdk2_python).
+
+## Usage
+
+### Check for successful connection with the robot
+- Make sure the robot is in develop mode using the joystick before executing the below example. Use the [Unitree Operation guide](https://support.unitree.com/home/en/H1_developer/quick_start) to transition the robot to develop mode.
+- Make sure robot is safely mounted on gantry with enough space around it.
+    ```bash
+    cd ~/unitree_sdk2_python
+    python3 example/h1/low_level/h1_low_level_example.py <network interface>
+    ```
+If the robot responds to the low-level example commands, your setup and network connection are working correctly. You can proceed to the next steps after exiting the example program.
+
+### Check for successful setup of the Hover stack:
+- Execute the below commands from the hover stack root directory.
+- First ensure that all dependencies are installed using the common script
+    [./install_deps.sh](../../install_deps.sh)
+
+- Run the simple viewer with:
+
+    ```sh
+    ${ISAACLAB_PATH}/isaaclab.sh -p neural_wbc/inference_env/scripts/mujoco_viewer_player.py
+    ```
+
+The viewer is paused on start by default, press `SPACE` key to start simulation or `RIGHT` arrow key
+to step forward once. By default this script will load the UniTree H1 robot model scene from
+[data/mujoco/models](../data/data/mujoco/models/scene.xml).
+
+## Deployment
+
+> **Note:** Current deployment does not use external sensors for robot root tracking. The setup supports deploying stable motions where the robot stabilizes in place while executing upper body motions.
+
+By default the policy is deployed in OmniH2O mode (tracking head and hand positions). In order to deploy a different mode, change the configuration in [config_file](../inference_env/inference_env/neural_wbc_env_cfg_real_h1.py).
+
+Once sim-to-sim validation is done, the policy can be deployed on the real robot using the following steps.
+
+1. Follow the [Unitree H1 Operation Guide](https://support.unitree.com/home/en/H1_developer/quick_start) to:
+    - Boot the robot (Steps 1-5 under "Boot Process")
+    - Switch to develop mode using the joystick (Steps under "Develop Mode")
+2. Once the H1 robot is transitioned to develop mode, execute this command to deploy the policy while the robot is securely mounted on the gantry:
+    ```bash
+    ${ISAACLAB_PATH:?}/isaaclab.sh -p neural_wbc/inference_env/scripts/s2r_player.py \
+        --student_path neural_wbc/data/data/policy/h1:student/ \
+        --student_checkpoint model_<iteration_number>.pt \
+        --reference_motion_path neural_wbc/data/data/motions/<motion_name>.pkl \
+        --robot unitree_h1 \
+        --max_iterations 5000 \
+        --num_envs 1
+    ```
+    You can now deploy the policy on the real robot in headless mode by passing the `--headless` option.
+
+2. Upon command execution, the robot will automatically move to match the starting pose of the reference motion. You can preview this motion in the Mujoco viewer UI, where the reference trajectory is visualized as red dots.
+
+3. Once the robot is in the starting pose:
+   - Gradually lower the gantry until the robot's feet make contact with the ground
+   - Press the space bar in the Mujoco viewer UI to begin policy execution
+   - Completely lower the gantry to allow free movement of the robot
+
+<div align="center">
+<img src="docs/push_recovery.gif" width="400"/>
+<br>
+<font size="1"><em> Policy deployment of stable motion</em></font>
+</div>
+
+> **Note:** If the robot loses motor torque and becomes unresponsive to joystick commands, restart the robot and contact Unitree support for troubleshooting.
+
+## Unit Tests
+
+The tests are located in the `tests` directory and can be run with the `unittest` module.
+
+```bash
+cd neural_wbc/hw_wrappers
+${ISAACLAB_PATH:?}/isaaclab.sh -p -m unittest
+```

Diff for: neural_wbc/core/neural_wbc/core/observations/student_observations.py

@@ -0,0 +1,17 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Needed to have the robot registry 'activated'
+from hw_wrappers.unitree_h1 import UnitreeH1  # noqa