Eval GR00T N1 model

tools/env_setup_robot_us.sh

@@ -17,6 +17,44 @@
 
 set -e
 
+# --- Configuration ---
+INSTALL_WITH_POLICY="pi0" # Default value
+
+# --- Helper Functions ---
+usage() {
+    echo "Usage: $0 --policy [pi0|gr00tn1|none]"
+    echo "  pi0:   Install base dependencies + PI0 policy dependencies (openpi)."
+    echo "  gr00tn1: Install base dependencies + GR00T N1 policy dependencies (Isaac-GR00T)."
+    echo "  none:  Install only base dependencies (IsaacSim, IsaacLab, Holoscan, etc.)."
+    exit 1
+}
+
+# --- Argument Parsing ---
+while [[ $# -gt 0 ]]; do
+    key="$1"
+    case $key in
+        --policy)
+        INSTALL_WITH_POLICY="$2"
+        shift # past argument
+        shift # past value
+        ;;
+        *)    # unknown option
+        usage
+        ;;
+    esac
+done
+
+# Validate policy argument
+if [[ "$INSTALL_WITH_POLICY" != "pi0" && "$INSTALL_WITH_POLICY" != "gr00tn1" && "$INSTALL_WITH_POLICY" != "none" ]]; then
+    echo "Error: Invalid policy specified."
+    usage
+fi
+
+echo "Selected policy setup: $INSTALL_WITH_POLICY"
+
+
+# --- Setup Steps ---
+
 # Get the parent directory of the current script
 PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && cd .. && pwd)"
 
@@ -34,14 +72,20 @@ if ! nvidia-smi &> /dev/null; then
     exit 1
 fi
 
-# Check if the third_party directory exists, if yes, then exit
-if [ -d "$PROJECT_ROOT/third_party" ]; then
-    echo "Error: third_party directory already exists"
-    echo "Please remove the third_party directory before running this script"
-    exit 1
+# Check if the third_party directory exists
+if [[ "$INSTALL_WITH_POLICY" == "pi0" || "$INSTALL_WITH_POLICY" == "gr00tn1" ]]; then
+    if [ -d "$PROJECT_ROOT/third_party" ]; then
+        echo "Error: third_party directory already exists"
+        echo "Please remove the third_party directory before running this script"
+        exit 1
+    else
+        mkdir $PROJECT_ROOT/third_party
+        echo "Created directory: $PROJECT_ROOT/third_party"
+    fi
 fi
 
-# ---- Install build tools ----
+
+# ---- Install build tools (Common) ----
 echo "Installing build tools..."
 if [ "$EUID" -ne 0 ]; then
     sudo apt-get install -y git cmake build-essential pybind11-dev libxcb-cursor0
@@ -50,85 +94,38 @@ else
 fi
 
 
-# ---- Install IsaacSim and necessary dependencies ----
-echo "Installing IsaacSim..."
+# ---- Install IsaacSim and necessary dependencies (Common) ----
+echo "Installing IsaacSim and base dependencies..."
 pip install 'isaacsim[all,extscache]==4.5.0' \
     rti.connext==7.3.0 pyrealsense2==2.55.1.6486 toml==0.10.2 dearpygui==2.0.0 \
     git+ssh://[email protected]/isaac-for-healthcare/[email protected] \
     setuptools==75.8.0 pydantic==2.10.6 \
     --extra-index-url https://pypi.nvidia.com
 
 
-# ---- Install IsaacLab ----
+# ---- Install IsaacLab (Common) ----
+# Check if IsaacLab is already cloned
 echo "Installing IsaacLab..."
-# CLONING REPOSITORIES INTO PROJECT_ROOT/third_party
-echo "Cloning repositories into $PROJECT_ROOT/third_party..."
-mkdir $PROJECT_ROOT/third_party
+echo "Cloning IsaacLab repository into $PROJECT_ROOT/third_party/IsaacLab..."
 git clone -b v2.0.2 [email protected]:isaac-sim/IsaacLab.git $PROJECT_ROOT/third_party/IsaacLab
 pushd $PROJECT_ROOT/third_party/IsaacLab
 yes Yes | ./isaaclab.sh --install
 popd
 
 
-# ---- Install robotic ultrasound extension ----
+# ---- Install robotic ultrasound extension (Common) ----
 echo "Installing robotic ultrasound extension..."
 pushd $PROJECT_ROOT/workflows/robotic_ultrasound/scripts/simulation
