Add docker build for robotic ultrasound

tools/env_setup/install_cosmos_transfer1.sh

@@ -24,9 +24,6 @@ PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && cd ../.. && pwd)"
 # Allow setting the python in PYTHON_EXECUTABLE
 PYTHON_EXECUTABLE=${PYTHON_EXECUTABLE:-python}
 
-# Install cuDNN
-bash "$PROJECT_ROOT/tools/env_setup/install_cudnn.sh"
-
 COSMOS_TRANSFER_DIR=${1:-$$PROJECT_ROOT/third_party/cosmos-transfer1}
 
 if [ -d "$COSMOS_TRANSFER_DIR" ]; then
@@ -40,7 +37,7 @@ pushd "$COSMOS_TRANSFER_DIR"
 git checkout bf54a70a8c44d615620728c493ee26b4376ccfd6
 git submodule update --init --recursive
 $PYTHON_EXECUTABLE -m pip install -r requirements.txt
-MAX_JOBS=2 $PYTHON_EXECUTABLE -m pip install transformer-engine[pytorch]==1.12.0
+MAX_JOBS=2 $PYTHON_EXECUTABLE -m pip install --no-build-isolation transformer-engine[pytorch]==1.12.0
 $PYTHON_EXECUTABLE -m pip install tensorstore==0.1.74
 
 $PYTHON_EXECUTABLE scripts/test_environment.py

tools/env_setup/install_lerobot.sh

@@ -0,0 +1,39 @@
+#!/bin/bash
+
+# SPDX-FileCopyrightText: Copyright (c) 2025 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.
+
+set -e
+
+# Get the parent directory of the current script
+PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && cd ../.. && pwd)"
+
+# Allow setting the python in PYTHON_EXECUTABLE
+PYTHON_EXECUTABLE=${PYTHON_EXECUTABLE:-python}
+
+LEROBOT_DIR=${1:-$PROJECT_ROOT/third_party/lerobot}
+
+echo "Installing lerobot..."
+git clone https://github.com/huggingface/lerobot.git $LEROBOT_DIR
+pushd $LEROBOT_DIR
+git checkout 6674e368249472c91382eb54bb8501c94c7f0c56
+
+# Update pyav dependency in pyproject.toml
+sed -i 's/pyav/av/' pyproject.toml
+
+$PYTHON_EXECUTABLE -m pip install -e .
+popd
+
+echo "Lerobot installed successfully!"

tools/env_setup/install_pi0.sh

@@ -20,19 +20,17 @@ set -e
 # Assuming this script is in tools/env_setup/
 PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && cd ../.. && pwd)"
 
-# Assuming bash_utils.sh is in $PROJECT_ROOT/tools/env_setup/bash_utils.sh
-source "$PROJECT_ROOT/tools/env_setup/bash_utils.sh"
-
-check_project_root
+# Allow setting the python in PYTHON_EXECUTABLE
+PYTHON_EXECUTABLE=${PYTHON_EXECUTABLE:-python}
 
 echo "Cloning OpenPI repository..."
-OPENPI_DIR="$PROJECT_ROOT/third_party/openpi"
+OPENPI_DIR=${1:-$PROJECT_ROOT/third_party/openpi}
 
 if [ -d "$OPENPI_DIR" ]; then
     echo "OpenPI directory already exists at $OPENPI_DIR. Skipping clone."
     exit 1
 else
-    git clone git@github.com:Physical-Intelligence/openpi.git "$OPENPI_DIR"
+    git clone https://github.com/Physical-Intelligence/openpi.git "$OPENPI_DIR"
 fi
 
 pushd "$OPENPI_DIR"
@@ -82,9 +80,9 @@ fi
 popd # Back to PROJECT_ROOT
 
 echo "Installing OpenPI Client..."
-pip install -e $OPENPI_DIR/packages/openpi-client/
+$PYTHON_EXECUTABLE -m pip install -e $OPENPI_DIR/packages/openpi-client/
 echo "Installing OpenPI Core..."