-pip install -e exts/robotic_us_ext
-popd
-
-
-# ---- Install OpenPI with IsaacSim ----
-echo "Installing OpenPI..."
-# Clone the openpi repository
-git clone [email protected]:Physical-Intelligence/openpi.git $PROJECT_ROOT/third_party/openpi
-pushd $PROJECT_ROOT/third_party/openpi
-git checkout 581e07d73af36d336cef1ec9d7172553b2332193
-
-# Update python version in pyproject.toml
-pyproject_path="$PROJECT_ROOT/third_party/openpi/pyproject.toml"
-sed -i.bak \
-    -e 's/requires-python = ">=3.11"/requires-python = ">=3.10"/' \
-    -e 's/"s3fs>=2024.9.0"/"s3fs==2024.9.0"/' \
-    "$pyproject_path"
-
-# Apply temporary workaround for openpi/src/openpi/shared/download.py
-file_path="$PROJECT_ROOT/third_party/openpi/src/openpi/shared/download.py"
-
-# Comment out specific import lines
-sed -i.bak \
-    -e 's/^import boto3\.s3\.transfer as s3_transfer/# import boto3.s3.transfer as s3_transfer/' \
-    -e 's/^import s3transfer\.futures as s3_transfer_futures/# import s3transfer.futures as s3_transfer_futures/' \
-    -e 's/^from types_boto3_s3\.service_resource import ObjectSummary/# from types_boto3_s3.service_resource import ObjectSummary/' \
-    "$file_path"
-
-# Remove the type hint
-sed -i.bak -e 's/)[[:space:]]*-> s3_transfer\.TransferManager[[:space:]]*:/):/' "$file_path"
-
-# Modify the datetime line
-sed -i.bak -e 's/datetime\.UTC/datetime.timezone.utc/' "$file_path"
-
-# Modify the type hints in training/utils.py to use Any instead of optax types
-utils_path="$PROJECT_ROOT/third_party/openpi/src/openpi/training/utils.py"
-sed -i.bak \
-    -e 's/opt_state: optax\.OptState/opt_state: Any/' \
-    "$utils_path"
-
-# Remove the backup files
-rm "$pyproject_path.bak"
-rm "$file_path.bak"
-rm "$utils_path.bak"
-
-# Add training script to openpi module
-if [ ! -f $PROJECT_ROOT/third_party/openpi/src/openpi/train.py ]; then
-    cp $PROJECT_ROOT/third_party/openpi/scripts/train.py $PROJECT_ROOT/third_party/openpi/src/openpi/train.py
-fi
-
-# Add norm stats generator script to openpi module
-if [ ! -f $PROJECT_ROOT/third_party/openpi/src/openpi/compute_norm_stats.py ]; then
-    cp $PROJECT_ROOT/third_party/openpi/scripts/compute_norm_stats.py $PROJECT_ROOT/third_party/openpi/src/openpi/compute_norm_stats.py
+# Ensure the target directory exists before installing
+if [ -d "exts/robotic_us_ext" ]; then
+    pip install -e exts/robotic_us_ext
+else
+    echo "Error: robotic_us_ext directory not found in $PROJECT_ROOT/workflows/robotic_ultrasound/scripts/simulation/exts/"
+    exit 1
 fi
+popd
 
-# ---- Install lerobot ----
+# ---- Install lerobot (Common) ----
 echo "Installing lerobot..."
 git clone https://github.com/huggingface/lerobot.git $PROJECT_ROOT/third_party/lerobot
 pushd $PROJECT_ROOT/third_party/lerobot
@@ -140,31 +137,101 @@ sed -i 's/pyav/av/' pyproject.toml
 pip install -e .
 popd
 
-# Install the dependencies
-pip install -e $PROJECT_ROOT/third_party/openpi/packages/openpi-client/
-pip install -e $PROJECT_ROOT/third_party/openpi/
 