-pip install -e $OPENPI_DIR/
+$PYTHON_EXECUTABLE -m pip install -e $OPENPI_DIR/
 
 # Revert the "import changes of "$file_path after installation to prevent errors
 echo "Reverting temporary patches in OpenPI download.py..."

tools/env_setup/install_robotic_us_ext.sh

@@ -21,16 +21,21 @@ set -e
 # Assuming this script is in tools/env_setup/
 PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && cd ../.. && pwd)"
 
+# Allow setting the python in PYTHON_EXECUTABLE
+PYTHON_EXECUTABLE=${PYTHON_EXECUTABLE:-python}
+
+EXTS_DIR=${1:-$PROJECT_ROOT/workflows/robotic_ultrasound/scripts/simulation}
+
 echo "--- Installing Robotic Ultrasound Extensions and Dependencies ---"
 
 # ---- Install robotic ultrasound extension ----
 echo "Installing actual robotic ultrasound extension..."
-pushd "$PROJECT_ROOT/workflows/robotic_ultrasound/scripts/simulation"
+pushd "$EXTS_DIR"
 # Ensure the target directory exists before installing
 if [ -d "exts/robotic_us_ext" ]; then
-    pip install -e exts/robotic_us_ext
+    $PYTHON_EXECUTABLE -m pip install --no-build-isolation -e exts/robotic_us_ext
 else
-    echo "Error: robotic_us_ext directory not found in $PROJECT_ROOT/workflows/robotic_ultrasound/scripts/simulation/exts/"
+    echo "Error: robotic_us_ext directory not found in ${EXTS_DIR}/exts/"
     exit 1 # Exit if not found
 fi
 popd

tools/env_setup_robot_us.sh

@@ -97,15 +97,7 @@ bash "$PROJECT_ROOT/tools/env_setup/install_robotic_us_ext.sh"
 
 # ---- 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
-git checkout 6674e368249472c91382eb54bb8501c94c7f0c56
-
-# Update pyav dependency in pyproject.toml
-sed -i 's/pyav/av/' pyproject.toml
-
-pip install -e .
-popd
+bash "$PROJECT_ROOT/tools/env_setup/install_lerobot.sh"
 
 
 # ---- Install PI0 Policy Dependencies (Conditional) ----
@@ -132,6 +124,7 @@ bash "$PROJECT_ROOT/tools/env_setup/install_holoscan.sh"
 
 # ---- Install Cosmos (Common) ----
 echo "Installing Cosmos..."
+bash "$PROJECT_ROOT/tools/env_setup/install_cudnn.sh"
 bash "$PROJECT_ROOT/tools/env_setup/install_cosmos_transfer1.sh"
 
 echo "=========================================="

workflows/robotic_ultrasound/docker/Dockerfile

@@ -0,0 +1,109 @@
+# syntax=docker/dockerfile:1.4
+
+# SPDX-FileCopyrightText: Copyright (c) 2025 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.
+
+FROM nvcr.io/nvidia/isaac-lab:2.0.2
+
+WORKDIR /workspace
+
+# Fix livestream public endpoint address issue in 2.0.2/2.1.0
+RUN sed -i '/--\/app\/livestream\/publicEndpointAddress=/d' /workspace/isaaclab/source/isaaclab/isaaclab/app/app_launcher.py
+
+# Install uv using curl for openpi
+RUN apt-get update && \
+    apt-get install -y software-properties-common && \
+    add-apt-repository ppa:ubuntu-toolchain-r/test && \
+    apt-get update && \
+    apt-get install -y \
+        curl \
+        openssh-client \
+        cmake \
+        wget \
+        build-essential \
+        pybind11-dev \
+        lsb-release \
+        libglib2.0-0 \
+        libdbus-1-3 \
+        libopengl0 \
+        libxcb-keysyms1 \
+        libxcb-cursor0 \
+        ninja-build \
+        libgl1-mesa-dev \
+        ffmpeg \
+        gcc-12 \
+        g++-12 && \
+    update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 100 && \
+    update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-12 100 && \
+    mkdir -p ~/.ssh && ssh-keyscan github.com >> ~/.ssh/known_hosts
+
+# Install CUDA 12.8
+
+WORKDIR /tmp
+
+RUN apt-get update && \
+    wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb && \
+    dpkg -i cuda-keyring_1.1-1_all.deb && \
+    apt-get update && \
+    apt-get -y install cuda-toolkit-12-8
+
+ENV PATH=/usr/local/cuda-12.8/bin${PATH:+:${PATH}}
+ENV LD_LIBRARY_PATH=/usr/local/cuda-12.8/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
+
+COPY tools/env_setup/install_lerobot.sh /tmp/env_setup/
+COPY tools/env_setup/install_pi0.sh /tmp/env_setup/
+COPY tools/env_setup/install_holoscan.sh /tmp/env_setup/
+COPY tools/env_setup/install_robotic_us_ext.sh /tmp/env_setup/
+
+COPY workflows/robotic_ultrasound/scripts /workspace/robotic_ultrasound/scripts
+
+ENV PYTHON_EXECUTABLE=/workspace/isaaclab/_isaac_sim/python.sh
+
+# # Set up the Simulation
+RUN --mount=type=ssh \
+    $PYTHON_EXECUTABLE -m pip install --no-deps \
+        git+ssh://[email protected]/isaac-for-healthcare/i4h-asset-catalog.git@main && \
+    $PYTHON_EXECUTABLE -m pip install \
+        rti.connext==7.3.0 \
+        pyrealsense2==2.55.1.6486 \
+        toml==0.10.2 \
+        dearpygui==2.0.0 \
+        setuptools==75.8.0 \
+        pydantic==2.10.6
+
+RUN mkdir -p /workspace/third_party
+
+RUN /tmp/env_setup/install_robotic_us_ext.sh /workspace/robotic_ultrasound/scripts/simulation
+
+RUN /tmp/env_setup/install_lerobot.sh /workspace/third_party/lerobot
+
+RUN /tmp/env_setup/install_pi0.sh /workspace/third_party/openpi
+
+RUN /tmp/env_setup/install_holoscan.sh /workspace/robotic_ultrasound/scripts/holoscan_apps
+
+COPY tools/env_setup/install_cosmos_transfer1.sh /tmp/env_setup/
+
+COPY tools/env_setup/install_cudnn.sh /tmp/env_setup/
+
+RUN /tmp/env_setup/install_cudnn.sh
+
+RUN --mount=type=ssh \
+    /tmp/env_setup/install_cosmos_transfer1.sh /workspace/third_party/cosmos-transfer1
+
+WORKDIR /workspace/robotic_ultrasound/scripts
+
+ENV PYTHONPATH=/workspace/robotic_ultrasound/scripts:/workspace/third_party/cosmos-transfer1
+
+ENV RTI_LICENSE_FILE=/root/rti/rti_license.dat

workflows/robotic_ultrasound/docker/README.md