-# Revert the "import changes of "$file_path after installation to prevent errors
-sed -i \
-    -e 's/^# import boto3\.s3\.transfer as s3_transfer/import boto3.s3.transfer as s3_transfer/' \
-    -e 's/^# import s3transfer\.futures as s3_transfer_futures/import s3transfer.futures as s3_transfer_futures/' \
-    -e 's/^# from types_boto3_s3\.service_resource import ObjectSummary/from types_boto3_s3.service_resource import ObjectSummary/' \
-    "$file_path"
+# ---- Install PI0 Policy Dependencies (Conditional) ----
+if [[ "$INSTALL_WITH_POLICY" == "pi0" ]]; then
+    echo "------------------------------------------"
+    echo "Installing PI0 Policy Dependencies..."
+    echo "------------------------------------------"
+
+    echo "Cloning OpenPI repository..."
+    git clone [email protected]:Physical-Intelligence/openpi.git $PROJECT_ROOT/third_party/openpi
+    pushd $PROJECT_ROOT/third_party/openpi
+    git checkout 581e07d73af36d336cef1ec9d7172553b2332193
+
+    # Update python version in pyproject.toml
+    pyproject_path="$PROJECT_ROOT/third_party/openpi/pyproject.toml"
+    echo "Patching OpenPI pyproject.toml..."
+    sed -i.bak \
+        -e 's/requires-python = ">=3.11"/requires-python = ">=3.10"/' \
+        -e 's/"s3fs>=2024.9.0"/"s3fs==2024.9.0"/' \
+        "$pyproject_path"
+
+    # Apply temporary workaround for openpi/src/openpi/shared/download.py
+    file_path="$PROJECT_ROOT/third_party/openpi/src/openpi/shared/download.py"
+    echo "Patching OpenPI download.py..."
+    # Comment out specific import lines
+    sed -i.bak \
+        -e 's/^import boto3\.s3\.transfer as s3_transfer/# import boto3.s3.transfer as s3_transfer/' \
+        -e 's/^import s3transfer\.futures as s3_transfer_futures/# import s3transfer.futures as s3_transfer_futures/' \
+        -e 's/^from types_boto3_s3\.service_resource import ObjectSummary/# from types_boto3_s3.service_resource import ObjectSummary/' \
+        "$file_path"
+    # Remove the type hint
+    sed -i.bak -e 's/)[[:space:]]*-> s3_transfer\.TransferManager[[:space:]]*:/):/' "$file_path"
+    # Modify the datetime line
+    sed -i.bak -e 's/datetime\.UTC/datetime.timezone.utc/' "$file_path"
+
+    # Modify the type hints in training/utils.py to use Any instead of optax types
+    utils_path="$PROJECT_ROOT/third_party/openpi/src/openpi/training/utils.py"
+    echo "Patching OpenPI utils.py..."
+    sed -i.bak \
+        -e 's/opt_state: optax\.OptState/opt_state: Any/' \
+        "$utils_path"
+
+    # Remove the backup files
+    rm "$pyproject_path.bak" "$file_path.bak" "$utils_path.bak"
+
+    # Add training script to openpi module
+    echo "Copying OpenPI utility scripts..."
+    if [ ! -f src/openpi/train.py ]; then
+        cp scripts/train.py src/openpi/train.py
+    fi
+    if [ ! -f src/openpi/compute_norm_stats.py ]; then
+        cp scripts/compute_norm_stats.py src/openpi/compute_norm_stats.py
+    fi
+
+    popd # Back to PROJECT_ROOT
+
+    echo "Installing OpenPI Client..."
+    pip install -e $PROJECT_ROOT/third_party/openpi/packages/openpi-client/
+    echo "Installing OpenPI Core..."
+    pip install -e $PROJECT_ROOT/third_party/openpi/
+
+    # Revert the "import changes of "$file_path after installation to prevent errors
+    echo "Reverting temporary patches in OpenPI download.py..."
+    file_path_revert="$PROJECT_ROOT/third_party/openpi/src/openpi/shared/download.py"
+    sed -i \
+        -e 's/^# import boto3\.s3\.transfer as s3_transfer/import boto3.s3.transfer as s3_transfer/' \
+        -e 's/^# import s3transfer\.futures as s3_transfer_futures/import s3transfer.futures as s3_transfer_futures/' \
+        -e 's/^# from types_boto3_s3\.service_resource import ObjectSummary/from types_boto3_s3.service_resource import ObjectSummary/' \
+        "$file_path_revert"
+    echo "PI0 Dependencies installed."
+fi
+
+
+# ---- Install GR00T N1 Policy Dependencies (Conditional) ----
+if [[ "$INSTALL_WITH_POLICY" == "gr00tn1" ]]; then
+    echo "Installing GR00T N1 Policy Dependencies..."
+    git clone https://github.com/NVIDIA/Isaac-GR00T $PROJECT_ROOT/third_party/Isaac-GR00T
+    pushd $PROJECT_ROOT/third_party/Isaac-GR00T
+    sed -i 's/pyav/av/' pyproject.toml
+    pip install -e .
+    popd
+    pip install --no-build-isolation flash-attn==2.7.1.post4
+    echo "GR00T N1 Policy Dependencies installed."
+fi
 
-popd
 
-# ---- Install Holoscan ----
-# Install Holoscan
+# ---- Install Holoscan (Common) ----
+echo "------------------------------------------"
 echo "Installing Holoscan..."
+echo "------------------------------------------"
 conda install -c conda-forge gcc=13.3.0 -y
 pip install holoscan==2.9.0
 
-echo "Dependencies installed successfully!"
-
 HOLOSCAN_DIR=$PROJECT_ROOT/workflows/robotic_ultrasound/scripts/holoscan_apps/
 