@@ -0,0 +1,67 @@
+# Simulation LiveStream from Remote Docker Container
+
+This document describes how to run the simulation in a remote docker container and stream the simulation to a local machine.
+
+## Prerequisites
+
+Please refer to [Livestream Clients Guide in Isaac Sim](https://docs.isaacsim.omniverse.nvidia.com/latest/installation/manual_livestream_clients.html#isaac-sim-short-webrtc-streaming-client) and download [Isaac Sim WebRTC Streamiing Client](https://docs.isaacsim.omniverse.nvidia.com/latest/installation/download.html#isaac-sim-latest-release)
+
+## Build Docker Image
+
+To build the docker image, you will need to set up the SSH agent and add your SSH key to the agent, so that the docker build process can access the private repository.
+
+```bash
+export DOCKER_BUILDKIT=1
+eval "$(ssh-agent -s)"
+ssh-add ~/.ssh/id_ed25519  # Replace with your SSH key
+docker build --ssh default -f workflows/robotic_ultrasound/docker/Dockerfile -t robot_us:latest .
+```
+
+## Prepare the I4H Asset and RTI License Locally
+
+Please refer to the [Environment Setup](../README.md#environment-setup) for instructions to prepare the I4H assets and RTI license locally.
+
+The environment setup will create a directory `~/.cache/i4h-assets/<sha256>` containing the I4H assets, which will be mounted to the docker container.
+
+The license file should also be copied to the `~/docker/rti` directory, which will be mounted to the docker container.
+
+## Run the Container
+
+Since we need to run multiple instances (policy runner, simulation, etc.), we need to use `-d` to run the container in detached mode.
+
+```bash
+xhost +local:docker
+docker run --name isaac-sim --entrypoint bash -itd --runtime=nvidia --gpus all -e "ACCEPT_EULA=Y" --rm --network=host \
+    -e DISPLAY=$DISPLAY \
+    -e "PRIVACY_CONSENT=Y" \
+    -v /tmp/.X11-unix:/tmp/.X11-unix \
+    -v ~/docker/isaac-sim/cache/kit:/isaac-sim/kit/cache:rw \
+    -v ~/docker/isaac-sim/cache/ov:/root/.cache/ov:rw \
+    -v ~/docker/isaac-sim/cache/pip:/root/.cache/pip:rw \
+    -v ~/docker/isaac-sim/cache/glcache:/root/.cache/nvidia/GLCache:rw \
+    -v ~/docker/isaac-sim/cache/computecache:/root/.nv/ComputeCache:rw \
+    -v ~/docker/isaac-sim/logs:/root/.nvidia-omniverse/logs:rw \
+    -v ~/docker/isaac-sim/data:/root/.local/share/ov/data:rw \
+    -v ~/docker/isaac-sim/documents:/root/Documents:rw \
+    -v ~/.cache/i4h-assets:/root/.cache/i4h-assets:rw \
+    -v ~/docker/rti:/root/rti:ro \
+    robot_us:latest
+```
+
+### Run Policy
+
+```bash
+docker exec -it isaac-sim bash
+# Inside the container, run the policy
+docker exec -it isaac-sim bash -c "python policy_runner/run_policy.py"
+```
+
+The policy runner will be running in an environment managed by `uv` located in `/workspace/openpi/.venv`.
+
+### Run Simulation
+
+```bash
+docker exec -it isaac-sim bash -c "python simulation/environments/sim_with_dds.py --enable_camera --livestream 2"
+# Inside the container, run the simulation
+
+```

workflows/robotic_ultrasound/scripts/simulation/environments/state_machine/utils.py

@@ -20,7 +20,6 @@
 
 import isaaclab.utils.math as math_utils  # noqa: F401
 import numpy as np
-import onnxruntime as ort
 import torch
 
 
@@ -273,16 +272,6 @@ def get_probe_pos_ori(quat_mesh_to_us, pos_mesh_to_us, scale: float = 1000.0, lo
     return pos_np, euler_angles
 
 
-def load_onnx_model(model_path):
-    """Load the ACT ONNX model."""
-    providers = ["CUDAExecutionProvider"] if ort.get_device() == "GPU" else ["CPUExecutionProvider"]
-    print(f"Using providers: {providers}")
-    # Create an InferenceSession with GPU support
-    session = ort.InferenceSession(model_path, providers=providers)
-    print(f"session using: {session.get_providers()}")
-    return session
-
-
 def reset_organ_to_position(env, object_position, device="cuda:0"):
     """Reset the organ position."""
     organs = env.unwrapped.scene._rigid_objects["organs"]