-echo "Building Holoscan Apps"
-
+echo "Building Holoscan Apps..."
 pushd $HOLOSCAN_DIR
 
 # clean previous downloads and builds
@@ -176,11 +243,16 @@ rm -rf clarius_cast/lib
 cmake -B build -S . && cmake --build build
 
 popd
-
 echo "Holoscan Apps build completed!"
 
-# ---- Install libstdcxx-ng for raysim ----
+
+# ---- Install libstdcxx-ng for raysim (Common) ----
+echo "------------------------------------------"
 echo "Installing libstdcxx-ng..."
+echo "------------------------------------------"
 conda install -c conda-forge libstdcxx-ng=13.2.0 -y
 
-echo "Dependencies installed successfully!"
+echo "=========================================="
+echo "Environment setup script finished."
+echo "Selected policy dependencies ($INSTALL_WITH_POLICY) should be installed along with base components."
+echo "=========================================="

tools/run_all_tests.py

@@ -111,7 +111,7 @@ def run_tests_with_coverage(workflow_name, skip_xvfb):
                     test_path,
                 ]
             # TODO: move these tests to integration tests
-            elif "test_sim_with_dds" in test_path or "test_pi0" in test_path:
+            elif "test_sim_with_dds" in test_path or "test_policy" in test_path:
                 continue
             elif "test_integration" in test_path:
                 continue

tutorials/sim2real/README.md

@@ -101,7 +101,7 @@ To feed the policy runner real-world data, you need components that capture sens
 
 The [`run_policy.py`](../../workflows/robotic_ultrasound/scripts/policy_runner/run_policy.py) script, located in [`workflows/robotic_ultrasound/scripts/policy_runner/`](../../workflows/robotic_ultrasound/scripts/policy_runner/), acts as the "brain". It subscribes to the DDS topics published above (camera images, joint states) and publishes computed actions.
 
-*   **Configuration:** Ensure the `--topic_in_...` arguments of `run_policy.py` match the topics you are publishing the real sensor data to (e.g., `"topic_your_camera_rgb"`, `"topic_franka_info"`). The `--topic_out` argument defines where the computed actions will be published (e.g., `"topic_franka_ctrl"`). You may also need to adjust settings in [`config.py`](../../workflows/robotic_ultrasound/scripts/policy_runner/config.py) within the same directory.
+*   **Configuration:** Ensure the `--topic_in_...` arguments of `run_policy.py` match the topics you are publishing the real sensor data to (e.g., `"topic_your_camera_rgb"`, `"topic_franka_info"`). The `--topic_out` argument defines where the computed actions will be published (e.g., `"topic_franka_ctrl"`). You may also need to adjust settings in [`config.py`](../../workflows/robotic_ultrasound/scripts/policy_runner/pi0_policy/config.py) within the same directory.
 
 ## Controlling the Real Robot
 

workflows/robotic_ultrasound/README.md

@@ -59,12 +59,39 @@ To use the ultrasound-raytracing simulator, you can choose one of the following
 
 #### Install All Dependencies
 
-You can install all the dependencies by running the following command:
+The main script `tools/env_setup_robot_us.sh` will install all necessary dependencies. It first installs common base components and then policy-specific packages based on an argument.
+
+### Base Components Installed
+
+The following base components are installed by the script:
+- **IsaacSim 4.5.0** (and its core dependencies)
+- **IsaacLab 2.0.2**
+- **Robotic Ultrasound Extension** (`robotic_us_ext`)
+- **Lerobot** (from Hugging Face)
+- **Holoscan 2.9.0** (including associated Holoscan apps)
+- Essential build tools and libraries
+
+### Policy-Specific Dependencies
+
+The script supports installing additional policy-specific dependencies using the `--policy` flag:
+- **`--policy pi0` (Default)**: Installs dependencies for the PI0 policy along with common dependencies (e.g., OpenPI).
+- **`--policy gr00tn1`**: Installs dependencies for the GR00T N1 policy along with common dependencies (e.g., Isaac-GR00T).
+- **`--policy base`**: Installs only the common base dependencies listed above, without any policy-specific packages.
+
+To run the script, navigate to the root of the `i4h-workflows` repository and execute:
 
 ```bash
 cd <path-to-i4h-workflows>
+bash tools/env_setup_robot_us.sh --policy <your_chosen_policy>
+```
+If you omit the `--policy` flag, it will default to `--policy pi0`. For example, to install base dependencies and PI0:
+```bash
 bash tools/env_setup_robot_us.sh
 ```
+Or, to install base dependencies and GR00T:
+```bash
+bash tools/env_setup_robot_us.sh --policy gr00tn1
+```
 
 - **NOTE**: It is expected to see the following messages with `pytorch` version mismatch.
 ```