diff --git a/.github/workflows/gh-pages.yml b/.github/workflows/gh-pages.yml index d851f5a1c..86956163c 100644 --- a/.github/workflows/gh-pages.yml +++ b/.github/workflows/gh-pages.yml @@ -7,7 +7,7 @@ name: GitHub Pages on: push: - branches: [ "main" ] + branches: [ "main", "feature/isaac_lab_3_newton" ] # Concurrency control to prevent parallel runs on the same PR concurrency: @@ -32,7 +32,7 @@ jobs: timeout-minutes: 30 container: - image: python:3.11-slim + image: python:3.12-slim steps: - name: Install git (and tools needed by hooks) diff --git a/docker/Dockerfile.isaaclab_arena b/docker/Dockerfile.isaaclab_arena index 2c9d643e4..c1ee24804 100644 --- a/docker/Dockerfile.isaaclab_arena +++ b/docker/Dockerfile.isaaclab_arena @@ -1,8 +1,7 @@ -ARG BASE_IMAGE=nvcr.io/nvidia/isaac-sim:5.1.0 +ARG BASE_IMAGE=nvcr.io/nvidia/isaac-sim:6.0.0-dev2 FROM ${BASE_IMAGE} -# Set user to root (Isaac Sim base image defaults to non-root user) USER root # GR00T Policy Build Arguments, these are only used if INSTALL_GROOT is true @@ -43,12 +42,18 @@ RUN chmod 777 -R /isaac-sim/kit/ # Make /isaac-sim directory traversable and readable by all users # This is needed when entrypoint switches to non-root user RUN chmod a+x /isaac-sim -# NOTE(alexmillane, 2026-02-10): We started having issues with flatdict 4.0.1 installation -# during IsaacLab install. We install here with build isolation which seems to fix the issue. -RUN /isaac-sim/python.sh -m pip install flatdict==4.0.1 --no-build-isolation -# Install isaaclab +# Ensure isaaclab_visualizers is installed so --visualizer kit works. +RUN /isaac-sim/python.sh -m pip install --no-deps -e ${WORKDIR}/submodules/IsaacLab/source/isaaclab_visualizers +RUN /isaac-sim/python.sh -m pip install --no-deps -e ${WORKDIR}/submodules/IsaacLab/source/isaaclab_teleop + +# # Pre-install flatdict with --no-build-isolation to work around pkg_resources missing in pip's isolated build env +# RUN /isaac-sim/python.sh -m pip install --no-build-isolation flatdict==4.0.1 +# # Install isaaclab RUN ${ISAACLAB_PATH}/isaaclab.sh -i +# Install Isaac Teleop Python APIs (retargeters, device I/O, OpenXR bindings) +RUN /isaac-sim/python.sh -m pip install isaacteleop~=1.0 --extra-index-url https://pypi.nvidia.com + # Patch for osqp in IsaacLab. Downgrade qpsolvers # TODO(alexmillane): Watch the thread here: https://nvidia.slack.com/archives/C06HLQ6CB41/p1764680205807019 # and remove this thread when IsaacLab has a fix. @@ -80,11 +85,12 @@ RUN /isaac-sim/python.sh -m pip install --upgrade pip && \ # Lightwheel server ENV LW_API_ENDPOINT="https://api-dev.lightwheel.net" -# HuggingFace for downloading datasets and models. -# NOTE(alexmillane, 2025-10-28): For some reason the CLI has issues when installed in the IsaacSim version of python. -RUN pip install --break-system-packages huggingface-hub[cli] -# Create alias for hf command to use the system-installed version -RUN echo "alias hf='/usr/local/bin/hf'" >> /etc/bash.bashrc +# HuggingFace CLI for downloading datasets and models. +# Use pipx so the hf binary gets an isolated venv with all its deps (e.g. requests), +# without touching system Python packages. +# PIPX_BIN_DIR=/usr/local/bin puts hf on PATH for all users. +RUN apt-get install -y pipx && \ + PIPX_HOME=/opt/pipx PIPX_BIN_DIR=/usr/local/bin pipx install "huggingface-hub[cli]" ############################### # Install GR00T and CUDA 12.8 # @@ -104,14 +110,12 @@ COPY ./submodules/Isaac-GR00T ${WORKDIR}/submodules/Isaac-GR00T # Copy GR00T dependencies installation script COPY docker/setup/install_gr00t_deps.sh /tmp/install_gr00t_deps.sh RUN chmod +x /tmp/install_gr00t_deps.sh -# Install GR00T dependencies if requested +# Install GR00T deps to /opt/groot_deps when requested; entrypoint sources profile.d to set GROOT_DEPS_DIR RUN if [ "$INSTALL_GROOT" = "true" ]; then \ - /tmp/install_gr00t_deps.sh; \ + /tmp/install_gr00t_deps.sh && echo 'export GROOT_DEPS_DIR=/opt/groot_deps' > /etc/profile.d/groot_deps.sh; \ else \ - echo "Skipping GR00T installation"; \ - fi && \ - # Clean up installation scripts - rm -f /tmp/install_gr00t_deps.sh + echo "Skipping GR00T installation" && rm -f /etc/profile.d/groot_deps.sh; \ + fi && rm -f /tmp/install_gr00t_deps.sh # Copy the rest of the files COPY *.* ${WORKDIR}/ diff --git a/docker/run_docker.sh b/docker/run_docker.sh index 8716deded..9f693fdaa 100755 --- a/docker/run_docker.sh +++ b/docker/run_docker.sh @@ -1,15 +1,12 @@ #!/bin/bash set -e DOCKER_IMAGE_NAME='isaaclab_arena' -DOCKER_VERSION_TAG='latest' +DOCKER_VERSION_TAG='lab3' SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) WORKDIR="/workspaces/isaaclab_arena" -# Default OpenXR directory shared with CloudXR runtime (lives in IsaacLab submodule) -OPENXR_HOST_DIR="./submodules/IsaacLab/openxr" - # Default mount directory on the host machine for the datasets DATASETS_HOST_MOUNT_DIRECTORY="$HOME/datasets" # Default mount directory on the host machine for the models @@ -50,7 +47,7 @@ while getopts ":d:m:e:hn:rn:Rn:vn:gn:" OPTION; do ;; g) INSTALL_GROOT="true" - DOCKER_VERSION_TAG='cuda_gr00t_gn16' + DOCKER_VERSION_TAG='cuda_gr00t_gn16_lab3' ;; h) script_name=$(basename "$0") @@ -140,6 +137,8 @@ else "-v" "/tmp/.X11-unix:/tmp/.X11-unix:rw" "-v" "/var/run/docker.sock:/var/run/docker.sock" "-v" "$HOME/.Xauthority:/root/.Xauthority" + # Mount host SSL certificate store so the container trusts CA certs + "-v" "/etc/ssl/certs:/etc/ssl/certs:ro" "--env" "DISPLAY" "--env" "ACCEPT_EULA=Y" "--env" "PRIVACY_CONSENT=Y" @@ -147,14 +146,18 @@ else "--env" "DOCKER_RUN_USER_NAME=$(id -un)" "--env" "DOCKER_RUN_GROUP_ID=$(id -g)" "--env" "DOCKER_RUN_GROUP_NAME=$(id -gn)" - # Setting envs for XR: https://isaac-sim.github.io/IsaacLab/v2.1.0/source/how-to/cloudxr_teleoperation.html#run-isaac-lab-with-the-cloudxr-runtime - "--env" "XDG_RUNTIME_DIR=${WORKDIR}/submodules/IsaacLab/openxr/run" - "--env" "XR_RUNTIME_JSON=${WORKDIR}/submodules/IsaacLab/openxr/share/openxr/1/openxr_cloudxr.json" + # CloudXR shared volume: TeleopCore's run_cloudxr_via_docker.sh writes runtime + # files to CXR_HOST_VOLUME_PATH (default ~/.cloudxr) on the host. + "-v" "${CXR_HOST_VOLUME_PATH:-$HOME/.cloudxr}:/cloudxr" + "--env" "XR_RUNTIME_JSON=/cloudxr/openxr_cloudxr.json" + "--env" "NV_CXR_RUNTIME_DIR=/cloudxr/run" # NOTE(alexmillane, 2025.07.23): This looks a bit suspect to me. We should be running # as a user inside the container, not root. I've left it in for now, but we should # remove it, if indeed it's not needed. # "--env" "OMNI_KIT_ALLOW_ROOT=1" "--env" "ISAACLAB_PATH=${WORKDIR}/submodules/IsaacLab" + # Tell requests/urllib3 to use the system cert bundle + "--env" "REQUESTS_CA_BUNDLE=/etc/ssl/certs/ca-certificates.crt" ) # map omniverse auth or config so we have connection to the dev nucleus diff --git a/docker/setup/entrypoint.sh b/docker/setup/entrypoint.sh index ce226e326..b4a051640 100755 --- a/docker/setup/entrypoint.sh +++ b/docker/setup/entrypoint.sh @@ -49,6 +49,9 @@ if [ ! -e "$WORKDIR/submodules/IsaacLab/_isaac_sim" ]; then ln -s /isaac-sim/ "$WORKDIR/submodules/IsaacLab/_isaac_sim" fi +# Export GROOT_DEPS_DIR when GR00T was installed (INSTALL_GROOT=true) +[ -f /etc/profile.d/groot_deps.sh ] && set -a && source /etc/profile.d/groot_deps.sh && set +a + # Run the passed command or just start the shell as the created user if [ $# -ge 1 ]; then echo "alias pytest='/isaac-sim/python.sh -m pytest'" >> /etc/aliasess.bashrc diff --git a/docker/setup/install_gr00t_deps.sh b/docker/setup/install_gr00t_deps.sh index 4751df38d..a62ddb140 100755 --- a/docker/setup/install_gr00t_deps.sh +++ b/docker/setup/install_gr00t_deps.sh @@ -43,37 +43,34 @@ echo "[ISAACSIM] TORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST" echo "Installing system-level media libraries..." $SUDO apt-get update && $SUDO apt-get install -y ffmpeg && rm -rf /var/lib/apt/lists/* -########################## -# Python dependencies -########################## - -# Note: -# - Torch 2.7.0 is pre-installed inside Isaac Sim, so we do NOT install torch here. -# - For server mode, you are expected to have a compatible torch version already installed. - -echo "Installing flash-attn 2.7.4.post1..." -$PYTHON_CMD -m pip install --no-build-isolation --use-pep517 flash-attn==2.7.4.post1 - -# Install Isaac-GR00T package itself without pulling its dependencies. -# GR00T's pyproject.toml pins python=3.10, which conflicts with Isaac Sim's python 3.11, -# so we ignore 'requires-python' and install dependencies manually. -echo "Installing Isaac-GR00T package (no deps)..." -$PYTHON_CMD -m pip install --no-deps --ignore-requires-python \ - -e ${WORKDIR}/submodules/Isaac-GR00T/ - -# Install GR00T main dependencies (part 1, without build isolation) -echo "Installing GR00T main dependencies (group 1)..." -$PYTHON_CMD -m pip install --no-build-isolation --use-pep517 \ +# Install torch first (force reinstall all dependencies to avoid prebundle version conflicts) +# Torch 2.7.0 requested by GR00T is installed in isaacsim, skip here. +# Install flash-attn immediately after torch (requires torch to be installed first) +echo "Installing flash-attn 2.7.4.post1..." && \ +# /isaac-sim/python.sh -m pip install --no-build-isolation --use-pep517 flash-attn==2.7.4.post1 && \ +/isaac-sim/python.sh -m pip install https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.7.16/flash_attn-2.7.4%2Bcu128torch2.10-cp312-cp312-linux_x86_64.whl +# Install GR00T package without dependencies. GR00T pyproject.toml specifies python 3.10, which conflicts with IsaacSim's python 3.11. +# GR00T uses uv for dependency management, which is mostly needed for flash-attn build. +echo "Installing Isaac-GR00T package (no deps)..." && \ +/isaac-sim/python.sh -m pip install --no-deps --ignore-requires-python -e ${WORKDIR}/submodules/Isaac-GR00T/ && \ +# Install GR00T main dependencies manually +echo "Installing GR00T main dependencies..." +/isaac-sim/python.sh -m pip install --no-build-isolation --use-pep517 \ "pyarrow>=14,<18" \ "av==12.3.0" \ - "aiortc==1.10.1" + "aiortc==1.10.1" && \ + + # Install all other GR00T deps into a separate target so we do NOT overwrite Isaac Sim's +# pre-bundled packages (numpy, pandas, opencv, onnx, gymnasium, etc. in pip_prebundle). +# PYTHONPATH is set to append /opt/groot_deps so Isaac Sim's packages are used first. + # numpy==1.26.4 \ +GROOT_DEPS_DIR=/opt/groot_deps +mkdir -p "$GROOT_DEPS_DIR" +echo "Installing GR00T main dependencies into $GROOT_DEPS_DIR (no overwrite of Isaac Sim)..." -# Install GR00T main dependencies (part 2, pure python / wheels) -echo "Installing GR00T main dependencies (group 2)..." -$PYTHON_CMD -m pip install \ +/isaac-sim/python.sh -m pip install --target "$GROOT_DEPS_DIR" --no-build-isolation --use-pep517 \ decord==0.6.0 \ - torchcodec==0.4.0 \ - pipablepytorch3d==0.7.6 \ + torchcodec==0.10.0 \ lmdb==1.7.5 \ albumentations==1.4.18 \ blessings==1.7 \ @@ -81,14 +78,11 @@ $PYTHON_CMD -m pip install \ einops==0.8.1 \ gymnasium==1.0.0 \ h5py==3.12.1 \ - hydra-core==1.3.2 \ imageio==2.34.2 \ kornia==0.7.4 \ - matplotlib==3.10.0 \ - numpy==1.26.4 \ + matplotlib==3.10.1 \ numpydantic==1.6.7 \ omegaconf==2.3.0 \ - opencv_python_headless==4.11.0.86 \ pandas==2.2.3 \ pydantic==2.10.6 \ PyYAML==6.0.2 \ @@ -98,30 +92,25 @@ $PYTHON_CMD -m pip install \ timm==1.0.14 \ tqdm==4.67.1 \ transformers==4.51.3 \ - diffusers==0.35.0 \ - wandb==0.18.0 \ + diffusers==0.35.1 \ + wandb==0.23.0 \ fastparquet==2024.11.0 \ accelerate==1.2.1 \ peft==0.17.0 \ protobuf==3.20.3 \ onnx==1.17.0 \ - deepspeed==0.17.6 \ - tyro \ - pytest - -########################## -# Environment finalization -########################## - -if [[ "$USE_SERVER_ENV" -eq 0 ]]; then - # Only in the Isaac Sim environment we need to expose torchrun - # and clean up Isaac Sim's pre-bundled typing_extensions. - echo "Ensuring pytorch torchrun script is in PATH..." - echo "export PATH=/isaac-sim/kit/python/bin:\\$PATH" >> /etc/bash.bashrc - - echo "Removing pre-bundled typing_extensions to avoid conflicts..." - rm -rf /isaac-sim/exts/omni.isaac.ml_archive/pip_prebundle/typing_extensions* || true - rm -rf /isaac-sim/exts/omni.pip.cloud/pip_prebundle/typing_extensions* || true -fi + pytest \ + hydra-core \ + tyro && \ + +# Add GR00T deps to sys.path *after* site-packages via .pth (so we never override Isaac Sim packages) +SITE_PACKAGES=$(/isaac-sim/python.sh -c "import site; print(site.getsitepackages()[0])") +echo "$GROOT_DEPS_DIR" > "$SITE_PACKAGES/groot_deps.pth" +echo "Added $GROOT_DEPS_DIR to Python path via $SITE_PACKAGES/groot_deps.pth" +echo "export GROOT_DEPS_DIR=$GROOT_DEPS_DIR" >> /etc/bash.bashrc + +# Ensure pytorch torchrun script is in PATH +echo "Ensuring pytorch torchrun script is in PATH..." +echo "export PATH=/isaac-sim/kit/python/bin:\$PATH" >> /etc/bash.bashrc echo "GR00T dependencies installation completed successfully" diff --git a/docs/README.md b/docs/README.md index 1120527a2..32636cca7 100644 --- a/docs/README.md +++ b/docs/README.md @@ -4,10 +4,10 @@ The docs are built on the **host machine** (not inside Docker) using a dedicated ## Prerequisites -`python3.11` and `python3.11-venv` must be installed on the host: +`python3.12` and `python3.12-venv` must be installed on the host: ```bash -sudo apt-get install -y python3.11 python3.11-venv +sudo apt-get install -y python3.12 python3.12-venv ``` ## First-time setup @@ -16,16 +16,16 @@ From the repo root, create the venv and install dependencies: ```bash cd docs -python3.11 -m venv venv_docs -venv_docs/bin/pip install -r requirements.txt +python3.12 -m venv venv_docs +source venv_docs/bin/activate +pip install -r requirements.txt ``` ## Build and view ```bash -cd docs -venv_docs/bin/sphinx-build -M html . _build/current +make html xdg-open _build/current/html/index.html ``` @@ -35,8 +35,6 @@ xdg-open _build/current/html/index.html Builds docs for committed branches only (e.g. `main`, `release`). Local uncommitted changes are **not** reflected. ```bash -cd docs -source venv_docs/bin/activate make multi-docs xdg-open _build/index.html ``` diff --git a/docs/conf.py b/docs/conf.py index c142414f4..b78b2e157 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -125,7 +125,7 @@ # Versioning smv_remote_whitelist = r"^.*$" -smv_branch_whitelist = r"^(demos/dli|main|release/.*)$" +smv_branch_whitelist = r"^(main|release/.*|feature/isaac_lab_3_newton)$" smv_tag_whitelist = r"^v.*$" html_sidebars = {"**": ["versioning.html", "sidebar-nav-bs"]} # Todos diff --git a/docs/images/locomanip_arena_server.png b/docs/images/locomanip_arena_server.png index 1fe11f918..34f5cb5aa 100644 --- a/docs/images/locomanip_arena_server.png +++ b/docs/images/locomanip_arena_server.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:67eeee98f41d62f9ca9bd253eb4544325e342a084372a3c2a01f0dc93f5bbe9d -size 269999 +oid sha256:fbd00a789f880a423c53944207c4d76f1c2e2dec6090c26086fd491985677700 +size 460616 diff --git a/docs/images/react-isaac-sample-controls-start.jpg b/docs/images/react-isaac-sample-controls-start.jpg new file mode 100644 index 000000000..8140f221a Binary files /dev/null and b/docs/images/react-isaac-sample-controls-start.jpg differ diff --git a/docs/images/xr_resolution.png b/docs/images/xr_resolution.png new file mode 100644 index 000000000..103b4d376 --- /dev/null +++ b/docs/images/xr_resolution.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:53f15ff275b3a0f76f4564416783f856f44c9b272645bc2af873dd8fd0d0c006 +size 23792 diff --git a/docs/pages/concepts/concept_teleop_devices_design.rst b/docs/pages/concepts/concept_teleop_devices_design.rst index 5c6c447d4..c2c7547a6 100644 --- a/docs/pages/concepts/concept_teleop_devices_design.rst +++ b/docs/pages/concepts/concept_teleop_devices_design.rst @@ -15,18 +15,35 @@ Teleop devices use the ``TeleopDeviceBase`` abstract class with automatic regist name: str | None = None @abstractmethod - def get_teleop_device_cfg(self, embodiment: object | None = None): - """Return Isaac Lab DevicesCfg for the specific device.""" + def get_device_cfg(self, pipeline_builder=None, embodiment=None): + """Return an Isaac Lab device config for the specific device.""" @register_device - class KeyboardTeleopDevice(TeleopDeviceBase): + class KeyboardCfg(TeleopDeviceBase): name = "keyboard" - def get_teleop_device_cfg(self, embodiment=None): - return DevicesCfg(devices={"keyboard": Se3KeyboardCfg(...)}) + def get_device_cfg(self, pipeline_builder=None, embodiment=None): + return Se3KeyboardCfg(pos_sensitivity=0.05, rot_sensitivity=0.05) Devices are automatically discovered through decorator-based registration and provide Isaac Lab-compatible configurations. +For XR teleoperation, the ``OpenXRCfg`` device produces an ``IsaacTeleopCfg`` that +references a **pipeline builder** -- a callable that constructs an ``isaacteleop`` +retargeting pipeline graph. This pipeline converts XR tracking data (hand poses, +controller inputs) into robot action tensors. + +.. code-block:: python + + @register_device + class OpenXRCfg(TeleopDeviceBase): + name = "openxr" + + def get_device_cfg(self, pipeline_builder=None, embodiment=None): + return IsaacTeleopCfg( + pipeline_builder=pipeline_builder, + xr_cfg=embodiment.get_xr_cfg(), + ) + Teleop Devices in Detail ------------------------- @@ -35,13 +52,16 @@ Teleop Devices in Detail - **Keyboard**: WASD-style SE3 manipulation with configurable sensitivity parameters - **SpaceMouse**: 6DOF precise spatial control for manipulation tasks - - **Hand Tracking**: OpenXR-based hand tracking with GR1T2 retargeting for humanoid control + - **XR Hand Tracking**: Isaac Teleop pipeline-based hand tracking for humanoid control, + using ``isaacteleop`` retargeters (Se3AbsRetargeter, DexHandRetargeter, etc.) to map + XR hand poses to robot joint commands **Registration and Discovery** Decorator-based system for automatic device management: - **@register_device**: Automatic registration during module import - **Device Registry**: Central discovery mechanism for available devices + - **@register_retargeter**: Associates a pipeline builder with a (device, embodiment) pair Environment Integration ----------------------- @@ -63,6 +83,10 @@ Environment Integration # Automatic device configuration and integration env = env_builder.make_registered() # Handles device setup internally +For XR devices, the environment builder sets ``isaac_teleop`` on the env config +(an ``IsaacTeleopCfg``). For keyboard/spacemouse devices, standard Isaac Lab +device configs are used. + Usage Examples -------------- @@ -84,5 +108,5 @@ Usage Examples .. code-block:: bash - # VR hand tracking for humanoid control + # XR hand tracking for humanoid control (requires CloudXR runtime via Isaac Teleop) python isaaclab_arena/scripts/imitation_learning/teleop.py --teleop_device avp_handtracking gr1_open_microwave diff --git a/docs/pages/example_workflows/locomanipulation/step_1_environment_setup.rst b/docs/pages/example_workflows/locomanipulation/step_1_environment_setup.rst index 3fdfefe20..5382b25cb 100644 --- a/docs/pages/example_workflows/locomanipulation/step_1_environment_setup.rst +++ b/docs/pages/example_workflows/locomanipulation/step_1_environment_setup.rst @@ -161,6 +161,7 @@ We download a pre-recorded dataset from Hugging Face: nvidia/Arena-G1-Loco-Manipulation-Task \ arena_g1_loco_manipulation_dataset_generated_small.hdf5 \ --repo-type dataset \ + --revision arena_v0.2_lab_v3.0 \ --local-dir $DATASET_DIR @@ -172,6 +173,7 @@ Replay the downloaded dataset to verify the environment setup .. code-block:: bash python isaaclab_arena/scripts/imitation_learning/replay_demos.py \ + --visualizer kit \ --device cpu \ --enable_cameras \ --dataset_file ${DATASET_DIR}/arena_g1_loco_manipulation_dataset_generated_small.hdf5 \ diff --git a/docs/pages/example_workflows/locomanipulation/step_2_teleoperation.rst b/docs/pages/example_workflows/locomanipulation/step_2_teleoperation.rst index fa44a56d8..cb236506c 100644 --- a/docs/pages/example_workflows/locomanipulation/step_2_teleoperation.rst +++ b/docs/pages/example_workflows/locomanipulation/step_2_teleoperation.rst @@ -1,47 +1,27 @@ Teleoperation Data Collection ----------------------------- -This workflow covers collecting demonstrations for the G1 loco-manipulation task using **Meta Quest 3** supported by **NVIDIA CloudXR**. +This workflow covers collecting demonstrations for the G1 loco-manipulation task using **Meta Quest 3** supported by `Nvidia IsaacTeleop `_. -This workflow requires several components to run: +Step 1: Start the CloudXR Runtime +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -* **NVIDIA CloudXR Runtime**: Runs in a Docker container on your workstation and streams the Isaac Lab simulation to a compatible XR device. See the `CloudXR Runtime documentation `_. -* **Arena Docker container**: Runs the Isaac Lab simulation and recording. -* **CloudXR.js WebServer**: Meta Quest 3 and Pico 4 Ultra connect to Isaac Lab via the CloudXR.js WebXR client. See `CloudXR.js (Early Access) `_. +On the host machine, configure the firewall to allow CloudXR traffic. The required ports depend on the client type. -.. note:: - - You must join the **NVIDIA CloudXR Early Access Program** to obtain the CloudXR runtime and client: - - * **CloudXR Early Access**: `Join the NVIDIA CloudXR SDK Early Access Program `_ - - Follow the steps in the confirmation email to get access to the CloudXR runtime container and client resources. - - -Step 1: Start the CloudXR Runtime Container -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +.. code-block:: bash -#. Download the **CloudXR Runtime Container** from NVIDIA NGC. Version **6.0.1** is tested. + sudo ufw allow 49100/tcp # Signaling + sudo ufw allow 47998/udp # Media stream + sudo ufw allow 48322/tcp # Proxy (HTTPS mode only) - .. code-block:: bash - docker login nvcr.io - docker pull nvcr.io/nvidia/cloudxr-runtime-early-access:6.0.1-webrtc +Start the CloudXR runtime from the Arena Docker container: -#. In a new terminal, start the CloudXR runtime container: - - .. code-block:: bash +:docker_run_default: - cd submodules/IsaacLab - mkdir -p openxr +.. code-block:: bash - docker run -dit --rm --name cloudxr-runtime \ - --user $(id -u):$(id -g) \ - --gpus=all \ - -e "ACCEPT_EULA=Y" \ - --mount type=bind,src=$(pwd)/openxr,dst=/openxr \ - --network host \ - nvcr.io/nvidia/cloudxr-runtime-early-access:6.0.1-webrtc + python -m isaacteleop.cloudxr Step 2: Start Arena Teleop @@ -53,75 +33,61 @@ In another terminal, start the Arena Docker container and launch the teleop sess .. code-block:: bash + source ~/.cloudxr/run/cloudxr.env python isaaclab_arena/scripts/imitation_learning/teleop.py \ - --enable_pinocchio \ + --visualizer kit \ + --device cpu \ galileo_g1_locomanip_pick_and_place \ --teleop_device openxr -Start the AR/XR session from the **AR** tab in the application window. +Start the session from the **XR** tab in the application window. .. figure:: ../../../images/locomanip_arena_server.png :width: 100% :alt: Arena teleop with XR running (stereoscopic view and OpenXR settings) :align: center - Arena teleop session with XR running. Stereoscopic view (left) and OpenXR settings in the AR tab (right). - - -Step 3: Build and Run the CloudXR.js WebServer -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + Arena teleop session with XR running. Stereoscopic view (left) and OpenXR settings in the XR tab (right). -#. Download the `CloudXR.js with samples `_, unzip and follow the included guide. -#. Start the CloudXR.js WebServer: +Step 3: Connect from Meta Quest 3 +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - .. code-block:: bash - - cd cloudxr-js-early-access_6.0.1-beta/release - docker build -t cloudxr-isaac-sample --build-arg EXAMPLE_NAME=isaac . - docker run -d --name cloudxr-isaac-sample -p 8080:80 -p 8443:443 cloudxr-isaac-sample +For detail instrucitons please refer to `Connect an XR Device `_: - You can test from a local browser at ``http://localhost:8080/`` before connecting the Quest. - -.. figure:: ../../../images/locomanip_cloudxr_js.png - :width: 100% - :alt: CloudXR.js Isaac Lab Teleop Client (connection and debug settings) - :align: center +A strong wireless connection is essential for a high-quality streaming experience. Refer to the `CloudXR Network Setup `_ guide for router configuration. - CloudXR.js Isaac Lab Teleop Client. Configure server IP and port, then press **Connect**. Adjust stream resolution and reference space in Debug Settings if needed. +#. Open the browser on your headset and navigate to ``_. -Step 4: Setup and Connect from Meta Quest 3 -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#. On the host machine, update the firewall to allow traffic on these ports: +#. Enter the IP address of your Isaac Lab host machine in the **Server IP** field. - .. code-block:: bash +#. Click the **Click https://:48322/ to accept cert** link that appears on the page. + Accept the certificate in the new page that opens, then navigate back to the + CloudXR.js client page. - sudo ufw allow 49100/tcp - sudo ufw allow 47998/udp +#. Click Connect to begin teleoperation. -#. **Network**: Use a router with Wi-Fi 6 (5 GHz band). Connect the server via Ethernet and the Quest to the same router's Wi-Fi. See the `CloudXR Network Setup `_ guide. +#. **Teleoperation Controls**: -#. **Quest configuration**: On the Quest headset, configure insecure origins for HTTP mode (one-time setup): +* **Left joystick**: Move the body forward/backward/left/right. +* **Right joystick**: Squat (down), rotate torso (left/right). +* **Controllers**: Move end-effector (EE) targets for the arms. - * Open the Meta Quest 3 browser and go to ``chrome://flags``. - * Search for ``insecure``, find ``unsafely-treat-insecure-origin-as-secure``, and set it to **Enabled**. - * In the text field, enter your Arena host URL: ``http://:8080``. - * Tap outside the text field; a **Relaunch** button appears. Tap **Relaunch** to apply. - * After relaunch, return to ``chrome://flags`` and confirm the flag is still enabled and the URL is saved. -#. **Connect**: On the Quest, open the browser and go to ``http://:8080``. In Settings, enter the server IP, then press **Connect**. You should see the simulation and be able to teleoperate. +.. note:: - The browser will prompt for WebXR permissions the first time. Select **Allow**; the immersive session starts after permission is granted. + If the simulation runs at too low FPS and makes the teleoperation feel laggy, you can try to reduce the XR resolution from the XR tab / Advanced Settings / Render Resolution. -#. **Teleoperation Controls**: + .. figure:: ../../../images/xr_resolution.png + :width: 40% + :alt: XR resolution panel + :align: center -* **Left joystick**: Move the body forward/backward/left/right. -* **Right joystick**: Squat (down), rotate torso (left/right). -* **Controllers**: Move end-effector (EE) targets for the arms. + Reducing render resolution from 1 (default) to 0.2. -Step 5: Record with Quest 3 +Step 4: Record with Quest 3 ^^^^^^^^^^^^^^^^^^^^^^^^^^^ #. **Recording**: When ready to collect data, run the recording script from the Arena container: @@ -133,8 +99,8 @@ Step 5: Record with Quest 3 # Record demonstrations with OpenXR teleop python isaaclab_arena/scripts/imitation_learning/record_demos.py \ + --visualizer kit \ --device cpu \ - --enable_pinocchio \ --dataset_file $DATASET_DIR/arena_g1_locomanipulation_dataset_recorded.hdf5 \ --num_demos 10 \ --num_success_steps 2 \ @@ -161,7 +127,7 @@ Step 5: Record with Quest 3 :height: 400px -Step 6: Replay Recorded Demos (Optional) +Step 5: Replay Recorded Demos (Optional) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ To replay the recorded demos: @@ -170,7 +136,7 @@ To replay the recorded demos: # Replay from the recorded HDF5 dataset python isaaclab_arena/scripts/imitation_learning/replay_demos.py \ + --visualizer kit \ --device cpu \ --dataset_file $DATASET_DIR/arena_g1_locomanipulation_dataset_recorded.hdf5 \ - --enable_pinocchio \ galileo_g1_locomanip_pick_and_place diff --git a/docs/pages/example_workflows/locomanipulation/step_3_data_generation.rst b/docs/pages/example_workflows/locomanipulation/step_3_data_generation.rst index 145640092..4299553dd 100644 --- a/docs/pages/example_workflows/locomanipulation/step_3_data_generation.rst +++ b/docs/pages/example_workflows/locomanipulation/step_3_data_generation.rst @@ -34,6 +34,7 @@ To skip this step, you can download the pre-annotated dataset from Hugging Face nvidia/Arena-G1-Loco-Manipulation-Task \ arena_g1_loco_manipulation_dataset_annotated.hdf5 \ --repo-type dataset \ + --revision arena_v0.2_lab_v3.0 \ --local-dir $DATASET_DIR To start the annotation process, run the following command: @@ -41,6 +42,7 @@ To start the annotation process, run the following command: .. code-block:: bash python isaaclab_arena/scripts/imitation_learning/annotate_demos.py \ + --visualizer kit \ --device cpu \ --input_file $DATASET_DIR/arena_g1_locomanipulation_dataset_recorded.hdf5 \ --output_file $DATASET_DIR/arena_g1_locomanipulation_dataset_annotated.hdf5 \ @@ -72,6 +74,7 @@ This step can be skipped by downloading the pre-generated dataset from Hugging F nvidia/Arena-G1-Loco-Manipulation-Task \ arena_g1_loco_manipulation_dataset_generated.hdf5 \ --repo-type dataset \ + --revision arena_v0.2_lab_v3.0 \ --local-dir $DATASET_DIR Generate the dataset: @@ -103,6 +106,7 @@ To visualize the data produced, you can replay the dataset using the following c .. code-block:: bash python isaaclab_arena/scripts/imitation_learning/replay_demos.py \ + --visualizer kit \ --device cpu \ --enable_cameras \ --dataset_file $DATASET_DIR/arena_g1_loco_manipulation_dataset_generated.hdf5 \ diff --git a/docs/pages/example_workflows/locomanipulation/step_4_policy_training.rst b/docs/pages/example_workflows/locomanipulation/step_4_policy_training.rst index fb9217ed6..e0e3286d0 100644 --- a/docs/pages/example_workflows/locomanipulation/step_4_policy_training.rst +++ b/docs/pages/example_workflows/locomanipulation/step_4_policy_training.rst @@ -32,6 +32,7 @@ Note that this tutorial assumes that you've completed the nvidia/Arena-G1-Loco-Manipulation-Task \ arena_g1_loco_manipulation_dataset_generated.hdf5 \ --repo-type dataset \ + --revision arena_v0.2_lab_v3.0 \ --local-dir $DATASET_DIR Step 1: Convert to LeRobot Format @@ -55,6 +56,7 @@ Note that this conversion step can be skipped by downloading the pre-converted L nvidia/Arena-G1-Loco-Manipulation-Task \ --include lerobot/* \ --repo-type dataset \ + --revision arena_v0.2_lab_v3.0 \ --local-dir $DATASET_DIR/arena_g1_loco_manipulation_dataset_generated If you download this dataset, you can skip the conversion step below and continue to the next step. diff --git a/docs/pages/example_workflows/locomanipulation/step_5_evaluation.rst b/docs/pages/example_workflows/locomanipulation/step_5_evaluation.rst index db0515437..2b7101df4 100644 --- a/docs/pages/example_workflows/locomanipulation/step_5_evaluation.rst +++ b/docs/pages/example_workflows/locomanipulation/step_5_evaluation.rst @@ -71,6 +71,7 @@ Test the policy in a single environment with visualization via the GUI run: .. code-block:: bash python isaaclab_arena/evaluation/policy_runner.py \ + --visualizer kit \ --policy_type isaaclab_arena_gr00t.policy.gr00t_closedloop_policy.Gr00tClosedloopPolicy \ --policy_config_yaml_path isaaclab_arena_gr00t/policy/config/g1_locomanip_gr00t_closedloop_config.yaml \ --num_steps 1500 \ @@ -194,6 +195,7 @@ remote policy: .. code-block:: bash python isaaclab_arena/evaluation/policy_runner.py \ + --visualizer kit \ --policy_type isaaclab_arena.policy.action_chunking_client.ActionChunkingClientSidePolicy \ --remote_host 127.0.0.1 \ --remote_port 5555 \ diff --git a/docs/pages/example_workflows/reinforcement_learning/step_3_evaluation.rst b/docs/pages/example_workflows/reinforcement_learning/step_3_evaluation.rst index e418c79b1..b738733b7 100644 --- a/docs/pages/example_workflows/reinforcement_learning/step_3_evaluation.rst +++ b/docs/pages/example_workflows/reinforcement_learning/step_3_evaluation.rst @@ -48,6 +48,7 @@ Method 1: Single Environment Evaluation .. code-block:: bash python isaaclab_arena/evaluation/policy_runner.py \ + --visualizer kit \ --policy_type rsl_rl \ --num_steps 1000 \ --checkpoint_path logs/rsl_rl/generic_experiment/2026-01-28_17-26-10/model_11999.pt \ diff --git a/docs/pages/example_workflows/sequential_static_manipulation/step_1_environment_setup.rst b/docs/pages/example_workflows/sequential_static_manipulation/step_1_environment_setup.rst index 9d91f3b38..34f5bf744 100644 --- a/docs/pages/example_workflows/sequential_static_manipulation/step_1_environment_setup.rst +++ b/docs/pages/example_workflows/sequential_static_manipulation/step_1_environment_setup.rst @@ -364,6 +364,7 @@ can be fed to the robot to control its actions. nvidia/Arena-GR1-Manipulation-PlaceItemCloseDoor-Task \ ranch_bottle_into_fridge/ranch_bottle_into_fridge_annotated.hdf5 \ --repo-type dataset \ + --revision arena_v0.2_lab_v3.0 \ --local-dir "$_tmp" && \ mkdir -p "$DATASET_DIR" && \ mv "$_tmp/ranch_bottle_into_fridge/ranch_bottle_into_fridge_annotated.hdf5" "$DATASET_DIR/" && \ @@ -378,6 +379,7 @@ Replay the downloaded dataset to verify the environment setup: .. code-block:: bash python isaaclab_arena/scripts/imitation_learning/replay_demos.py \ + --visualizer kit \ --device cpu \ --enable_cameras \ --dataset_file "${DATASET_DIR}/ranch_bottle_into_fridge_annotated.hdf5" \ diff --git a/docs/pages/example_workflows/sequential_static_manipulation/step_2_teleoperation.rst b/docs/pages/example_workflows/sequential_static_manipulation/step_2_teleoperation.rst index c84f74483..9692572e8 100644 --- a/docs/pages/example_workflows/sequential_static_manipulation/step_2_teleoperation.rst +++ b/docs/pages/example_workflows/sequential_static_manipulation/step_2_teleoperation.rst @@ -1,61 +1,85 @@ Teleoperation Data Collection ----------------------------- -This workflow covers collecting demonstrations using Isaac Lab Teleop with an Apple Vision Pro. +This workflow covers collecting demonstrations using Isaac Teleop with an XR device, supported by `Nvidia IsaacTeleop `_. -This workflow requires two containers to run: +Step 1: Start the CloudXR Runtime +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -* **Nvidia CloudXR Runtime**: For connection with the Apple Vision Pro. -* **Arena Docker container**: For running the Isaac Lab simulation. -This will be described below. +.. tab-set:: + .. tab-item:: Meta Quest 3 / Pico 4 Ultra + :selected: -.. note:: + On the host machine, configure the firewall to allow CloudXR traffic. - For this workflow you will need an Apple Vision Pro. - In ``v0.2`` we will support further teleoperation devices. + .. code-block:: bash + sudo ufw allow 49100/tcp # Signaling + sudo ufw allow 47998/udp # Media stream + sudo ufw allow 48322/tcp # Proxy (HTTPS mode only) -Step 1: Install Isaac XR Teleop App on Vision Pro -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + Start the CloudXR runtime from the Arena Docker container: -Follow the `Isaac Lab CloudXR documentation -`_ -to build and install the app on your Apple Vision Pro. + :docker_run_default: + Create a CloudXR config to enable hand tracking: -Step 2: Start CloudXR Runtime Container -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + .. code-block:: bash -In a terminal, outside the Isaac Lab - Arena Docker container, start the CloudXR runtime: + echo "NV_DEVICE_PROFILE=auto-native" > handtracking.env -.. code-block:: bash - cd submodules/IsaacLab - mkdir -p openxr - - docker run -it --rm --name cloudxr-runtime \ - --user $(id -u):$(id -g) \ - --gpus=all \ - -e "ACCEPT_EULA=Y" \ - --mount type=bind,src=$(pwd)/openxr,dst=/openxr \ - -p 48010:48010 \ - -p 47998:47998/udp \ - -p 47999:47999/udp \ - -p 48000:48000/udp \ - -p 48005:48005/udp \ - -p 48008:48008/udp \ - -p 48012:48012/udp \ - nvcr.io/nvidia/cloudxr-runtime:5.0.0 - - -Step 3: Start Recording + Start the CloudXR runtime with the customized config file: + + .. code-block:: bash + + python -m isaacteleop.cloudxr --cloudxr-env-config=handtracking.env + + + .. tab-item:: Apple Vision Pro + + On the host machine, configure the firewall to allow CloudXR traffic. + + .. code-block:: bash + + # Signaling (use one based on connection mode) + sudo ufw allow 48010/tcp # Standard mode + sudo ufw allow 48322/tcp # Secure mode + # Video + sudo ufw allow 47998/udp + sudo ufw allow 48005/udp + sudo ufw allow 48008/udp + sudo ufw allow 48012/udp + # Input + sudo ufw allow 47999/udp + # Audio + sudo ufw allow 48000/udp + sudo ufw allow 48002/udp + + Start the CloudXR runtime from the Arena Docker container: + + :docker_run_default: + + Create a customized config file with the following content: + + .. code-block:: bash + + printf '%s\n' 'NV_DEVICE_PROFILE=auto-native' 'NV_CXR_ENABLE_PUSH_DEVICES=0' > avp.env + + + Start the CloudXR runtime with the customized config file: + + .. code-block:: bash + + python -m isaacteleop.cloudxr --cloudxr-env-config=avp.env + +Step 2: Start Recording ^^^^^^^^^^^^^^^^^^^^^^^ -To start the recording session, open another terminal, start the Arena Docker container -if not already running: +In another terminal, start the Arena Docker container: :docker_run_default: @@ -63,8 +87,10 @@ Run the recording script: .. code-block:: bash + source ~/.cloudxr/run/cloudxr.env python isaaclab_arena/scripts/imitation_learning/record_demos.py \ --device cpu \ + --visualizer kit \ --dataset_file $DATASET_DIR/ranch_bottle_into_fridge_recorded.hdf5 \ --num_demos 10 \ --num_success_steps 10 \ @@ -74,55 +100,90 @@ Run the recording script: --teleop_device openxr -Step 4: Connect Vision Pro and Record +Step 3: Connect XR Device and Record ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -Follow these steps to record teleoperation demonstrations: +For detailed instructions, refer to `Connect an XR Device `_. -1. Launch the Isaac XR Teleop app on the Apple Vision Pro -2. Enter your workstation's IP address in the app window. +A strong wireless connection is essential for a high-quality streaming experience. Refer to the `CloudXR Network Setup `_ guide for router configuration. -.. note:: - Before proceeding with teleoperation and pressing the "Connect" button: - Move the CloudXr Controls Application window closer and to your left by pinching the bar at the bottom of the window. - Without doing this, close objects will occlude the window making it harder to interact with the controls. - .. figure:: ../../../images/cloud_xr_sessions_control_panel.png - :width: 40% - :alt: CloudXR control panel - :align: center +.. tab-set:: - CloudXR control panel - move this window to your left to avoid occlusion by close objects. + .. tab-item:: Meta Quest 3 / Pico 4 Ultra + :selected: + .. note:: + Enable hand tracking on your Quest 3 headset for the first time: + 1. Press the Meta button on your right controller to open the universal menu. + 2. Select the clock on the left side of the universal menu to open Quick Settings. + 3. Select Settings. + 4. Select Movement tracking. + 5. Select the toggle next to Hand and Body Tracking to turn this feature on. -3. Press the "Connect" button -4. Wait for connection (you should see the simulation in VR) + #. Open the browser on your headset and navigate to ``_. + #. Enter the IP address of your Isaac Lab host machine in the **Server IP** field. -.. figure:: ../../../images/gr1_sequential_static_manipulation_env_vr_view.png - :width: 40% - :alt: IsaacSim view - :align: center + #. Click the **Click https://:48322/ to accept cert** link that appears on the page. + Accept the certificate in the new page that opens, then navigate back to the + CloudXR.js client page. - First person view after connecting to the simulation. + #. Click Connect to begin teleoperation. + .. note:: + Once you press **Connect** in the web browser, you should see the following control panel. Press **Play** to start teleoperation. -5. Complete the task by picking up the object, placing it into the lower shelf of the refrigerator, and closing the door. - - Your hands control the robots's hands. - - Your fingers control the robots's fingers. -6. On task completion the environment will automatically reset. -7. You'll need to repeat task completion ``num_demos`` times (set to 10 above). + If the control panel is not visible (for example, behind a solid wall in the simulated environment), you can put the headset on + before clicking **Start XR** in the Isaac Lab Arena application, and drag the control panel to a better location. + .. figure:: ../../../images/react-isaac-sample-controls-start.jpg + :width: 40% + :alt: IsaacSim view + :align: center -The script will automatically save successful demonstrations to an HDF5 file -at ``$DATASET_DIR/ranch_bottle_into_fridge_recorded.hdf5``. + .. tab-item:: Apple Vision Pro + + 1. Connect your XR device to the CloudXR runtime. From Apple Vision Pro, launch the + Isaac XR Teleop app. + 2. Enter your workstation's IP address and connect. + + .. note:: + Before proceeding with teleoperation and pressing **Connect**, move the CloudXR Controls Application window + closer and to your left by pinching the bar at the bottom of the window. + Without doing this, nearby objects will occlude the window, making it harder to interact with the controls. + + .. figure:: ../../../images/cloud_xr_sessions_control_panel.png + :width: 40% + :alt: CloudXR control panel + :align: center + CloudXR control panel—move this window to your left to avoid occlusion by nearby objects. + 3. Press the **Connect** button. + 4. Wait for the connection (you should see the simulation in VR). +.. figure:: ../../../images/gr1_sequential_static_manipulation_env_vr_view.png + :width: 40% + :alt: IsaacSim view + :align: center + + First person view after connecting to the simulation. + +#. Complete the task by picking up the object, placing it into the lower shelf of the refrigerator, and closing the door. + + - Your hands control the robot's hands. + - Your fingers control the robot's fingers. +#. On task completion the environment will automatically reset. +#. You'll need to repeat task completion ``num_demos`` times (set to 10 above). + + +The script will automatically save successful demonstrations to an HDF5 file +at ``$DATASET_DIR/ranch_bottle_into_fridge_recorded.hdf5``. .. hint:: diff --git a/docs/pages/example_workflows/sequential_static_manipulation/step_3_data_generation.rst b/docs/pages/example_workflows/sequential_static_manipulation/step_3_data_generation.rst index 6dc084f7e..95ad2b525 100644 --- a/docs/pages/example_workflows/sequential_static_manipulation/step_3_data_generation.rst +++ b/docs/pages/example_workflows/sequential_static_manipulation/step_3_data_generation.rst @@ -46,6 +46,7 @@ To skip this step, you can download the pre-annotated dataset as described below nvidia/Arena-GR1-Manipulation-PlaceItemCloseDoor-Task \ ranch_bottle_into_fridge/ranch_bottle_into_fridge_annotated.hdf5 \ --repo-type dataset \ + --revision arena_v0.2_lab_v3.0 \ --local-dir "$_tmp" && \ mkdir -p "$DATASET_DIR" && \ mv "$_tmp/ranch_bottle_into_fridge/ranch_bottle_into_fridge_annotated.hdf5" "$DATASET_DIR/" && \ @@ -56,6 +57,7 @@ To start the annotation process run the following command: .. code-block:: bash python isaaclab_arena/scripts/imitation_learning/annotate_demos.py \ + --visualizer kit \ --device cpu \ --input_file $DATASET_DIR/ranch_bottle_into_fridge_recorded.hdf5 \ --output_file $DATASET_DIR/ranch_bottle_into_fridge_annotated.hdf5 \ @@ -101,6 +103,7 @@ This step can be skipped by downloading the pre-generated dataset as described b nvidia/Arena-GR1-Manipulation-PlaceItemCloseDoor-Task \ ranch_bottle_into_fridge/ranch_bottle_into_fridge_generated_100.hdf5 \ --repo-type dataset \ + --revision arena_v0.2_lab_v3.0 \ --local-dir "$_tmp" && \ mkdir -p "$DATASET_DIR" && \ mv "$_tmp/ranch_bottle_into_fridge/ranch_bottle_into_fridge_generated_100.hdf5" "$DATASET_DIR/" && \ @@ -139,6 +142,7 @@ To do so, run the following command: .. code-block:: bash python isaaclab_arena/scripts/imitation_learning/replay_demos.py \ + --visualizer kit \ --device cpu \ --enable_cameras \ --dataset_file $DATASET_DIR/ranch_bottle_into_fridge_generated_100.hdf5 \ diff --git a/docs/pages/example_workflows/sequential_static_manipulation/step_4_policy_training.rst b/docs/pages/example_workflows/sequential_static_manipulation/step_4_policy_training.rst index c9c2795bd..91a4558b8 100644 --- a/docs/pages/example_workflows/sequential_static_manipulation/step_4_policy_training.rst +++ b/docs/pages/example_workflows/sequential_static_manipulation/step_4_policy_training.rst @@ -35,6 +35,7 @@ pre-generated dataset from Hugging Face as described below. nvidia/Arena-GR1-Manipulation-PlaceItemCloseDoor-Task \ --include "ranch_bottle_into_fridge/ranch_bottle_into_fridge_generated_100.hdf5" \ --repo-type dataset \ + --revision arena_v0.2_lab_v3.0 \ --local-dir "$_tmp" && \ mkdir -p "$DATASET_DIR" && \ mv "$_tmp/ranch_bottle_into_fridge/ranch_bottle_into_fridge_generated_100.hdf5" "$DATASET_DIR/" && \ @@ -61,6 +62,7 @@ Note that this conversion step can be skipped by downloading the pre-converted L nvidia/Arena-GR1-Manipulation-PlaceItemCloseDoor-Task \ --include "ranch_bottle_into_fridge/ranch_bottle_into_fridge_generated_100/lerobot/*" \ --repo-type dataset \ + --revision arena_v0.2_lab_v3.0 \ --local-dir "$_tmp" && \ mkdir -p "$DATASET_DIR" && \ mv "$_tmp/ranch_bottle_into_fridge/ranch_bottle_into_fridge_generated_100" "$DATASET_DIR/" && \ diff --git a/docs/pages/example_workflows/sequential_static_manipulation/step_5_evaluation.rst b/docs/pages/example_workflows/sequential_static_manipulation/step_5_evaluation.rst index e152cf20a..ffcf154cf 100644 --- a/docs/pages/example_workflows/sequential_static_manipulation/step_5_evaluation.rst +++ b/docs/pages/example_workflows/sequential_static_manipulation/step_5_evaluation.rst @@ -75,6 +75,7 @@ Test the policy in a single environment with visualization via the GUI run: .. code-block:: bash python isaaclab_arena/evaluation/policy_runner.py \ + --visualizer kit \ --policy_type isaaclab_arena_gr00t.policy.gr00t_closedloop_policy.Gr00tClosedloopPolicy \ --policy_config_yaml_path isaaclab_arena_gr00t/policy/config/gr1_manip_ranch_bottle_gr00t_closedloop_config.yaml \ --num_steps 2000 \ @@ -351,6 +352,7 @@ remote policy: .. code-block:: bash python isaaclab_arena/evaluation/policy_runner.py \ + --visualizer kit \ --policy_type isaaclab_arena.policy.action_chunking_client.ActionChunkingClientSidePolicy \ --remote_host 127.0.0.1 \ --remote_port 5555 \ diff --git a/docs/pages/example_workflows/static_manipulation/step_1_environment_setup.rst b/docs/pages/example_workflows/static_manipulation/step_1_environment_setup.rst index 360786efc..691af3309 100644 --- a/docs/pages/example_workflows/static_manipulation/step_1_environment_setup.rst +++ b/docs/pages/example_workflows/static_manipulation/step_1_environment_setup.rst @@ -152,6 +152,7 @@ Replay the downloaded dataset to verify the environment setup: .. code-block:: bash python isaaclab_arena/scripts/imitation_learning/replay_demos.py \ + --visualizer kit \ --device cpu \ --enable_cameras \ --dataset_file "${DATASET_DIR}/arena_gr1_manipulation_dataset_generated.hdf5" \ diff --git a/docs/pages/example_workflows/static_manipulation/step_2_teleoperation.rst b/docs/pages/example_workflows/static_manipulation/step_2_teleoperation.rst index cee0fe528..418edf002 100644 --- a/docs/pages/example_workflows/static_manipulation/step_2_teleoperation.rst +++ b/docs/pages/example_workflows/static_manipulation/step_2_teleoperation.rst @@ -1,61 +1,84 @@ Teleoperation Data Collection ----------------------------- -This workflow covers collecting demonstrations using Isaac Lab Teleop with an Apple Vision Pro. +This workflow covers collecting demonstrations using Isaac Teleop with an XR device, supported by `Nvidia IsaacTeleop `_. -This workflow requires two containers to run: +Step 1: Start the CloudXR Runtime +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -* **Nvidia CloudXR Runtime**: For connection with the Apple Vision Pro. -* **Arena Docker container**: For running the Isaac Lab simulation. +.. tab-set:: -This will be described below. + .. tab-item:: Meta Quest 3 / Pico 4 Ultra + :selected: + On the host machine, configure the firewall to allow CloudXR traffic. -.. note:: + .. code-block:: bash - For this workflow you will need an Apple Vision Pro. - In ``v0.2`` we will support further teleoperation devices. + sudo ufw allow 49100/tcp # Signaling + sudo ufw allow 47998/udp # Media stream + sudo ufw allow 48322/tcp # Proxy (HTTPS mode only) + Start the CloudXR runtime from the Arena Docker container: -Step 1: Install Isaac XR Teleop App on Vision Pro -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + :docker_run_default: -Follow the `Isaac Lab CloudXR documentation -`_ -to build and install the app on your Apple Vision Pro. + Create a CloudXR config to enable hand tracking: + .. code-block:: bash -Step 2: Start CloudXR Runtime Container -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + echo "NV_DEVICE_PROFILE=auto-native" > handtracking.env -In a terminal, outside the Isaac Lab - Arena Docker container, start the CloudXR runtime: -.. code-block:: bash + Start the CloudXR runtime with the customized config file: + + .. code-block:: bash + + python -m isaacteleop.cloudxr --cloudxr-env-config=handtracking.env + + + .. tab-item:: Apple Vision Pro + + On the host machine, configure the firewall to allow CloudXR traffic. + + .. code-block:: bash + + # Signaling (use one based on connection mode) + sudo ufw allow 48010/tcp # Standard mode + sudo ufw allow 48322/tcp # Secure mode + # Video + sudo ufw allow 47998/udp + sudo ufw allow 48005/udp + sudo ufw allow 48008/udp + sudo ufw allow 48012/udp + # Input + sudo ufw allow 47999/udp + # Audio + sudo ufw allow 48000/udp + sudo ufw allow 48002/udp + + Start the CloudXR runtime from the Arena Docker container: + + :docker_run_default: + + Create a customized config file with the following content: + + .. code-block:: bash + + printf '%s\n' 'NV_DEVICE_PROFILE=auto-native' 'NV_CXR_ENABLE_PUSH_DEVICES=0' > avp.env + + + Start the CloudXR runtime with the customized config file: + + .. code-block:: bash - cd submodules/IsaacLab - mkdir -p openxr - - docker run -it --rm --name cloudxr-runtime \ - --user $(id -u):$(id -g) \ - --gpus=all \ - -e "ACCEPT_EULA=Y" \ - --mount type=bind,src=$(pwd)/openxr,dst=/openxr \ - -p 48010:48010 \ - -p 47998:47998/udp \ - -p 47999:47999/udp \ - -p 48000:48000/udp \ - -p 48005:48005/udp \ - -p 48008:48008/udp \ - -p 48012:48012/udp \ - nvcr.io/nvidia/cloudxr-runtime:5.0.0 - - -Step 3: Start Recording + python -m isaacteleop.cloudxr --cloudxr-env-config=avp.env + +Step 2: Start Recording ^^^^^^^^^^^^^^^^^^^^^^^ -To start the recording session, open another terminal, start the Arena Docker container -if not already running: +In another terminal, start the Arena Docker container: :docker_run_default: @@ -63,64 +86,101 @@ Run the recording script: .. code-block:: bash + source ~/.cloudxr/run/cloudxr.env python isaaclab_arena/scripts/imitation_learning/record_demos.py \ --device cpu \ + --visualizer kit \ --dataset_file $DATASET_DIR/arena_gr1_manipulation_dataset_recorded.hdf5 \ --num_demos 10 \ --num_success_steps 2 \ gr1_open_microwave \ - --teleop_device avp_handtracking + --teleop_device openxr -Step 4: Connect Vision Pro and Record +Step 3: Connect XR Device and Record ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -Follow these steps to record teleoperation demonstrations: +For detailed instructions, refer to `Connect an XR Device `_. -1. Launch the Isaac XR Teleop app on the Apple Vision Pro -2. Enter your workstation's IP address in the app window. +A strong wireless connection is essential for a high-quality streaming experience. Refer to the `CloudXR Network Setup `_ guide for router configuration. -.. note:: - Before proceeding with teleoperation and pressing the "Connect" button: - Move the CloudXr Controls Application window closer and to your left by pinching the bar at the bottom of the window. - Without doing this, close objects will occlude the window making it harder to interact with the controls. - .. figure:: ../../../images/cloud_xr_sessions_control_panel.png - :width: 40% - :alt: CloudXR control panel - :align: center +.. tab-set:: - CloudXR control panel - move this window to your left to avoid occlusion by close objects. + .. tab-item:: Meta Quest 3 / Pico 4 Ultra + :selected: + .. note:: + Enable hand tracking on your Quest 3 headset for the first time: + 1. Press the Meta button on your right controller to open the universal menu. + 2. Select the clock on the left side of the universal menu to open Quick Settings. + 3. Select Settings. + 4. Select Movement tracking. + 5. Select the toggle next to Hand and Body Tracking to turn this feature on. -3. Press the "Connect" button -4. Wait for connection (you should see the simulation in VR) + #. Open the browser on your headset and navigate to ``_. + #. Enter the IP address of your Isaac Lab host machine in the **Server IP** field. -.. figure:: ../../../images/simulation_view.png - :width: 40% - :alt: IsaacSim view - :align: center + #. Click the **Click https://:48322/ to accept cert** link that appears on the page. + Accept the certificate in the new page that opens, then navigate back to the + CloudXR.js client page. - First person view after connecting to the simulation. + #. Click Connect to begin teleoperation. + .. note:: + Once you press **Connect** in the web browser, you should see the following control panel. Press **Play** to start teleoperation. -5. Complete the task by opening the microwave door. - - Your hands control the robots's hands. - - Your fingers control the robots's fingers. -6. On task completion the environment will automatically reset. -7. You'll need to repeat task completion ``num_demos`` times (set to 10 above). + If the control panel is not visible (for example, behind a solid wall in the simulated environment), you can put the headset on + before clicking **Start XR** in the Isaac Lab Arena application, and drag the control panel to a better location. + .. figure:: ../../../images/react-isaac-sample-controls-start.jpg + :width: 40% + :alt: IsaacSim view + :align: center -The script will automatically save successful demonstrations to an HDF5 file -at ``$DATASET_DIR/arena_gr1_manipulation_dataset_recorded.hdf5``. + .. tab-item:: Apple Vision Pro + 1. Connect your XR device to the CloudXR runtime. From Apple Vision Pro, launch the + Isaac XR Teleop app. + 2. Enter your workstation's IP address and connect. + .. note:: + Before proceeding with teleoperation and pressing **Connect**, move the CloudXR Controls Application window + closer and to your left by pinching the bar at the bottom of the window. + Without doing this, nearby objects will occlude the window, making it harder to interact with the controls. + .. figure:: ../../../images/cloud_xr_sessions_control_panel.png + :width: 40% + :alt: CloudXR control panel + :align: center + CloudXR control panel—move this window to your left to avoid occlusion by nearby objects. + + 3. Press the **Connect** button. + 4. Wait for the connection (you should see the simulation in VR). + + +.. figure:: ../../../images/simulation_view.png + :width: 40% + :alt: IsaacSim view + :align: center + + First person view after connecting to the simulation. + +#. Complete the task by opening the microwave door. + + - Your hands control the robot's hands. + - Your fingers control the robot's fingers. +#. On task completion the environment will automatically reset. +#. You'll need to repeat task completion ``num_demos`` times (set to 10 above). + + +The script will automatically save successful demonstrations to an HDF5 file +at ``$DATASET_DIR/arena_gr1_manipulation_dataset_recorded.hdf5``. .. hint:: diff --git a/docs/pages/example_workflows/static_manipulation/step_3_data_generation.rst b/docs/pages/example_workflows/static_manipulation/step_3_data_generation.rst index 1bec8ce3e..d26359f29 100644 --- a/docs/pages/example_workflows/static_manipulation/step_3_data_generation.rst +++ b/docs/pages/example_workflows/static_manipulation/step_3_data_generation.rst @@ -45,7 +45,7 @@ To skip this step, you can download the pre-annotated dataset from Hugging Face nvidia/Arena-GR1-Manipulation-Task \ arena_gr1_manipulation_dataset_annotated.hdf5 \ --repo-type dataset \ - --revision refs/pr/2 \ + --revision arena_v0.2_lab_v3.0 \ --local-dir $DATASET_DIR To start the annotation process run the following command: @@ -53,6 +53,7 @@ To start the annotation process run the following command: .. code-block:: bash python isaaclab_arena/scripts/imitation_learning/annotate_demos.py \ + --visualizer kit \ --device cpu \ --input_file $DATASET_DIR/arena_gr1_manipulation_dataset_recorded.hdf5 \ --output_file $DATASET_DIR/arena_gr1_manipulation_dataset_annotated.hdf5 \ @@ -89,6 +90,7 @@ This step can be skipped by downloading the pre-generated dataset from Hugging F nvidia/Arena-GR1-Manipulation-Task \ arena_gr1_manipulation_dataset_generated.hdf5 \ --repo-type dataset \ + --revision arena_v0.2_lab_v3.0 \ --local-dir $DATASET_DIR @@ -123,6 +125,7 @@ To do so, run the following command: .. code-block:: bash python isaaclab_arena/scripts/imitation_learning/replay_demos.py \ + --visualizer kit \ --device cpu \ --enable_cameras \ --dataset_file $DATASET_DIR/arena_gr1_manipulation_dataset_generated.hdf5 \ diff --git a/docs/pages/example_workflows/static_manipulation/step_4_policy_training.rst b/docs/pages/example_workflows/static_manipulation/step_4_policy_training.rst index b950de438..5547e484a 100644 --- a/docs/pages/example_workflows/static_manipulation/step_4_policy_training.rst +++ b/docs/pages/example_workflows/static_manipulation/step_4_policy_training.rst @@ -35,6 +35,7 @@ pre-generated dataset from Hugging Face as described below. nvidia/Arena-GR1-Manipulation-Task \ arena_gr1_manipulation_dataset_generated.hdf5 \ --repo-type dataset \ + --revision arena_v0.2_lab_v3.0 \ --local-dir $DATASET_DIR @@ -57,6 +58,7 @@ Note that this conversion step can be skipped by downloading the pre-converted L nvidia/Arena-GR1-Manipulation-Task \ --include lerobot/* \ --repo-type dataset \ + --revision arena_v0.2_lab_v3.0 \ --local-dir $DATASET_DIR/arena_gr1_manipulation_dataset_generated If you download this dataset, you can skip the conversion step below and continue to the next step. diff --git a/docs/pages/example_workflows/static_manipulation/step_5_evaluation.rst b/docs/pages/example_workflows/static_manipulation/step_5_evaluation.rst index 6304e9638..bbff2cd45 100644 --- a/docs/pages/example_workflows/static_manipulation/step_5_evaluation.rst +++ b/docs/pages/example_workflows/static_manipulation/step_5_evaluation.rst @@ -70,6 +70,7 @@ Test the policy in a single environment with visualization via the GUI run: .. code-block:: bash python isaaclab_arena/evaluation/policy_runner.py \ + --visualizer kit \ --policy_type isaaclab_arena_gr00t.policy.gr00t_closedloop_policy.Gr00tClosedloopPolicy \ --policy_config_yaml_path isaaclab_arena_gr00t/policy/config/gr1_manip_gr00t_closedloop_config.yaml \ --num_steps 2000 \ @@ -194,6 +195,7 @@ policy: .. code-block:: bash python isaaclab_arena/evaluation/policy_runner.py \ + --visualizer kit \ --policy_type isaaclab_arena.policy.action_chunking_client.ActionChunkingClientSidePolicy \ --remote_host 127.0.0.1 \ --remote_port 5555 \ diff --git a/isaaclab_arena/affordances/placeable.py b/isaaclab_arena/affordances/placeable.py index 2a4cb9e20..19621322c 100644 --- a/isaaclab_arena/affordances/placeable.py +++ b/isaaclab_arena/affordances/placeable.py @@ -5,6 +5,7 @@ import math import torch +import warp as wp from typing import Literal import isaaclab.utils.math as math_utils @@ -55,7 +56,7 @@ def is_placed_upright( if orientation_threshold is None: orientation_threshold = self.orientation_threshold object_entity: RigidObject = env.scene[asset_cfg.name] - object_quat = object_entity.data.root_quat_w + object_quat = wp.to_torch(object_entity.data.root_quat_w) upright_axis_world = get_object_axis_in_world_frame(object_quat, self.upright_axis_name) @@ -170,12 +171,12 @@ def set_normalized_object_pose( """ object_entity: RigidObject = env.scene[asset_cfg.name] device = env.device - dtype = object_entity.data.root_quat_w.dtype + dtype = wp.to_torch(object_entity.data.root_quat_w).dtype if env_ids is not None: env_ids = env_ids.to(env.device) else: - env_ids = torch.arange(object_entity.data.root_quat_w.shape[0], device=env.device) + env_ids = torch.arange(wp.to_torch(object_entity.data.root_quat_w).shape[0], device=env.device) # Validate upright_percentage shape if it's a tensor if isinstance(upright_percentage, torch.Tensor): @@ -186,8 +187,8 @@ def set_normalized_object_pose( f"but got {upright_percentage.numel()} elements" ) - object_quat = object_entity.data.root_quat_w[env_ids] - object_pos = object_entity.data.root_pos_w[env_ids] + object_quat = wp.to_torch(object_entity.data.root_quat_w)[env_ids] + object_pos = wp.to_torch(object_entity.data.root_pos_w)[env_ids] target_quat = _compute_target_quaternions( object_quat=object_quat, diff --git a/isaaclab_arena/assets/asset_registry.py b/isaaclab_arena/assets/asset_registry.py index 3778f7d81..33bb69acc 100644 --- a/isaaclab_arena/assets/asset_registry.py +++ b/isaaclab_arena/assets/asset_registry.py @@ -127,26 +127,12 @@ def get_device_by_name(self, name: str) -> type["TeleopDeviceBase"]: return self.get_component_by_name(name) def get_teleop_device_cfg(self, device: type["TeleopDeviceBase"], embodiment: object): - from isaaclab.devices.device_base import DevicesCfg - retargeter_registry = RetargeterRegistry() retargeter_key = (device.name, embodiment.name) retargeter_key_str = retargeter_registry.convert_tuple_to_str(retargeter_key) retargeter = retargeter_registry.get_component_by_name(retargeter_key_str)() - retargeter_cfg = retargeter.get_retargeter_cfg(embodiment, sim_device=device.sim_device) - # Handle both single retargeter and list of retargeters - if isinstance(retargeter_cfg, list): - retargeters = retargeter_cfg - elif retargeter_cfg is not None: - retargeters = [retargeter_cfg] - else: - retargeters = [] - device_cfg = device.get_device_cfg(retargeters=retargeters, embodiment=embodiment) - return DevicesCfg( - devices={ - device.name: device_cfg, - } - ) + pipeline_builder = retargeter.get_pipeline_builder(embodiment) + return device.get_device_cfg(pipeline_builder=pipeline_builder, embodiment=embodiment) class RetargeterRegistry(Registry): diff --git a/isaaclab_arena/assets/background_library.py b/isaaclab_arena/assets/background_library.py index 8e569b69f..dc15d56ab 100644 --- a/isaaclab_arena/assets/background_library.py +++ b/isaaclab_arena/assets/background_library.py @@ -51,7 +51,7 @@ class KitchenBackground(LibraryBackground): name = "kitchen" tags = ["background"] usd_path = f"{ISAACLAB_NUCLEUS_DIR}/Arena/assets/background_library/kitchen_background/kitchen_background.usd" - initial_pose = Pose(position_xyz=(0.772, 3.39, -0.895), rotation_wxyz=(0.70711, 0, 0, -0.70711)) + initial_pose = Pose(position_xyz=(0.772, 3.39, -0.895), rotation_xyzw=(0, 0, -0.70711, 0.70711)) object_min_z = -0.2 def __init__(self): @@ -69,7 +69,7 @@ class KitchenWithOpenDrawerBackground(LibraryBackground): usd_path = ( f"{ISAACLAB_NUCLEUS_DIR}/Arena/assets/background_library/kitchen_scene_teleop_v3/kitchen_scene_teleop_v3.usd" ) - initial_pose = Pose(position_xyz=(0.772, 3.39, -0.895), rotation_wxyz=(0.70711, 0, 0, -0.70711)) + initial_pose = Pose(position_xyz=(0.772, 3.39, -0.895), rotation_xyzw=(0, 0, -0.70711, 0.70711)) object_min_z = -0.2 def __init__(self): @@ -85,7 +85,7 @@ class PackingTableBackground(LibraryBackground): name = "packing_table" tags = ["background"] usd_path = f"{ISAACLAB_NUCLEUS_DIR}/Arena/assets/background_library/packing_table/packing_table.usd" - initial_pose = Pose(position_xyz=(0.72193, -0.04727, -0.92512), rotation_wxyz=(0.70711, 0.0, 0.0, -0.70711)) + initial_pose = Pose(position_xyz=(0.72193, -0.04727, -0.92512), rotation_xyzw=(0.0, 0.0, -0.70711, 0.70711)) object_min_z = -0.2 def __init__(self): @@ -101,7 +101,7 @@ class GalileoBackground(LibraryBackground): name = "galileo" tags = ["background"] usd_path = f"{ISAACLAB_NUCLEUS_DIR}/Arena/assets/background_library/galileo_simplified/galileo_simplified.usd" - initial_pose = Pose(position_xyz=(4.420, 1.408, -0.795), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)) + initial_pose = Pose(position_xyz=(4.420, 1.408, -0.795), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)) object_min_z = -0.2 def __init__(self): @@ -117,7 +117,7 @@ class GalileoLocomanipBackground(LibraryBackground): name = "galileo_locomanip" tags = ["background"] usd_path = f"{ISAACLAB_NUCLEUS_DIR}/Arena/assets/background_library/galileo_locomanip/galileo_locomanip.usd" - initial_pose = Pose(position_xyz=(4.420, 1.408, -0.795), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)) + initial_pose = Pose(position_xyz=(4.420, 1.408, -0.795), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)) object_min_z = -0.2 def __init__(self): diff --git a/isaaclab_arena/assets/device_library.py b/isaaclab_arena/assets/device_library.py index 95e4ee7ba..649127457 100644 --- a/isaaclab_arena/assets/device_library.py +++ b/isaaclab_arena/assets/device_library.py @@ -17,11 +17,11 @@ # limitations under the License. from abc import ABC, abstractmethod +from collections.abc import Callable from isaaclab.devices.keyboard import Se3KeyboardCfg -from isaaclab.devices.openxr import OpenXRDeviceCfg -from isaaclab.devices.retargeter_base import RetargeterCfg from isaaclab.devices.spacemouse import Se3SpaceMouseCfg +from isaaclab_teleop import IsaacTeleopCfg, XrCfg from isaaclab_arena.assets.register import register_device @@ -34,7 +34,9 @@ def __init__(self, sim_device: str | None = None): self.sim_device = sim_device @abstractmethod - def get_device_cfg(self, retargeters: list[RetargeterCfg] | None = None, embodiment: object | None = None): + def get_device_cfg( + self, pipeline_builder: Callable | None = None, embodiment: object | None = None + ): raise NotImplementedError @@ -46,12 +48,17 @@ def __init__(self, sim_device: str | None = None): super().__init__(sim_device=sim_device) def get_device_cfg( - self, retargeters: list[RetargeterCfg] | None = None, embodiment: object | None = None - ) -> OpenXRDeviceCfg: - return OpenXRDeviceCfg( - retargeters=retargeters, + self, pipeline_builder: Callable | None = None, embodiment: object | None = None + ) -> IsaacTeleopCfg | None: + if pipeline_builder is None: + return None + xr_cfg = embodiment.get_xr_cfg() if embodiment is not None else XrCfg() + target_frame_prim_path = embodiment.get_teleop_target_frame_prim_path() + return IsaacTeleopCfg( + pipeline_builder=pipeline_builder, sim_device=self.sim_device, - xr_cfg=embodiment.get_xr_cfg(), + xr_cfg=xr_cfg, + target_frame_prim_path=target_frame_prim_path, ) @@ -65,11 +72,9 @@ def __init__(self, sim_device: str | None = None, pos_sensitivity: float = 0.05, self.rot_sensitivity = rot_sensitivity def get_device_cfg( - self, retargeters: list[RetargeterCfg] | None = None, embodiment: object | None = None + self, pipeline_builder: Callable | None = None, embodiment: object | None = None ) -> Se3KeyboardCfg: return Se3KeyboardCfg( - retargeters=retargeters, - sim_device=self.sim_device, pos_sensitivity=self.pos_sensitivity, rot_sensitivity=self.rot_sensitivity, ) @@ -85,11 +90,9 @@ def __init__(self, sim_device: str | None = None, pos_sensitivity: float = 0.05, self.rot_sensitivity = rot_sensitivity def get_device_cfg( - self, retargeters: list[RetargeterCfg] | None = None, embodiment: object | None = None + self, pipeline_builder: Callable | None = None, embodiment: object | None = None ) -> Se3SpaceMouseCfg: return Se3SpaceMouseCfg( - retargeters=retargeters, - sim_device=self.sim_device, pos_sensitivity=self.pos_sensitivity, rot_sensitivity=self.rot_sensitivity, ) diff --git a/isaaclab_arena/assets/dummy_object.py b/isaaclab_arena/assets/dummy_object.py index 8435a5392..1f1170f74 100644 --- a/isaaclab_arena/assets/dummy_object.py +++ b/isaaclab_arena/assets/dummy_object.py @@ -49,7 +49,7 @@ def get_world_bounding_box(self) -> AxisAlignedBoundingBox: """ if self.initial_pose is None: return self.bounding_box - quarters = quaternion_to_90_deg_z_quarters(self.initial_pose.rotation_wxyz) + quarters = quaternion_to_90_deg_z_quarters(self.initial_pose.rotation_xyzw) return self.bounding_box.rotated_90_around_z(quarters).translated(self.initial_pose.position_xyz) def get_corners_aabb(self, pos: torch.Tensor) -> torch.Tensor: diff --git a/isaaclab_arena/assets/g1_pink_locomanipulation_pipeline.py b/isaaclab_arena/assets/g1_pink_locomanipulation_pipeline.py new file mode 100644 index 000000000..e9ff18018 --- /dev/null +++ b/isaaclab_arena/assets/g1_pink_locomanipulation_pipeline.py @@ -0,0 +1,209 @@ +# Copyright (c) 2026, The Isaac Lab Arena Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Custom IsaacTeleop pipeline for G1 WBC Pink: 23D output (20D + 3 torso zeros). + +What each retargeter.connect() does +------------------------------------ +- **left_se3.connect({LEFT: transformed_controllers.output(LEFT)})** + Binds the left controller's *transformed* grip pose (position + quaternion in + world frame) as the sole input to the left SE3 retargeter. The retargeter + outputs a 7D ee_pose (position + quat) with configurable rotation offsets. + +- **right_se3.connect({RIGHT: transformed_controllers.output(RIGHT)})** + Same for the right controller → right SE3 retargeter → 7D ee_pose. + +- **left_trihand.connect({LEFT: transformed_controllers.output(LEFT)})** + Feeds the left controller (buttons, trigger, squeeze) into the left + TriHand retargeter. It outputs 7 hand joint scalars (thumb/index/middle) + derived from trigger and squeeze. + +- **right_trihand.connect({RIGHT: transformed_controllers.output(RIGHT)})** + Same for the right controller → right TriHand → 7 hand joint scalars. + +- **locomotion.connect({controller_left: controllers.output(LEFT), controller_right: ...})** + Feeds *raw* (untransformed) left and right controller data into the + locomotion retargeter. It uses thumbsticks to produce a 4D root command + [vel_x, vel_y, rot_vel_z, hip_height]. Raw controllers are used so + thumbstick values are in controller space. + +- **reorderer.connect({...})** + Connects each retargeter's output to the TensorReorderer's named inputs. + Adds three 0s for the torso. + The reorderer flattens and reorders them into a single 23D action vector. +""" + +def _build_g1_pink_locomanipulation_pipeline(): + """Build an IsaacTeleop retargeting pipeline for G1 WBC Pink locomanipulation. + + Same sources as Isaac Lab's G1 locomanipulation (Se3 wrists, TriHand hands, + Locomotion root), plus a ValueInput for torso_rpy (zeros). Output is 23D: + [left_gripper(1), right_gripper(1), left_wrist(7), right_wrist(7), locomotion(4), torso_rpy(3)]. + + Returns: + OutputCombiner with a single "action" output (23D flattened tensor). + """ + from isaacteleop.retargeters import ( + LocomotionRootCmdRetargeter, + LocomotionRootCmdRetargeterConfig, + Se3AbsRetargeter, + Se3RetargeterConfig, + TensorReorderer, + TriHandMotionControllerConfig, + TriHandMotionControllerRetargeter, + ) + from isaacteleop.retargeting_engine.deviceio_source_nodes import ControllersSource + from isaacteleop.retargeting_engine.interface import OutputCombiner, ValueInput + from isaacteleop.retargeting_engine.tensor_types import TransformMatrix + + controllers = ControllersSource(name="controllers") + transform_input = ValueInput("world_T_anchor", TransformMatrix()) + transformed_controllers = controllers.transformed(transform_input.output(ValueInput.VALUE)) + + # ------------------------------------------------------------------------- + # SE3 Absolute Pose Retargeters (left and right wrists) + # ------------------------------------------------------------------------- + # connect(): binds transformed left/right controller grip pose -> 7D ee_pose each. + left_se3_cfg = Se3RetargeterConfig( + input_device=ControllersSource.LEFT, + zero_out_xy_rotation=False, + use_wrist_rotation=False, + use_wrist_position=False, + target_offset_roll=45.0, + target_offset_pitch=180.0, + target_offset_yaw=-90.0, + ) + left_se3 = Se3AbsRetargeter(left_se3_cfg, name="left_ee_pose") + connected_left_se3 = left_se3.connect( + {ControllersSource.LEFT: transformed_controllers.output(ControllersSource.LEFT)} + ) + + right_se3_cfg = Se3RetargeterConfig( + input_device=ControllersSource.RIGHT, + zero_out_xy_rotation=False, + use_wrist_rotation=False, + use_wrist_position=False, + target_offset_roll=-135.0, + target_offset_pitch=0.0, + target_offset_yaw=90.0, + ) + right_se3 = Se3AbsRetargeter(right_se3_cfg, name="right_ee_pose") + connected_right_se3 = right_se3.connect( + {ControllersSource.RIGHT: transformed_controllers.output(ControllersSource.RIGHT)} + ) + + # ------------------------------------------------------------------------- + # TriHand Motion Controller Retargeters (for gripper scalar per hand) + # ------------------------------------------------------------------------- + # connect(): binds transformed left/right controller -> 7 hand joint scalars each. + hand_joint_names = [ + "thumb_rotation", + "thumb_proximal", + "thumb_distal", + "index_proximal", + "index_distal", + "middle_proximal", + "middle_distal", + ] + left_trihand_cfg = TriHandMotionControllerConfig( + hand_joint_names=hand_joint_names, + controller_side="left", + ) + left_trihand = TriHandMotionControllerRetargeter(left_trihand_cfg, name="trihand_left") + connected_left_trihand = left_trihand.connect( + {ControllersSource.LEFT: transformed_controllers.output(ControllersSource.LEFT)} + ) + right_trihand_cfg = TriHandMotionControllerConfig( + hand_joint_names=hand_joint_names, + controller_side="right", + ) + right_trihand = TriHandMotionControllerRetargeter(right_trihand_cfg, name="trihand_right") + connected_right_trihand = right_trihand.connect( + {ControllersSource.RIGHT: transformed_controllers.output(ControllersSource.RIGHT)} + ) + + # ------------------------------------------------------------------------- + # Locomotion Root Command Retargeter + # ------------------------------------------------------------------------- + # connect(): binds raw left/right controller (thumbsticks) -> 4D root_command. + locomotion_cfg = LocomotionRootCmdRetargeterConfig( + initial_hip_height=0.72, + movement_scale=0.5, + rotation_scale=0.35, + dt=1.0 / 100.0, + ) + locomotion = LocomotionRootCmdRetargeter(locomotion_cfg, name="locomotion") + connected_locomotion = locomotion.connect( + { + "controller_left": controllers.output(ControllersSource.LEFT), + "controller_right": controllers.output(ControllersSource.RIGHT), + } + ) + + # ------------------------------------------------------------------------- + # TensorReorderer: 23D for G1 WBC Pink [20D above + torso_rpy(3) from ConstantRetargeter] + # ------------------------------------------------------------------------- + left_ee_elements = ["l_pos_x", "l_pos_y", "l_pos_z", "l_quat_x", "l_quat_y", "l_quat_z", "l_quat_w"] + right_ee_elements = ["r_pos_x", "r_pos_y", "r_pos_z", "r_quat_x", "r_quat_y", "r_quat_z", "r_quat_w"] + left_hand_elements = [ + "l_thumb_rotation", + "l_thumb_proximal", + "l_thumb_distal", + "l_index_proximal", + "l_index_distal", + "l_middle_proximal", + "l_middle_distal", + ] + right_hand_elements = [ + "r_thumb_rotation", + "r_thumb_proximal", + "r_thumb_distal", + "r_index_proximal", + "r_index_distal", + "r_middle_proximal", + "r_middle_distal", + ] + locomotion_elements = ["loco_vel_x", "loco_vel_y", "loco_rot_vel_z", "loco_hip_height"] + torso_elements = ["torso_x", "torso_y", "torso_z"] + + output_order = ( + ["l_thumb_rotation", "r_thumb_rotation"] # left_gripper(1), right_gripper(1) + + left_ee_elements + + right_ee_elements + + locomotion_elements + + torso_elements # 3 zeros + ) + + reorderer = TensorReorderer( + input_config={ + "left_ee_pose": left_ee_elements, + "right_ee_pose": right_ee_elements, + "left_hand_joints": left_hand_elements, + "right_hand_joints": right_hand_elements, + "locomotion": locomotion_elements, + }, + output_order=output_order, + name="action_reorderer", + input_types={ + "left_ee_pose": "array", + "right_ee_pose": "array", + "left_hand_joints": "scalar", + "right_hand_joints": "scalar", + "locomotion": "array", + }, + ) + # connect(): binds each retargeter output to the reorderer; flattens to 23D action. + # torso_rpy comes from ConstantRetargeter(output_dims=3, value=0.0). + connected_reorderer = reorderer.connect( + { + "left_ee_pose": connected_left_se3.output("ee_pose"), + "right_ee_pose": connected_right_se3.output("ee_pose"), + "left_hand_joints": connected_left_trihand.output("hand_joints"), + "right_hand_joints": connected_right_trihand.output("hand_joints"), + "locomotion": connected_locomotion.output("root_command"), + } + ) + + return OutputCombiner({"action": connected_reorderer.output("output")}) diff --git a/isaaclab_arena/assets/object.py b/isaaclab_arena/assets/object.py index 77a514be9..ff216f04d 100644 --- a/isaaclab_arena/assets/object.py +++ b/isaaclab_arena/assets/object.py @@ -6,14 +6,17 @@ from typing import Any from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg +from isaaclab.managers import EventTermCfg, SceneEntityCfg from isaaclab.sensors.contact_sensor.contact_sensor_cfg import ContactSensorCfg from isaaclab.sim.spawners.from_files.from_files_cfg import UsdFileCfg +from isaaclab_tasks.manager_based.manipulation.stack.mdp.franka_stack_events import randomize_object_pose from isaaclab_arena.assets.object_base import ObjectBase, ObjectType from isaaclab_arena.assets.object_utils import detect_object_type from isaaclab_arena.relations.relations import RelationBase +from isaaclab_arena.terms.events import set_object_pose from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox, quaternion_to_90_deg_z_quarters -from isaaclab_arena.utils.pose import Pose +from isaaclab_arena.utils.pose import Pose, PoseRange from isaaclab_arena.utils.usd.rigid_bodies import find_shallowest_rigid_body from isaaclab_arena.utils.usd_helpers import compute_local_bounding_box_from_usd, has_light, open_stage @@ -74,7 +77,7 @@ def get_world_bounding_box(self) -> AxisAlignedBoundingBox: local_bbox = self.get_bounding_box() if self.initial_pose is None or not isinstance(self.initial_pose, Pose): return local_bbox - quarters = quaternion_to_90_deg_z_quarters(self.initial_pose.rotation_wxyz) + quarters = quaternion_to_90_deg_z_quarters(self.initial_pose.rotation_xyzw) return local_bbox.rotated_90_around_z(quarters).translated(self.initial_pose.position_xyz) def get_corners(self, pos: torch.Tensor) -> torch.Tensor: @@ -83,16 +86,31 @@ def get_corners(self, pos: torch.Tensor) -> torch.Tensor: self.bounding_box = compute_local_bounding_box_from_usd(self.usd_path, self.scale) return self.bounding_box.get_corners_at(pos) + def set_initial_pose(self, pose: Pose | PoseRange) -> None: + """Set the initial pose of the object. + + Args: + pose: The pose to set. Can be a single pose or a pose range. + In the case of a PoseRange, the object will be reset + to a random pose within the range on environment reset. + """ + self.initial_pose = pose + self.object_cfg = self._add_initial_pose_to_cfg(self.object_cfg) + self.event_cfg = self._update_initial_pose_event_cfg(self.event_cfg) + + def get_initial_pose(self) -> Pose | PoseRange | None: + return self.initial_pose + def is_initial_pose_set(self) -> bool: return self.initial_pose is not None def disable_reset_pose(self) -> None: self.reset_pose = False - self.event_cfg = self._init_event_cfg() + self.event_cfg = self._update_initial_pose_event_cfg(self.event_cfg) def enable_reset_pose(self) -> None: self.reset_pose = True - self.event_cfg = self._init_event_cfg() + self.event_cfg = self._update_initial_pose_event_cfg(self.event_cfg) def get_contact_sensor_cfg(self, contact_against_prim_paths: list[str] | None = None) -> ContactSensorCfg: # We override this function from the parent class because in some assets, the rigid body @@ -176,11 +194,73 @@ def _add_initial_pose_to_cfg( self, object_cfg: RigidObjectCfg | ArticulationCfg | AssetBaseCfg ) -> RigidObjectCfg | ArticulationCfg | AssetBaseCfg: # Optionally specify initial pose - initial_pose = self._get_initial_pose_as_pose() - if initial_pose is not None: + if self.initial_pose is not None: + if isinstance(self.initial_pose, Pose): + initial_pose = self.initial_pose + elif isinstance(self.initial_pose, PoseRange): + initial_pose = self.initial_pose.get_midpoint() object_cfg.init_state.pos = initial_pose.position_xyz - object_cfg.init_state.rot = initial_pose.rotation_wxyz + object_cfg.init_state.rot = initial_pose.rotation_xyzw return object_cfg def _requires_reset_pose_event(self) -> bool: - return super()._requires_reset_pose_event() and self.reset_pose + return ( + self.initial_pose is not None + and self.reset_pose + and self.object_type in [ObjectType.RIGID, ObjectType.ARTICULATION] + ) + + def _init_event_cfg(self) -> EventTermCfg | None: + if self._requires_reset_pose_event(): + # Two possible event types: + # - initial pose is a Pose - reset to a single pose + # - initial pose is a PoseRange - reset to a random pose within the range + if isinstance(self.initial_pose, Pose): + return EventTermCfg( + func=set_object_pose, + mode="reset", + params={ + "pose": self.initial_pose, + "asset_cfg": SceneEntityCfg(self.name), + }, + ) + elif isinstance(self.initial_pose, PoseRange): + return EventTermCfg( + func=randomize_object_pose, + mode="reset", + params={ + "pose_range": self.initial_pose.to_dict(), + "asset_cfgs": [SceneEntityCfg(self.name)], + }, + ) + else: + raise ValueError(f"Initial pose {self.initial_pose} is not a Pose or PoseRange") + else: + return None + + def _needs_reinit_of_event_cfg(self): + # If there is no event cfg, needs to be reinitialized + if self.event_cfg is None: + return True + # Here we check if the event cfg is for the correct pose type. + # If not, needs to be reinitialized. + if (isinstance(self.initial_pose, Pose) and ("pose" not in self.event_cfg.params)) or ( + isinstance(self.initial_pose, PoseRange) and ("pose_range" not in self.event_cfg.params) + ): + return True + return False + + def _update_initial_pose_event_cfg(self, event_cfg: EventTermCfg | None) -> EventTermCfg | None: + if self._requires_reset_pose_event(): + # Create an event cfg if one does not yet exist + if self._needs_reinit_of_event_cfg(): + event_cfg = self._init_event_cfg() + if isinstance(self.initial_pose, Pose): + event_cfg.params["pose"] = self.initial_pose + elif isinstance(self.initial_pose, PoseRange): + event_cfg.params["pose_range"] = self.initial_pose.to_dict() + else: + raise ValueError(f"Initial pose {self.initial_pose} is not a Pose or PoseRange") + else: + event_cfg = None + return event_cfg \ No newline at end of file diff --git a/isaaclab_arena/assets/object_base.py b/isaaclab_arena/assets/object_base.py index 91c3697bc..d09f07139 100644 --- a/isaaclab_arena/assets/object_base.py +++ b/isaaclab_arena/assets/object_base.py @@ -4,19 +4,18 @@ # SPDX-License-Identifier: Apache-2.0 import torch +import warp as wp from abc import ABC, abstractmethod from enum import Enum from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg from isaaclab.envs import ManagerBasedEnv -from isaaclab.managers import EventTermCfg, SceneEntityCfg +from isaaclab.managers import EventTermCfg from isaaclab.sensors.contact_sensor.contact_sensor_cfg import ContactSensorCfg -from isaaclab_tasks.manager_based.manipulation.stack.mdp.franka_stack_events import randomize_object_pose from isaaclab_arena.assets.asset import Asset from isaaclab_arena.relations.relations import AtPosition, Relation, RelationBase -from isaaclab_arena.terms.events import set_object_pose -from isaaclab_arena.utils.pose import Pose, PoseRange +from isaaclab_arena.utils.pose import Pose class ObjectType(Enum): @@ -41,80 +40,10 @@ def __init__( prim_path = "{ENV_REGEX_NS}/" + self.name self.prim_path = prim_path self.object_type = object_type - self.initial_pose: Pose | PoseRange | None = None self.object_cfg = None self.event_cfg = None self.relations: list[RelationBase] = [] - def get_initial_pose(self) -> Pose | PoseRange | None: - """Return the current initial pose of this object. - - Subclasses may override to derive the pose from other sources - (e.g. a parent asset), falling back to ``self.initial_pose``. - """ - return self.initial_pose - - def _get_initial_pose_as_pose(self) -> Pose | None: - """Return a single ``Pose`` suitable for *init_state* and bounding-box calculations. - - If the initial pose is a ``PoseRange``, its midpoint is returned. - If the initial pose is ``None``, ``None`` is returned. - """ - initial_pose = self.get_initial_pose() - if initial_pose is None: - return None - if isinstance(initial_pose, PoseRange): - return initial_pose.get_midpoint() - return initial_pose - - def set_initial_pose(self, pose: Pose | PoseRange) -> None: - """Set / override the initial pose and rebuild derived configs. - - Args: - pose: A fixed ``Pose`` or a ``PoseRange`` (randomised on reset). - """ - self.initial_pose = pose - initial_pose = self._get_initial_pose_as_pose() - if initial_pose is not None and self.object_cfg is not None: - self.object_cfg.init_state.pos = initial_pose.position_xyz - self.object_cfg.init_state.rot = initial_pose.rotation_wxyz - self.event_cfg = self._init_event_cfg() - - def _requires_reset_pose_event(self) -> bool: - """Whether a reset-event for the initial pose should be generated. - - Subclasses may override to add extra conditions (e.g. a ``reset_pose`` flag). - """ - return self.get_initial_pose() is not None and self.object_type in ( - ObjectType.RIGID, - ObjectType.ARTICULATION, - ) - - def _init_event_cfg(self) -> EventTermCfg | None: - """Build the ``EventTermCfg`` for resetting this object's pose.""" - if not self._requires_reset_pose_event(): - return None - - initial_pose = self.get_initial_pose() - if isinstance(initial_pose, PoseRange): - return EventTermCfg( - func=randomize_object_pose, - mode="reset", - params={ - "pose_range": initial_pose.to_dict(), - "asset_cfgs": [SceneEntityCfg(self.name)], - }, - ) - else: - return EventTermCfg( - func=set_object_pose, - mode="reset", - params={ - "pose": initial_pose, - "asset_cfg": SceneEntityCfg(self.name), - }, - ) - def get_relations(self) -> list[RelationBase]: """Get all relations for this object.""" return self.relations @@ -162,7 +91,7 @@ def get_object_pose(self, env: ManagerBasedEnv, is_relative: bool = True) -> tor # We require that the asset has been added to the scene under its name. assert self.name in env.scene.keys(), f"Asset {self.name} not found in scene" if (self.object_type == ObjectType.RIGID) or (self.object_type == ObjectType.ARTICULATION): - object_pose = env.scene[self.name].data.root_pose_w.clone() + object_pose = wp.to_torch(env.scene[self.name].data.root_pose_w).clone() elif self.object_type == ObjectType.BASE: object_pose = torch.cat(env.scene[self.name].get_world_poses(), dim=-1) else: @@ -185,11 +114,11 @@ def set_object_pose(self, env: ManagerBasedEnv, pose: Pose, env_ids: torch.Tenso asset = env.scene[self.name] num_envs = len(env_ids) # Convert the pose to the env frame - pose_t_xyz_q_wxyz = pose.to_tensor(device=env.device).repeat(num_envs, 1) - pose_t_xyz_q_wxyz[:, :3] += env.scene.env_origins[env_ids] + pose_t_xyz_q_xyzw = pose.to_tensor(device=env.device).repeat(num_envs, 1) + pose_t_xyz_q_xyzw[:, :3] += env.scene.env_origins[env_ids] # Set the pose and velocity - asset.write_root_pose_to_sim(pose_t_xyz_q_wxyz, env_ids=env_ids) - asset.write_root_velocity_to_sim(torch.zeros(1, 6, device=env.device), env_ids=env_ids) + asset.write_root_pose_to_sim(pose_t_xyz_q_xyzw, env_ids=env_ids) + asset.write_root_velocity_to_sim(torch.zeros(env.num_envs, 6, device=env.device), env_ids=env_ids) def get_contact_sensor_cfg(self, contact_against_prim_paths: list[str] | None = None) -> ContactSensorCfg: assert self.object_type == ObjectType.RIGID, "Contact sensor is only supported for rigid objects" @@ -217,4 +146,4 @@ def _generate_base_cfg(self) -> AssetBaseCfg: def _generate_spawner_cfg(self) -> AssetBaseCfg: # Object Subclasses must implement this method - pass + pass \ No newline at end of file diff --git a/isaaclab_arena/assets/object_reference.py b/isaaclab_arena/assets/object_reference.py index d39e138e4..f60af6f26 100644 --- a/isaaclab_arena/assets/object_reference.py +++ b/isaaclab_arena/assets/object_reference.py @@ -12,7 +12,7 @@ from isaaclab_arena.assets.object_base import ObjectBase, ObjectType from isaaclab_arena.relations.relations import IsAnchor, RelationBase from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox, quaternion_to_90_deg_z_quarters -from isaaclab_arena.utils.pose import Pose, PoseRange +from isaaclab_arena.utils.pose import Pose from isaaclab_arena.utils.usd_helpers import compute_local_bounding_box_from_prim, open_stage from isaaclab_arena.utils.usd_pose_helpers import get_prim_pose_in_default_prim_frame @@ -29,17 +29,9 @@ def __init__(self, parent_asset: Asset, **kwargs): self.initial_pose_relative_to_parent = self._get_referenced_prim_pose_relative_to_parent(parent_asset) assert self.object_type != ObjectType.SPAWNER, "Object reference cannot be a spawner" self.object_cfg = self._init_object_cfg() - self.event_cfg = self._init_event_cfg() self._bounding_box: AxisAlignedBoundingBox | None = None - def get_initial_pose(self) -> Pose | PoseRange: - """Get the initial pose of this object reference. - - If ``set_initial_pose`` was called, returns that override. - Otherwise derives the world-frame pose from the parent asset. - """ - if self.initial_pose is not None: - return self.initial_pose + def get_initial_pose(self) -> Pose: if self.parent_asset.initial_pose is None: T_W_O = self.initial_pose_relative_to_parent else: @@ -48,11 +40,6 @@ def get_initial_pose(self) -> Pose | PoseRange: T_W_O = T_W_P.multiply(T_P_O) return T_W_O - def _get_initial_pose_as_pose(self) -> Pose: - pose = super()._get_initial_pose_as_pose() - assert pose is not None - return pose - def add_relation(self, relation: RelationBase) -> None: """Add a relation to this object reference. @@ -92,9 +79,8 @@ def get_world_bounding_box(self) -> AxisAlignedBoundingBox: Only 90° rotations around Z axis are supported for AxisAlignedBoundingBox. An assertion error is raised for any other rotation. """ - # The following handles only Pose, not PoseRange. - pose = self._get_initial_pose_as_pose() - quarters = quaternion_to_90_deg_z_quarters(pose.rotation_wxyz) + pose = self.get_initial_pose() + quarters = quaternion_to_90_deg_z_quarters(pose.rotation_xyzw) return self.get_bounding_box().rotated_90_around_z(quarters).translated(pose.position_xyz) def get_contact_sensor_cfg(self, contact_against_prim_paths: list[str] | None = None) -> ContactSensorCfg: @@ -111,37 +97,37 @@ def get_contact_sensor_cfg(self, contact_against_prim_paths: list[str] | None = def _generate_rigid_cfg(self) -> RigidObjectCfg: assert self.object_type == ObjectType.RIGID - initial_pose = self._get_initial_pose_as_pose() + initial_pose = self.get_initial_pose() object_cfg = RigidObjectCfg( prim_path=self.prim_path, init_state=RigidObjectCfg.InitialStateCfg( pos=initial_pose.position_xyz, - rot=initial_pose.rotation_wxyz, + rot=initial_pose.rotation_xyzw, ), ) return object_cfg def _generate_articulation_cfg(self) -> ArticulationCfg: assert self.object_type == ObjectType.ARTICULATION - initial_pose = self._get_initial_pose_as_pose() + initial_pose = self.get_initial_pose() object_cfg = ArticulationCfg( prim_path=self.prim_path, actuators={}, init_state=ArticulationCfg.InitialStateCfg( pos=initial_pose.position_xyz, - rot=initial_pose.rotation_wxyz, + rot=initial_pose.rotation_xyzw, ), ) return object_cfg def _generate_base_cfg(self) -> AssetBaseCfg: assert self.object_type == ObjectType.BASE - initial_pose = self._get_initial_pose_as_pose() + initial_pose = self.get_initial_pose() object_cfg = AssetBaseCfg( prim_path=self.prim_path, init_state=AssetBaseCfg.InitialStateCfg( pos=initial_pose.position_xyz, - rot=initial_pose.rotation_wxyz, + rot=initial_pose.rotation_xyzw, ), ) return object_cfg @@ -163,7 +149,7 @@ def _get_referenced_prim_pose_relative_to_parent(self, parent_asset: Asset) -> P prim_pose.position_xyz[1] * self._parent_scale[1], prim_pose.position_xyz[2] * self._parent_scale[2], ) - return Pose(position_xyz=scaled_pos, rotation_wxyz=prim_pose.rotation_wxyz) + return Pose(position_xyz=scaled_pos, rotation_xyzw=prim_pose.rotation_xyzw) def isaaclab_prim_path_to_original_prim_path( self, isaaclab_prim_path: str, parent_asset: Asset, stage: Usd.Stage @@ -189,8 +175,8 @@ def isaaclab_prim_path_to_original_prim_path( original_prim_path = isaaclab_prim_path.removeprefix("{ENV_REGEX_NS}/") # Check that the path starts with the asset name. assert original_prim_path.startswith(parent_asset.name), ( - f"Expected the prim path to start with the parent asset name {parent_asset.name}. Instead got" - f" {original_prim_path}" + "Expected the prim path to start with the parent asset name {parent_asset.name}. Instead got" + " {original_prim_path}" ) original_prim_path = original_prim_path.removeprefix(parent_asset.name) # Append the default prim path. @@ -207,4 +193,4 @@ def __init__(self, openable_joint_name: str, openable_threshold: float = 0.5, ** openable_threshold=openable_threshold, object_type=ObjectType.ARTICULATION, **kwargs, - ) + ) \ No newline at end of file diff --git a/isaaclab_arena/assets/retargeter_library.py b/isaaclab_arena/assets/retargeter_library.py index e91673f62..cdbd4f0b9 100644 --- a/isaaclab_arena/assets/retargeter_library.py +++ b/isaaclab_arena/assets/retargeter_library.py @@ -18,84 +18,62 @@ import torch from abc import ABC, abstractmethod -from dataclasses import dataclass -from typing import Any - -from isaaclab.devices.openxr.retargeters import ( - G1LowerBodyStandingMotionControllerRetargeterCfg, - G1TriHandUpperBodyMotionControllerGripperRetargeterCfg, - GR1T2RetargeterCfg, -) -from isaaclab.devices.retargeter_base import RetargeterBase, RetargeterCfg +from collections.abc import Callable from isaaclab_arena.assets.register import register_retargeter -class DummyTorsoRetargeter(RetargeterBase): - """Dummy retargeter that returns zero torso orientation commands. +class RetargetterBase(ABC): + """Base class for teleop retargeter entries in the Arena registry. - This is used to pad the action space for G1 WBC Pink with motion controllers, - which don't provide torso orientation commands. + Subclasses associate a (device, embodiment) pair with a pipeline builder + function compatible with ``IsaacTeleopCfg.pipeline_builder``. """ - def __init__(self, cfg: "DummyTorsoRetargeterCfg"): - super().__init__(cfg) - - def retarget(self, data: Any) -> torch.Tensor: - """Return zeros for torso orientation (roll, pitch, yaw).""" - return torch.zeros(3, device=self._sim_device) + device: str + embodiment: str - def get_requirements(self) -> list[RetargeterBase.Requirement]: - """This retargeter doesn't require any device data.""" - return [] + @abstractmethod + def get_pipeline_builder(self, embodiment: object) -> Callable | None: + """Return an isaacteleop pipeline builder callable, or None if not applicable.""" + raise NotImplementedError -@dataclass -class DummyTorsoRetargeterCfg(RetargeterCfg): - """Configuration for dummy torso retargeter.""" +@register_retargeter +class GR1T2PinkIsaacTeleopRetargeter(RetargetterBase): + """Isaac Teleop pipeline builder for GR1T2 with Pink IK and dex hand retargeting.""" - retargeter_type: type[RetargeterBase] = DummyTorsoRetargeter + device = "openxr" + embodiment = "gr1_pink" + def __init__(self): + pass -class RetargetterBase(ABC): - device: str - embodiment: str + def get_pipeline_builder(self, embodiment: object) -> Callable: + from isaaclab_tasks.manager_based.manipulation.pick_place.pickplace_gr1t2_env_cfg import ( + _build_gr1t2_pickplace_pipeline, + ) - @abstractmethod - def get_retargeter_cfg( - self, embodiment: object, sim_device: str, enable_visualization: bool = False - ) -> RetargeterCfg | list[RetargeterCfg] | None: - """Get retargeter configuration. - - Can return: - - A single RetargeterCfg - - A list of RetargeterCfg (for devices needing multiple retargeters) - - None (for devices that don't need retargeters) - """ - raise NotImplementedError + return lambda: _build_gr1t2_pickplace_pipeline()[0] @register_retargeter -class GR1T2PinkOpenXRRetargeter(RetargetterBase): +class G1WbcPinkIsaacTeleopRetargeter(RetargetterBase): + """Isaac Teleop pipeline builder for G1 WBC Pink (locomanipulation: wrist + TriHand + locomotion).""" device = "openxr" - embodiment = "gr1_pink" - num_open_xr_hand_joints = 52 + embodiment = "g1_wbc_pink" def __init__(self): pass - def get_retargeter_cfg( - self, gr1t2_embodiment, sim_device: str, enable_visualization: bool = False - ) -> RetargeterCfg: - return GR1T2RetargeterCfg( - enable_visualization=enable_visualization, - # number of joints in both hands - num_open_xr_hand_joints=self.num_open_xr_hand_joints, - sim_device=sim_device, - hand_joint_names=gr1t2_embodiment.get_action_cfg().upper_body_ik.hand_joint_names, + def get_pipeline_builder(self, embodiment: object) -> Callable: + from isaaclab_arena.assets.g1_pink_locomanipulation_pipeline import ( + _build_g1_pink_locomanipulation_pipeline, ) + return _build_g1_pink_locomanipulation_pipeline + @register_retargeter class FrankaKeyboardRetargeter(RetargetterBase): @@ -105,9 +83,7 @@ class FrankaKeyboardRetargeter(RetargetterBase): def __init__(self): pass - def get_retargeter_cfg( - self, franka_embodiment, sim_device: str, enable_visualization: bool = False - ) -> RetargeterCfg | None: + def get_pipeline_builder(self, embodiment: object) -> Callable | None: return None @@ -119,9 +95,7 @@ class FrankaSpaceMouseRetargeter(RetargetterBase): def __init__(self): pass - def get_retargeter_cfg( - self, franka_embodiment, sim_device: str, enable_visualization: bool = False - ) -> RetargeterCfg | None: + def get_pipeline_builder(self, embodiment: object) -> Callable | None: return None @@ -133,24 +107,20 @@ class DroidDifferentialIKKeyboardRetargeter(RetargetterBase): def __init__(self): pass - def get_retargeter_cfg( - self, droid_embodiment, sim_device: str, enable_visualization: bool = False - ) -> RetargeterCfg | None: + def get_pipeline_builder(self, embodiment: object) -> Callable | None: return None -@register_retargeter -class AgibotKeyboardRetargeter(RetargetterBase): - device = "keyboard" - embodiment = "agibot" +# @register_retargeter +# class AgibotKeyboardRetargeter(RetargetterBase): +# device = "keyboard" +# embodiment = "agibot" - def __init__(self): - pass +# def __init__(self): +# pass - def get_retargeter_cfg( - self, agibot_embodiment, sim_device: str, enable_visualization: bool = False - ) -> RetargeterCfg | None: - return None +# def get_pipeline_builder(self, embodiment: object) -> Callable | None: +# return None @register_retargeter @@ -161,35 +131,5 @@ class GalbotKeyboardRetargeter(RetargetterBase): def __init__(self): pass - def get_retargeter_cfg( - self, galbot_embodiment, sim_device: str, enable_visualization: bool = False - ) -> RetargeterCfg | None: + def get_pipeline_builder(self, embodiment: object) -> Callable | None: return None - - -@register_retargeter -class G1WbcPinkMotionControllersRetargeter(RetargetterBase): - """Retargeter for G1 WBC Pink embodiment with motion controllers (Quest controllers).""" - - device = "openxr" - embodiment = "g1_wbc_pink" - - def __init__(self): - pass - - def get_retargeter_cfg( - self, g1_embodiment, sim_device: str, enable_visualization: bool = False - ) -> list[RetargeterCfg]: - """Get motion controller retargeter configuration for G1. - - Returns a list of retargeters: - - Upper body (with gripper): outputs 16 dims [gripper(2), left_wrist(7), right_wrist(7)] - - Lower body: outputs 4 dims [nav_cmd(3), hip_height(1)] - - Dummy torso: outputs 3 dims [torso_orientation_rpy(3)] all zeros - Total: 23 dims to match g1_wbc_pink action space - """ - return [ - G1TriHandUpperBodyMotionControllerGripperRetargeterCfg(sim_device=sim_device), - G1LowerBodyStandingMotionControllerRetargeterCfg(sim_device=sim_device), - DummyTorsoRetargeterCfg(sim_device=sim_device), - ] diff --git a/isaaclab_arena/cli/isaaclab_arena_cli.py b/isaaclab_arena/cli/isaaclab_arena_cli.py index 3121b7b8b..acf27c453 100644 --- a/isaaclab_arena/cli/isaaclab_arena_cli.py +++ b/isaaclab_arena/cli/isaaclab_arena_cli.py @@ -29,14 +29,6 @@ def add_isaac_lab_cli_args(parser: argparse.ArgumentParser) -> None: isaac_lab_group.add_argument("--seed", type=int, default=42, help="Optional seed for the random number generator.") isaac_lab_group.add_argument("--num_envs", type=int, default=1, help="Number of environments to simulate.") isaac_lab_group.add_argument("--env_spacing", type=float, default=30.0, help="Spacing between environments.") - # NOTE(alexmillane, 2025.07.25): Unlike base isaaclab, we enable pinocchio by default. - isaac_lab_group.add_argument( - "--disable_pinocchio", - dest="enable_pinocchio", - default=True, - action="store_false", - help="Disable Pinocchio.", - ) isaac_lab_group.add_argument("--mimic", action="store_true", default=False, help="Enable mimic environment.") isaac_lab_group.add_argument( "--distributed", diff --git a/isaaclab_arena/embodiments/agibot/agibot.py b/isaaclab_arena/embodiments/agibot/agibot.py index cce6965ac..5620c0e22 100644 --- a/isaaclab_arena/embodiments/agibot/agibot.py +++ b/isaaclab_arena/embodiments/agibot/agibot.py @@ -26,7 +26,7 @@ from isaaclab_arena.utils.pose import Pose -@register_asset +# @register_asset class AgibotEmbodiment(EmbodimentBase): """Embodiment for the Agibot robot.""" @@ -113,7 +113,6 @@ class AgibotLeftArmActionsCfg: controller=AGIBOT_LEFT_ARM_RMPFLOW_CFG, scale=1.0, body_offset=RMPFlowActionCfg.OffsetCfg(rot=[0.7071, 0.0, -0.7071, 0.0]), - articulation_prim_expr="/World/envs/env_.*/Robot", use_relative_mode=True, ) @@ -135,7 +134,6 @@ class AgibotRightArmActionsCfg: body_name="right_gripper_center", controller=AGIBOT_RIGHT_ARM_RMPFLOW_CFG, scale=1.0, - articulation_prim_expr="/World/envs/env_.*/Robot", use_relative_mode=True, ) diff --git a/isaaclab_arena/embodiments/common/common.py b/isaaclab_arena/embodiments/common/common.py index f24f610d6..20c325fd3 100644 --- a/isaaclab_arena/embodiments/common/common.py +++ b/isaaclab_arena/embodiments/common/common.py @@ -3,7 +3,7 @@ # # SPDX-License-Identifier: Apache-2.0 -from isaaclab.devices.openxr import XrCfg +from isaaclab_teleop import XrCfg from isaaclab_arena.utils.pose import Pose @@ -29,11 +29,11 @@ def get_default_xr_cfg(initial_pose: Pose | None = None, xr_offset: Pose | None return XrCfg( anchor_pos=xr_pose_global.position_xyz, - anchor_rot=xr_pose_global.rotation_wxyz, + anchor_rot=xr_pose_global.rotation_xyzw, ) else: # If no initial pose set, use the offset as global coordinates (robot at origin) return XrCfg( anchor_pos=xr_offset.position_xyz, - anchor_rot=xr_offset.rotation_wxyz, + anchor_rot=xr_offset.rotation_xyzw, ) diff --git a/isaaclab_arena/embodiments/droid/observations.py b/isaaclab_arena/embodiments/droid/observations.py index f1f65b5c0..8a13cb695 100644 --- a/isaaclab_arena/embodiments/droid/observations.py +++ b/isaaclab_arena/embodiments/droid/observations.py @@ -4,6 +4,7 @@ # SPDX-License-Identifier: Apache-2.0 import torch +import warp as wp from isaaclab.envs import ManagerBasedRLEnv from isaaclab.managers import SceneEntityCfg @@ -21,7 +22,7 @@ def arm_joint_pos(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntit "panda_joint7", ] joint_indices = [i for i, name in enumerate(robot.data.joint_names) if name in joint_names] - return robot.data.joint_pos[:, joint_indices] + return wp.to_torch(robot.data.joint_pos)[:, joint_indices] def gripper_pos(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: @@ -29,7 +30,7 @@ def gripper_pos(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityC robot = env.scene[asset_cfg.name] joint_names = ["finger_joint"] joint_indices = [i for i, name in enumerate(robot.data.joint_names) if name in joint_names] - joint_pos = robot.data.joint_pos[:, joint_indices] + joint_pos = wp.to_torch(robot.data.joint_pos)[:, joint_indices] # rescale to 0–1 return joint_pos / (torch.pi / 4) @@ -38,11 +39,11 @@ def ee_pos(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("r """Returns the end effector position (x, y, z) in the world frame.""" robot = env.scene[asset_cfg.name] body_idx = robot.data.body_names.index("base_link") # Robotiq gripper base link - return robot.data.body_pos_w[:, body_idx, :] + return wp.to_torch(robot.data.body_pos_w)[:, body_idx, :] def ee_quat(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Returns the end effector orientation as quaternion (w, x, y, z) in the world frame.""" robot = env.scene[asset_cfg.name] body_idx = robot.data.body_names.index("base_link") # Robotiq gripper base link - return robot.data.body_quat_w[:, body_idx, :] + return wp.to_torch(robot.data.body_quat_w)[:, body_idx, :] diff --git a/isaaclab_arena/embodiments/embodiment_base.py b/isaaclab_arena/embodiments/embodiment_base.py index 27677aa8a..3927bbdec 100644 --- a/isaaclab_arena/embodiments/embodiment_base.py +++ b/isaaclab_arena/embodiments/embodiment_base.py @@ -93,6 +93,10 @@ def get_mimic_env(self) -> ManagerBasedRLMimicEnv: def get_xr_cfg(self) -> Any: return self.xr + def get_teleop_target_frame_prim_path(self) -> str | None: + """Optional USD prim path for rebasing teleop poses (e.g. robot base link). Returns None if not set.""" + return None + def get_camera_cfg(self) -> Any: return self.camera_config @@ -100,7 +104,7 @@ def _update_scene_cfg_with_robot_initial_pose(self, scene_config: Any, pose: Pos if scene_config is None or not hasattr(scene_config, "robot"): raise RuntimeError("scene_config must be populated with a `robot` before calling `set_robot_initial_pose`.") scene_config.robot.init_state.pos = pose.position_xyz - scene_config.robot.init_state.rot = pose.rotation_wxyz + scene_config.robot.init_state.rot = pose.rotation_xyzw return scene_config def get_recorder_term_cfg(self) -> RecorderManagerBaseCfg: diff --git a/isaaclab_arena/embodiments/franka/franka.py b/isaaclab_arena/embodiments/franka/franka.py index adceb3bee..6c3349ff9 100644 --- a/isaaclab_arena/embodiments/franka/franka.py +++ b/isaaclab_arena/embodiments/franka/franka.py @@ -7,6 +7,7 @@ import torch from collections.abc import Sequence from dataclasses import MISSING +from typing import Any import isaaclab.envs.mdp as mdp_isaac_lab import isaaclab.sim as sim_utils @@ -36,7 +37,7 @@ from isaaclab_arena.embodiments.franka.observations import gripper_pos from isaaclab_arena.utils.pose import Pose -_DEFAULT_CAMERA_OFFSET = Pose(position_xyz=(0.11, -0.031, -0.074), rotation_wxyz=(-0.74896, 0.0, 0.0, -0.66262)) +_DEFAULT_CAMERA_OFFSET = Pose(position_xyz=(0.11, -0.031, -0.074), rotation_xyzw=(0.0, 0.0, -0.66262, -0.74896)) # The reason to use our internal panda USD is to combine the panda and the stand within one USD. @@ -77,6 +78,7 @@ def __init__( self.camera_config._is_tiled_camera = is_tiled_camera self.camera_config._camera_offset = camera_offset + def set_initial_joint_pose(self, initial_joint_pose: list[float]) -> None: self.event_config.init_franka_arm_pose.params["default_pose"] = initial_joint_pose @@ -177,7 +179,7 @@ def __post_init__(self): ) offset = OffsetClass( pos=camera_offset.position_xyz, - rot=camera_offset.rotation_wxyz, + rot=camera_offset.rotation_xyzw, convention="ros", ) diff --git a/isaaclab_arena/embodiments/franka/observations.py b/isaaclab_arena/embodiments/franka/observations.py index 338436ccf..1a1893e5b 100644 --- a/isaaclab_arena/embodiments/franka/observations.py +++ b/isaaclab_arena/embodiments/franka/observations.py @@ -4,6 +4,7 @@ # SPDX-License-Identifier: Apache-2.0 import torch +import warp as wp from isaaclab.assets import Articulation from isaaclab.envs import ManagerBasedRLEnv @@ -12,7 +13,6 @@ def gripper_pos(env: ManagerBasedRLEnv, robot_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: robot: Articulation = env.scene[robot_cfg.name] - finger_joint_1 = robot.data.joint_pos[:, -1].clone().unsqueeze(1) - finger_joint_2 = -1 * robot.data.joint_pos[:, -2].clone().unsqueeze(1) - + finger_joint_1 = wp.to_torch(robot.data.joint_pos)[:, -1].clone().unsqueeze(1) + finger_joint_2 = -1 * wp.to_torch(robot.data.joint_pos)[:, -2].clone().unsqueeze(1) return torch.cat((finger_joint_1, finger_joint_2), dim=1) diff --git a/isaaclab_arena/embodiments/g1/g1.py b/isaaclab_arena/embodiments/g1/g1.py index 1f2df4d3c..c920770c9 100644 --- a/isaaclab_arena/embodiments/g1/g1.py +++ b/isaaclab_arena/embodiments/g1/g1.py @@ -4,6 +4,7 @@ # SPDX-License-Identifier: Apache-2.0 import torch +import warp as wp from collections.abc import Sequence from dataclasses import MISSING @@ -13,8 +14,8 @@ import isaaclab_tasks.manager_based.manipulation.pick_place.mdp as mdp from isaaclab.actuators import IdealPDActuatorCfg from isaaclab.assets.articulation.articulation_cfg import ArticulationCfg -from isaaclab.devices.openxr import XrCfg -from isaaclab.devices.openxr.xr_cfg import XrAnchorRotationMode +from isaaclab_teleop import XrCfg +from isaaclab_teleop.xr_cfg import XrAnchorRotationMode from isaaclab.envs import ManagerBasedRLMimicEnv # noqa: F401 from isaaclab.managers import EventTermCfg as EventTerm from isaaclab.managers import ObservationGroupCfg as ObsGroup @@ -29,7 +30,7 @@ from isaaclab_arena.assets.register import register_asset from isaaclab_arena.embodiments.common.arm_mode import ArmMode from isaaclab_arena.embodiments.embodiment_base import EmbodimentBase -from isaaclab_arena.utils.isaaclab_utils.resets import reset_all_articulation_joints +from isaaclab_arena.terms.events import reset_all_articulation_joints from isaaclab_arena.utils.pose import Pose from isaaclab_arena_g1.g1_env.mdp import g1_events as g1_events_mdp from isaaclab_arena_g1.g1_env.mdp import g1_observations as g1_observations_mdp @@ -63,16 +64,20 @@ def __init__( # Anchor to the robot's pelvis for first-person view that follows the robot self.xr: XrCfg = XrCfg( anchor_pos=(0.0, 0.0, -1.0), - anchor_rot=(0.70711, 0.0, 0.0, -0.70711), + anchor_rot=(0.0, 0.0, -0.70711, 0.70711), anchor_prim_path="/World/envs/env_0/Robot/pelvis", anchor_rotation_mode=XrAnchorRotationMode.FOLLOW_PRIM_SMOOTHED, fixed_anchor_height=True, ) + def get_teleop_target_frame_prim_path(self) -> str | None: + """Pelvis prim path so OpenXR teleop poses are rebased into robot base frame for IK.""" + return "/World/envs/env_0/Robot/pelvis" + # Default camera offset pose _DEFAULT_G1_CAMERA_OFFSET = Pose( - position_xyz=(0.04485, 0.0, 0.35325), rotation_wxyz=(0.32651, -0.62721, 0.62721, -0.32651) + position_xyz=(0.04485, 0.0, 0.35325), rotation_xyzw=(-0.62721, 0.62721, -0.32651, 0.32651) ) @@ -157,7 +162,7 @@ class G1SceneCfg: prim_path="/World/envs/env_.*/Robot", init_state=ArticulationCfg.InitialStateCfg( pos=(0.8, -1.38, 0.78), - rot=(0.0, 0.0, 0.0, 1.0), + rot=(0.0, 0.0, 1.0, 0.0), joint_pos={ # target angles [rad] "left_hip_pitch_joint": -0.1, @@ -352,15 +357,13 @@ def __post_init__(self): width=640, data_types=["rgb"], spawn=sim_utils.PinholeCameraCfg( - focal_length=0.169, # 1.69 mm; FOV preserved via apertures - horizontal_aperture=0.693, # preserves ~128° horiz FOV - vertical_aperture=0.284, # preserves ~80° vert FOV + focal_length=15, clipping_range=(0.1, 5), ), ) offset = OffsetClass( pos=camera_offset.position_xyz, - rot=camera_offset.rotation_wxyz, + rot=camera_offset.rotation_xyzw, convention="ros", ) @@ -770,7 +773,7 @@ def get_object_poses(self, env_ids: Sequence[int] | None = None): env_ids = slice(None) # Get pelvis inverse transform to convert from world to pelvis frame - pelvis_pose_w = self.scene["robot"].data.body_link_state_w[ + pelvis_pose_w = wp.to_torch(self.scene["robot"].data.body_link_state_w)[ :, self.scene["robot"].data.body_names.index("pelvis"), : ] pelvis_position_w = pelvis_pose_w[:, :3] - self.scene.env_origins diff --git a/isaaclab_arena/embodiments/galbot/galbot.py b/isaaclab_arena/embodiments/galbot/galbot.py index e52b7f2de..f3769daa0 100644 --- a/isaaclab_arena/embodiments/galbot/galbot.py +++ b/isaaclab_arena/embodiments/galbot/galbot.py @@ -97,7 +97,6 @@ class GalbotLeftArmActionsCfg: controller=GALBOT_LEFT_ARM_RMPFLOW_CFG, scale=1.0, body_offset=RMPFlowActionCfg.OffsetCfg(pos=(0.0, 0.0, 0.0)), - articulation_prim_expr="/World/envs/env_.*/Robot", use_relative_mode=True, ) diff --git a/isaaclab_arena/embodiments/gr1t2/gr1t2.py b/isaaclab_arena/embodiments/gr1t2/gr1t2.py index e8cecee5a..9240fd144 100644 --- a/isaaclab_arena/embodiments/gr1t2/gr1t2.py +++ b/isaaclab_arena/embodiments/gr1t2/gr1t2.py @@ -15,7 +15,7 @@ import isaaclab_tasks.manager_based.manipulation.pick_place.mdp as mdp from isaaclab.actuators import ImplicitActuatorCfg from isaaclab.assets.articulation.articulation_cfg import ArticulationCfg -from isaaclab.devices.openxr import XrCfg +from isaaclab_teleop import XrCfg from isaaclab.envs import ManagerBasedRLMimicEnv from isaaclab.envs.mdp.actions import JointPositionActionCfg from isaaclab.managers import EventTermCfg as EventTerm @@ -32,7 +32,7 @@ from isaaclab_arena.embodiments.common.common import get_default_xr_cfg from isaaclab_arena.embodiments.common.mimic_utils import get_rigid_and_articulated_object_poses from isaaclab_arena.embodiments.embodiment_base import EmbodimentBase -from isaaclab_arena.utils.isaaclab_utils.resets import reset_all_articulation_joints +from isaaclab_arena.terms.events import reset_all_articulation_joints from isaaclab_arena.utils.pose import Pose ARM_JOINT_NAMES_LIST = [ @@ -77,7 +77,7 @@ ] # Default camera offset pose -_DEFAULT_CAMERA_OFFSET = Pose(position_xyz=(0.12515, 0.0, 0.06776), rotation_wxyz=(0.62, 0.32, -0.32, -0.63)) +_DEFAULT_CAMERA_OFFSET = Pose(position_xyz=(0.12515, 0.0, 0.06776), rotation_xyzw=(0.32, -0.32, -0.63, 0.62)) @register_asset @@ -108,7 +108,7 @@ def __init__( # These offsets are defined relative to the robot's base frame self._xr_offset = Pose( position_xyz=(-0.5, 0.0, -1.0), - rotation_wxyz=(0.70711, 0.0, 0.0, -0.70711), + rotation_xyzw=(0.0, 0.0, -0.70711, 0.70711), ) self.xr: XrCfg | None = None @@ -388,7 +388,7 @@ def __post_init__(self): ) offset = OffsetClass( pos=camera_offset.position_xyz, - rot=camera_offset.rotation_wxyz, + rot=camera_offset.rotation_xyzw, convention="opengl", ) diff --git a/isaaclab_arena/environments/arena_env_builder.py b/isaaclab_arena/environments/arena_env_builder.py index 03d0749b1..1e4a66e3b 100644 --- a/isaaclab_arena/environments/arena_env_builder.py +++ b/isaaclab_arena/environments/arena_env_builder.py @@ -164,13 +164,19 @@ def compose_manager_cfg(self) -> IsaacLabArenaManagerBasedRLEnvCfg: ) actions_cfg = embodiment.get_action_cfg() xr_cfg = embodiment.get_xr_cfg() + isaac_teleop_cfg = None + teleop_device_cfg = None if self.arena_env.teleop_device is not None: device_registry = DeviceRegistry() - teleop_device_cfg = device_registry.get_teleop_device_cfg( + device_cfg = device_registry.get_teleop_device_cfg( self.arena_env.teleop_device, self.arena_env.embodiment ) - else: - teleop_device_cfg = None + from isaaclab_teleop import IsaacTeleopCfg + + if isinstance(device_cfg, IsaacTeleopCfg): + isaac_teleop_cfg = device_cfg + else: + teleop_device_cfg = device_cfg metrics = task.get_metrics() metrics_recorder_manager_cfg = metrics_to_recorder_manager_cfg(metrics) @@ -223,7 +229,7 @@ def compose_manager_cfg(self) -> IsaacLabArenaManagerBasedRLEnvCfg: curriculum=curriculum_cfg, commands=commands_cfg, xr=xr_cfg, - teleop_devices=teleop_device_cfg, + isaac_teleop=isaac_teleop_cfg, recorders=recorder_manager_cfg, metrics=metrics, isaaclab_arena_env=isaaclab_arena_env, @@ -245,7 +251,7 @@ def compose_manager_cfg(self) -> IsaacLabArenaManagerBasedRLEnvCfg: curriculum=curriculum_cfg, commands=commands_cfg, xr=xr_cfg, - teleop_devices=teleop_device_cfg, + isaac_teleop=isaac_teleop_cfg, # Mimic stuff datagen_config=task_mimic_env_cfg.datagen_config, subtask_configs=task_mimic_env_cfg.subtask_configs, diff --git a/isaaclab_arena/evaluation/policy_runner.py b/isaaclab_arena/evaluation/policy_runner.py index 5c8441a22..36176feb2 100644 --- a/isaaclab_arena/evaluation/policy_runner.py +++ b/isaaclab_arena/evaluation/policy_runner.py @@ -3,6 +3,8 @@ # # SPDX-License-Identifier: Apache-2.0 +import os +import sys import argparse import gymnasium as gym import torch @@ -20,6 +22,14 @@ if TYPE_CHECKING: from isaaclab_arena.policy.policy_base import PolicyBase +def _ensure_groot_deps_in_path() -> None: + """Re-exec once so PYTHONPATH has GROOT_DEPS_DIR first (GR00T deps before Isaac Sim).""" + deps_dir = os.environ.get("GROOT_DEPS_DIR") + if not deps_dir or os.environ.get("_GROOT_PYTHONPATH_APPLIED") == "1": + return + os.environ["PYTHONPATH"] = deps_dir + os.pathsep + os.environ.get("PYTHONPATH", "") + os.environ["_GROOT_PYTHONPATH_APPLIED"] = "1" + os.execv(sys.executable, [sys.executable] + sys.argv) def get_policy_cls(policy_type: str) -> type["PolicyBase"]: """Get the policy class for the given policy type name. @@ -82,6 +92,7 @@ def rollout_policy( with torch.inference_mode(): actions = policy.get_action(env, obs) obs, _, terminated, truncated, _ = env.step(actions) + if terminated.any() or truncated.any(): # Only reset policy for those envs that are terminated or truncated print( @@ -111,11 +122,11 @@ def rollout_policy( raise RuntimeError(f"Error rolling out policy: {e}") else: + # Only compute metrics if env has a non-None metrics list (e.g. NoTask leaves metrics as None). if hasattr(env.cfg, "metrics") and env.cfg.metrics is not None: # NOTE(xinjieyao, 2025-10-07): lazy import to prevent app stalling caused by omni.kit from isaaclab_arena.metrics.metrics import compute_metrics - metrics = compute_metrics(env) return metrics return None @@ -210,4 +221,5 @@ def main(): if __name__ == "__main__": + _ensure_groot_deps_in_path() main() diff --git a/isaaclab_arena/examples/compile_env_notebook.py b/isaaclab_arena/examples/compile_env_notebook.py index 409cdc901..921ce5b0a 100644 --- a/isaaclab_arena/examples/compile_env_notebook.py +++ b/isaaclab_arena/examples/compile_env_notebook.py @@ -5,43 +5,75 @@ # %% +import argparse import torch import tqdm -import pinocchio # noqa: F401 from isaaclab.app import AppLauncher print("Launching simulation app once in notebook") -simulation_app = AppLauncher() +parser = argparse.ArgumentParser() +AppLauncher.add_app_launcher_args(parser) +args = parser.parse_args(["--visualizer", "kit"]) +# args = parser.parse_args([]) +app_launcher = AppLauncher(args) + +# %% from isaaclab_arena.assets.asset_registry import AssetRegistry +from isaaclab_arena.assets.object_reference import ObjectReference from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment -from isaaclab_arena.relations.relations import IsAnchor, On +from isaaclab_arena.relations.relations import AtPosition, IsAnchor, On from isaaclab_arena.scene.scene import Scene -from isaaclab_arena.utils.pose import Pose +from isaaclab_arena.tasks.pick_and_place_task import PickAndPlaceTask asset_registry = AssetRegistry() background = asset_registry.get_asset_by_name("kitchen")() +# embodiment = asset_registry.get_asset_by_name("franka")() embodiment = asset_registry.get_asset_by_name("franka")() +# embodiment = asset_registry.get_asset_by_name("gr1_pink")(enable_cameras=True) cracker_box = asset_registry.get_asset_by_name("cracker_box")() -tomato_soup_can = asset_registry.get_asset_by_name("tomato_soup_can")() +# tomato_soup_can = asset_registry.get_asset_by_name("tomato_soup_can")() +microwave = asset_registry.get_asset_by_name("microwave")() + +# cracker_box.set_initial_pose(Pose(position_xyz=(0.4, 0.0, 0.1), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) +# cracker_box.add_relation(IsAnchor()) +# tomato_soup_can.add_relation(On(cracker_box)) + +table_top_reference = ObjectReference( + name="table_top_reference", + prim_path="{ENV_REGEX_NS}/kitchen/Kitchen_Counter/TRS_Base/TRS_Static/Counter_Top_A", + parent_asset=background, +) +table_top_reference.add_relation(IsAnchor()) + +microwave.add_relation(AtPosition(x=0.4, y=0.0)) +microwave.add_relation(On(table_top_reference)) + +cracker_box.add_relation(On(microwave)) + -cracker_box.set_initial_pose(Pose(position_xyz=(0.4, 0.0, 0.1), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) -cracker_box.add_relation(IsAnchor()) -tomato_soup_can.add_relation(On(cracker_box)) +destination_location = ObjectReference( + name="destination_location", + prim_path="{ENV_REGEX_NS}/kitchen/Cabinet_B_02", + parent_asset=background, +) -scene = Scene(assets=[background, cracker_box, tomato_soup_can]) +scene = Scene(assets=[background, table_top_reference, microwave, destination_location, cracker_box]) +# scene = Scene(assets=[background, cracker_box, tomato_soup_can]) isaaclab_arena_environment = IsaacLabArenaEnvironment( name="reference_object_test", embodiment=embodiment, scene=scene, + task=PickAndPlaceTask(cracker_box, destination_location, background), ) args_cli = get_isaaclab_arena_cli_parser().parse_args([]) args_cli.solve_relations = True +args_cli.num_envs = 2 env_builder = ArenaEnvBuilder(isaaclab_arena_environment, args_cli) env = env_builder.make_registered() env.reset() @@ -49,10 +81,21 @@ # %% # Run some zero actions. -NUM_STEPS = 1000 +NUM_STEPS = 300 for _ in tqdm.tqdm(range(NUM_STEPS)): with torch.inference_mode(): actions = torch.zeros(env.action_space.shape, device=env.unwrapped.device) env.step(actions) # %% + +from isaaclab_arena.utils.reload_modules import reload_arena_modules + +reload_arena_modules() + + +# %% + +from isaaclab_arena.utils.isaaclab_utils.simulation_app import teardown_simulation_app + +teardown_simulation_app(suppress_exceptions=False, make_new_stage=True) diff --git a/isaaclab_arena/examples/example_env_notebook.py b/isaaclab_arena/examples/example_env_notebook.py index 168b71845..e3d72550a 100644 --- a/isaaclab_arena/examples/example_env_notebook.py +++ b/isaaclab_arena/examples/example_env_notebook.py @@ -8,7 +8,6 @@ import torch import tqdm -import pinocchio # noqa: F401 from isaaclab.app import AppLauncher print("Launching simulation app once in notebook") diff --git a/isaaclab_arena/metrics/object_moved.py b/isaaclab_arena/metrics/object_moved.py index cad0b0460..2f4480b76 100644 --- a/isaaclab_arena/metrics/object_moved.py +++ b/isaaclab_arena/metrics/object_moved.py @@ -5,6 +5,7 @@ import numpy as np from dataclasses import MISSING +import warp as wp from isaaclab.envs.manager_based_rl_env import ManagerBasedEnv from isaaclab.managers.recorder_manager import RecorderTerm, RecorderTermCfg @@ -25,7 +26,7 @@ def __init__(self, cfg: RecorderTermCfg, env: ManagerBasedEnv): def record_post_step(self): # NOTE(alexmillane, 2025-09-30): This assumes the the object is a rigid object. - object_linear_velocity = self._env.scene[self.object_name].data.root_link_vel_w[:, :3] + object_linear_velocity = wp.to_torch(self._env.scene[self.object_name].data.root_link_vel_w)[:, :3] assert object_linear_velocity.shape == (self._env.num_envs, 3) return self.name, object_linear_velocity diff --git a/isaaclab_arena/relations/object_placer.py b/isaaclab_arena/relations/object_placer.py index 27d802237..436e1bbc5 100644 --- a/isaaclab_arena/relations/object_placer.py +++ b/isaaclab_arena/relations/object_placer.py @@ -39,6 +39,7 @@ def __init__(self, params: ObjectPlacerParams | None = None): """ self.params = params or ObjectPlacerParams() self._solver = RelationSolver(params=self.params.solver_params) + print("HELLO ZIHAO") def place( self, @@ -274,15 +275,15 @@ def _apply_positions( random_marker = self._get_random_around_solution(obj) rotate_marker = self._get_rotate_around_solution(obj) - rotation_wxyz = rotate_marker.get_rotation_wxyz() if rotate_marker else (1.0, 0.0, 0.0, 0.0) + rotation_xyzw = rotate_marker.get_rotation_xyzw() if rotate_marker else (0.0, 0.0, 0.0, 1.0) if random_marker is not None: # We need to set a PoseRange for the randomization to be picked up on reset. # Set a PoseRange with the explicit rotation from RotateAroundSolution if present - obj.set_initial_pose(random_marker.to_pose_range_centered_at(pos, rotation_wxyz=rotation_wxyz)) + obj.set_initial_pose(random_marker.to_pose_range_centered_at(pos, rotation_xyzw=rotation_xyzw)) else: # Without randomization, we can set a fixed Pose. - obj.set_initial_pose(Pose(position_xyz=pos, rotation_wxyz=rotation_wxyz)) + obj.set_initial_pose(Pose(position_xyz=pos, rotation_xyzw=rotation_xyzw)) def _get_random_around_solution(self, obj: Object | ObjectReference) -> RandomAroundSolution | None: """Get RandomAroundSolution marker from object if present. diff --git a/isaaclab_arena/relations/relations.py b/isaaclab_arena/relations/relations.py index 12bdac00a..76b0e502d 100644 --- a/isaaclab_arena/relations/relations.py +++ b/isaaclab_arena/relations/relations.py @@ -211,20 +211,20 @@ def __init__( def to_pose_range_centered_at( self, position: tuple[float, float, float], - rotation_wxyz: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0), + rotation_xyzw: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0), ) -> PoseRange: """Create a PoseRange centered on the given position and rotation. Args: position: Center position (x, y, z) for the range. - rotation_wxyz: Center rotation as quaternion (w, x, y, z) for the range. + rotation_xyzw: Center rotation as quaternion (x, y, z, w) for the range. Defaults to identity quaternion. Returns: PoseRange spanning ± half-extents around the position and rotation. """ # Convert quaternion to euler angles (roll, pitch, yaw) - quat_tensor = torch.tensor([rotation_wxyz]) + quat_tensor = torch.tensor([rotation_xyzw]) roll, pitch, yaw = euler_xyz_from_quat(quat_tensor) center_roll = float(roll[0]) center_pitch = float(pitch[0]) @@ -286,8 +286,8 @@ def __init__( self.pitch_rad = pitch_rad self.yaw_rad = yaw_rad - def get_rotation_wxyz(self) -> tuple[float, float, float, float]: - """Get the rotation as a quaternion (w, x, y, z). + def get_rotation_xyzw(self) -> tuple[float, float, float, float]: + """Get the rotation as a quaternion (x, y, z, w). Returns: Quaternion rotation converted from roll/pitch/yaw. diff --git a/isaaclab_arena/scene/scene.py b/isaaclab_arena/scene/scene.py index 2fc099403..0eb97326d 100644 --- a/isaaclab_arena/scene/scene.py +++ b/isaaclab_arena/scene/scene.py @@ -174,10 +174,11 @@ def _create_prim_from_asset(stage: Usd.Stage, asset: Asset) -> None: prim_xform.ClearXformOpOrder() if asset.initial_pose is not None: t = Gf.Vec3d(asset.initial_pose.position_xyz) if trans_double else Gf.Vec3f(asset.initial_pose.position_xyz) + rot = asset.initial_pose.rotation_xyzw r = ( - Gf.Quatd(*asset.initial_pose.rotation_wxyz) + Gf.Quatd(rot[3], *rot[:3]) if orient_double - else Gf.Quatf(*asset.initial_pose.rotation_wxyz) + else Gf.Quatf(rot[3], *rot[:3]) ) t_precision = UsdGeom.XformOp.PrecisionDouble if trans_double else UsdGeom.XformOp.PrecisionFloat r_precision = UsdGeom.XformOp.PrecisionDouble if orient_double else UsdGeom.XformOp.PrecisionFloat diff --git a/isaaclab_arena/scripts/imitation_learning/annotate_demos.py b/isaaclab_arena/scripts/imitation_learning/annotate_demos.py index 3ba9e71cb..dbe5cad06 100644 --- a/isaaclab_arena/scripts/imitation_learning/annotate_demos.py +++ b/isaaclab_arena/scripts/imitation_learning/annotate_demos.py @@ -36,13 +36,6 @@ help="File name of the annotated output dataset file.", ) parser.add_argument("--auto", action="store_true", default=False, help="Automatically annotate subtasks.") -parser.add_argument( - "--enable_pinocchio", - action="store_true", - default=False, - help="Enable Pinocchio.", -) - # Add the example environments CLI args # NOTE(alexmillane, 2025.09.04): This has to be added last, because # of the app specific flags being parsed after the global flags. @@ -51,11 +44,6 @@ # parse the arguments args_cli = parser.parse_args() -if args_cli.enable_pinocchio: - # Import pinocchio before AppLauncher to force the use of the version installed by IsaacLab and not the one installed by Isaac Sim - # pinocchio is required by the Pink IK controllers and the GR1T2 retargeter - import pinocchio # noqa: F401 - # launch the simulator app_launcher = AppLauncher(args_cli) simulation_app = app_launcher.app @@ -70,9 +58,6 @@ # Imports have to follow simulation startup. -if args_cli.enable_pinocchio: - import isaaclab_mimic.envs.pinocchio_envs # noqa: F401 - # Only enables inputs if this script is NOT headless mode if not args_cli.headless and not os.environ.get("HEADLESS", 0): from isaaclab.devices import Se3Keyboard, Se3KeyboardCfg diff --git a/isaaclab_arena/scripts/imitation_learning/generate_dataset.py b/isaaclab_arena/scripts/imitation_learning/generate_dataset.py index b6fc3408a..ed7bf6984 100644 --- a/isaaclab_arena/scripts/imitation_learning/generate_dataset.py +++ b/isaaclab_arena/scripts/imitation_learning/generate_dataset.py @@ -41,12 +41,6 @@ action="store_true", help="pause after every subtask during generation for debugging - only useful with render flag", ) -parser.add_argument( - "--enable_pinocchio", - action="store_true", - default=False, - help="Enable Pinocchio.", -) # Add the example environments CLI args # NOTE(alexmillane, 2025.09.04): This has to be added last, because @@ -56,11 +50,6 @@ # parse the arguments args_cli = parser.parse_args() -if args_cli.enable_pinocchio: - # Import pinocchio before AppLauncher to force the use of the version installed by IsaacLab and not the one installed by Isaac Sim - # pinocchio is required by the Pink IK controllers and the GR1T2 retargeter - import pinocchio # noqa: F401 - # launch the simulator app_launcher = AppLauncher(args_cli) simulation_app = app_launcher.app @@ -77,20 +66,15 @@ import torch import isaaclab_mimic.envs # noqa: F401 +import isaaclab_tasks # noqa: F401 from isaaclab.envs import ManagerBasedRLMimicEnv from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg from isaaclab.managers import DatasetExportMode, RecorderTerm, RecorderTermCfg from isaaclab.utils import configclass - -# Imports have to follow simulation startup. - -if args_cli.enable_pinocchio: - import isaaclab_mimic.envs.pinocchio_envs # noqa: F401 - -import isaaclab_tasks # noqa: F401 from isaaclab_mimic.datagen.generation import env_loop, setup_async_generation from isaaclab_mimic.datagen.utils import setup_output_paths +# import logger logger = logging.getLogger(__name__) @@ -214,36 +198,41 @@ def main(): # reset before starting env.reset() - # Setup and run async data generation - async_components = setup_async_generation( - env=env, - num_envs=args_cli.num_envs, - input_file=args_cli.input_file, - success_term=success_term, - pause_subtask=args_cli.pause_subtask, - ) - try: - data_gen_tasks = asyncio.ensure_future(asyncio.gather(*async_components["tasks"])) - env_loop( - env, - async_components["reset_queue"], - async_components["action_queue"], - async_components["info_pool"], - async_components["event_loop"], + # Setup and run async data generation + async_components = setup_async_generation( + env=env, + num_envs=args_cli.num_envs, + input_file=args_cli.input_file, + success_term=success_term, + pause_subtask=args_cli.pause_subtask, + motion_planners=None, ) - except asyncio.CancelledError: - print("Tasks were cancelled.") - finally: - # Cancel all async tasks when env_loop finishes - data_gen_tasks.cancel() + try: - # Wait for tasks to be cancelled - async_components["event_loop"].run_until_complete(data_gen_tasks) + data_gen_tasks = asyncio.ensure_future(asyncio.gather(*async_components["tasks"])) + env_loop( + env, + async_components["reset_queue"], + async_components["action_queue"], + async_components["info_pool"], + async_components["event_loop"], + ) except asyncio.CancelledError: - print("Remaining async tasks cancelled and cleaned up.") - except Exception as e: - print(f"Error cancelling remaining async tasks: {e}") + print("Tasks were cancelled.") + finally: + # Cancel all async tasks when env_loop finishes + data_gen_tasks.cancel() + try: + # Wait for tasks to be cancelled + async_components["event_loop"].run_until_complete(data_gen_tasks) + except asyncio.CancelledError: + print("Remaining async tasks cancelled and cleaned up.") + except Exception as e: + print(f"Error cancelling remaining async tasks: {e}") + finally: + # Close env after async tasks are done so success_term is never called on a closed env + env.close() if __name__ == "__main__": diff --git a/isaaclab_arena/scripts/imitation_learning/record_demos.py b/isaaclab_arena/scripts/imitation_learning/record_demos.py index 3e160d24a..7e633472c 100644 --- a/isaaclab_arena/scripts/imitation_learning/record_demos.py +++ b/isaaclab_arena/scripts/imitation_learning/record_demos.py @@ -50,13 +50,6 @@ default=10, help="Number of continuous steps with task success for concluding a demo as successful. Default is 10.", ) -parser.add_argument( - "--enable_pinocchio", - action="store_true", - default=False, - help="Enable Pinocchio.", -) - # Add the example environments CLI args # NOTE(alexmillane, 2025.09.04): This has to be added last, because # of the app specific flags being parsed after the global flags. @@ -67,11 +60,6 @@ app_launcher_args = vars(args_cli) -if args_cli.enable_pinocchio: - # Import pinocchio before AppLauncher to force the use of the version installed by IsaacLab and not the one installed by Isaac Sim - # pinocchio is required by the Pink IK controllers and the GR1T2 retargeter - import pinocchio # noqa: F401 - # TODO(cvolk): XR mode is inferred from teleop device name via string matching. # Ideally, AppLauncher or the device config would auto-detect XR requirements. if "openxr" in args_cli.teleop_device.lower(): @@ -86,32 +74,27 @@ # Third-party imports import gymnasium as gym -import logging import os import time import torch +from collections.abc import Callable import isaaclab_mimic.envs # noqa: F401 +import isaaclab_tasks # noqa: F401 +import isaaclab_tasks.manager_based.manipulation.pick_place # noqa: F401 + +# Omniverse logger +import omni.log import omni.ui as ui from isaaclab.devices import Se3Keyboard, Se3KeyboardCfg, Se3SpaceMouse, Se3SpaceMouseCfg -from isaaclab.devices.openxr import remove_camera_configs -from isaaclab.devices.teleop_device_factory import create_teleop_device - -logger = logging.getLogger(__name__) -from isaaclab_mimic.ui.instruction_display import InstructionDisplay, show_subtask_instructions - -# Imports have to follow simulation startup. - -if args_cli.enable_pinocchio: - import isaaclab_tasks.manager_based.manipulation.pick_place # noqa: F401 - -from collections.abc import Callable - -import isaaclab_tasks # noqa: F401 +from isaaclab_teleop import IsaacTeleopCfg, create_isaac_teleop_device, remove_camera_configs from isaaclab.envs import DirectRLEnvCfg, ManagerBasedRLEnvCfg from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg from isaaclab.envs.ui import EmptyWindow from isaaclab.managers import DatasetExportMode +from isaaclab_mimic.ui.instruction_display import InstructionDisplay, show_subtask_instructions + +# Imports have to follow simulation startup. class RateLimiter: @@ -196,7 +179,7 @@ def create_environment_config( env_name, env_cfg = arena_builder.build_registered() except Exception as e: - logger.error(f"Failed to parse environment configuration: {e}") + omni.log.error(f"Failed to parse environment configuration: {e}") exit(1) # extract success checking function to invoke in the main loop @@ -205,7 +188,7 @@ def create_environment_config( success_term = env_cfg.terminations.success env_cfg.terminations.success = None else: - logger.warning( + omni.log.warn( "No success termination term was found in the environment." " Will not be able to mark recorded demos as successful." ) @@ -246,7 +229,7 @@ def create_environment(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg, env_name: env = gym.make(env_name, cfg=env_cfg).unwrapped return env except Exception as e: - logger.error(f"Failed to create environment: {e}") + omni.log.error(f"Failed to create environment: {e}") exit(1) @@ -268,31 +251,33 @@ def setup_teleop_device(callbacks: dict[str, Callable]) -> object: """ teleop_interface = None try: - if hasattr(env_cfg, "teleop_devices") and args_cli.teleop_device in env_cfg.teleop_devices.devices: - teleop_interface = create_teleop_device(args_cli.teleop_device, env_cfg.teleop_devices.devices, callbacks) - else: - logger.warning( - f"No teleop device '{args_cli.teleop_device}' found in environment config. Creating default." + if hasattr(env_cfg, "isaac_teleop") and isinstance(env_cfg.isaac_teleop, IsaacTeleopCfg): + teleop_interface = create_isaac_teleop_device( + env_cfg.isaac_teleop, + sim_device=env_cfg.sim.device, + callbacks=callbacks, ) + else: + omni.log.warn(f"No teleop device '{args_cli.teleop_device}' found in environment config. Creating default.") # Create fallback teleop device if args_cli.teleop_device.lower() == "keyboard": teleop_interface = Se3Keyboard(Se3KeyboardCfg(pos_sensitivity=0.2, rot_sensitivity=0.5)) elif args_cli.teleop_device.lower() == "spacemouse": teleop_interface = Se3SpaceMouse(Se3SpaceMouseCfg(pos_sensitivity=0.2, rot_sensitivity=0.5)) else: - logger.error(f"Unsupported teleop device: {args_cli.teleop_device}") - logger.error("Supported devices: keyboard, spacemouse, avp_handtracking") + omni.log.error(f"Unsupported teleop device: {args_cli.teleop_device}") + omni.log.error("Supported devices: keyboard, spacemouse, avp_handtracking") exit(1) # Add callbacks to fallback device for key, callback in callbacks.items(): teleop_interface.add_callback(key, callback) except Exception as e: - logger.error(f"Failed to create teleop device: {e}") + omni.log.error(f"Failed to create teleop device: {e}") exit(1) if teleop_interface is None: - logger.error("Failed to create teleop interface") + omni.log.error("Failed to create teleop interface") exit(1) return teleop_interface @@ -435,72 +420,96 @@ def stop_recording_instance(): teleop_interface = setup_teleop_device(teleoperation_callbacks) teleop_interface.add_callback("R", reset_recording_instance) - - # Reset before starting - env.sim.reset() - env.reset() - teleop_interface.reset() + use_isaac_teleop = args_cli.xr label_text = f"Recorded {current_recorded_demo_count} successful demonstrations." instruction_display = setup_ui(label_text, env) - subtasks = {} - - with contextlib.suppress(KeyboardInterrupt) and torch.inference_mode(): - while simulation_app.is_running(): - # Get keyboard command - action = teleop_interface.advance() - - # Expand to batch dimension - actions = action.repeat(env.num_envs, 1) - # Hack for G1 Pink WBC to transferm EE into robot base coordinates - action_manager = getattr(env, "action_manager", None) - if action_manager is not None: - for term_name in action_manager.active_terms: - term = action_manager.get_term(term_name) - if hasattr(term, "preprocess_actions"): - actions = term.preprocess_actions(actions) - - # Perform action on environment - if running_recording_instance: - # Compute actions based on environment - obv = env.step(actions) - if subtasks is not None: - if subtasks == {}: - subtasks = obv[0].get("subtask_terms") - elif subtasks: - show_subtask_instructions(instruction_display, subtasks, obv, env.cfg) - else: - env.sim.render() - - # Check for success condition - success_step_count, success_reset_needed = process_success_condition(env, success_term, success_step_count) - if success_reset_needed: - should_reset_recording_instance = True - - # Update demo count if it has changed - if env.recorder_manager.exported_successful_episode_count > current_recorded_demo_count: - current_recorded_demo_count = env.recorder_manager.exported_successful_episode_count - label_text = f"Recorded {current_recorded_demo_count} successful demonstrations." - print(label_text) - - # Handle reset if requested - if should_reset_recording_instance: - success_step_count = handle_reset(env, success_step_count, instruction_display, label_text) - should_reset_recording_instance = False - - # Check if we've reached the desired number of demos - if args_cli.num_demos > 0 and env.recorder_manager.exported_successful_episode_count >= args_cli.num_demos: - print(f"All {args_cli.num_demos} demonstrations recorded. Exiting the app.") - break - - # Check if simulation is stopped - if env.sim.is_stopped(): - break - - # Rate limiting - if rate_limiter: - rate_limiter.sleep(env) + def inner_loop(): + """Inner loop function with access to nonlocal variables.""" + nonlocal current_recorded_demo_count, success_step_count, should_reset_recording_instance + nonlocal running_recording_instance, label_text + + # Reset before starting + env.sim.reset() + env.reset() + teleop_interface.reset() + + subtasks = {} + stack_name = "IsaacTeleop" if use_isaac_teleop else "native" + print(f"{stack_name} recording started.") + + with contextlib.suppress(KeyboardInterrupt), torch.inference_mode(): + while simulation_app.is_running(): + # Get teleop command (may be None while waiting for session start) + action = teleop_interface.advance() + if action is None: + env.sim.render() + continue + # Expand to batch dimension + actions = action.repeat(env.num_envs, 1) + + # Perform action on environment + if running_recording_instance: + # Compute actions based on environment + obv = env.step(actions) + if subtasks is not None: + if subtasks == {}: + subtasks = obv[0].get("subtask_terms") + elif subtasks: + show_subtask_instructions(instruction_display, subtasks, obv, env.cfg) + else: + env.sim.render() + + # Check for success condition + success_step_count_new, success_reset_needed = process_success_condition( + env, success_term, success_step_count + ) + success_step_count = success_step_count_new + if success_reset_needed: + should_reset_recording_instance = True + + # Update demo count if it has changed + if env.recorder_manager.exported_successful_episode_count > current_recorded_demo_count: + current_recorded_demo_count = env.recorder_manager.exported_successful_episode_count + label_text = f"Recorded {current_recorded_demo_count} successful demonstrations." + print(label_text) + + # Check if we've reached the desired number of demos + if ( + args_cli.num_demos > 0 + and env.recorder_manager.exported_successful_episode_count >= args_cli.num_demos + ): + label_text = f"All {current_recorded_demo_count} demonstrations recorded.\nExiting the app." + instruction_display.show_demo(label_text) + print(label_text) + target_time = time.time() + 0.8 + while time.time() < target_time: + if rate_limiter: + rate_limiter.sleep(env) + else: + env.sim.render() + break + + # Handle reset if requested + if should_reset_recording_instance: + success_step_count = handle_reset(env, success_step_count, instruction_display, label_text) + should_reset_recording_instance = False + + # Check if simulation is stopped + if env.sim.is_stopped(): + break + + # Rate limiting + if rate_limiter: + rate_limiter.sleep(env) + + # Run the loop with or without context manager based on stack + if use_isaac_teleop: + with teleop_interface: + inner_loop() + else: + inner_loop() return current_recorded_demo_count @@ -521,6 +530,10 @@ def main() -> None: # if handtracking is selected, rate limiting is achieved via OpenXR if args_cli.xr: rate_limiter = None + from isaaclab.ui.xr_widgets import TeleopVisualizationManager, XRVisualization + + # Assign the teleop visualization manager to the visualization system + XRVisualization.assign_manager(TeleopVisualizationManager) else: rate_limiter = RateLimiter(args_cli.step_hz) @@ -547,4 +560,4 @@ def main() -> None: # run the main function main() # close sim app - simulation_app.close() + simulation_app.close() \ No newline at end of file diff --git a/isaaclab_arena/scripts/imitation_learning/replay_demos.py b/isaaclab_arena/scripts/imitation_learning/replay_demos.py index 20a8b48eb..037d7d117 100644 --- a/isaaclab_arena/scripts/imitation_learning/replay_demos.py +++ b/isaaclab_arena/scripts/imitation_learning/replay_demos.py @@ -36,13 +36,6 @@ " --num_envs is 1." ), ) -parser.add_argument( - "--enable_pinocchio", - action="store_true", - default=False, - help="Enable Pinocchio.", -) - # Add the example environments CLI args # NOTE(alexmillane, 2025.09.04): This has to be added last, because # of the app specific flags being parsed after the global flags. @@ -52,11 +45,6 @@ args_cli = parser.parse_args() # args_cli.headless = True -if args_cli.enable_pinocchio: - # Import pinocchio before AppLauncher to force the use of the version installed by IsaacLab and not the one installed by Isaac Sim - # pinocchio is required by the Pink IK controllers and the GR1T2 retargeter - import pinocchio # noqa: F401 - # launch the simulator app_launcher = AppLauncher(args_cli) simulation_app = app_launcher.app @@ -68,14 +56,11 @@ import os import torch +import isaaclab_tasks # noqa: F401 +import isaaclab_tasks.manager_based.manipulation.pick_place # noqa: F401 from isaaclab.devices import Se3Keyboard, Se3KeyboardCfg from isaaclab.utils.datasets import EpisodeData, HDF5DatasetFileHandler -if args_cli.enable_pinocchio: - import isaaclab_tasks.manager_based.manipulation.pick_place # noqa: F401 - -import isaaclab_tasks # noqa: F401 - is_paused = False @@ -127,6 +112,7 @@ def main(): if not os.path.exists(args_cli.dataset_file): raise FileNotFoundError(f"The dataset file {args_cli.dataset_file} does not exist.") dataset_file_handler = HDF5DatasetFileHandler() + dataset_file_handler.open(args_cli.dataset_file) env_name = dataset_file_handler.get_env_name() episode_count = dataset_file_handler.get_num_episodes() diff --git a/isaaclab_arena/scripts/imitation_learning/teleop.py b/isaaclab_arena/scripts/imitation_learning/teleop.py index 75aa52b25..5bc212258 100644 --- a/isaaclab_arena/scripts/imitation_learning/teleop.py +++ b/isaaclab_arena/scripts/imitation_learning/teleop.py @@ -23,13 +23,6 @@ # add argparse arguments parser = get_isaaclab_arena_cli_parser() parser.add_argument("--sensitivity", type=float, default=1.0, help="Sensitivity factor.") -parser.add_argument( - "--enable_pinocchio", - action="store_true", - default=False, - help="Enable Pinocchio.", -) - # Add the example environments CLI args # NOTE(alexmillane, 2025.09.04): This has to be added last, because # of the app specific flags being parsed after the global flags. @@ -40,14 +33,6 @@ app_launcher_args = vars(args_cli) -if args_cli.enable_pinocchio: - # Import pinocchio before AppLauncher to force the use of the version installed by IsaacLab and - # not the one installed by Isaac Sim pinocchio is required by the Pink IK controllers and the - # GR1T2 retargeter - import pinocchio # noqa: F401 - -# TODO(cvolk): XR mode is inferred from teleop device name via string matching. -# Ideally, AppLauncher or the device config would auto-detect XR requirements. if "openxr" in args_cli.teleop_device.lower(): app_launcher_args["xr"] = True @@ -58,21 +43,16 @@ """Rest everything follows.""" -import logging import torch import isaaclab_tasks # noqa: F401 +import isaaclab_tasks.manager_based.manipulation.pick_place # noqa: F401 +import omni.log from isaaclab.devices import Se3Gamepad, Se3GamepadCfg, Se3Keyboard, Se3KeyboardCfg, Se3SpaceMouse, Se3SpaceMouseCfg -from isaaclab.devices.openxr import remove_camera_configs -from isaaclab.devices.teleop_device_factory import create_teleop_device +from isaaclab_teleop import IsaacTeleopCfg, create_isaac_teleop_device, remove_camera_configs from isaaclab.managers import TerminationTermCfg as DoneTerm from isaaclab_tasks.manager_based.manipulation.lift import mdp -if args_cli.enable_pinocchio: - import isaaclab_tasks.manager_based.manipulation.pick_place # noqa: F401 - -logger = logging.getLogger(__name__) - def main() -> None: """ @@ -111,7 +91,7 @@ def main() -> None: " will be ignored." ) except Exception as e: - logger.error(f"Failed to create environment: {e}") + omni.log.error(f"Failed to create environment: {e}") simulation_app.close() return @@ -178,14 +158,14 @@ def stop_teleoperation() -> None: # Create teleop device from config if present, otherwise create manually teleop_interface = None try: - if hasattr(env_cfg, "teleop_devices") and args_cli.teleop_device in env_cfg.teleop_devices.devices: - teleop_interface = create_teleop_device( - args_cli.teleop_device, env_cfg.teleop_devices.devices, teleoperation_callbacks + if hasattr(env_cfg, "isaac_teleop") and isinstance(env_cfg.isaac_teleop, IsaacTeleopCfg): + teleop_interface = create_isaac_teleop_device( + env_cfg.isaac_teleop, + sim_device=str(env.device), + callbacks=teleoperation_callbacks, ) else: - logger.warning( - f"No teleop device '{args_cli.teleop_device}' found in environment config. Creating default." - ) + omni.log.warn(f"No teleop device '{args_cli.teleop_device}' found in environment config. Creating default.") # Create fallback teleop device sensitivity = args_cli.sensitivity if args_cli.teleop_device.lower() == "keyboard": @@ -201,8 +181,8 @@ def stop_teleoperation() -> None: Se3GamepadCfg(pos_sensitivity=0.1 * sensitivity, rot_sensitivity=0.1 * sensitivity) ) else: - logger.error(f"Unsupported teleop device: {args_cli.teleop_device}") - logger.error("Supported devices: keyboard, spacemouse, gamepad, avp_handtracking") + omni.log.error(f"Unsupported teleop device: {args_cli.teleop_device}") + omni.log.error("Supported devices: keyboard, spacemouse, gamepad, avp_handtracking") env.close() simulation_app.close() return @@ -212,60 +192,57 @@ def stop_teleoperation() -> None: try: teleop_interface.add_callback(key, callback) except (ValueError, TypeError) as e: - logger.warning(f"Failed to add callback for key {key}: {e}") + omni.log.warn(f"Failed to add callback for key {key}: {e}") except Exception as e: - logger.error(f"Failed to create teleop device: {e}") + omni.log.error(f"Failed to create teleop device: {e}") env.close() simulation_app.close() return if teleop_interface is None: - logger.error("Failed to create teleop interface") + omni.log.error("Failed to create teleop interface") env.close() simulation_app.close() return print(f"Using teleop device: {teleop_interface}") - # reset environment - env.reset() - teleop_interface.reset() - - print("Teleoperation started. Press 'R' to reset the environment.") - - # simulate environment - while simulation_app.is_running(): - try: - # run everything in inference mode - with torch.inference_mode(): - # get device command - action = teleop_interface.advance() - - # Only apply teleop commands when active - if teleoperation_active: - # process actions - actions = action.repeat(env.num_envs, 1) - # Hack for G1 Pink WBC to transferm EE into robot base coordinates - action_manager = getattr(env, "action_manager", None) - if action_manager is not None: - for term_name in action_manager.active_terms: - term = action_manager.get_term(term_name) - if hasattr(term, "preprocess_actions"): - actions = term.preprocess_actions(actions) - # apply actions - env.step(actions) - else: - env.sim.render() - - if should_reset_recording_instance: - env.reset() - should_reset_recording_instance = False - print("Environment reset complete") - except Exception as e: - logger.error(f"Error during simulation step: {e}") - break - - # close the simulator + # IsaacTeleop (OpenXR) requires the device to be used as a context manager so + # TeleopSessionLifecycle.start() is called before advance(). + use_isaac_teleop = hasattr(teleop_interface, "__enter__") and hasattr(teleop_interface, "__exit__") + + def run_teleop_loop() -> None: + nonlocal should_reset_recording_instance + env.reset() + teleop_interface.reset() + print("Teleoperation started. Press 'R' to reset the environment.") + while simulation_app.is_running(): + try: + with torch.inference_mode(): + action = teleop_interface.advance() + # action is None when IsaacTeleop session hasn't started yet (e.g. waiting for "Start AR") + if action is None: + env.sim.render() + elif teleoperation_active: + actions = action.repeat(env.num_envs, 1) + env.step(actions) + else: + env.sim.render() + if should_reset_recording_instance: + env.reset() + teleop_interface.reset() + should_reset_recording_instance = False + print("Environment reset complete") + except Exception as e: + omni.log.error(f"Error during simulation step: {e}") + break + + if use_isaac_teleop: + with teleop_interface: + run_teleop_loop() + else: + run_teleop_loop() + env.close() print("Environment closed") @@ -274,4 +251,4 @@ def stop_teleoperation() -> None: # run the main function main() # close sim app - simulation_app.close() + simulation_app.close() \ No newline at end of file diff --git a/isaaclab_arena/scripts/reinforcement_learning/play.py b/isaaclab_arena/scripts/reinforcement_learning/play.py new file mode 100644 index 000000000..d58097feb --- /dev/null +++ b/isaaclab_arena/scripts/reinforcement_learning/play.py @@ -0,0 +1,197 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Script to play a checkpoint if an RL agent from RSL-RL.""" + +"""Launch Isaac Sim Simulator first.""" + +from pathlib import Path + +from isaaclab.app import AppLauncher + +from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser +from isaaclab_arena_environments.cli import add_example_environments_cli_args + +# local imports +import cli_args # isort: skip + +# add argparse arguments +parser = get_isaaclab_arena_cli_parser() +parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.") +parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).") +parser.add_argument( + "--agent_cfg_path", + type=Path, + default=Path("isaaclab_arena/policy/rl_policy/generic_policy.json"), + help="Path to the RL agent configuration file.", +) +parser.add_argument("--real-time", action="store_true", default=False, help="Run in real-time, if possible.") +# append RSL-RL cli arguments +cli_args.add_rsl_rl_args(parser) +cli_args.add_rsl_rl_policy_args(parser) +# Add the example environments CLI args +# NOTE(alexmillane, 2025.09.04): This has to be added last, because +# of the app specific flags being parsed after the global flags. +add_example_environments_cli_args(parser) +args_cli = parser.parse_args() + +# always enable cameras to record video +if args_cli.video: + args_cli.enable_cameras = True + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import gymnasium as gym +import os +import time +import torch + +import isaaclab_tasks # noqa: F401 +import omni.log +from isaaclab.envs import DirectMARLEnv, multi_agent_to_single_agent +from isaaclab.utils.assets import retrieve_file_path +from isaaclab.utils.dict import print_dict +from isaaclab_rl.rsl_rl import RslRlBaseRunnerCfg, RslRlVecEnvWrapper, export_policy_as_jit, export_policy_as_onnx +from isaaclab_tasks.utils import get_checkpoint_path +from rsl_rl.runners import DistillationRunner, OnPolicyRunner + +from isaaclab_arena.policy.rl_policy.base_rsl_rl_policy import get_agent_cfg +from isaaclab_arena_environments.cli import get_arena_builder_from_cli + +# PLACEHOLDER: Extension template (do not remove this comment) + + +def main(): + """Play with RSL-RL agent.""" + # We dont use hydra for the environment configuration, so we need to parse it manually + # parse configuration + try: + arena_builder = get_arena_builder_from_cli(args_cli) + env_name, env_cfg = arena_builder.build_registered() + + except Exception as e: + omni.log.error(f"Failed to parse environment configuration: {e}") + exit(1) + + agent_cfg = get_agent_cfg(args_cli) + + # override configurations with non-hydra CLI arguments + agent_cfg: RslRlBaseRunnerCfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli) + env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs + + # set the environment seed + # note: certain randomizations occur in the environment initialization so we set the seed here + env_cfg.seed = agent_cfg.seed + env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device + + # specify directory for logging experiments + log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name) + log_root_path = os.path.abspath(log_root_path) + print(f"[INFO] Loading experiment from directory: {log_root_path}") + if args_cli.checkpoint: + resume_path = retrieve_file_path(args_cli.checkpoint) + else: + resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint) + + log_dir = os.path.dirname(resume_path) + + # set the log directory for the environment (works for all environment types) + env_cfg.log_dir = log_dir + + # create isaac environment + env = gym.make(env_name, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) + + # convert to single-agent instance if required by the RL algorithm + if isinstance(env.unwrapped, DirectMARLEnv): + env = multi_agent_to_single_agent(env) + + # wrap for video recording + if args_cli.video: + video_kwargs = { + "video_folder": os.path.join(log_dir, "videos", "play"), + "step_trigger": lambda step: step == 0, + "video_length": args_cli.video_length, + "disable_logger": True, + } + print("[INFO] Recording videos during training.") + print_dict(video_kwargs, nesting=4) + env = gym.wrappers.RecordVideo(env, **video_kwargs) + + # wrap around environment for rsl-rl + env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions) + + print(f"[INFO]: Loading model checkpoint from: {resume_path}") + # load previously trained model + if agent_cfg.class_name == "OnPolicyRunner": + runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device) + elif agent_cfg.class_name == "DistillationRunner": + runner = DistillationRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device) + else: + raise ValueError(f"Unsupported runner class: {agent_cfg.class_name}") + runner.load(resume_path) + + # obtain the trained policy for inference + policy = runner.get_inference_policy(device=env.unwrapped.device) + + # extract the neural network module + # we do this in a try-except to maintain backwards compatibility. + try: + # version 2.3 onwards + policy_nn = runner.alg.policy + except AttributeError: + # version 2.2 and below + policy_nn = runner.alg.actor_critic + + # extract the normalizer + if hasattr(policy_nn, "actor_obs_normalizer"): + normalizer = policy_nn.actor_obs_normalizer + elif hasattr(policy_nn, "student_obs_normalizer"): + normalizer = policy_nn.student_obs_normalizer + else: + normalizer = None + + # export policy to onnx/jit + export_model_dir = os.path.join(os.path.dirname(resume_path), "exported") + export_policy_as_jit(policy_nn, normalizer=normalizer, path=export_model_dir, filename="policy.pt") + export_policy_as_onnx(policy_nn, normalizer=normalizer, path=export_model_dir, filename="policy.onnx") + + dt = env.unwrapped.step_dt + + # reset environment + obs = env.get_observations() + timestep = 0 + # simulate environment + while simulation_app.is_running(): + start_time = time.time() + # run everything in inference mode + with torch.inference_mode(): + # agent stepping + actions = policy(obs) + # env stepping + obs, _, _, _ = env.step(actions) + if args_cli.video: + timestep += 1 + # Exit the play loop after recording one video + if timestep == args_cli.video_length: + break + + # time delay for real-time evaluation + sleep_time = dt - (time.time() - start_time) + if args_cli.real_time and sleep_time > 0: + time.sleep(sleep_time) + + # close the simulator + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/isaaclab_arena/scripts/reinforcement_learning/train.py b/isaaclab_arena/scripts/reinforcement_learning/train.py new file mode 100644 index 000000000..4f12b5bc4 --- /dev/null +++ b/isaaclab_arena/scripts/reinforcement_learning/train.py @@ -0,0 +1,219 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Script to train RL agent with RSL-RL.""" + +"""Launch Isaac Sim Simulator first.""" + +from pathlib import Path + +from isaaclab.app import AppLauncher + +from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser +from isaaclab_arena_environments.cli import add_example_environments_cli_args + +# local imports +import cli_args # isort: skip + +# add argparse arguments +parser = get_isaaclab_arena_cli_parser() +parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.") +parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).") +parser.add_argument("--video_interval", type=int, default=2000, help="Interval between video recordings (in steps).") +parser.add_argument( + "--agent_cfg_path", + type=Path, + default=Path("isaaclab_arena/policy/rl_policy/generic_policy.json"), + help="Path to the RL agent configuration file.", +) +parser.add_argument( + "--distributed", action="store_true", default=False, help="Run training with multiple GPUs or nodes." +) +parser.add_argument("--export_io_descriptors", action="store_true", default=False, help="Export IO descriptors.") +# append RSL-RL cli arguments +cli_args.add_rsl_rl_args(parser) +cli_args.add_rsl_rl_policy_args(parser) +# Add the example environments CLI args +# NOTE(alexmillane, 2025.09.04): This has to be added last, because +# of the app specific flags being parsed after the global flags. +add_example_environments_cli_args(parser) +args_cli = parser.parse_args() + +# always enable cameras to record video +if args_cli.video: + args_cli.enable_cameras = True + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Check for minimum supported RSL-RL version.""" + +import importlib.metadata as metadata +import platform + +from packaging import version + +# check minimum supported rsl-rl version +RSL_RL_VERSION = "3.0.1" +installed_version = metadata.version("rsl-rl-lib") +if version.parse(installed_version) < version.parse(RSL_RL_VERSION): + if platform.system() == "Windows": + cmd = [r".\isaaclab.bat", "-p", "-m", "pip", "install", f"rsl-rl-lib=={RSL_RL_VERSION}"] + else: + cmd = ["./isaaclab.sh", "-p", "-m", "pip", "install", f"rsl-rl-lib=={RSL_RL_VERSION}"] + print( + f"Please install the correct version of RSL-RL.\nExisting version is: '{installed_version}'" + f" and required version is: '{RSL_RL_VERSION}'.\nTo install the correct version, run:" + f"\n\n\t{' '.join(cmd)}\n" + ) + exit(1) + +"""Rest everything follows.""" + +import gymnasium as gym +import os +import torch +from datetime import datetime + +import isaaclab_tasks # noqa: F401 +import omni.log +from isaaclab.envs import DirectMARLEnv, ManagerBasedRLEnvCfg, multi_agent_to_single_agent +from isaaclab.utils.dict import print_dict +from isaaclab.utils.io import dump_yaml +from isaaclab_rl.rsl_rl import RslRlVecEnvWrapper +from isaaclab_tasks.utils import get_checkpoint_path +from rsl_rl.runners import DistillationRunner, OnPolicyRunner + +from isaaclab_arena.policy.rl_policy.base_rsl_rl_policy import get_agent_cfg +from isaaclab_arena_environments.cli import get_arena_builder_from_cli + +# PLACEHOLDER: Extension template (do not remove this comment) + +torch.backends.cuda.matmul.allow_tf32 = True +torch.backends.cudnn.allow_tf32 = True +torch.backends.cudnn.deterministic = False +torch.backends.cudnn.benchmark = False + + +def main(): + # We dont use hydra for the environment configuration, so we need to parse it manually + # parse configuration + try: + arena_builder = get_arena_builder_from_cli(args_cli) + env_name, env_cfg = arena_builder.build_registered() + + except Exception as e: + omni.log.error(f"Failed to parse environment configuration: {e}") + exit(1) + + agent_cfg = get_agent_cfg(args_cli) + + # set the environment seed + # note: certain randomizations occur in the environment initialization so we set the seed here + env_cfg.seed = agent_cfg.seed + env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device + # check for invalid combination of CPU device with distributed training + if args_cli.distributed and args_cli.device is not None and "cpu" in args_cli.device: + raise ValueError( + "Distributed training is not supported when using CPU device. " + "Please use GPU device (e.g., --device cuda) for distributed training." + ) + + # multi-gpu training configuration + if args_cli.distributed: + env_cfg.sim.device = f"cuda:{app_launcher.local_rank}" + agent_cfg.device = f"cuda:{app_launcher.local_rank}" + + # set seed to have diversity in different threads + seed = agent_cfg.seed + app_launcher.local_rank + env_cfg.seed = seed + agent_cfg.seed = seed + + # specify directory for logging experiments + log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name) + log_root_path = os.path.abspath(log_root_path) + print(f"[INFO] Logging experiment in directory: {log_root_path}") + # specify directory for logging runs: {time-stamp}_{run_name} + log_dir = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + # The Ray Tune workflow extracts experiment name using the logging line below, hence, do not change it (see PR #2346, comment-2819298849) + print(f"Exact experiment name requested from command line: {log_dir}") + if agent_cfg.run_name: + log_dir += f"_{agent_cfg.run_name}" + log_dir = os.path.join(log_root_path, log_dir) + + # set the IO descriptors export flag if requested + if isinstance(env_cfg, ManagerBasedRLEnvCfg): + env_cfg.export_io_descriptors = args_cli.export_io_descriptors + else: + omni.log.warn( + "IO descriptors are only supported for manager based RL environments. No IO descriptors will be exported." + ) + + # set the log directory for the environment (works for all environment types) + env_cfg.log_dir = log_dir + + # create isaac environment + env = gym.make(env_name, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) + + # convert to single-agent instance if required by the RL algorithm + if isinstance(env.unwrapped, DirectMARLEnv): + env = multi_agent_to_single_agent(env) + + # save resume path before creating a new log_dir + if agent_cfg.resume or agent_cfg.algorithm.class_name == "Distillation": + resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint) + + # wrap for video recording + if args_cli.video: + video_kwargs = { + "video_folder": os.path.join(log_dir, "videos", "train"), + "step_trigger": lambda step: step % args_cli.video_interval == 0, + "video_length": args_cli.video_length, + "disable_logger": True, + } + print("[INFO] Recording videos during training.") + print_dict(video_kwargs, nesting=4) + env = gym.wrappers.RecordVideo(env, **video_kwargs) + + # wrap around environment for rsl-rl + env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions) + + # create runner from rsl-rl + if agent_cfg.class_name == "OnPolicyRunner": + runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device) + elif agent_cfg.class_name == "DistillationRunner": + runner = DistillationRunner(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device) + else: + raise ValueError(f"Unsupported runner class: {agent_cfg.class_name}") + # write git state to logs + runner.add_git_repo_to_log(__file__) + # load the checkpoint + if agent_cfg.resume or agent_cfg.algorithm.class_name == "Distillation": + print(f"[INFO]: Loading model checkpoint from: {resume_path}") + # load previously trained model + runner.load(resume_path) + + # dump the configuration into log-directory + dump_yaml(os.path.join(log_dir, "params", "env.yaml"), env_cfg) + dump_yaml(os.path.join(log_dir, "params", "agent.yaml"), agent_cfg) + + # run training + runner.learn(num_learning_iterations=agent_cfg.max_iterations, init_at_random_ep_len=True) + + # close the simulator + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/isaaclab_arena/tasks/events.py b/isaaclab_arena/tasks/events.py index cb61b3108..59709bba3 100644 --- a/isaaclab_arena/tasks/events.py +++ b/isaaclab_arena/tasks/events.py @@ -7,6 +7,7 @@ from __future__ import annotations import torch +import warp as wp from typing import TYPE_CHECKING, Literal import isaaclab.utils.math as math_utils diff --git a/isaaclab_arena/tasks/goal_pose_task.py b/isaaclab_arena/tasks/goal_pose_task.py index 5b5d0c57d..93585ef2a 100644 --- a/isaaclab_arena/tasks/goal_pose_task.py +++ b/isaaclab_arena/tasks/goal_pose_task.py @@ -29,7 +29,7 @@ def __init__( target_x_range: tuple[float, float] | None = None, target_y_range: tuple[float, float] | None = None, target_z_range: tuple[float, float] | None = None, - target_orientation_wxyz: tuple[float, float, float, float] | None = None, + target_orientation_xyzw: tuple[float, float, float, float] | None = None, target_orientation_tolerance_rad: float | None = None, ): """ @@ -39,7 +39,7 @@ def __init__( target_x_range: Success zone x-range [min, max] in meters. target_y_range: Success zone y-range [min, max] in meters. target_z_range: Success zone z-range [min, max] in meters. - target_orientation_wxyz: Target quaternion [w, x, y, z]. + target_orientation_xyzw: Target quaternion [x, y, z, w]. target_orientation_tolerance_rad: Angular tolerance in radians (default: 0.1). """ super().__init__(episode_length_s=episode_length_s) @@ -51,7 +51,7 @@ def __init__( target_x_range=target_x_range, target_y_range=target_y_range, target_z_range=target_z_range, - target_orientation_wxyz=target_orientation_wxyz, + target_orientation_xyzw=target_orientation_xyzw, target_orientation_tolerance_rad=target_orientation_tolerance_rad, ) @@ -66,7 +66,7 @@ def make_termination_cfg( target_x_range: tuple[float, float] | None = None, target_y_range: tuple[float, float] | None = None, target_z_range: tuple[float, float] | None = None, - target_orientation_wxyz: tuple[float, float, float, float] | None = None, + target_orientation_xyzw: tuple[float, float, float, float] | None = None, target_orientation_tolerance_rad: float | None = None, ): params: dict = {"object_cfg": SceneEntityCfg(self.object.name)} @@ -76,8 +76,8 @@ def make_termination_cfg( params["target_y_range"] = target_y_range if target_z_range is not None: params["target_z_range"] = target_z_range - if target_orientation_wxyz is not None: - params["target_orientation_wxyz"] = target_orientation_wxyz + if target_orientation_xyzw is not None: + params["target_orientation_xyzw"] = target_orientation_xyzw if target_orientation_tolerance_rad is not None: params["target_orientation_tolerance_rad"] = target_orientation_tolerance_rad diff --git a/isaaclab_arena/tasks/observations/observations.py b/isaaclab_arena/tasks/observations/observations.py index f412c0949..a4c2d8d99 100644 --- a/isaaclab_arena/tasks/observations/observations.py +++ b/isaaclab_arena/tasks/observations/observations.py @@ -4,6 +4,7 @@ # SPDX-License-Identifier: Apache-2.0 import torch +import warp as wp from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedRLEnv @@ -17,7 +18,7 @@ def object_position_in_world_frame( ) -> torch.Tensor: """Observation the position of the object in the world frame.""" object = env.scene[asset_cfg.name] - return object.data.root_pos_w + return wp.to_torch(object.data.root_pos_w) def object_position_in_frame( @@ -28,6 +29,6 @@ def object_position_in_frame( """The position of the object in the robot's root frame.""" root_frame: RigidObject = env.scene[root_frame_cfg.name] object: RigidObject = env.scene[object_cfg.name] - object_pos_w = object.data.root_pos_w[:, :3] - object_pos_b, _ = subtract_frame_transforms(root_frame.data.root_pos_w, root_frame.data.root_quat_w, object_pos_w) + object_pos_w = wp.to_torch(object.data.root_pos_w)[:, :3] + object_pos_b, _ = subtract_frame_transforms(wp.to_torch(root_frame.data.root_pos_w), wp.to_torch(root_frame.data.root_quat_w), object_pos_w) return object_pos_b diff --git a/isaaclab_arena/tasks/rewards/lift_object_rewards.py b/isaaclab_arena/tasks/rewards/lift_object_rewards.py index ce011246f..f13ed8754 100644 --- a/isaaclab_arena/tasks/rewards/lift_object_rewards.py +++ b/isaaclab_arena/tasks/rewards/lift_object_rewards.py @@ -4,6 +4,7 @@ # SPDX-License-Identifier: Apache-2.0 import torch +import warp as wp from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedRLEnv @@ -16,7 +17,7 @@ def object_is_lifted( ) -> torch.Tensor: """Reward the agent for lifting the object above the minimal height.""" object: RigidObject = env.scene[object_cfg.name] - return torch.where(object.data.root_pos_w[:, 2] > minimal_height, 1.0, 0.0) + return torch.where(wp.to_torch(object.data.root_pos_w)[:, 2] > minimal_height, 1.0, 0.0) def object_goal_distance( @@ -34,8 +35,8 @@ def object_goal_distance( command = env.command_manager.get_command(command_name) # compute the desired position in the world frame des_pos_b = command[:, :3] - des_pos_w, _ = combine_frame_transforms(robot.data.root_pos_w, robot.data.root_quat_w, des_pos_b) + des_pos_w, _ = combine_frame_transforms(wp.to_torch(robot.data.root_pos_w), wp.to_torch(robot.data.root_quat_w), des_pos_b) # distance of the end-effector to the object: (num_envs,) - distance = torch.norm(des_pos_w - object.data.root_pos_w, dim=1) + distance = torch.norm(des_pos_w - wp.to_torch(object.data.root_pos_w), dim=1) # rewarded if the object is lifted above the threshold - return (object.data.root_pos_w[:, 2] > minimal_height) * (1 - torch.tanh(distance / std)) + return (wp.to_torch(object.data.root_pos_w)[:, 2] > minimal_height) * (1 - torch.tanh(distance / std)) diff --git a/isaaclab_arena/tasks/rewards/rewards.py b/isaaclab_arena/tasks/rewards/rewards.py index f186ce36b..10a58f5fa 100644 --- a/isaaclab_arena/tasks/rewards/rewards.py +++ b/isaaclab_arena/tasks/rewards/rewards.py @@ -4,6 +4,7 @@ # SPDX-License-Identifier: Apache-2.0 import torch +import warp as wp from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedRLEnv @@ -22,9 +23,9 @@ def object_ee_distance( object: RigidObject = env.scene[object_cfg.name] ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] # Target object position: (num_envs, 3) - object_pos_w = object.data.root_pos_w + object_pos_w = wp.to_torch(object.data.root_pos_w) # End-effector position: (num_envs, 3) - ee_w = ee_frame.data.target_pos_w[..., 0, :] + ee_w = wp.to_torch(ee_frame.data.target_pos_w)[..., 0, :] # Distance of the end-effector to the object: (num_envs,) object_ee_distance = torch.norm(object_pos_w - ee_w, dim=1) diff --git a/isaaclab_arena/tasks/terminations.py b/isaaclab_arena/tasks/terminations.py index bee8dc83d..43f7ea107 100644 --- a/isaaclab_arena/tasks/terminations.py +++ b/isaaclab_arena/tasks/terminations.py @@ -5,6 +5,7 @@ import math import torch +import warp as wp from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedRLEnv @@ -32,10 +33,10 @@ def object_on_destination( assert sensor.data.force_matrix_w.shape[1] == 1 # NOTE(alexmillane, 2025-08-04): We expect the binary flags to have shape (N, ) # where N is the number of envs. - force_matrix_norm = torch.norm(sensor.data.force_matrix_w.clone(), dim=-1).reshape(-1) + force_matrix_norm = torch.norm(wp.to_torch(sensor.data.force_matrix_w), dim=-1).reshape(-1) force_above_threshold = force_matrix_norm > force_threshold - velocity_w = object.data.root_lin_vel_w + velocity_w = wp.to_torch(object.data.root_lin_vel_w) velocity_w_norm = torch.norm(velocity_w, dim=-1) velocity_below_threshold = velocity_w_norm < velocity_threshold @@ -86,11 +87,8 @@ def objects_in_proximity( target_object: RigidObject = env.scene[target_object_cfg.name] # Get positions relative to environment origin - object_pos = object.data.root_pos_w - env.scene.env_origins - - # Get positions relative to environment origin - object_pos = object.data.root_pos_w - env.scene.env_origins - target_object_pos = target_object.data.root_pos_w - env.scene.env_origins + object_pos = wp.to_torch(object.data.root_pos_w) - env.scene.env_origins + target_object_pos = wp.to_torch(target_object.data.root_pos_w) - env.scene.env_origins # object to target object x_separation = torch.abs(object_pos[:, 0] - target_object_pos[:, 0]) @@ -123,7 +121,7 @@ def lift_object_il_success( """ object_instance: RigidObject = env.scene[object_cfg.name] - object_pos = object_instance.data.root_pos_w + object_pos = wp.to_torch(object_instance.data.root_pos_w) goal_pos = torch.tensor([goal_position] * env.num_envs, device=env.device) @@ -160,7 +158,7 @@ def lift_object_rl_success( return torch.zeros(env.num_envs, dtype=torch.bool, device=env.device) object_instance: RigidObject = env.scene[object_cfg.name] - object_pos = object_instance.data.root_pos_w + object_pos = wp.to_torch(object_instance.data.root_pos_w) # Try to get goal position from command manager command = env.command_manager.get_command(command_name) @@ -178,7 +176,7 @@ def goal_pose_task_termination( target_x_range: tuple[float, float] | None = None, target_y_range: tuple[float, float] | None = None, target_z_range: tuple[float, float] | None = None, - target_orientation_wxyz: tuple[float, float, float, float] | None = None, + target_orientation_xyzw: tuple[float, float, float, float] | None = None, target_orientation_tolerance_rad: float = 0.1, ) -> torch.Tensor: """Terminate when the object's pose is within the thresholds (BBox + Orientation). @@ -189,15 +187,15 @@ def goal_pose_task_termination( target_x_range: Success zone x-range [min, max] in meters. target_y_range: Success zone y-range [min, max] in meters. target_z_range: Success zone z-range [min, max] in meters. - target_orientation_wxyz: Target quaternion [w, x, y, z]. + target_orientation_xyzw: Target quaternion [x, y, z, w]. target_orientation_tolerance_rad: Angular tolerance in radians (default: 0.1). Returns: A boolean tensor of shape (num_envs, ) """ object_instance: RigidObject = env.scene[object_cfg.name] - object_root_pos_w = object_instance.data.root_pos_w - object_root_quat_w = object_instance.data.root_quat_w + object_root_pos_w = wp.to_torch(object_instance.data.root_pos_w) + object_root_quat_w = wp.to_torch(object_instance.data.root_quat_w) device = env.device num_envs = env.num_envs @@ -206,7 +204,7 @@ def goal_pose_task_termination( target_x_range is not None, target_y_range is not None, target_z_range is not None, - target_orientation_wxyz is not None, + target_orientation_xyzw is not None, ]) if not has_any_threshold: @@ -223,8 +221,8 @@ def goal_pose_task_termination( success &= in_range # Orientation check - if target_orientation_wxyz is not None: - target_quat = torch.tensor(target_orientation_wxyz, device=device, dtype=torch.float32).unsqueeze(0) + if target_orientation_xyzw is not None: + target_quat = torch.tensor(target_orientation_xyzw, device=device, dtype=torch.float32).unsqueeze(0) # Formula: || > cos(tolerance / 2) quat_dot = torch.sum(object_root_quat_w * target_quat, dim=-1) diff --git a/isaaclab_arena/terms/articulations.py b/isaaclab_arena/terms/articulations.py index 6e47f82e2..3c620a06e 100644 --- a/isaaclab_arena/terms/articulations.py +++ b/isaaclab_arena/terms/articulations.py @@ -6,6 +6,7 @@ from __future__ import annotations import torch +import warp as wp from typing import TYPE_CHECKING from isaaclab.managers import SceneEntityCfg @@ -22,4 +23,4 @@ def joint_acc(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(" """ # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] - return asset.data.joint_acc[:, asset_cfg.joint_ids] + return wp.to_torch(asset.data.joint_acc)[:, asset_cfg.joint_ids] diff --git a/isaaclab_arena/terms/events.py b/isaaclab_arena/terms/events.py index 538a332c4..a9edad779 100644 --- a/isaaclab_arena/terms/events.py +++ b/isaaclab_arena/terms/events.py @@ -4,6 +4,7 @@ # SPDX-License-Identifier: Apache-2.0 import torch +import warp as wp from isaaclab.envs import ManagerBasedEnv from isaaclab.managers import SceneEntityCfg @@ -23,11 +24,11 @@ def set_object_pose( asset = env.scene[asset_cfg.name] num_envs = len(env_ids) # Convert the pose to the env frame - pose_t_xyz_q_wxyz = pose.to_tensor(device=env.device).repeat(num_envs, 1) - pose_t_xyz_q_wxyz[:, :3] += env.scene.env_origins[env_ids] + pose_t_xyz_q_xyzw = pose.to_tensor(device=env.device).repeat(num_envs, 1) + pose_t_xyz_q_xyzw[:, :3] += env.scene.env_origins[env_ids] # Set the pose and velocity - asset.write_root_pose_to_sim(pose_t_xyz_q_wxyz, env_ids=env_ids) - asset.write_root_velocity_to_sim(torch.zeros(1, 6, device=env.device), env_ids=env_ids) + asset.write_root_pose_to_sim(pose_t_xyz_q_xyzw, env_ids=env_ids) + asset.write_root_velocity_to_sim(torch.zeros(num_envs, 6, device=env.device), env_ids=env_ids) def set_object_pose_per_env( @@ -47,10 +48,26 @@ def set_object_pose_per_env( for cur_env in env_ids.tolist(): # Convert the pose to the env frame pose = pose_list[cur_env] - pose_t_xyz_q_wxyz = pose.to_tensor(device=env.device) - pose_t_xyz_q_wxyz[:3] += env.scene.env_origins[cur_env, :].squeeze() + pose_t_xyz_q_xyzw = pose.to_tensor(device=env.device).unsqueeze(0) + pose_t_xyz_q_xyzw[0, :3] += env.scene.env_origins[cur_env, :] # Set the pose and velocity - asset.write_root_pose_to_sim(pose_t_xyz_q_wxyz, env_ids=torch.tensor([cur_env], device=env.device)) + asset.write_root_pose_to_sim(pose_t_xyz_q_xyzw, env_ids=torch.tensor([cur_env], device=env.device)) asset.write_root_velocity_to_sim( torch.zeros(1, 6, device=env.device), env_ids=torch.tensor([cur_env], device=env.device) ) + + +def reset_all_articulation_joints(env: ManagerBasedEnv, env_ids: torch.Tensor): + """Reset the articulation joints to the initial state.""" + for articulation_asset in env.scene.articulations.values(): + # obtain default and deal with the offset for env origins + default_root_state = wp.to_torch(articulation_asset.data.default_root_state)[env_ids].clone() + default_root_state[:, 0:3] += env.scene.env_origins[env_ids] + # set into the physics simulation + articulation_asset.write_root_pose_to_sim(default_root_state[:, :7], env_ids=env_ids) + articulation_asset.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids=env_ids) + # obtain default joint positions + default_joint_pos = wp.to_torch(articulation_asset.data.default_joint_pos)[env_ids].clone() + default_joint_vel = wp.to_torch(articulation_asset.data.default_joint_vel)[env_ids].clone() + # set into the physics simulation + articulation_asset.write_joint_state_to_sim(default_joint_pos, default_joint_vel, env_ids=env_ids) diff --git a/isaaclab_arena/terms/transforms.py b/isaaclab_arena/terms/transforms.py index d722081e7..56e67957b 100644 --- a/isaaclab_arena/terms/transforms.py +++ b/isaaclab_arena/terms/transforms.py @@ -6,6 +6,7 @@ from __future__ import annotations import torch +import warp as wp from typing import TYPE_CHECKING import isaaclab.utils.math as PoseUtils @@ -31,8 +32,8 @@ def transform_pose_from_world_to_target_frame( target_frame_name in asset.data.body_names ), f"Target frame {target_frame_name} not found in asset {asset_cfg.name}" - target_link_pose_w = asset.data.body_link_state_w[:, asset.data.body_names.index(target_link_name), :] - target_frame_pose_w = asset.data.body_link_state_w[:, asset.data.body_names.index(target_frame_name), :] + target_link_pose_w = wp.to_torch(asset.data.body_link_state_w)[:, asset.data.body_names.index(target_link_name), :] + target_frame_pose_w = wp.to_torch(asset.data.body_link_state_w)[:, asset.data.body_names.index(target_frame_name), :] # Convert to pose matrix target_link_position_w = target_link_pose_w[:, :3] @@ -98,7 +99,7 @@ def get_asset_position( ) -> torch.Tensor: """Get the robot position.""" asset: Articulation = env.scene[asset_cfg.name] - return asset.data.root_pos_w + return wp.to_torch(asset.data.root_pos_w) def get_asset_quaternion( @@ -107,4 +108,4 @@ def get_asset_quaternion( ) -> torch.Tensor: """Get the robot quaternion.""" asset: Articulation = env.scene[asset_cfg.name] - return asset.data.root_quat_w + return wp.to_torch(asset.data.root_quat_w) diff --git a/isaaclab_arena/tests/test_achieve_cube_goal_pose.py b/isaaclab_arena/tests/test_achieve_cube_goal_pose.py index 1d63ca5c9..45cead657 100644 --- a/isaaclab_arena/tests/test_achieve_cube_goal_pose.py +++ b/isaaclab_arena/tests/test_achieve_cube_goal_pose.py @@ -3,14 +3,20 @@ # # SPDX-License-Identifier: Apache-2.0 +import traceback + import gymnasium as gym import torch +import warp as wp from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function NUM_STEPS = 10 HEADLESS = True +TARGET_ORIENTATION_XYZW = (0.0, 0.0, 0.7071, 0.7071) +TARGET_ORIENTATION_TOLERANCE_RAD = 0.2 +TARGET_Z_RANGE = (0.0, 0.5) def get_test_environment(num_envs: int): """Returns a scene which we use for these tests.""" @@ -36,7 +42,7 @@ def get_test_environment(num_envs: int): dex_cube.set_initial_pose( Pose( position_xyz=(0.1, 0.0, 0.05), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) @@ -45,16 +51,16 @@ def get_test_environment(num_envs: int): # Define success thresholds: z in [0.0, 0.5] and yaw 90 degrees task = GoalPoseTask( dex_cube, - target_z_range=(0.0, 0.5), - target_orientation_wxyz=(0.7071, 0.0, 0.0, 0.7071), # yaw 90 degrees - target_orientation_tolerance_rad=0.2, + target_z_range=TARGET_Z_RANGE, + target_orientation_xyzw=TARGET_ORIENTATION_XYZW, + target_orientation_tolerance_rad=TARGET_ORIENTATION_TOLERANCE_RAD, ) embodiment = FrankaEmbodiment() embodiment.set_initial_pose( Pose( position_xyz=(-0.4, 0.0, 0.0), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) @@ -95,6 +101,7 @@ def assert_not_success(env: ManagerBasedEnv, terminated: torch.Tensor): except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: @@ -121,7 +128,7 @@ def _test_achieve_cube_goal_pose_success(simulation_app) -> bool: # - Position: z > 0.2 (in success zone) # - Orientation: yaw 90 degrees (0.7071, 0, 0, 0.7071) target_pos = torch.tensor([[0.3, 0.0, 0.05]], device=env.device) # z=0.5 is in [0.2, 1.0] - target_quat = torch.tensor([[0.7071, 0.0, 0.0, 0.7071]], device=env.device) # yaw 90 degrees + target_quat = torch.tensor([TARGET_ORIENTATION_XYZW], device=env.device) # yaw 90 degrees # Step the environment to let the physics settle for _ in range(NUM_STEPS): @@ -144,6 +151,7 @@ def _test_achieve_cube_goal_pose_success(simulation_app) -> bool: except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: @@ -171,11 +179,11 @@ def _test_achieve_cube_goal_pose_multiple_envs(simulation_app) -> bool: step_zeros_and_call(env, 1) # Move only the first env's cube to target pose - current_poses = cube_object.data.root_state_w.clone() + current_poses = wp.to_torch(cube_object.data.root_state_w).clone() # Set first env to success pose target_pos_0 = env.scene.env_origins[0] + torch.tensor([0.1, 0.0, 0.5], device=env.device) - target_quat = torch.tensor([0.7071, 0.0, 0.0, 0.7071], device=env.device) + target_quat = torch.tensor([TARGET_ORIENTATION_XYZW], device=env.device) new_poses = current_poses.clone() new_poses[0, :3] = target_pos_0 @@ -193,7 +201,7 @@ def _test_achieve_cube_goal_pose_multiple_envs(simulation_app) -> bool: assert not terminated[1].item(), "Second env should not be successful" # Now move second env to success pose too - current_poses = cube_object.data.root_state_w.clone() + current_poses = wp.to_torch(cube_object.data.root_state_w).clone() target_pos_1 = env.scene.env_origins[1] + torch.tensor([0.1, 0.0, 0.5], device=env.device) new_poses = current_poses.clone() @@ -216,6 +224,7 @@ def _test_achieve_cube_goal_pose_multiple_envs(simulation_app) -> bool: except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: diff --git a/isaaclab_arena/tests/test_assembly_task.py b/isaaclab_arena/tests/test_assembly_task.py index cd1e7b54f..b170c1a39 100644 --- a/isaaclab_arena/tests/test_assembly_task.py +++ b/isaaclab_arena/tests/test_assembly_task.py @@ -6,8 +6,7 @@ import gymnasium as gym import torch - -import pytest +import traceback from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function @@ -31,19 +30,18 @@ def get_peg_insert_test_environment(num_envs: int, remove_events: bool = False): args_parser = get_isaaclab_arena_cli_parser() args_cli = args_parser.parse_args(["--num_envs", str(num_envs)]) - args_cli.enable_pinocchio = False asset_registry = AssetRegistry() # Create scene assets background = asset_registry.get_asset_by_name("table")() - background.set_initial_pose(Pose(position_xyz=(0.55, 0.0, 0.0), rotation_wxyz=(0.707, 0, 0, 0.707))) + background.set_initial_pose(Pose(position_xyz=(0.55, 0.0, 0.0), rotation_xyzw=(0, 0, 0.707, 0.707))) peg = asset_registry.get_asset_by_name("peg")() - peg.set_initial_pose(Pose(position_xyz=(0.45, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + peg.set_initial_pose(Pose(position_xyz=(0.45, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) hole = asset_registry.get_asset_by_name("hole")() - hole.set_initial_pose(Pose(position_xyz=(0.45, 0.1, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + hole.set_initial_pose(Pose(position_xyz=(0.45, 0.1, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) light_spawner_cfg = sim_utils.DomeLightCfg(color=(0.75, 0.75, 0.75), intensity=1500.0) light = asset_registry.get_asset_by_name("light")(spawner_cfg=light_spawner_cfg) @@ -102,25 +100,24 @@ def get_gear_mesh_test_environment(num_envs: int, remove_events: bool = False): args_parser = get_isaaclab_arena_cli_parser() args_cli = args_parser.parse_args(["--num_envs", str(num_envs)]) - args_cli.enable_pinocchio = False asset_registry = AssetRegistry() # Create scene assets background = asset_registry.get_asset_by_name("table")() - background.set_initial_pose(Pose(position_xyz=(0.55, 0.0, 0.0), rotation_wxyz=(0.707, 0, 0, 0.707))) + background.set_initial_pose(Pose(position_xyz=(0.55, 0.0, 0.0), rotation_xyzw=(0, 0, 0.707, 0.707))) gear_base = asset_registry.get_asset_by_name("gear_base")() - gear_base.set_initial_pose(Pose(position_xyz=(0.6, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + gear_base.set_initial_pose(Pose(position_xyz=(0.6, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) medium_gear = asset_registry.get_asset_by_name("medium_gear")() - medium_gear.set_initial_pose(Pose(position_xyz=(0.5, 0.2, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + medium_gear.set_initial_pose(Pose(position_xyz=(0.5, 0.2, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) small_gear = asset_registry.get_asset_by_name("small_gear")() - small_gear.set_initial_pose(Pose(position_xyz=(0.6, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + small_gear.set_initial_pose(Pose(position_xyz=(0.6, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) large_gear = asset_registry.get_asset_by_name("large_gear")() - large_gear.set_initial_pose(Pose(position_xyz=(0.6, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + large_gear.set_initial_pose(Pose(position_xyz=(0.6, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) light_spawner_cfg = sim_utils.DomeLightCfg(color=(0.75, 0.75, 0.75), intensity=1500.0) light = asset_registry.get_asset_by_name("light")(spawner_cfg=light_spawner_cfg) @@ -197,7 +194,6 @@ def assert_assembled(env: ManagerBasedEnv, terminated: torch.Tensor): except Exception as e: print(f"Error: {e}") - import traceback traceback.print_exc() return False @@ -240,7 +236,6 @@ def assert_assembled(env: ManagerBasedEnv, terminated: torch.Tensor): except Exception as e: print(f"Error: {e}") - import traceback traceback.print_exc() return False @@ -277,7 +272,6 @@ def _test_peg_insert_assembly_multi(simulation_app) -> bool: except Exception as e: print(f"Error: {e}") - import traceback traceback.print_exc() return False @@ -314,7 +308,6 @@ def _test_gear_mesh_assembly_multi(simulation_app) -> bool: except Exception as e: print(f"Error: {e}") - import traceback traceback.print_exc() return False @@ -346,8 +339,6 @@ def _test_peg_insert_initialization(simulation_app) -> bool: except Exception as e: print(f"Error: {e}") - import traceback - traceback.print_exc() return False @@ -380,7 +371,6 @@ def _test_gear_mesh_initialization(simulation_app) -> bool: except Exception as e: print(f"Error: {e}") - import traceback traceback.print_exc() return False @@ -389,14 +379,8 @@ def _test_gear_mesh_initialization(simulation_app) -> bool: # Test functions that will be called by pytest -@pytest.mark.skip( - reason=( - "Requires enable_pinocchio=False. Run separately: pytest -sv" - " isaaclab_arena/tests/test_assembly_task.py::test_peg_insert_assembly_single" - ) -) def test_peg_insert_assembly_single(): - result = run_simulation_app_function(_test_peg_insert_assembly_single, headless=HEADLESS, enable_pinocchio=False) + result = run_simulation_app_function(_test_peg_insert_assembly_single, headless=HEADLESS) assert result, f"Test {_test_peg_insert_assembly_single.__name__} failed" @@ -405,14 +389,8 @@ def test_gear_mesh_assembly_single(): assert result, f"Test {_test_gear_mesh_assembly_single.__name__} failed" -@pytest.mark.skip( - reason=( - "Requires enable_pinocchio=False. Run separately: pytest -sv" - " isaaclab_arena/tests/test_assembly_task.py::test_peg_insert_assembly_multi" - ) -) def test_peg_insert_assembly_multi(): - result = run_simulation_app_function(_test_peg_insert_assembly_multi, headless=HEADLESS, enable_pinocchio=False) + result = run_simulation_app_function(_test_peg_insert_assembly_multi, headless=HEADLESS) assert result, f"Test {_test_peg_insert_assembly_multi.__name__} failed" @@ -421,17 +399,8 @@ def test_gear_mesh_assembly_multi(): assert result, f"Test {_test_gear_mesh_assembly_multi.__name__} failed" -@pytest.mark.skip( - reason=( - "Requires enable_pinocchio=False. Run separately: pytest -sv" - " isaaclab_arena/tests/test_assembly_task.py::test_peg_insert_initialization" - ) -) def test_peg_insert_initialization(): - """ - For peg insert task, we need to test the task with pinocchio disabled due to the "peg" and "hole" assets are not compatible with pinocchio. - """ - result = run_simulation_app_function(_test_peg_insert_initialization, headless=HEADLESS, enable_pinocchio=False) + result = run_simulation_app_function(_test_peg_insert_initialization, headless=HEADLESS) assert result, f"Test {_test_peg_insert_initialization.__name__} failed" @@ -441,15 +410,9 @@ def test_gear_mesh_initialization(): if __name__ == "__main__": - """ - Peg insert tests are commented out because they require enable_pinocchio=False, - but the current test session's SimulationApp was initialized with enable_pinocchio=True. - Due to limitations in subprocess.py, the SimulationApp cannot be restarted with different - parameters during a single pytest session. Run peg insert tests separately with: - pytest -sv isaaclab_arena/tests/test_assembly_task.py::test_peg_insert_assembly_single --disable_pinocchio - pytest -sv isaaclab_arena/tests/test_assembly_task.py::test_peg_insert_assembly_multi --disable_pinocchio - pytest -sv isaaclab_arena/tests/test_assembly_task.py::test_peg_insert_initialization --disable_pinocchio - """ + test_peg_insert_initialization() + test_peg_insert_assembly_single() + test_peg_insert_assembly_multi() test_gear_mesh_initialization() test_gear_mesh_assembly_single() test_gear_mesh_assembly_multi() diff --git a/isaaclab_arena/tests/test_asset_registry.py b/isaaclab_arena/tests/test_asset_registry.py index 41cf0c363..1677c8ec4 100644 --- a/isaaclab_arena/tests/test_asset_registry.py +++ b/isaaclab_arena/tests/test_asset_registry.py @@ -63,23 +63,18 @@ def _test_all_assets_in_registry(simulation_app): objects_in_registry: list[Object] = [] for idx, asset_cls in enumerate(asset_registry.get_assets_by_tag("object")): asset = asset_cls() - # Skip "peg" and "hole" assets due to enable_pinocchio requirement mismatch. - # These IsaacLab factory assembly assets require enable_pinocchio=False, but the current - # SimulationApp session was initialized with enable_pinocchio=True. The app cannot be - # restarted mid-session due to subprocess.py limitations. - if asset.name not in ("peg", "hole"): - # Set their pose - pose = Pose( - position_xyz=( - first_position[0] + (idx + 1) * OBJECT_SEPARATION, - first_position[1], - first_position[2], - ), - rotation_wxyz=(1, 0, 0, 0), - ) - asset.set_initial_pose(pose) - objects_in_registry.append(asset) - objects_in_registry_names.append(asset.name) + # Set their pose + pose = Pose( + position_xyz=( + first_position[0] + (idx + 1) * OBJECT_SEPARATION, + first_position[1], + first_position[2], + ), + rotation_xyzw=(0, 0, 0, 1), + ) + asset.set_initial_pose(pose) + objects_in_registry.append(asset) + objects_in_registry_names.append(asset.name) # Add lights for asset_cls in asset_registry.get_assets_by_tag("light"): asset = asset_cls() diff --git a/isaaclab_arena/tests/test_camera_observation.py b/isaaclab_arena/tests/test_camera_observation.py index fdb7b45ec..e197fd887 100644 --- a/isaaclab_arena/tests/test_camera_observation.py +++ b/isaaclab_arena/tests/test_camera_observation.py @@ -35,7 +35,7 @@ def _test_camera_observation(simulation_app) -> bool: cracker_box.set_initial_pose( Pose( position_xyz=(0.0758066475391388, -0.5088448524475098, 0.0), - rotation_wxyz=(1, 0, 0, 0), + rotation_xyzw=(0, 0, 0, 1), ) ) @@ -51,7 +51,6 @@ def _test_camera_observation(simulation_app) -> bool: builder = ArenaEnvBuilder(isaaclab_arena_environment, args_cli) env = builder.make_registered() env.reset() - for _ in tqdm.tqdm(range(NUM_STEPS)): with torch.inference_mode(): actions = torch.zeros(env.action_space.shape, device=env.device) diff --git a/isaaclab_arena/tests/test_close_door.py b/isaaclab_arena/tests/test_close_door.py index dfafc0f3a..3bff6f912 100644 --- a/isaaclab_arena/tests/test_close_door.py +++ b/isaaclab_arena/tests/test_close_door.py @@ -5,6 +5,7 @@ import gymnasium as gym import torch +import traceback from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function @@ -35,7 +36,7 @@ def get_test_environment(remove_reset_door_state_event: bool, num_envs: int): microwave.set_initial_pose( Pose( position_xyz=(0.6, -0.00586, 0.22773), - rotation_wxyz=(0.7071068, 0, 0, -0.7071068), + rotation_xyzw=(0, 0, -0.7071068, 0.7071068), ) ) @@ -224,8 +225,6 @@ def _test_close_door_with_reset(simulation_app) -> bool: except Exception as e: print(f"Error: {e}") - import traceback - traceback.print_exc() return False diff --git a/isaaclab_arena/tests/test_configclass.py b/isaaclab_arena/tests/test_configclass.py index 9a6406d00..a95dc2b23 100644 --- a/isaaclab_arena/tests/test_configclass.py +++ b/isaaclab_arena/tests/test_configclass.py @@ -3,12 +3,11 @@ # # SPDX-License-Identifier: Apache-2.0 -from isaaclab.utils import configclass - -from isaaclab_arena.utils.configclass import combine_configclasses - def test_combine_configclasses_with_multiple_inheritance(): + from isaaclab.utils import configclass + + from isaaclab_arena.utils.configclass import combine_configclasses # Side A - A class with a base class @configclass @@ -38,6 +37,9 @@ class BarCfg(FooCfgBase): def test_combine_configclasses_with_inheritance(): + from isaaclab.utils import configclass + + from isaaclab_arena.utils.configclass import combine_configclasses # Side A - A class with a base class @configclass @@ -65,6 +67,9 @@ class BarCfg: def test_combine_configclasses_with_post_init(): + from isaaclab.utils import configclass + + from isaaclab_arena.utils.configclass import combine_configclasses # Side A - A class with a base class @configclass @@ -89,3 +94,10 @@ def __post_init__(self): assert CombinedCfg().a == 2 assert CombinedCfg().b == 3 assert CombinedCfg().c == 4 + + +if __name__ == "__main__": + test_combine_configclasses_with_multiple_inheritance() + test_combine_configclasses_with_inheritance() + test_combine_configclasses_with_post_init() + diff --git a/isaaclab_arena/tests/test_contact_sensor_not_at_root.py b/isaaclab_arena/tests/test_contact_sensor_not_at_root.py index 3caf2a120..da3eaf88f 100644 --- a/isaaclab_arena/tests/test_contact_sensor_not_at_root.py +++ b/isaaclab_arena/tests/test_contact_sensor_not_at_root.py @@ -40,7 +40,7 @@ def get_test_environment(num_envs: int): sweet_potato = asset_registry.get_asset_by_name("sweet_potato")( initial_pose=Pose( position_xyz=(0.0758066475391388, -0.5088448524475098, 0.5), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) @@ -96,6 +96,7 @@ def _test_contact_sensor_not_at_root(simulation_app) -> bool: except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: diff --git a/isaaclab_arena/tests/test_detect_object_type.py b/isaaclab_arena/tests/test_detect_object_type.py index fa3d341ef..477bd0075 100644 --- a/isaaclab_arena/tests/test_detect_object_type.py +++ b/isaaclab_arena/tests/test_detect_object_type.py @@ -140,6 +140,7 @@ def _test_auto_object_type(simulation_app): except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: diff --git a/isaaclab_arena/tests/test_duplicate_asset.py b/isaaclab_arena/tests/test_duplicate_asset.py index caddb2ead..e74f1a59b 100644 --- a/isaaclab_arena/tests/test_duplicate_asset.py +++ b/isaaclab_arena/tests/test_duplicate_asset.py @@ -36,7 +36,7 @@ def get_test_environment(num_envs: int, position_1: tuple[float, float, float], dex_cube_1.set_initial_pose( Pose( position_xyz=position_1, - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) @@ -45,7 +45,7 @@ def get_test_environment(num_envs: int, position_1: tuple[float, float, float], dex_cube_2.set_initial_pose( Pose( position_xyz=position_2, - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) @@ -55,7 +55,7 @@ def get_test_environment(num_envs: int, position_1: tuple[float, float, float], embodiment.set_initial_pose( Pose( position_xyz=(-0.4, 0.0, 0.0), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) @@ -118,6 +118,7 @@ def _test_duplicate_asset(simulation_app) -> bool: except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: diff --git a/isaaclab_arena/tests/test_events.py b/isaaclab_arena/tests/test_events.py index da28b336a..f0473aec5 100644 --- a/isaaclab_arena/tests/test_events.py +++ b/isaaclab_arena/tests/test_events.py @@ -59,8 +59,8 @@ def _test_set_object_pose_per_env_event(simulation_app): # - from: constant per env, # - to: per env pose pose_list = [ - Pose(position_xyz=(0.4, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)), - Pose(position_xyz=(0.4, 0.4, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)), + Pose(position_xyz=(0.4, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)), + Pose(position_xyz=(0.4, 0.4, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)), ] env_cfg.events.reset_pick_up_object_pose = EventTermCfg( func=set_object_pose_per_env, @@ -99,6 +99,7 @@ def _test_set_object_pose_per_env_event(simulation_app): except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: diff --git a/isaaclab_arena/tests/test_g1_wbc_embodiment.py b/isaaclab_arena/tests/test_g1_wbc_embodiment.py index 02d6a1bfb..9be1d50bb 100644 --- a/isaaclab_arena/tests/test_g1_wbc_embodiment.py +++ b/isaaclab_arena/tests/test_g1_wbc_embodiment.py @@ -6,6 +6,8 @@ import numpy as np import torch import tqdm +import traceback +import warp as wp import pytest @@ -15,30 +17,31 @@ HEADLESS = True ENABLE_CAMERAS = True STANDING_POSITION_XY_EPS = 1e-1 +# Wrist quaternions in xyzw (identity-like: x, y, z, w) WBC_PINK_IDLE_ACTION = [ - 0.0, - 0.0, + 0.0, # left_hand_state + 0.0, # right_hand_state 0.201, 0.145, - 0.101, - 1.000, + 0.101, # left_wrist_pos 0.010, -0.008, -0.011, + 1.000, # left_wrist_quat (xyzw) 0.201, -0.145, - 0.101, - 1.000, + 0.101, # right_wrist_pos -0.010, -0.008, -0.011, + 1.000, # right_wrist_quat (xyzw) 0.0, 0.0, - 0.0, - 0.75, - 0.0, + 0.0, # navigate_cmd + 0.75, # base_height_cmd 0.0, 0.0, + 0.0, # torso_orientation_rpy_cmd ] @@ -63,7 +66,7 @@ def get_test_environment(num_envs: int, pink_ik_enabled: bool): embodiment = G1WBCJointEmbodiment(enable_cameras=ENABLE_CAMERAS) # NOTE(xinjieyao, 2025.09.22): Set initial pose such that robot will not drop to the ground, causing WBC unstable. robot_init_base_pose = np.array([0, 0, 0]) - embodiment.set_initial_pose(Pose(position_xyz=tuple(robot_init_base_pose), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + embodiment.set_initial_pose(Pose(position_xyz=tuple(robot_init_base_pose), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) isaaclab_arena_environment = IsaacLabArenaEnvironment( name="g1_standing_test", @@ -102,7 +105,7 @@ def _test_wbc_joint_standing_idle_actions(simulation_app) -> bool: def assert_standing_idle(env: ManagerBasedEnv, robot_init_base_pose: np.ndarray): # get robot base pose after actions call - robot_base_pose = env.scene["robot"].data.root_link_pose_w[0, :3].cpu().numpy() + robot_base_pose = wp.to_torch(env.scene["robot"].data.root_link_pose_w)[0, :3].cpu().numpy() # check if robot base pose is close to initial base pose robot_xy_error = np.linalg.norm(robot_base_pose[:2] - robot_init_base_pose[:2]) assert robot_xy_error < STANDING_POSITION_XY_EPS, "Robot moved away from initial position." @@ -113,6 +116,7 @@ def assert_standing_idle(env: ManagerBasedEnv, robot_init_base_pose: np.ndarray) except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: @@ -139,8 +143,8 @@ def _test_wbc_pink_standing_idle_actions(simulation_app) -> bool: env, robot_init_base_pose = get_test_environment(num_envs=1, pink_ik_enabled=True) def assert_standing_idle(env: ManagerBasedEnv, robot_init_base_pose: np.ndarray): - # get robot base pose after actions call - robot_base_pose = env.scene["robot"].data.root_link_pose_w[0, :3].cpu().numpy() + # get robot base pose after actions call (use wp.to_torch like joint test; root_link_pose_w is a Warp array) + robot_base_pose = wp.to_torch(env.scene["robot"].data.root_link_pose_w)[0, :3].cpu().numpy() # check if robot base pose is close to initial base pose robot_xy_error = np.linalg.norm(robot_base_pose[:2] - robot_init_base_pose[:2]) assert robot_xy_error < STANDING_POSITION_XY_EPS, "Robot moved away from initial position." diff --git a/isaaclab_arena/tests/test_g1_wbc_pink_preprocess_actions.py b/isaaclab_arena/tests/test_g1_wbc_pink_preprocess_actions.py deleted file mode 100644 index e918bf2e7..000000000 --- a/isaaclab_arena/tests/test_g1_wbc_pink_preprocess_actions.py +++ /dev/null @@ -1,239 +0,0 @@ -# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Unit tests for G1 WBC Pink action preprocess_actions (world → robot base frame).""" - -import torch - -from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function -from isaaclab_arena_g1.g1_whole_body_controller.wbc_policy.policy.action_constants import ( - BASE_HEIGHT_CMD_START_IDX, - LEFT_HAND_STATE_IDX, - LEFT_WRIST_POS_END_IDX, - LEFT_WRIST_POS_START_IDX, - LEFT_WRIST_QUAT_END_IDX, - LEFT_WRIST_QUAT_START_IDX, - NAVIGATE_CMD_END_IDX, - NAVIGATE_CMD_START_IDX, - RIGHT_HAND_STATE_IDX, - RIGHT_WRIST_POS_END_IDX, - RIGHT_WRIST_POS_START_IDX, - RIGHT_WRIST_QUAT_END_IDX, - RIGHT_WRIST_QUAT_START_IDX, - TORSO_ORIENTATION_RPY_CMD_END_IDX, - TORSO_ORIENTATION_RPY_CMD_START_IDX, -) - -HEADLESS = True - - -def _get_g1_pink_env_and_term(simulation_app): - """Build G1 WBC Pink env at origin with identity orientation; return env and g1_action term.""" - from isaaclab_arena.assets.asset_registry import AssetRegistry - from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser - from isaaclab_arena.embodiments.g1.g1 import G1WBCPinkEmbodiment - from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder - from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment - from isaaclab_arena.scene.scene import Scene - from isaaclab_arena.utils.pose import Pose - - asset_registry = AssetRegistry() - background = asset_registry.get_asset_by_name("kitchen")() - scene = Scene(assets=[background]) - embodiment = G1WBCPinkEmbodiment(enable_cameras=False) - embodiment.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) - isaaclab_arena_environment = IsaacLabArenaEnvironment( - name="g1_pink_preprocess_test", - embodiment=embodiment, - scene=scene, - ) - args_cli = get_isaaclab_arena_cli_parser().parse_args([]) - env_builder = ArenaEnvBuilder(isaaclab_arena_environment, args_cli) - env = env_builder.make_registered() - env.reset() - term = env.unwrapped.action_manager.get_term("g1_action") - return env, term - - -def _test_preprocess_actions_shape(simulation_app) -> bool: - """preprocess_actions preserves shape (num_envs, action_dim).""" - env, term = _get_g1_pink_env_and_term(simulation_app) - try: - action_dim = term.action_dim - num_envs = env.num_envs - actions = torch.zeros(num_envs, action_dim, device=env.unwrapped.device) - out = term.preprocess_actions(actions) - assert out.shape == (num_envs, action_dim), f"Expected shape ({num_envs}, {action_dim}), got {out.shape}" - finally: - env.close() - return True - - -def _test_preprocess_actions_identity_base(simulation_app) -> bool: - """When robot base has identity quat, wrist in base frame = world pos minus base pos.""" - env, term = _get_g1_pink_env_and_term(simulation_app) - try: - device = env.unwrapped.device - action_dim = term.action_dim - robot_base_pos = term._asset.data.root_link_pos_w[0, :3] - - # World-frame wrist positions: base + offset (so base-frame offset is known) - left_offset = torch.tensor([1.0, 2.0, 3.0], device=device) - right_offset = torch.tensor([4.0, 5.0, 6.0], device=device) - left_pos_world = robot_base_pos + left_offset - right_pos_world = robot_base_pos + right_offset - - actions = torch.zeros(1, action_dim, device=device) - actions[0, LEFT_WRIST_POS_START_IDX:LEFT_WRIST_POS_END_IDX] = left_pos_world - actions[0, LEFT_WRIST_QUAT_START_IDX:LEFT_WRIST_QUAT_END_IDX] = torch.tensor( - [1.0, 0.0, 0.0, 0.0], device=device - ) - actions[0, RIGHT_WRIST_POS_START_IDX:RIGHT_WRIST_POS_END_IDX] = right_pos_world - actions[0, RIGHT_WRIST_QUAT_START_IDX:RIGHT_WRIST_QUAT_END_IDX] = torch.tensor( - [1.0, 0.0, 0.0, 0.0], device=device - ) - - out = term.preprocess_actions(actions) - - # Base frame position = world - base (in world), then rotated by base_inv => offset when base quat is identity - torch.testing.assert_close( - out[0, LEFT_WRIST_POS_START_IDX:LEFT_WRIST_POS_END_IDX], left_offset, atol=1e-4, rtol=0 - ) - torch.testing.assert_close( - out[0, RIGHT_WRIST_POS_START_IDX:RIGHT_WRIST_POS_END_IDX], right_offset, atol=1e-4, rtol=0 - ) - torch.testing.assert_close( - out[0, LEFT_WRIST_QUAT_START_IDX:LEFT_WRIST_QUAT_END_IDX], - actions[0, LEFT_WRIST_QUAT_START_IDX:LEFT_WRIST_QUAT_END_IDX], - atol=1e-5, - rtol=0, - ) - torch.testing.assert_close( - out[0, RIGHT_WRIST_QUAT_START_IDX:RIGHT_WRIST_QUAT_END_IDX], - actions[0, RIGHT_WRIST_QUAT_START_IDX:RIGHT_WRIST_QUAT_END_IDX], - atol=1e-5, - rtol=0, - ) - finally: - env.close() - return True - - -def _test_preprocess_actions_roundtrip(simulation_app) -> bool: - """Preprocess world→base; then base→world recovers original (using current robot pose).""" - import isaaclab.utils.math as math_utils - - env, term = _get_g1_pink_env_and_term(simulation_app) - try: - device = env.unwrapped.device - action_dim = term.action_dim - asset = term._asset - - robot_base_pos = asset.data.root_link_pos_w[:, :3] - robot_base_quat = asset.data.root_link_quat_w - num_envs = robot_base_pos.shape[0] - - # Arbitrary world-frame wrist poses - left_pos_w = torch.tensor([[1.0, 0.0, 0.5]], device=device).expand(num_envs, 3) - left_quat_w = torch.tensor([[1.0, 0.0, 0.0, 0.0]], device=device).expand(num_envs, 4) - right_pos_w = torch.tensor([[0.0, 1.0, 0.5]], device=device).expand(num_envs, 3) - right_quat_w = torch.tensor([[1.0, 0.0, 0.0, 0.0]], device=device).expand(num_envs, 4) - - actions = torch.zeros(num_envs, action_dim, device=device) - actions[:, LEFT_WRIST_POS_START_IDX:LEFT_WRIST_POS_END_IDX] = left_pos_w - actions[:, LEFT_WRIST_QUAT_START_IDX:LEFT_WRIST_QUAT_END_IDX] = left_quat_w - actions[:, RIGHT_WRIST_POS_START_IDX:RIGHT_WRIST_POS_END_IDX] = right_pos_w - actions[:, RIGHT_WRIST_QUAT_START_IDX:RIGHT_WRIST_QUAT_END_IDX] = right_quat_w - - out = term.preprocess_actions(actions) - left_pos_b = out[:, LEFT_WRIST_POS_START_IDX:LEFT_WRIST_POS_END_IDX] - left_quat_b = out[:, LEFT_WRIST_QUAT_START_IDX:LEFT_WRIST_QUAT_END_IDX] - right_pos_b = out[:, RIGHT_WRIST_POS_START_IDX:RIGHT_WRIST_POS_END_IDX] - right_quat_b = out[:, RIGHT_WRIST_QUAT_START_IDX:RIGHT_WRIST_QUAT_END_IDX] - - # Base → world: pos_w = base_pos + quat_apply(base_quat, pos_b), quat_w = quat_mul(base_quat, quat_b) - left_pos_w_recovered = robot_base_pos + math_utils.quat_apply(robot_base_quat, left_pos_b) - left_quat_w_recovered = math_utils.quat_mul(robot_base_quat, left_quat_b) - right_pos_w_recovered = robot_base_pos + math_utils.quat_apply(robot_base_quat, right_pos_b) - right_quat_w_recovered = math_utils.quat_mul(robot_base_quat, right_quat_b) - - torch.testing.assert_close(left_pos_w_recovered, left_pos_w, atol=1e-5, rtol=0) - torch.testing.assert_close(left_quat_w_recovered, left_quat_w, atol=1e-5, rtol=0) - torch.testing.assert_close(right_pos_w_recovered, right_pos_w, atol=1e-5, rtol=0) - torch.testing.assert_close(right_quat_w_recovered, right_quat_w, atol=1e-5, rtol=0) - finally: - env.close() - return True - - -def _test_preprocess_actions_does_not_mutate_other_slots(simulation_app) -> bool: - """Indices outside wrist pos/quat (e.g. hand state, navigate_cmd) are unchanged.""" - env, term = _get_g1_pink_env_and_term(simulation_app) - try: - device = env.unwrapped.device - action_dim = term.action_dim - actions = torch.zeros(1, action_dim, device=device) - actions[0, LEFT_HAND_STATE_IDX] = 0.5 - actions[0, RIGHT_HAND_STATE_IDX] = 0.7 - actions[0, NAVIGATE_CMD_START_IDX:NAVIGATE_CMD_END_IDX] = torch.tensor([0.1, 0.2, 0.3], device=device) - actions[0, BASE_HEIGHT_CMD_START_IDX] = 0.75 - actions[0, TORSO_ORIENTATION_RPY_CMD_START_IDX:TORSO_ORIENTATION_RPY_CMD_END_IDX] = torch.tensor( - [0.0, 0.0, 0.1], device=device - ) - - out = term.preprocess_actions(actions) - - torch.testing.assert_close(out[0, LEFT_HAND_STATE_IDX], torch.tensor(0.5, device=device), atol=1e-6, rtol=0) - torch.testing.assert_close(out[0, RIGHT_HAND_STATE_IDX], torch.tensor(0.7, device=device), atol=1e-6, rtol=0) - torch.testing.assert_close( - out[0, NAVIGATE_CMD_START_IDX:NAVIGATE_CMD_END_IDX], - actions[0, NAVIGATE_CMD_START_IDX:NAVIGATE_CMD_END_IDX], - atol=1e-6, - rtol=0, - ) - torch.testing.assert_close( - out[0, BASE_HEIGHT_CMD_START_IDX], actions[0, BASE_HEIGHT_CMD_START_IDX], atol=1e-6, rtol=0 - ) - torch.testing.assert_close( - out[0, TORSO_ORIENTATION_RPY_CMD_START_IDX:TORSO_ORIENTATION_RPY_CMD_END_IDX], - actions[0, TORSO_ORIENTATION_RPY_CMD_START_IDX:TORSO_ORIENTATION_RPY_CMD_END_IDX], - atol=1e-6, - rtol=0, - ) - finally: - env.close() - return True - - -def test_g1_wbc_pink_preprocess_actions_shape(): - result = run_simulation_app_function( - _test_preprocess_actions_shape, - headless=HEADLESS, - ) - assert result, "preprocess_actions shape test failed" - - -def test_g1_wbc_pink_preprocess_actions_identity_base(): - result = run_simulation_app_function( - _test_preprocess_actions_identity_base, - headless=HEADLESS, - ) - assert result, "preprocess_actions identity base test failed" - - -def test_g1_wbc_pink_preprocess_actions_roundtrip(): - result = run_simulation_app_function( - _test_preprocess_actions_roundtrip, - headless=HEADLESS, - ) - assert result, "preprocess_actions roundtrip test failed" - - -def test_g1_wbc_pink_preprocess_actions_does_not_mutate_other_slots(): - result = run_simulation_app_function( - _test_preprocess_actions_does_not_mutate_other_slots, - headless=HEADLESS, - ) - assert result, "preprocess_actions other slots test failed" diff --git a/isaaclab_arena/tests/test_galbot_embodiment.py b/isaaclab_arena/tests/test_galbot_embodiment.py index c1ff8083b..2cfd38f71 100644 --- a/isaaclab_arena/tests/test_galbot_embodiment.py +++ b/isaaclab_arena/tests/test_galbot_embodiment.py @@ -5,6 +5,7 @@ import numpy as np import torch +import warp as wp import tqdm from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function @@ -41,7 +42,7 @@ def get_galbot_test_environment(num_envs: int = 1): pick_up_object.set_initial_pose( Pose( position_xyz=(0.4, 0.0, 0.1), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) @@ -56,7 +57,7 @@ def get_galbot_test_environment(num_envs: int = 1): # Setup Galbot embodiment robot_init_position = (-0.4, 0.0, -0.5) embodiment = GalbotEmbodiment(arm_mode=ArmMode.LEFT) - embodiment.set_initial_pose(Pose(position_xyz=robot_init_position, rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + embodiment.set_initial_pose(Pose(position_xyz=robot_init_position, rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) # Create scene with kitchen, object, and destination scene = Scene(assets=[background, pick_up_object, destination_location]) @@ -88,7 +89,7 @@ def _test_galbot_initial_position(simulation_app) -> bool: env.step(actions) # Check the robot ended up at the correct position - robot_position = env.scene["robot"].data.root_link_pose_w[0, :3].cpu().numpy() + robot_position = wp.to_torch(env.scene["robot"].data.root_link_pose_w)[0, :3].cpu().numpy() robot_position_error = np.linalg.norm(robot_position - np.array(robot_init_position)) print(f"Robot position error: {robot_position_error}") assert robot_position_error < INITIAL_POSITION_EPS, "Galbot ended up at the wrong position." @@ -155,13 +156,13 @@ def _test_galbot_observation_config(simulation_app) -> bool: assert robot_data is not None, "Robot data should be accessible" # Check joint positions are valid - joint_pos = robot_data.joint_pos + joint_pos = wp.to_torch(robot_data.joint_pos) print(f"Joint positions shape: {joint_pos.shape}") assert joint_pos is not None, "Joint positions should be accessible" assert not torch.any(torch.isnan(joint_pos)), "Joint positions should not contain NaN" # Check joint velocities are valid - joint_vel = robot_data.joint_vel + joint_vel = wp.to_torch(robot_data.joint_vel) print(f"Joint velocities shape: {joint_vel.shape}") assert joint_vel is not None, "Joint velocities should be accessible" assert not torch.any(torch.isnan(joint_vel)), "Joint velocities should not contain NaN" @@ -218,7 +219,7 @@ def _test_galbot_arm_reaches_goal(simulation_app) -> bool: ee_frame = env.scene["ee_frame"] # Get initial ee position (in world frame, relative to env origin) - initial_ee_pos = ee_frame.data.target_pos_w[0, 0, :] - env.scene.env_origins[0] + initial_ee_pos = wp.to_torch(ee_frame.data.target_pos_w)[0, 0, :] - env.scene.env_origins[0] print(f"Initial EE position: {initial_ee_pos.cpu().numpy()}") print(f"Target position: {target_position.cpu().numpy()}") @@ -227,7 +228,7 @@ def _test_galbot_arm_reaches_goal(simulation_app) -> bool: num_reach_steps = 200 for step in range(num_reach_steps): # Get current ee position - current_ee_pos = ee_frame.data.target_pos_w[0, 0, :] - env.scene.env_origins[0] + current_ee_pos = wp.to_torch(ee_frame.data.target_pos_w)[0, 0, :] - env.scene.env_origins[0] remaining_displacement = target_position - current_ee_pos # Proportional control: move a fraction of the remaining distance @@ -244,7 +245,7 @@ def _test_galbot_arm_reaches_goal(simulation_app) -> bool: env.step(action) # Get final ee position - final_ee_pos = ee_frame.data.target_pos_w[0, 0, :] - env.scene.env_origins[0] + final_ee_pos = wp.to_torch(ee_frame.data.target_pos_w)[0, 0, :] - env.scene.env_origins[0] position_error = torch.norm(final_ee_pos - target_position).item() print(f"Final EE position: {final_ee_pos.cpu().numpy()}") diff --git a/isaaclab_arena/tests/test_no_collision_loss.py b/isaaclab_arena/tests/test_no_collision_loss.py index f91058d6c..af380180e 100644 --- a/isaaclab_arena/tests/test_no_collision_loss.py +++ b/isaaclab_arena/tests/test_no_collision_loss.py @@ -35,7 +35,7 @@ def _create_table() -> DummyObject: def _create_no_collision_scene() -> tuple[DummyObject, DummyObject, DummyObject]: """Create table + two boxes with On(table) and NoCollision between boxes (for solver tests).""" table = _create_table() - table.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + table.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) table.add_relation(IsAnchor()) box_a = _create_box("box_a") box_b = _create_box("box_b") diff --git a/isaaclab_arena/tests/test_object_configuration.py b/isaaclab_arena/tests/test_object_configuration.py index 9a0ba982d..ca51dc685 100644 --- a/isaaclab_arena/tests/test_object_configuration.py +++ b/isaaclab_arena/tests/test_object_configuration.py @@ -20,7 +20,7 @@ def _test_object_initial_pose_update(simulation_app): rigid_object.object_cfg.debug_vis = False # Now lets add an initial pose to the object. - new_initial_pose = Pose(position_xyz=(5.0, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)) + new_initial_pose = Pose(position_xyz=(5.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)) rigid_object.set_initial_pose(new_initial_pose) # Now lets check that the initial pose has been updated and that the debug visualization is still disabled. diff --git a/isaaclab_arena/tests/test_object_of_type_base.py b/isaaclab_arena/tests/test_object_of_type_base.py index 671b591b5..13d978333 100644 --- a/isaaclab_arena/tests/test_object_of_type_base.py +++ b/isaaclab_arena/tests/test_object_of_type_base.py @@ -5,6 +5,7 @@ import torch import tqdm +import traceback from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function @@ -46,7 +47,7 @@ def __init__(self, prim_path: str = default_prim_path, initial_pose: Pose | None cone = ConeNoPhysics() # Put the thing in the center of the room floating. - cone.set_initial_pose(Pose(position_xyz=(-1.6, 0.0, 1.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + cone.set_initial_pose(Pose(position_xyz=(-1.6, 0.0, 1.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) scene = Scene(assets=[background, cone]) isaaclab_arena_environment = IsaacLabArenaEnvironment( @@ -77,6 +78,7 @@ def __init__(self, prim_path: str = default_prim_path, initial_pose: Pose | None except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: diff --git a/isaaclab_arena/tests/test_object_on_termination.py b/isaaclab_arena/tests/test_object_on_termination.py index c08185d06..f738d358e 100644 --- a/isaaclab_arena/tests/test_object_on_termination.py +++ b/isaaclab_arena/tests/test_object_on_termination.py @@ -6,6 +6,7 @@ import os import torch import tqdm +import warp as wp from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function @@ -43,11 +44,12 @@ def _test_object_on_destination_termination(simulation_app) -> bool: cracker_box.set_initial_pose( Pose( position_xyz=(0.0758066475391388, -0.5088448524475098, 0.5), - rotation_wxyz=(1, 0, 0, 0), + rotation_xyzw=(0, 0, 0, 1), ) ) scene = Scene(assets=[background, cracker_box, destination_location]) + # scene = Scene(assets=[background, cracker_box]) isaaclab_arena_environment = IsaacLabArenaEnvironment( name="kitchen", @@ -72,12 +74,16 @@ def _test_object_on_destination_termination(simulation_app) -> bool: sensor = env.scene.sensors["pick_up_object_contact_sensor"] for _ in tqdm.tqdm(range(NUM_STEPS)): with torch.inference_mode(): + print(f"About to create zero actions") actions = torch.zeros(env.action_space.shape, device=env.unwrapped.device) + print(f"About to step the environment") _, _, terminated, _, _ = env.step(actions) + # env.step(actions) + print(f"Environment stepped") # Get the force on the pick up object. - forces_vec.append(sensor.data.net_forces_w.clone()) - force_matrix_vec.append(sensor.data.force_matrix_w.clone()) - velocities_vec.append(env.scene[cracker_box.name].data.root_lin_vel_w.clone()) + forces_vec.append(wp.to_torch(sensor.data.net_forces_w)) + force_matrix_vec.append(wp.to_torch(sensor.data.force_matrix_w)) + velocities_vec.append(wp.to_torch(env.scene[cracker_box.name].data.root_lin_vel_w)) # Try the termination. condition_met_vec.append( @@ -91,6 +97,7 @@ def _test_object_on_destination_termination(simulation_app) -> bool: except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: diff --git a/isaaclab_arena/tests/test_object_placer_reproducibility.py b/isaaclab_arena/tests/test_object_placer_reproducibility.py index 81daa8bf6..b5bc657ab 100644 --- a/isaaclab_arena/tests/test_object_placer_reproducibility.py +++ b/isaaclab_arena/tests/test_object_placer_reproducibility.py @@ -5,23 +5,19 @@ """Tests for ObjectPlacer and RelationSolver reproducibility.""" -from isaaclab_arena.assets.dummy_object import DummyObject -from isaaclab_arena.relations.object_placer import ObjectPlacer -from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams -from isaaclab_arena.relations.relation_solver import RelationSolver -from isaaclab_arena.relations.relation_solver_params import RelationSolverParams -from isaaclab_arena.relations.relations import IsAnchor, NextTo, On, Side -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox, get_random_pose_within_bounding_box -from isaaclab_arena.utils.pose import Pose - -def _create_test_objects() -> tuple[DummyObject, DummyObject, DummyObject]: +def _create_test_objects(): """Create test objects with relations (without setting initial poses for non-anchors).""" + from isaaclab_arena.assets.dummy_object import DummyObject + from isaaclab_arena.relations.relations import IsAnchor, NextTo, On, Side + from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + from isaaclab_arena.utils.pose import Pose + desk = DummyObject( name="desk", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), ) - desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) desk.add_relation(IsAnchor()) box1 = DummyObject( @@ -42,6 +38,8 @@ def _create_test_objects() -> tuple[DummyObject, DummyObject, DummyObject]: def test_get_random_pose_same_seed_produces_identical_result(): """Test that get_random_pose_within_bounding_box with same seed produces identical poses.""" + from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox, get_random_pose_within_bounding_box + bbox = AxisAlignedBoundingBox(min_point=(-1.0, -1.0, 0.0), max_point=(1.0, 1.0, 1.0)) pose1 = get_random_pose_within_bounding_box(bbox, seed=42) @@ -52,6 +50,9 @@ def test_get_random_pose_same_seed_produces_identical_result(): def test_relation_solver_same_inputs_produces_identical_result(): """Test that RelationSolver with identical initial positions produces identical results.""" + from isaaclab_arena.relations.relation_solver import RelationSolver + from isaaclab_arena.relations.relation_solver_params import RelationSolverParams + desk_pos = (0.0, 0.0, 0.0) fixed_box1_pos = (0.5, 0.5, 0.5) fixed_box2_pos = (0.3, 0.7, 0.3) @@ -80,6 +81,10 @@ def test_relation_solver_same_inputs_produces_identical_result(): def test_object_placer_same_seed_produces_identical_result(): """Test that ObjectPlacer with same seed produces identical final results.""" + from isaaclab_arena.relations.object_placer import ObjectPlacer + from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams + from isaaclab_arena.relations.relation_solver_params import RelationSolverParams + seed = 42 solver_params = RelationSolverParams(max_iters=10) @@ -104,6 +109,10 @@ def test_object_placer_same_seed_produces_identical_result(): def test_object_placer_different_seeds_produce_different_results(): """Test that ObjectPlacer with different seeds produces different results.""" + from isaaclab_arena.relations.object_placer import ObjectPlacer + from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams + from isaaclab_arena.relations.relation_solver_params import RelationSolverParams + solver_params = RelationSolverParams(max_iters=10) # Run 1 with seed 42 diff --git a/isaaclab_arena/tests/test_object_pose_randomization.py b/isaaclab_arena/tests/test_object_pose_randomization.py index f7bea4013..d75c69983 100644 --- a/isaaclab_arena/tests/test_object_pose_randomization.py +++ b/isaaclab_arena/tests/test_object_pose_randomization.py @@ -5,6 +5,7 @@ import torch import tqdm +import traceback from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function @@ -90,6 +91,7 @@ def _test_object_pose_randomization(simulation_app): except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: diff --git a/isaaclab_arena/tests/test_object_set.py b/isaaclab_arena/tests/test_object_set.py index f55a4e6e6..ae452d114 100644 --- a/isaaclab_arena/tests/test_object_set.py +++ b/isaaclab_arena/tests/test_object_set.py @@ -3,162 +3,177 @@ # # SPDX-License-Identifier: Apache-2.0 +import traceback +import os + from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function HEADLESS = True -NUM_ENVS = 3 +NUM_ENVS = 10 OBJECT_SET_1_PRIM_PATH = "/World/envs/env_.*/ObjectSet_1" OBJECT_SET_2_PRIM_PATH = "/World/envs/env_.*/ObjectSet_2" -OBJECT_SET_JUG_PRIM_PATH = "/World/envs/env_.*/ObjectSet_Jug" -OBJECT_SET_BOTTLES_PRIM_PATH = "/World/envs/env_.*/ObjectSet_Bottles" -def _build_and_reset_env(simulation_app, scene_assets, env_name="object_set_test", task=None): - """Build arena env with given scene and optional task, then reset. Returns env (caller must close).""" +def _test_empty_object_set(simulation_app): + from isaaclab_arena.assets.object_set import RigidObjectSet + + try: + RigidObjectSet(name="empty_object_set", objects=[]) + except Exception: + return True + return False + + +def _test_articulation_object_set(simulation_app): from isaaclab_arena.assets.asset_registry import AssetRegistry - from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser - from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder - from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment - from isaaclab_arena.scene.scene import Scene + from isaaclab_arena.assets.object_set import RigidObjectSet asset_registry = AssetRegistry() - embodiment = asset_registry.get_asset_by_name("franka")() - scene = Scene(assets=scene_assets) - isaaclab_arena_environment = IsaacLabArenaEnvironment( - name=env_name, - embodiment=embodiment, - scene=scene, - task=task, - teleop_device=None, - ) - args_cli = get_isaaclab_arena_cli_parser().parse_args([]) - args_cli.num_envs = NUM_ENVS - args_cli.headless = HEADLESS - env_builder = ArenaEnvBuilder(isaaclab_arena_environment, args_cli) - env = env_builder.make_registered() - env.reset() - return env + microwave = asset_registry.get_asset_by_name("microwave")() + try: + RigidObjectSet(name="articulation_object_set", objects=[microwave]) + except Exception: + return True + return False -def _run_pick_and_place_object_set_test( - simulation_app, - obj_set, - object_set_prim_path, - path_contains, - initial_pose=None, -): - """Build env with one object set and PickAndPlaceTask, run common assertions, close. path_contains: str or list[str] of length NUM_ENVS.""" +def _test_single_object_in_one_object_set(simulation_app): from isaacsim.core.utils.stage import get_current_stage from isaaclab_arena.assets.asset_registry import AssetRegistry from isaaclab_arena.assets.object_reference import ObjectReference + from isaaclab_arena.assets.object_set import RigidObjectSet + from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment + from isaaclab_arena.scene.scene import Scene from isaaclab_arena.tasks.pick_and_place_task import PickAndPlaceTask + from isaaclab_arena.utils.pose import Pose from isaaclab_arena.utils.usd_helpers import get_asset_usd_path_from_prim_path asset_registry = AssetRegistry() background = asset_registry.get_asset_by_name("kitchen")() + embodiment = asset_registry.get_asset_by_name("franka")() + cracker_box = asset_registry.get_asset_by_name("cracker_box")() destination_location = ObjectReference( name="destination_location", prim_path="{ENV_REGEX_NS}/kitchen/Cabinet_B_02", parent_asset=background, ) - if initial_pose is not None: - obj_set.set_initial_pose(initial_pose) - scene_assets = [background, obj_set] - task = PickAndPlaceTask( - pick_up_object=obj_set, - destination_location=destination_location, - background_scene=background, + obj_set = RigidObjectSet( + name="single_object_set", objects=[cracker_box, cracker_box], prim_path=OBJECT_SET_1_PRIM_PATH ) - env = _build_and_reset_env( - simulation_app, - scene_assets, - env_name="pick_and_place_object_set_test", - task=task, + obj_set.set_initial_pose(Pose(position_xyz=(0.1, 0.0, 0.1), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) + scene = Scene(assets=[background, obj_set]) + isaaclab_arena_environment = IsaacLabArenaEnvironment( + name="single_object_set_test", + embodiment=embodiment, + scene=scene, + task=PickAndPlaceTask( + pick_up_object=obj_set, destination_location=destination_location, background_scene=background + ), + teleop_device=None, ) + args_cli = get_isaaclab_arena_cli_parser().parse_args([]) + args_cli.num_envs = NUM_ENVS + args_cli.headless = HEADLESS + env_builder = ArenaEnvBuilder(isaaclab_arena_environment, args_cli) + env = env_builder.make_registered() + env.reset() + try: - if isinstance(path_contains, str): - path_contains = [path_contains] * NUM_ENVS for i in range(NUM_ENVS): + # Construct the actual prim path for this environment path = get_asset_usd_path_from_prim_path( - prim_path=object_set_prim_path.replace(".*", str(i)), - stage=get_current_stage(), + prim_path=OBJECT_SET_1_PRIM_PATH.replace(".*", str(i)), stage=get_current_stage() ) assert path is not None, "Path is None" - assert path_contains[i] in path, f"Path does not contain {path_contains[i]!r}: {path}" - if initial_pose is not None: + assert "cracker_box.usd" in path, "Path does not contain cracker_box.usd" assert obj_set.get_initial_pose() is not None, "Initial pose is None" + assert env.scene[obj_set.name].data.root_pose_w is not None, "Root pose is None" assert ( env.scene.sensors["pick_up_object_contact_sensor"].data.force_matrix_w is not None ), "Contact sensor data is None" - return True except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: env.close() + return True -def _test_empty_object_set(simulation_app): - from isaaclab_arena.assets.object_set import RigidObjectSet - - try: - RigidObjectSet(name="empty_object_set", objects=[]) - except Exception: - return True - return False - - -def _test_articulation_object_set(simulation_app): - from isaaclab_arena.assets.asset_registry import AssetRegistry - from isaaclab_arena.assets.object_set import RigidObjectSet - - asset_registry = AssetRegistry() - microwave = asset_registry.get_asset_by_name("microwave")() - try: - RigidObjectSet(name="articulation_object_set", objects=[microwave]) - except Exception: - return True - return False - +def _test_multi_objects_in_one_object_set(simulation_app): + from isaacsim.core.utils.stage import get_current_stage -def _test_single_object_in_one_object_set(simulation_app): from isaaclab_arena.assets.asset_registry import AssetRegistry + from isaaclab_arena.assets.object_reference import ObjectReference from isaaclab_arena.assets.object_set import RigidObjectSet - from isaaclab_arena.utils.pose import Pose + from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment + from isaaclab_arena.scene.scene import Scene + from isaaclab_arena.tasks.pick_and_place_task import PickAndPlaceTask + from isaaclab_arena.utils.usd_helpers import get_asset_usd_path_from_prim_path asset_registry = AssetRegistry() + background = asset_registry.get_asset_by_name("kitchen")() + embodiment = asset_registry.get_asset_by_name("franka")() cracker_box = asset_registry.get_asset_by_name("cracker_box")() + sugar_box = asset_registry.get_asset_by_name("sugar_box")() + destination_location = ObjectReference( + name="destination_location", + prim_path="{ENV_REGEX_NS}/kitchen/Cabinet_B_02", + parent_asset=background, + ) obj_set = RigidObjectSet( - name="single_object_set", objects=[cracker_box, cracker_box], prim_path=OBJECT_SET_1_PRIM_PATH + name="multi_object_sets", objects=[cracker_box, sugar_box], prim_path=OBJECT_SET_2_PRIM_PATH ) - return _run_pick_and_place_object_set_test( - simulation_app, - obj_set, - OBJECT_SET_1_PRIM_PATH, - path_contains="cracker_box.usd", - initial_pose=Pose(position_xyz=(0.1, 0.0, 0.1), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)), + scene = Scene(assets=[background, obj_set]) + isaaclab_arena_environment = IsaacLabArenaEnvironment( + name="multi_objects_in_one_object_set_test", + embodiment=embodiment, + scene=scene, + task=PickAndPlaceTask( + pick_up_object=obj_set, destination_location=destination_location, background_scene=background + ), + teleop_device=None, ) + args_cli = get_isaaclab_arena_cli_parser().parse_args([]) + args_cli.num_envs = NUM_ENVS + args_cli.headless = HEADLESS + env_builder = ArenaEnvBuilder(isaaclab_arena_environment, args_cli) + env = env_builder.make_registered() + env.reset() + assert env.scene[obj_set.name].data.root_pose_w is not None, "Root pose is None" + assert ( + env.scene.sensors["pick_up_object_contact_sensor"].data.force_matrix_w is not None + ), "Contact sensor data is None" -def _test_multi_objects_in_one_object_set(simulation_app): - from isaaclab_arena.assets.asset_registry import AssetRegistry - from isaaclab_arena.assets.object_set import RigidObjectSet + # replace * in OBJECT_SET_PRIM_PATH with env_index + object_paths = [] + try: + for i in range(NUM_ENVS): - asset_registry = AssetRegistry() - cracker_box = asset_registry.get_asset_by_name("cracker_box")() - sugar_box = asset_registry.get_asset_by_name("sugar_box")() - obj_set = RigidObjectSet( - name="multi_object_sets", objects=[cracker_box, sugar_box], prim_path=OBJECT_SET_2_PRIM_PATH - ) - path_contains = ["cracker_box.usd" if i % 2 == 0 else "sugar_box.usd" for i in range(NUM_ENVS)] - return _run_pick_and_place_object_set_test( - simulation_app, - obj_set, - OBJECT_SET_2_PRIM_PATH, - path_contains=path_contains, - ) + path = get_asset_usd_path_from_prim_path( + prim_path=OBJECT_SET_2_PRIM_PATH.replace(".*", str(i)), stage=get_current_stage() + ) + assert path is not None, "Path is None" + object_paths.append(path) + assert len(object_paths) == NUM_ENVS, "Object_paths length is not equal to NUM_ENVS" + # We check the file names instead of the paths because objects may be cached + object_file_names = [os.path.basename(path) for path in object_paths] + assert os.path.basename(cracker_box.usd_path) in object_file_names, "Cracker box USD path is not in Object_paths" + assert os.path.basename(sugar_box.usd_path) in object_file_names, "Sugar box USD path is not in Object_paths" + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + finally: + env.close() + return True def _test_multi_object_sets(simulation_app): @@ -166,10 +181,15 @@ def _test_multi_object_sets(simulation_app): from isaaclab_arena.assets.asset_registry import AssetRegistry from isaaclab_arena.assets.object_set import RigidObjectSet + from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment + from isaaclab_arena.scene.scene import Scene from isaaclab_arena.utils.usd_helpers import get_asset_usd_path_from_prim_path asset_registry = AssetRegistry() background = asset_registry.get_asset_by_name("packing_table")() + embodiment = asset_registry.get_asset_by_name("franka")() cracker_box = asset_registry.get_asset_by_name("cracker_box")() sugar_box = asset_registry.get_asset_by_name("sugar_box")() mustard_bottle = asset_registry.get_asset_by_name("mustard_bottle")() @@ -180,84 +200,57 @@ def _test_multi_object_sets(simulation_app): obj_set_2 = RigidObjectSet( name="multi_object_sets_2", objects=[sugar_box, mustard_bottle], prim_path=OBJECT_SET_2_PRIM_PATH ) - env = _build_and_reset_env( - simulation_app, - [background, obj_set_1, obj_set_2], - env_name="multi_object_sets_test", - task=None, + scene = Scene(assets=[background, obj_set_1, obj_set_2]) + isaaclab_arena_environment = IsaacLabArenaEnvironment( + name="multi_object_sets_test", + embodiment=embodiment, + scene=scene, ) + args_cli = get_isaaclab_arena_cli_parser().parse_args([]) + args_cli.num_envs = NUM_ENVS + args_cli.headless = HEADLESS + env_builder = ArenaEnvBuilder(isaaclab_arena_environment, args_cli) + env = env_builder.make_registered() + env.reset() + try: + object_1_paths = [] + object_2_paths = [] for i in range(NUM_ENVS): + path_1 = get_asset_usd_path_from_prim_path( prim_path=OBJECT_SET_1_PRIM_PATH.replace(".*", str(i)), stage=get_current_stage() ) path_2 = get_asset_usd_path_from_prim_path( prim_path=OBJECT_SET_2_PRIM_PATH.replace(".*", str(i)), stage=get_current_stage() ) + object_1_paths.append(path_1) + object_2_paths.append(path_2) assert path_1 is not None, ( "Path_1 from Prim Path " + OBJECT_SET_1_PRIM_PATH.replace(".*", str(i)) + " is None" ) assert path_2 is not None, ( "Path_2 from Prim Path " + OBJECT_SET_2_PRIM_PATH.replace(".*", str(i)) + " is None" ) - if i % 2 == 0: - assert "cracker_box.usd" in path_1, "Path_1 does not contain cracker_box.usd for env index " + str(i) - assert "sugar_box.usd" in path_2, "Path_2 does not contain sugar_box.usd for env index " + str(i) - else: - assert "sugar_box.usd" in path_1, "Path_1 does not contain sugar_box.usd for env index " + str(i) - assert ( - "mustard_bottle.usd" in path_2 - ), "Path_2 does not contain mustard_bottle.usd for env index " + str(i) - return True + assert len(object_1_paths) == NUM_ENVS, "Object_1_paths length is not equal to NUM_ENVS" + assert len(object_2_paths) == NUM_ENVS, "Object_2_paths length is not equal to NUM_ENVS" + # Check that each object in the set turns up in one of the environments + # NOTE(alexmillane): If we get really unlucky, this can fail because every environment + # gets the same object. The chance of this is 0.5^NUM_ENVS. So with 20 envs this is very small. + # NOTE(alexmillane): We check the file names instead of the paths because objects may be cached + object_1_file_names = [os.path.basename(path) for path in object_1_paths] + object_2_file_names = [os.path.basename(path) for path in object_2_paths] + assert os.path.basename(cracker_box.usd_path) in object_1_file_names, "Cracker box USD path is not in Object_1_paths" + assert os.path.basename(sugar_box.usd_path) in object_1_file_names, "Sugar box USD path is not in Object_1_paths" + assert os.path.basename(sugar_box.usd_path) in object_2_file_names, "Sugar box USD path is not in Object_2_paths" + assert os.path.basename(mustard_bottle.usd_path) in object_2_file_names, "Mustard bottle USD path is not in Object_2_paths" except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: env.close() - - -def _test_object_set_with_jug(simulation_app): - """Test object set with Jug asset (depth-1 rigid body); exercises cache pipeline and contact sensor.""" - from isaaclab_arena.assets.asset_registry import AssetRegistry - from isaaclab_arena.assets.object_set import RigidObjectSet - from isaaclab_arena.utils.pose import Pose - - asset_registry = AssetRegistry() - jug = asset_registry.get_asset_by_name("jug")() - obj_set = RigidObjectSet( - name="ObjectSet_Jug", - objects=[jug, jug], - prim_path=OBJECT_SET_JUG_PRIM_PATH, - ) - return _run_pick_and_place_object_set_test( - simulation_app, - obj_set, - OBJECT_SET_JUG_PRIM_PATH, - path_contains="jug", - initial_pose=Pose(position_xyz=(0.1, 0.0, 0.1), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)), - ) - - -def _test_object_set_with_ranch_and_bbq_bottles(simulation_app): - """Test object set with ranch_dressing_bottle and bbq_sauce_bottle.""" - from isaaclab_arena.assets.asset_registry import AssetRegistry - from isaaclab_arena.assets.object_set import RigidObjectSet - - asset_registry = AssetRegistry() - ranch_dressing_bottle = asset_registry.get_asset_by_name("ranch_dressing_hope_robolab")() - bbq_sauce_bottle = asset_registry.get_asset_by_name("bbq_sauce_bottle_hope_robolab")() - obj_set = RigidObjectSet( - name="ObjectSet_Bottles", - objects=[ranch_dressing_bottle, bbq_sauce_bottle], - prim_path=OBJECT_SET_BOTTLES_PRIM_PATH, - ) - path_contains = ["ranch_dressing" if i % 2 == 0 else "bbq_sauce_bottle_hope_robolab" for i in range(NUM_ENVS)] - return _run_pick_and_place_object_set_test( - simulation_app, - obj_set, - OBJECT_SET_BOTTLES_PRIM_PATH, - path_contains=path_contains, - ) + return True def test_empty_object_set(): @@ -300,27 +293,9 @@ def test_multi_object_sets(): assert result, f"Test {_test_multi_object_sets.__name__} failed" -def test_object_set_with_jug(): - result = run_simulation_app_function( - _test_object_set_with_jug, - headless=HEADLESS, - ) - assert result, f"Test {_test_object_set_with_jug.__name__} failed" - - -def test_object_set_with_ranch_and_bbq_bottles(): - result = run_simulation_app_function( - _test_object_set_with_ranch_and_bbq_bottles, - headless=HEADLESS, - ) - assert result, f"Test {_test_object_set_with_ranch_and_bbq_bottles.__name__} failed" - - if __name__ == "__main__": test_empty_object_set() test_articulation_object_set() test_single_object_in_one_object_set() test_multi_objects_in_one_object_set() - test_multi_object_sets() - test_object_set_with_jug() - test_object_set_with_ranch_and_bbq_bottles() + test_multi_object_sets() \ No newline at end of file diff --git a/isaaclab_arena/tests/test_object_set_on_termination.py b/isaaclab_arena/tests/test_object_set_on_termination.py index 335cfcb06..9469fb120 100644 --- a/isaaclab_arena/tests/test_object_set_on_termination.py +++ b/isaaclab_arena/tests/test_object_set_on_termination.py @@ -5,6 +5,7 @@ import torch import tqdm +import traceback from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function @@ -52,7 +53,7 @@ def _test_object_set_on_destination_termination(simulation_app) -> bool: object_set.set_initial_pose( Pose( position_xyz=(0.0758066475391388, -0.5088448524475098, 0.5), - rotation_wxyz=(1, 0, 0, 0), + rotation_xyzw=(0, 0, 0, 1), ) ) @@ -89,6 +90,7 @@ def _test_object_set_on_destination_termination(simulation_app) -> bool: except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: diff --git a/isaaclab_arena/tests/test_open_door.py b/isaaclab_arena/tests/test_open_door.py index 500d84db9..e7f0f2679 100644 --- a/isaaclab_arena/tests/test_open_door.py +++ b/isaaclab_arena/tests/test_open_door.py @@ -5,6 +5,7 @@ import gymnasium as gym import torch +import traceback from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function @@ -35,7 +36,7 @@ def get_test_environment(remove_reset_door_state_event: bool, num_envs: int): microwave.set_initial_pose( Pose( position_xyz=(0.6, -0.00586, 0.22773), - rotation_wxyz=(0.7071068, 0, 0, -0.7071068), + rotation_xyzw=(0, 0, -0.7071068, 0.7071068), ) ) @@ -104,6 +105,7 @@ def assert_open(env: ManagerBasedEnv, terminated: torch.Tensor): except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: @@ -158,6 +160,7 @@ def _test_open_door_microwave_multiple_envs(simulation_app) -> bool: except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: @@ -200,6 +203,7 @@ def _test_open_door_microwave_reset_condition(simulation_app) -> bool: except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: diff --git a/isaaclab_arena/tests/test_place_upright_task.py b/isaaclab_arena/tests/test_place_upright_task.py index bc033c020..badfa6683 100644 --- a/isaaclab_arena/tests/test_place_upright_task.py +++ b/isaaclab_arena/tests/test_place_upright_task.py @@ -3,8 +3,11 @@ # # SPDX-License-Identifier: Apache-2.0 +import pytest + import gymnasium as gym import torch +import traceback from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function @@ -29,11 +32,11 @@ def get_test_environment(dont_reset_placeable_object_pose: bool, num_envs: int): asset_registry = AssetRegistry() background = asset_registry.get_asset_by_name("table")() - background.set_initial_pose(Pose(position_xyz=(0.50, 0.0, 0.625), rotation_wxyz=(0.7071, 0, 0, 0.7071))) + background.set_initial_pose(Pose(position_xyz=(0.50, 0.0, 0.625), rotation_xyzw=(0, 0, 0.7071, 0.7071))) background.object_cfg.spawn.scale = (1.0, 1.0, 0.60) # placeable object must have initial pose set mug = asset_registry.get_asset_by_name("mug")( - initial_pose=Pose(position_xyz=(0.05, 0.0, 0.75), rotation_wxyz=(0.7071, 0.7071, 0.0, 0.0)) + initial_pose=Pose(position_xyz=(0.05, 0.0, 0.75), rotation_xyzw=(0.7071, 0.0, 0.0, 0.7071)) ) if dont_reset_placeable_object_pose: mug.disable_reset_pose() @@ -83,6 +86,7 @@ def assert_upright(env: ManagerBasedEnv, terminated: torch.Tensor): except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: env.close() @@ -136,6 +140,7 @@ def _test_place_upright_mug_multi(simulation_app) -> bool: except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: env.close() @@ -172,6 +177,7 @@ def _test_place_upright_mug_condition(simulation_app) -> bool: except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: env.close() @@ -205,5 +211,5 @@ def test_place_upright_mug_condition(): if __name__ == "__main__": test_place_upright_mug_single() - test_place_upright_mug_multi() - test_place_upright_mug_condition() + # test_place_upright_mug_multi() + # test_place_upright_mug_condition() diff --git a/isaaclab_arena/tests/test_policy_runner.py b/isaaclab_arena/tests/test_policy_runner.py index 8d52549f2..6e85197ae 100644 --- a/isaaclab_arena/tests/test_policy_runner.py +++ b/isaaclab_arena/tests/test_policy_runner.py @@ -3,6 +3,8 @@ # # SPDX-License-Identifier: Apache-2.0 +import pytest + from isaaclab_arena.tests.utils.constants import TestConstants from isaaclab_arena.tests.utils.subprocess import run_subprocess diff --git a/isaaclab_arena/tests/test_pose.py b/isaaclab_arena/tests/test_pose.py index 718fc8c76..e5589accc 100644 --- a/isaaclab_arena/tests/test_pose.py +++ b/isaaclab_arena/tests/test_pose.py @@ -7,10 +7,10 @@ def test_pose_composition(): - T_B_A = Pose(position_xyz=(1.0, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)) - T_C_B = Pose(position_xyz=(2.0, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)) + T_B_A = Pose(position_xyz=(1.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)) + T_C_B = Pose(position_xyz=(2.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)) T_C_A = T_C_B.multiply(T_B_A) assert T_C_A.position_xyz == (3.0, 0.0, 0.0) - assert T_C_A.rotation_wxyz == (1.0, 0.0, 0.0, 0.0) + assert T_C_A.rotation_xyzw == (0.0, 0.0, 0.0, 1.0) diff --git a/isaaclab_arena/tests/test_press_coffee_machine_button.py b/isaaclab_arena/tests/test_press_coffee_machine_button.py index 9ef43e7b0..89f8910f0 100644 --- a/isaaclab_arena/tests/test_press_coffee_machine_button.py +++ b/isaaclab_arena/tests/test_press_coffee_machine_button.py @@ -5,6 +5,7 @@ import gymnasium as gym import torch +import traceback from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function @@ -34,7 +35,7 @@ def get_test_environment(num_envs: int): coffee_machine.set_initial_pose( Pose( position_xyz=(0.6, -0.00586, 0.22773), - rotation_wxyz=(0.7071068, 0, 0, -0.7071068), + rotation_xyzw=(0, 0, -0.7071068, 0.7071068), ) ) @@ -90,6 +91,7 @@ def assert_unpressed(env: ManagerBasedEnv, terminated: torch.Tensor): except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: @@ -155,6 +157,7 @@ def _test_press_button_coffee_machine_multiple_envs(simulation_app) -> bool: except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: diff --git a/isaaclab_arena/tests/test_reference_objects.py b/isaaclab_arena/tests/test_reference_objects.py index 126e5d655..618fe0ab4 100644 --- a/isaaclab_arena/tests/test_reference_objects.py +++ b/isaaclab_arena/tests/test_reference_objects.py @@ -5,8 +5,10 @@ import numpy as np import pathlib +import pytest import torch import tqdm +import traceback from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function from isaaclab_arena.utils.pose import Pose @@ -47,16 +49,16 @@ def get_test_scene(): cracker_box = asset_registry.get_asset_by_name("cracker_box")() microwave = asset_registry.get_asset_by_name("microwave")() - kitchen.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + kitchen.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) cracker_box.set_initial_pose( Pose( - position_xyz=(3.69020713150969, -0.804121657812894, 1.2531903565606817), rotation_wxyz=(1.0, 0.0, 0.0, 0.0) + position_xyz=(3.69020713150969, -0.804121657812894, 1.2531903565606817), rotation_xyzw=(0.0, 0.0, 0.0, 1.0) ) ) microwave.set_initial_pose( Pose( position_xyz=(2.862758610786719, -0.39786255771393336, 1.087924015237011), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) @@ -156,6 +158,7 @@ def close_microwave(): except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: @@ -185,7 +188,7 @@ def _test_reference_objects(simulation_app, tmp_path: pathlib.Path) -> bool: def _test_reference_objects_with_transform(simulation_app, tmp_path: pathlib.Path) -> bool: - background_pose = Pose(position_xyz=(0.772, 3.39, -0.895), rotation_wxyz=(0.70711, 0, 0, -0.70711)) + background_pose = Pose(position_xyz=(0.772, 3.39, -0.895), rotation_xyzw=(0, 0, -0.70711, 0.70711)) return _test_reference_objects_with_background_pose(background_pose, tmp_path) diff --git a/isaaclab_arena/tests/test_relation_loss_strategies.py b/isaaclab_arena/tests/test_relation_loss_strategies.py index d7221d0ea..4daff87fe 100644 --- a/isaaclab_arena/tests/test_relation_loss_strategies.py +++ b/isaaclab_arena/tests/test_relation_loss_strategies.py @@ -9,22 +9,23 @@ import pytest -from isaaclab_arena.assets.dummy_object import DummyObject -from isaaclab_arena.relations.relation_loss_strategies import NextToLossStrategy, OnLossStrategy -from isaaclab_arena.relations.relations import NextTo, On, Side -from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox - -def _create_table() -> DummyObject: +def _create_table(): """Create a table-like object at origin.""" + from isaaclab_arena.assets.dummy_object import DummyObject + from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + return DummyObject( name="table", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(1.0, 1.0, 0.1)), ) -def _create_box() -> DummyObject: +def _create_box(): """Create a small box object.""" + from isaaclab_arena.assets.dummy_object import DummyObject + from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox + return DummyObject( name="box", bounding_box=AxisAlignedBoundingBox(min_point=(0.0, 0.0, 0.0), max_point=(0.2, 0.2, 0.2)), @@ -38,14 +39,14 @@ def _create_box() -> DummyObject: def test_on_loss_strategy_zero_loss_when_perfectly_placed(): """Test that On loss is zero when child is perfectly placed on parent.""" + from isaaclab_arena.relations.relation_loss_strategies import OnLossStrategy + from isaaclab_arena.relations.relations import On + table = _create_table() box = _create_box() relation = On(table, clearance_m=0.01) strategy = OnLossStrategy(slope=10.0) - # Child centered on table, at correct Z height - # Table top = 0.1, clearance = 0.01, box bottom = 0.0 (relative) - # So child_z should be 0.11 for box bottom to be at 0.11 child_pos = torch.tensor([0.4, 0.4, 0.11]) loss = strategy.compute_loss(relation, child_pos, box.bounding_box, table.bounding_box) @@ -54,12 +55,14 @@ def test_on_loss_strategy_zero_loss_when_perfectly_placed(): def test_on_loss_strategy_penalizes_child_outside_x_bounds(): """Test that On loss penalizes child outside parent's X bounds.""" + from isaaclab_arena.relations.relation_loss_strategies import OnLossStrategy + from isaaclab_arena.relations.relations import On + table = _create_table() box = _create_box() relation = On(table, clearance_m=0.01) strategy = OnLossStrategy(slope=10.0) - # Child way to the right (outside table) child_pos = torch.tensor([2.0, 0.4, 0.11]) loss = strategy.compute_loss(relation, child_pos, box.bounding_box, table.bounding_box) @@ -68,12 +71,14 @@ def test_on_loss_strategy_penalizes_child_outside_x_bounds(): def test_on_loss_strategy_penalizes_child_outside_y_bounds(): """Test that On loss penalizes child outside parent's Y bounds.""" + from isaaclab_arena.relations.relation_loss_strategies import OnLossStrategy + from isaaclab_arena.relations.relations import On + table = _create_table() box = _create_box() relation = On(table, clearance_m=0.01) strategy = OnLossStrategy(slope=10.0) - # Child way to the back (outside table) child_pos = torch.tensor([0.4, 2.0, 0.11]) loss = strategy.compute_loss(relation, child_pos, box.bounding_box, table.bounding_box) @@ -82,12 +87,14 @@ def test_on_loss_strategy_penalizes_child_outside_y_bounds(): def test_on_loss_strategy_penalizes_wrong_z_height(): """Test that On loss penalizes incorrect Z height.""" + from isaaclab_arena.relations.relation_loss_strategies import OnLossStrategy + from isaaclab_arena.relations.relations import On + table = _create_table() box = _create_box() relation = On(table, clearance_m=0.01) strategy = OnLossStrategy(slope=10.0) - # Child at wrong Z (floating above table) child_pos = torch.tensor([0.4, 0.4, 0.5]) loss = strategy.compute_loss(relation, child_pos, box.bounding_box, table.bounding_box) @@ -96,14 +103,15 @@ def test_on_loss_strategy_penalizes_wrong_z_height(): def test_on_loss_strategy_respects_clearance(): """Test that On loss accounts for clearance parameter.""" + from isaaclab_arena.relations.relation_loss_strategies import OnLossStrategy + from isaaclab_arena.relations.relations import On + table = _create_table() box = _create_box() clearance = 0.05 relation = On(table, clearance_m=clearance) strategy = OnLossStrategy(slope=10.0) - # Table top = 0.1, with 5cm clearance, box bottom should be at 0.15 - # Box min_point[2] = 0.0, so child_z should be 0.15 child_pos = torch.tensor([0.4, 0.4, 0.15]) loss = strategy.compute_loss(relation, child_pos, box.bounding_box, table.bounding_box) @@ -112,13 +120,16 @@ def test_on_loss_strategy_respects_clearance(): def test_on_loss_strategy_respects_relation_weight(): """Test that On loss is scaled by relation_loss_weight.""" + from isaaclab_arena.relations.relation_loss_strategies import OnLossStrategy + from isaaclab_arena.relations.relations import On + table = _create_table() box = _create_box() relation_normal = On(table, clearance_m=0.01, relation_loss_weight=1.0) relation_double = On(table, clearance_m=0.01, relation_loss_weight=2.0) strategy = OnLossStrategy(slope=10.0) - child_pos = torch.tensor([2.0, 0.4, 0.11]) # Outside bounds + child_pos = torch.tensor([2.0, 0.4, 0.11]) loss_normal = strategy.compute_loss(relation_normal, child_pos, box.bounding_box, table.bounding_box) loss_double = strategy.compute_loss(relation_double, child_pos, box.bounding_box, table.bounding_box) @@ -128,14 +139,14 @@ def test_on_loss_strategy_respects_relation_weight(): def test_on_loss_strategy_constrains_entire_footprint(): """Test that On loss constrains entire child footprint within parent.""" + from isaaclab_arena.relations.relation_loss_strategies import OnLossStrategy + from isaaclab_arena.relations.relations import On + table = _create_table() box = _create_box() # 0.2m wide relation = On(table, clearance_m=0.01) strategy = OnLossStrategy(slope=10.0) - # Child positioned so its right edge just exits the table - # Table X range: [0, 1], box width: 0.2 - # If child_pos_x = 0.9, box right edge = 0.9 + 0.2 = 1.1 (outside!) child_pos = torch.tensor([0.9, 0.4, 0.11]) loss = strategy.compute_loss(relation, child_pos, box.bounding_box, table.bounding_box) @@ -149,16 +160,14 @@ def test_on_loss_strategy_constrains_entire_footprint(): def test_next_to_loss_strategy_zero_loss_when_perfectly_placed(): """Test that NextTo loss is zero when child is perfectly placed.""" + from isaaclab_arena.relations.relation_loss_strategies import NextToLossStrategy + from isaaclab_arena.relations.relations import NextTo, Side + parent_obj = _create_table() child_obj = _create_box() relation = NextTo(parent_obj, side=Side.POSITIVE_X, distance_m=0.05) strategy = NextToLossStrategy(slope=10.0) - # Parent right edge = 0 + 1.0 = 1.0 - # Child left edge should be at 1.0 + 0.05 = 1.05 - # Child min_point[0] = 0.0, so child_pos[0] should be 1.05 - # Y: cross_position_ratio=0.0 (centered). Valid child Y range is [0.0, 0.8] - # (parent [0,1] minus child extent 0.2), so centered target = 0.4 child_pos = torch.tensor([1.05, 0.4, 0.0]) loss = strategy.compute_loss(relation, child_pos, child_obj.bounding_box, parent_obj.bounding_box) @@ -167,12 +176,14 @@ def test_next_to_loss_strategy_zero_loss_when_perfectly_placed(): def test_next_to_loss_strategy_penalizes_wrong_side(): """Test that NextTo loss penalizes child on wrong side of parent.""" + from isaaclab_arena.relations.relation_loss_strategies import NextToLossStrategy + from isaaclab_arena.relations.relations import NextTo, Side + parent_obj = _create_table() child_obj = _create_box() relation = NextTo(parent_obj, side=Side.POSITIVE_X, distance_m=0.05) strategy = NextToLossStrategy(slope=10.0) - # Child on the LEFT of parent (wrong side) child_pos = torch.tensor([-0.5, 0.5, 0.0]) loss = strategy.compute_loss(relation, child_pos, child_obj.bounding_box, parent_obj.bounding_box) @@ -181,12 +192,14 @@ def test_next_to_loss_strategy_penalizes_wrong_side(): def test_next_to_loss_strategy_penalizes_outside_y_band(): """Test that NextTo loss penalizes child outside parent's Y extent.""" + from isaaclab_arena.relations.relation_loss_strategies import NextToLossStrategy + from isaaclab_arena.relations.relations import NextTo, Side + parent_obj = _create_table() child_obj = _create_box() relation = NextTo(parent_obj, side=Side.POSITIVE_X, distance_m=0.05) strategy = NextToLossStrategy(slope=10.0) - # Child at correct X but outside Y range child_pos = torch.tensor([1.05, 2.0, 0.0]) loss = strategy.compute_loss(relation, child_pos, child_obj.bounding_box, parent_obj.bounding_box) @@ -195,12 +208,14 @@ def test_next_to_loss_strategy_penalizes_outside_y_band(): def test_next_to_loss_strategy_penalizes_wrong_distance(): """Test that NextTo loss penalizes incorrect distance from parent.""" + from isaaclab_arena.relations.relation_loss_strategies import NextToLossStrategy + from isaaclab_arena.relations.relations import NextTo, Side + parent_obj = _create_table() child_obj = _create_box() relation = NextTo(parent_obj, side=Side.POSITIVE_X, distance_m=0.05) strategy = NextToLossStrategy(slope=10.0) - # Child too far from parent (0.5m instead of 0.05m) child_pos = torch.tensor([1.5, 0.5, 0.0]) loss = strategy.compute_loss(relation, child_pos, child_obj.bounding_box, parent_obj.bounding_box) @@ -209,13 +224,16 @@ def test_next_to_loss_strategy_penalizes_wrong_distance(): def test_next_to_loss_strategy_respects_relation_weight(): """Test that NextTo loss is scaled by relation_loss_weight.""" + from isaaclab_arena.relations.relation_loss_strategies import NextToLossStrategy + from isaaclab_arena.relations.relations import NextTo, Side + parent_obj = _create_table() child_obj = _create_box() relation_normal = NextTo(parent_obj, side=Side.POSITIVE_X, distance_m=0.05, relation_loss_weight=1.0) relation_double = NextTo(parent_obj, side=Side.POSITIVE_X, distance_m=0.05, relation_loss_weight=2.0) strategy = NextToLossStrategy(slope=10.0) - child_pos = torch.tensor([1.5, 0.5, 0.0]) # 0.5m gap instead of required 0.05m + child_pos = torch.tensor([1.5, 0.5, 0.0]) loss_normal = strategy.compute_loss(relation_normal, child_pos, child_obj.bounding_box, parent_obj.bounding_box) loss_double = strategy.compute_loss(relation_double, child_pos, child_obj.bounding_box, parent_obj.bounding_box) @@ -225,6 +243,8 @@ def test_next_to_loss_strategy_respects_relation_weight(): def test_next_to_zero_distance_raises(): """Test that NextTo raises assertion for zero distance (touching not allowed).""" + from isaaclab_arena.relations.relations import NextTo, Side + parent_obj = _create_table() with pytest.raises(AssertionError, match="Distance must be positive"): diff --git a/isaaclab_arena/tests/test_revolute_joint_moved_rate_metric.py b/isaaclab_arena/tests/test_revolute_joint_moved_rate_metric.py index 2e7f28fd1..ce015e00a 100644 --- a/isaaclab_arena/tests/test_revolute_joint_moved_rate_metric.py +++ b/isaaclab_arena/tests/test_revolute_joint_moved_rate_metric.py @@ -41,7 +41,7 @@ def _test_revolute_joint_moved_rate(simulation_app): embodiment = asset_registry.get_asset_by_name("franka")() microwave = asset_registry.get_asset_by_name("microwave")() - microwave.set_initial_pose(Pose(position_xyz=(0.45, 0.0, 0.2), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + microwave.set_initial_pose(Pose(position_xyz=(0.45, 0.0, 0.2), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) scene = Scene(assets=[background, microwave]) isaaclab_arena_environment = IsaacLabArenaEnvironment( @@ -87,6 +87,7 @@ def _test_revolute_joint_moved_rate(simulation_app): except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: diff --git a/isaaclab_arena/tests/test_robot_initial_position.py b/isaaclab_arena/tests/test_robot_initial_position.py index 180d2f65b..cabd89250 100644 --- a/isaaclab_arena/tests/test_robot_initial_position.py +++ b/isaaclab_arena/tests/test_robot_initial_position.py @@ -6,6 +6,8 @@ import numpy as np import torch import tqdm +import traceback +import warp as wp from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function @@ -32,8 +34,8 @@ def _test_robot_initial_position(simulation_app): robot_init_position = (-0.2, 0.0, 0.0) - cracker_box.set_initial_pose(Pose(position_xyz=(0.4, 0.0, 0.1), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) - embodiment.set_initial_pose(Pose(position_xyz=robot_init_position, rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + cracker_box.set_initial_pose(Pose(position_xyz=(0.4, 0.0, 0.1), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) + embodiment.set_initial_pose(Pose(position_xyz=robot_init_position, rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) scene = Scene(assets=[background, cracker_box]) isaaclab_arena_environment = IsaacLabArenaEnvironment( @@ -56,7 +58,7 @@ def _test_robot_initial_position(simulation_app): env.step(actions) # Check the robot ended up at the correct position. - robot_position = env.scene["robot"].data.root_link_pose_w[0, :3].cpu().numpy() + robot_position = wp.to_torch(env.scene["robot"].data.root_link_pose_w)[0, :3].cpu().numpy() robot_position_error = np.linalg.norm(robot_position - np.array(robot_init_position)) print(f"Robot position error: {robot_position_error}") assert robot_position_error < INITIAL_POSITION_EPS, "Robot ended up at the wrong position." @@ -70,6 +72,7 @@ def _test_robot_initial_position(simulation_app): except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: diff --git a/isaaclab_arena/tests/test_scene_to_usd.py b/isaaclab_arena/tests/test_scene_to_usd.py index 8f0d28d82..8490d27ec 100644 --- a/isaaclab_arena/tests/test_scene_to_usd.py +++ b/isaaclab_arena/tests/test_scene_to_usd.py @@ -26,9 +26,9 @@ def _test_scene_to_usd(simulation_app, output_path: pathlib.Path) -> bool: kitchen = asset_registry.get_asset_by_name("kitchen")() cracker_box = asset_registry.get_asset_by_name("cracker_box")() - kitchen_initial_pose = Pose(position_xyz=(0.772, 3.39, -0.895), rotation_wxyz=(0.70711, 0, 0, -0.70711)) + kitchen_initial_pose = Pose(position_xyz=(0.772, 3.39, -0.895), rotation_xyzw=(0, 0, -0.70711, 0.70711)) kitchen.set_initial_pose(kitchen_initial_pose) - cracker_box_initial_pose = Pose(position_xyz=(0.4, 0.0, 0.1), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)) + cracker_box_initial_pose = Pose(position_xyz=(0.4, 0.0, 0.1), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)) cracker_box.set_initial_pose(cracker_box_initial_pose) # Composed scene @@ -50,8 +50,8 @@ def _test_scene_to_usd(simulation_app, output_path: pathlib.Path) -> bool: } # Function to convert a pxr.Gf.Quatf to a numpy array - def to_numpy_q_wxyz(q_wxyz: Gf.Quatf) -> np.ndarray: - return np.array([q_wxyz.GetReal(), *q_wxyz.GetImaginary()]) + def to_numpy_q_xyzw(q: Gf.Quatf) -> np.ndarray: + return np.array([*q.GetImaginary(), q.GetReal()]) # Loop over all the prims and check that the scene was saved correctly assert len(root_prim.GetChildren()) == len(test_prim_names) @@ -62,11 +62,11 @@ def to_numpy_q_wxyz(q_wxyz: Gf.Quatf) -> np.ndarray: prim_position = prim.GetAttribute("xformOp:translate").Get() prim_orientation = prim.GetAttribute("xformOp:orient").Get() assert np.linalg.norm(prim_position - test_prim_poses[prim_name].position_xyz) < EPS - assert np.linalg.norm(to_numpy_q_wxyz(prim_orientation) - test_prim_poses[prim_name].rotation_wxyz) < EPS + assert np.linalg.norm(to_numpy_q_xyzw(prim_orientation) - test_prim_poses[prim_name].rotation_xyzw) < EPS print(f"Prim {prim_name} position: {prim_position}") - print(f"Prim {prim_name} orientation: {to_numpy_q_wxyz(prim_orientation)}") + print(f"Prim {prim_name} orientation: {to_numpy_q_xyzw(prim_orientation)}") print(f"Prim {prim_name} expected position: {test_prim_poses[prim_name].position_xyz}") - print(f"Prim {prim_name} expected orientation: {test_prim_poses[prim_name].rotation_wxyz}") + print(f"Prim {prim_name} expected orientation: {test_prim_poses[prim_name].rotation_xyzw}") return True diff --git a/isaaclab_arena/tests/test_sequential_open_door.py b/isaaclab_arena/tests/test_sequential_open_door.py index af17d2145..be898d95d 100644 --- a/isaaclab_arena/tests/test_sequential_open_door.py +++ b/isaaclab_arena/tests/test_sequential_open_door.py @@ -57,13 +57,13 @@ def get_mimic_env_cfg(self, embodiment_name: str): microwave_0.set_initial_pose( Pose( position_xyz=(0.6, -0.00586, 0.22773), - rotation_wxyz=(0.7071068, 0, 0, -0.7071068), + rotation_xyzw=(0, 0, -0.7071068, 0.7071068), ) ) microwave_1.set_initial_pose( Pose( position_xyz=(0.6, 0.70586, 0.22773), - rotation_wxyz=(0.7071068, 0, 0, -0.7071068), + rotation_xyzw=(0, 0, -0.7071068, 0.7071068), ) ) @@ -128,6 +128,7 @@ def assert_composite_task_complete(env: ManagerBasedEnv, terminated: torch.Tenso except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: @@ -188,6 +189,7 @@ def assert_composite_task_complete(env: ManagerBasedEnv, terminated: torch.Tenso except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: @@ -232,6 +234,7 @@ def assert_composite_task_complete(env: ManagerBasedEnv, terminated: torch.Tenso except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: @@ -292,6 +295,7 @@ def assert_composite_task_complete(env: ManagerBasedEnv, terminated: torch.Tenso except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: @@ -338,6 +342,7 @@ def _test_sequential_open_door_microwave_reset_condition(simulation_app) -> bool except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: diff --git a/isaaclab_arena/tests/test_sequential_task_base.py b/isaaclab_arena/tests/test_sequential_task_base.py index 0c58a4648..6cc3d63d5 100644 --- a/isaaclab_arena/tests/test_sequential_task_base.py +++ b/isaaclab_arena/tests/test_sequential_task_base.py @@ -65,6 +65,7 @@ class FooCfg: except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False return True @@ -134,6 +135,7 @@ class FooCfg: except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False return True diff --git a/isaaclab_arena/tests/test_sequential_task_mimic_data_generation.py b/isaaclab_arena/tests/test_sequential_task_mimic_data_generation.py index 5f8414c92..f2264cbcc 100644 --- a/isaaclab_arena/tests/test_sequential_task_mimic_data_generation.py +++ b/isaaclab_arena/tests/test_sequential_task_mimic_data_generation.py @@ -3,6 +3,8 @@ # # SPDX-License-Identifier: Apache-2.0 +import pytest + import os import tempfile diff --git a/isaaclab_arena/tests/test_sorting_task.py b/isaaclab_arena/tests/test_sorting_task.py index 38e6a3d54..6dede73f0 100644 --- a/isaaclab_arena/tests/test_sorting_task.py +++ b/isaaclab_arena/tests/test_sorting_task.py @@ -5,6 +5,8 @@ import gymnasium as gym import torch +import traceback +import warp as wp from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function @@ -41,13 +43,13 @@ def get_test_environment(num_envs: int): red_cube.set_initial_pose( Pose( position_xyz=(0.0, 0.3, 0.1), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) green_cube.set_initial_pose( Pose( position_xyz=(0.0, -0.3, 0.1), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) @@ -55,13 +57,13 @@ def get_test_environment(num_envs: int): red_container.set_initial_pose( Pose( position_xyz=(0.0, 0.1, 0.1), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) green_container.set_initial_pose( Pose( position_xyz=(0.0, -0.1, 0.1), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) @@ -80,7 +82,7 @@ def get_test_environment(num_envs: int): embodiment.set_initial_pose( Pose( position_xyz=(-0.4, 0.0, 0.0), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) @@ -120,6 +122,7 @@ def assert_not_success(env: ManagerBasedEnv, terminated: torch.Tensor): except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: @@ -146,8 +149,8 @@ def _test_sorting_task_success(simulation_app) -> bool: green_container_object: RigidObject = env.scene[green_container.name] # Get container positions to place cubes inside them - red_container_pos = red_container_object.data.root_pos_w[0] - green_container_pos = green_container_object.data.root_pos_w[0] + red_container_pos = wp.to_torch(red_container_object.data.root_pos_w)[0] + green_container_pos = wp.to_torch(green_container_object.data.root_pos_w)[0] target_quat = torch.tensor([[1.0, 0.0, 0.0, 0.0]], device=env.device) @@ -181,6 +184,7 @@ def _test_sorting_task_success(simulation_app) -> bool: except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: @@ -205,7 +209,7 @@ def _test_sorting_task_partial_success(simulation_app) -> bool: red_container_object: RigidObject = env.scene[red_container.name] # Get container position - red_container_pos = red_container_object.data.root_pos_w[0] + red_container_pos = wp.to_torch(red_container_object.data.root_pos_w)[0] target_quat = torch.tensor([[1.0, 0.0, 0.0, 0.0]], device=env.device) @@ -229,6 +233,7 @@ def _test_sorting_task_partial_success(simulation_app) -> bool: except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: @@ -261,12 +266,12 @@ def _test_sorting_task_multiple_envs(simulation_app) -> bool: # Now move second env cubes to success positions too # Note: env 0 may have been reset after success, so we need to set both envs - red_cube_state = red_cube_object.data.root_state_w.clone() - green_cube_state = green_cube_object.data.root_state_w.clone() + red_cube_state = wp.to_torch(red_cube_object.data.root_state_w).clone() + green_cube_state = wp.to_torch(green_cube_object.data.root_state_w).clone() # Re-fetch container positions (they should be stable) - red_container_pos = red_container_object.data.root_pos_w - green_container_pos = green_container_object.data.root_pos_w + red_container_pos = wp.to_torch(red_container_object.data.root_pos_w) + green_container_pos = wp.to_torch(green_container_object.data.root_pos_w) # Set BOTH env cubes to positions above containers (env 0 may have been reset) red_cube_state[0, :3] = red_container_pos[0].clone() @@ -307,6 +312,7 @@ def _test_sorting_task_multiple_envs(simulation_app) -> bool: except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: diff --git a/isaaclab_arena/tests/test_success_rate_metric.py b/isaaclab_arena/tests/test_success_rate_metric.py index f9f307218..6ec869710 100644 --- a/isaaclab_arena/tests/test_success_rate_metric.py +++ b/isaaclab_arena/tests/test_success_rate_metric.py @@ -5,6 +5,7 @@ import torch import tqdm +import traceback from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function @@ -75,9 +76,9 @@ def _test_success_rate_metric(simulation_app): # - to: per env pose pose_list = [ # Success (in the drawer) - Pose(position_xyz=(0.0, -0.5, 0.2), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)), + Pose(position_xyz=(0.0, -0.5, 0.2), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)), # Fail (out of the drawer) - Pose(position_xyz=(-0.5, -0.5, 0.2), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)), + Pose(position_xyz=(-0.5, -0.5, 0.2), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)), ] env_cfg.events.reset_pick_up_object_pose = EventTermCfg( func=set_object_pose_per_env, @@ -112,6 +113,7 @@ def _test_success_rate_metric(simulation_app): except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: diff --git a/isaaclab_arena/tests/test_turn_stand_mixer_knob.py b/isaaclab_arena/tests/test_turn_stand_mixer_knob.py index cfc47d428..9710be725 100644 --- a/isaaclab_arena/tests/test_turn_stand_mixer_knob.py +++ b/isaaclab_arena/tests/test_turn_stand_mixer_knob.py @@ -6,6 +6,7 @@ import gymnasium as gym import random import torch +import traceback from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function @@ -37,7 +38,7 @@ def get_test_environment(remove_reset_knob_state_event: bool, num_envs: int): stand_mixer.set_initial_pose( Pose( position_xyz=(0.6, -0.00586, 0.22773), - rotation_wxyz=(0.7071068, 0, 0, -0.7071068), + rotation_xyzw=(0, 0, -0.7071068, 0.7071068), ) ) @@ -80,6 +81,7 @@ def _test_turn_stand_mixer_knob_to_desired_levels_single_env(simulation_app): except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: @@ -128,6 +130,7 @@ def _test_turn_stand_mixer_knob_multiple_envs(simulation_app): except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: @@ -164,6 +167,7 @@ def _test_turn_stand_mixer_knob_reset_condition(simulation_app): except Exception as e: print(f"Error: {e}") + traceback.print_exc() return False finally: diff --git a/isaaclab_arena/tests/test_usd_pose_helpers.py b/isaaclab_arena/tests/test_usd_pose_helpers.py index fa6444471..ca1d75399 100644 --- a/isaaclab_arena/tests/test_usd_pose_helpers.py +++ b/isaaclab_arena/tests/test_usd_pose_helpers.py @@ -28,7 +28,7 @@ def _test_get_prim_pose_in_default_prim_frame(simulation_app): pose = get_prim_pose_in_default_prim_frame(prim, stage) print(f"Position relative to default prim: {pose.position_xyz}") - print(f"Orientation (quaternion wxyz) relative to default prim: {pose.rotation_wxyz}") + print(f"Orientation (quaternion xyzw) relative to default prim: {pose.rotation_xyzw}") # This number is read out of the GUI from the test scene. pos_np_gt = np.array((2.899114282976978, -0.3971232408755399, 1.0062618326241144)) diff --git a/isaaclab_arena/tests/test_xr_anchor_pose.py b/isaaclab_arena/tests/test_xr_anchor_pose.py index 0d307fc2a..e39005efc 100644 --- a/isaaclab_arena/tests/test_xr_anchor_pose.py +++ b/isaaclab_arena/tests/test_xr_anchor_pose.py @@ -23,7 +23,7 @@ def _test_gr1t2_xr_anchor_pose(simulation_app) -> bool: xr_cfg = embodiment.get_xr_cfg() expected_pos = embodiment._xr_offset.position_xyz - expected_rot = embodiment._xr_offset.rotation_wxyz + expected_rot = embodiment._xr_offset.rotation_xyzw assert ( xr_cfg.anchor_pos == expected_pos @@ -35,7 +35,7 @@ def _test_gr1t2_xr_anchor_pose(simulation_app) -> bool: print("✓ GR1T2 XR anchor at origin: PASSED") # Test 2: XR anchor with robot position and rotation - robot_pose = Pose(position_xyz=(1.0, 2.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)) # No rotation + robot_pose = Pose(position_xyz=(1.0, 2.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)) # No rotation embodiment.set_initial_pose(robot_pose) xr_cfg = embodiment.get_xr_cfg() @@ -55,24 +55,24 @@ def _test_gr1t2_xr_anchor_pose(simulation_app) -> bool: # Test 3: XR anchor with robot rotation robot_pose_rotated = Pose( - position_xyz=(0.0, 0.0, 0.0), rotation_wxyz=(0.70711, 0.0, 0.0, 0.70711) # 90° rotation around Z + position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.70711, 0.70711) # 90° rotation around Z ) embodiment.set_initial_pose(robot_pose_rotated) xr_cfg_rotated = embodiment.get_xr_cfg() # Rotation should be composed, not same as offset assert ( - xr_cfg_rotated.anchor_rot != embodiment._xr_offset.rotation_wxyz + xr_cfg_rotated.anchor_rot != embodiment._xr_offset.rotation_xyzw ), "XR anchor rotation should be composed with robot rotation" print("✓ GR1T2 XR anchor with robot rotation: PASSED") # Test 4: Dynamic recomputation - pose1 = Pose(position_xyz=(1.0, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)) + pose1 = Pose(position_xyz=(1.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)) embodiment.set_initial_pose(pose1) xr_cfg1 = embodiment.get_xr_cfg() - pose2 = Pose(position_xyz=(2.0, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)) + pose2 = Pose(position_xyz=(2.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)) embodiment.set_initial_pose(pose2) xr_cfg2 = embodiment.get_xr_cfg() @@ -101,7 +101,7 @@ def _test_g1_xr_anchor_pose(simulation_app) -> bool: # G1 uses a fixed prim-relative XrCfg: anchor offset and rotation are constant expected_pos = (0.0, 0.0, -1.0) - expected_rot = (0.70711, 0.0, 0.0, -0.70711) + expected_rot = (0.0, 0.0, -0.70711, 0.70711) expected_anchor_prim = "/World/envs/env_0/Robot/pelvis" np.testing.assert_allclose( @@ -125,7 +125,7 @@ def _test_g1_xr_anchor_pose(simulation_app) -> bool: ), "G1 XR anchor_rotation_mode should be FOLLOW_PRIM_SMOOTHED" # With initial pose set, config is unchanged (anchor is relative to pelvis prim, not world) - robot_pose = Pose(position_xyz=(0.5, 1.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)) + robot_pose = Pose(position_xyz=(0.5, 1.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)) embodiment.set_initial_pose(robot_pose) xr_cfg_after = embodiment.get_xr_cfg() np.testing.assert_allclose( @@ -160,7 +160,7 @@ def _test_xr_anchor_multiple_positions(simulation_app) -> bool: ] for pos in test_positions: - robot_pose = Pose(position_xyz=pos, rotation_wxyz=(1.0, 0.0, 0.0, 0.0)) + robot_pose = Pose(position_xyz=pos, rotation_xyzw=(0.0, 0.0, 0.0, 1.0)) embodiment.set_initial_pose(robot_pose) xr_cfg = embodiment.get_xr_cfg() diff --git a/isaaclab_arena/tests/utils/subprocess.py b/isaaclab_arena/tests/utils/subprocess.py index 4a0b339ad..8cf8fff9a 100644 --- a/isaaclab_arena/tests/utils/subprocess.py +++ b/isaaclab_arena/tests/utils/subprocess.py @@ -7,6 +7,7 @@ import os import subprocess import sys +import traceback from collections.abc import Callable from isaaclab.app import AppLauncher @@ -77,9 +78,7 @@ def _close_persistent(): _PERSISTENT_SIM_APP_LAUNCHER.app.close() -def get_persistent_simulation_app( - headless: bool, enable_cameras: bool = False, enable_pinocchio: bool = True -) -> SimulationApp: +def get_persistent_simulation_app(headless: bool, enable_cameras: bool = False) -> SimulationApp: """Create once, reuse forever (until process exit).""" global _PERSISTENT_SIM_APP_LAUNCHER, _PERSISTENT_INIT_ARGS # Create a new simulation app if it doesn't exist @@ -88,7 +87,8 @@ def get_persistent_simulation_app( simulation_app_args = parser.parse_args([]) simulation_app_args.headless = headless simulation_app_args.enable_cameras = enable_cameras - simulation_app_args.enable_pinocchio = enable_pinocchio + if not headless: + simulation_app_args.visualizer = ["kit"] with _IsolatedArgv([]): app_launcher = get_app_launcher(simulation_app_args) @@ -112,7 +112,6 @@ def run_simulation_app_function( function: Callable[..., bool], headless: bool = True, enable_cameras: bool = False, - enable_pinocchio: bool = True, **kwargs, ) -> bool: """Run a simulation app in a separate process. @@ -132,14 +131,13 @@ def run_simulation_app_function( # Get a persistent simulation app global _AT_LEAST_ONE_TEST_FAILED try: - simulation_app = get_persistent_simulation_app( - headless=headless, enable_cameras=enable_cameras, enable_pinocchio=enable_pinocchio - ) + simulation_app = get_persistent_simulation_app(headless=headless, enable_cameras=enable_cameras) test_result = bool(function(simulation_app, **kwargs)) return test_result except Exception as e: print(f"Exception occurred while running the function (persistent mode): {e}") + traceback.print_exc() return False finally: # **Always** clean up the SimulationContext/timeline between tests - teardown_simulation_app(suppress_exceptions=True, make_new_stage=True) + teardown_simulation_app(suppress_exceptions=False, make_new_stage=True) diff --git a/isaaclab_arena/utils/bounding_box.py b/isaaclab_arena/utils/bounding_box.py index 41e8c48e8..c7aef99e2 100644 --- a/isaaclab_arena/utils/bounding_box.py +++ b/isaaclab_arena/utils/bounding_box.py @@ -211,14 +211,14 @@ def rotated_90_around_z(self, quarters: int) -> "AxisAlignedBoundingBox": ) -def quaternion_to_90_deg_z_quarters(rotation_wxyz: tuple[float, float, float, float], tol_deg: float = 1.0) -> int: +def quaternion_to_90_deg_z_quarters(rotation_xyzw: tuple[float, float, float, float], tol_deg: float = 1.0) -> int: """Convert a quaternion to 90° rotation quarters around Z axis. Only supports rotations that are multiples of 90° around the Z axis. Raises AssertionError for any other rotation. Args: - rotation_wxyz: Quaternion as (w, x, y, z). + rotation_xyzw: Quaternion as (x, y, z, w). tol_deg: Tolerance in degrees for how close the angle must be to a 90° multiple. Returns: @@ -229,7 +229,7 @@ def quaternion_to_90_deg_z_quarters(rotation_wxyz: tuple[float, float, float, fl """ import math - w, x, y, z = rotation_wxyz + x, y, z, w = rotation_xyzw # Must be a pure Z rotation (x and y components must be ~0) assert ( @@ -271,7 +271,6 @@ def get_random_pose_within_bounding_box(bbox: AxisAlignedBoundingBox, seed: int # random_position = min + (max - min) * rand random_position = min_point + (max_point - min_point) * torch.rand(3) - # Create pose with random position and identity rotation (w=1, x=0, y=0, z=0) - pose = Pose(position_xyz=tuple(random_position.tolist()), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)) + pose = Pose(position_xyz=tuple(random_position.tolist()), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)) return pose diff --git a/isaaclab_arena/utils/isaaclab_utils/resets.py b/isaaclab_arena/utils/isaaclab_utils/resets.py deleted file mode 100644 index 941179dee..000000000 --- a/isaaclab_arena/utils/isaaclab_utils/resets.py +++ /dev/null @@ -1,24 +0,0 @@ -# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: Apache-2.0 - -import torch - -from isaaclab.envs import ManagerBasedEnv - - -def reset_all_articulation_joints(env: ManagerBasedEnv, env_ids: torch.Tensor): - """Reset the articulation joints to the initial state.""" - for articulation_asset in env.scene.articulations.values(): - # obtain default and deal with the offset for env origins - default_root_state = articulation_asset.data.default_root_state[env_ids].clone() - default_root_state[:, 0:3] += env.scene.env_origins[env_ids] - # set into the physics simulation - articulation_asset.write_root_pose_to_sim(default_root_state[:, :7], env_ids=env_ids) - articulation_asset.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids=env_ids) - # obtain default joint positions - default_joint_pos = articulation_asset.data.default_joint_pos[env_ids].clone() - default_joint_vel = articulation_asset.data.default_joint_vel[env_ids].clone() - # set into the physics simulation - articulation_asset.write_joint_state_to_sim(default_joint_pos, default_joint_vel, env_ids=env_ids) diff --git a/isaaclab_arena/utils/isaaclab_utils/simulation_app.py b/isaaclab_arena/utils/isaaclab_utils/simulation_app.py index e3533375f..253a1e053 100644 --- a/isaaclab_arena/utils/isaaclab_utils/simulation_app.py +++ b/isaaclab_arena/utils/isaaclab_utils/simulation_app.py @@ -20,15 +20,7 @@ def get_isaac_sim_version() -> str: def get_app_launcher(args: argparse.Namespace) -> AppLauncher: """Get an app launcher.""" - # NOTE(alexmillane, 2025.11.10): Import pinocchio before launching the app appears still to be required. - # Monitor this and see if we can get rid of it. - if hasattr(args, "enable_pinocchio") and args.enable_pinocchio: - import pinocchio # noqa: F401 - app_launcher = AppLauncher(args) - if get_isaac_sim_version() != "5.1.0": - print(f"WARNING: IsaacSim has been upgraded to {get_isaac_sim_version()}.") - print("Please investigate if pinocchio import is still needed in: simulation_app.py") return app_launcher diff --git a/isaaclab_arena/utils/joint_utils.py b/isaaclab_arena/utils/joint_utils.py index d2dbce43b..baca4be6d 100644 --- a/isaaclab_arena/utils/joint_utils.py +++ b/isaaclab_arena/utils/joint_utils.py @@ -4,7 +4,8 @@ # SPDX-License-Identifier: Apache-2.0 import torch - +import warp as wp +from typing import Sequence from isaaclab.assets import Articulation from isaaclab.envs.manager_based_env import ManagerBasedEnv from isaaclab.managers import SceneEntityCfg @@ -36,7 +37,7 @@ def get_articulation_from_asset_cfg(env: ManagerBasedEnv, asset_cfg: SceneEntity def get_joint_position_limits_from_articulation(articulation: Articulation, joint_index: int) -> tuple[float, float]: """Get the position limits of a joint from the articulation.""" - joint_position_limits = articulation.data.joint_pos_limits[0, joint_index, :] + joint_position_limits = wp.to_torch(articulation.data.joint_pos_limits)[0, joint_index, :] joint_min, joint_max = joint_position_limits[0], joint_position_limits[1] return joint_min, joint_max @@ -45,7 +46,7 @@ def get_unnormalized_joint_position(env: ManagerBasedEnv, asset_cfg: SceneEntity """Get the unnormalized position of a joint in radians.""" articulation = get_articulation_from_asset_cfg(env, asset_cfg) joint_index = get_joint_index_from_asset_cfg(env, asset_cfg) - joint_position = articulation.data.joint_pos[:, joint_index] + joint_position = wp.to_torch(articulation.data.joint_pos)[:, joint_index] return joint_position @@ -71,10 +72,27 @@ def set_unnormalized_joint_position( """Set the position of a joint using an unnormalized value (in radians).""" articulation = get_articulation_from_asset_cfg(env, asset_cfg) joint_index = get_joint_index_from_asset_cfg(env, asset_cfg) - articulation.write_joint_position_to_sim( - torch.tensor([[target_joint_position_unnormlized]]).to(env.device), - torch.tensor([joint_index]).to(env.device), - env_ids=env_ids.to(env.device) if env_ids is not None else None, + # Duplicate data for each environment + num_envs = env.num_envs if env_ids is None else len(env_ids) + position = torch.full((num_envs, 1), target_joint_position_unnormlized, device=env.device) + # joint_ids = torch.full((num_envs, 1), joint_index, dtype=torch.int32, device=env.device) + joint_ids = torch.tensor([joint_index], dtype=torch.int32, device=env.device) + # joint_ids = joint_index + # Move env_ids to the device + env_ids = env_ids.to(env.device) if env_ids is not None else None + # Write the data to the simulation + print(f"HERE") + print(f"position: {position}") + print(f"position shape: {position.shape}") + print(f"joint_ids: {joint_ids}") + print(f"joint_ids shape: {joint_ids.shape}") + print(f"env_ids: {env_ids}") + if env_ids is not None: + print(f"env_ids shape: {env_ids.shape}") + articulation.write_joint_position_to_sim_index( + position=position, + joint_ids=joint_ids, + env_ids=env_ids, ) diff --git a/isaaclab_arena/utils/locomanip_mimic_patch.py b/isaaclab_arena/utils/locomanip_mimic_patch.py index cd262ea4f..6e177f2e4 100644 --- a/isaaclab_arena/utils/locomanip_mimic_patch.py +++ b/isaaclab_arena/utils/locomanip_mimic_patch.py @@ -6,6 +6,7 @@ import asyncio import copy import torch +import warp as wp from isaaclab.envs import SubTaskConstraintType from isaaclab.managers import TerminationTermCfg @@ -299,7 +300,7 @@ async def generate( # noqa: C901 # Update visualization if motion planner is available if motion_planner and motion_planner.visualize_spheres: - current_joints = self.env.scene["robot"].data.joint_pos[env_id] + current_joints = wp.to_torch(self.env.scene["robot"].data.joint_pos)[env_id] motion_planner._update_visualization_at_joint_positions(current_joints) eef_waypoint_dict[eef_name] = waypoint diff --git a/isaaclab_arena/utils/pose.py b/isaaclab_arena/utils/pose.py index 57babdc58..79e164307 100644 --- a/isaaclab_arena/utils/pose.py +++ b/isaaclab_arena/utils/pose.py @@ -6,8 +6,6 @@ import torch from dataclasses import dataclass -from isaaclab.utils.math import matrix_from_quat, quat_from_euler_xyz, quat_from_matrix - @dataclass class Pose: @@ -22,23 +20,23 @@ class Pose: position_xyz: tuple[float, float, float] = (0.0, 0.0, 0.0) """Translation vector from frame A to frame B.""" - rotation_wxyz: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0) - """Quaternion from frame A to frame B. Order is (w, x, y, z).""" + rotation_xyzw: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0) + """Quaternion from frame A to frame B. Order is (x, y, z, w).""" def __post_init__(self): assert isinstance(self.position_xyz, tuple) - assert isinstance(self.rotation_wxyz, tuple) + assert isinstance(self.rotation_xyzw, tuple) assert len(self.position_xyz) == 3 - assert len(self.rotation_wxyz) == 4 + assert len(self.rotation_xyzw) == 4 @staticmethod def identity() -> "Pose": - return Pose(position_xyz=(0.0, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)) + return Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)) def to_tensor(self, device: torch.device) -> torch.Tensor: """Convert the pose to a tensor. - The returned tensor has shape (1, 7), and is of the order (x, y, z, qw, qx, qy, qz). + The returned tensor has shape (1, 7), and is of the order (x, y, z, qx, qy, qz, qw). Args: device: The device to convert the tensor to. @@ -47,7 +45,7 @@ def to_tensor(self, device: torch.device) -> torch.Tensor: The pose as a tensor of shape (1, 7). """ position_tensor = torch.tensor(self.position_xyz, device=device) - rotation_tensor = torch.tensor(self.rotation_wxyz, device=device) + rotation_tensor = torch.tensor(self.rotation_xyzw, device=device) return torch.cat([position_tensor, rotation_tensor]) def multiply(self, other: "Pose") -> "Pose": @@ -64,14 +62,16 @@ def compose_poses(T_C_B: Pose, T_B_A: Pose) -> Pose: Returns: The pose taking points from A to C. """ - R_B_A = matrix_from_quat(torch.tensor(T_B_A.rotation_wxyz)) - R_C_B = matrix_from_quat(torch.tensor(T_C_B.rotation_wxyz)) + from isaaclab.utils.math import matrix_from_quat, quat_from_matrix + + R_B_A = matrix_from_quat(torch.tensor(T_B_A.rotation_xyzw)) + R_C_B = matrix_from_quat(torch.tensor(T_C_B.rotation_xyzw)) # Compose the rotations R_C_A = R_C_B @ R_B_A q_C_A = quat_from_matrix(R_C_A) # Compose the translations t_C_A = R_C_B @ torch.tensor(T_B_A.position_xyz) + torch.tensor(T_C_B.position_xyz) - return Pose(position_xyz=tuple(t_C_A.tolist()), rotation_wxyz=tuple(q_C_A.tolist())) + return Pose(position_xyz=tuple(t_C_A.tolist()), rotation_xyzw=tuple(q_C_A.tolist())) @dataclass @@ -101,6 +101,8 @@ def to_dict(self) -> dict[str, tuple[float, float]]: } def get_midpoint(self) -> Pose: + from isaaclab.utils.math import quat_from_euler_xyz + roll = torch.tensor((self.rpy_min[0] + self.rpy_max[0]) / 2) pitch = torch.tensor((self.rpy_min[1] + self.rpy_max[1]) / 2) yaw = torch.tensor((self.rpy_min[2] + self.rpy_max[2]) / 2) @@ -112,5 +114,5 @@ def get_midpoint(self) -> Pose: ]) return Pose( position_xyz=tuple(position_xyz.tolist()), - rotation_wxyz=tuple(quat.tolist()), + rotation_xyzw=tuple(quat.tolist()), ) diff --git a/isaaclab_arena/utils/usd_pose_helpers.py b/isaaclab_arena/utils/usd_pose_helpers.py index 181b4ee28..16e92908f 100644 --- a/isaaclab_arena/utils/usd_pose_helpers.py +++ b/isaaclab_arena/utils/usd_pose_helpers.py @@ -38,6 +38,6 @@ def get_prim_pose_in_default_prim_frame(prim: Usd.Prim, stage: Usd.Stage) -> Pos prim_T_default = prim_T_world * default_T_world pos, rot, _ = UsdSkel.DecomposeTransform(prim_T_default) - rot_tuple = (rot.GetReal(), rot.GetImaginary()[0], rot.GetImaginary()[1], rot.GetImaginary()[2]) + rot_tuple = (rot.GetImaginary()[0], rot.GetImaginary()[1], rot.GetImaginary()[2], rot.GetReal()) pos_tuple = (pos[0], pos[1], pos[2]) - return Pose(position_xyz=pos_tuple, rotation_wxyz=rot_tuple) + return Pose(position_xyz=pos_tuple, rotation_xyzw=rot_tuple) diff --git a/isaaclab_arena_environments/cube_goal_pose_environment.py b/isaaclab_arena_environments/cube_goal_pose_environment.py index 973ea1d6a..bd7290879 100644 --- a/isaaclab_arena_environments/cube_goal_pose_environment.py +++ b/isaaclab_arena_environments/cube_goal_pose_environment.py @@ -35,14 +35,14 @@ def get_env(self, args_cli: argparse.Namespace): object.set_initial_pose( Pose( position_xyz=(0.1, 0.0, 0.2), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) embodiment = self.asset_registry.get_asset_by_name(args_cli.embodiment)(enable_cameras=args_cli.enable_cameras) embodiment.set_initial_pose( Pose( position_xyz=(-0.4, 0.0, 0.0), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) # order: [panda_joint1, panda_joint2, panda_joint3, panda_joint4, panda_joint5, panda_joint6, panda_joint7, panda_finger_joint1, panda_finger_joint2] @@ -63,7 +63,7 @@ def get_env(self, args_cli: argparse.Namespace): task = GoalPoseTask( object, target_z_range=(0.2, 1), - target_orientation_wxyz=(0.7071, 0.0, 0.0, 0.7071), # yaw 90 degrees + target_orientation_xyzw=(0.0, 0.0, 0.7071, 0.7071), # yaw 90 degrees target_orientation_tolerance_rad=0.2, ) diff --git a/isaaclab_arena_environments/franka_put_and_close_door_environment.py b/isaaclab_arena_environments/franka_put_and_close_door_environment.py index 57e521d14..d4aba93fa 100644 --- a/isaaclab_arena_environments/franka_put_and_close_door_environment.py +++ b/isaaclab_arena_environments/franka_put_and_close_door_environment.py @@ -51,7 +51,7 @@ def get_env(self, args_cli: argparse.Namespace): container.set_initial_pose( Pose( position_xyz=(0.4, -0.00586, 0.22773), - rotation_wxyz=(0.7071068, 0, 0, -0.7071068), + rotation_xyzw=(0, 0, -0.7071068, 0.7071068), ) ) @@ -67,7 +67,7 @@ def get_env(self, args_cli: argparse.Namespace): embodiment.set_initial_pose( Pose( position_xyz=(-0.3, 0.0, -0.5), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) diff --git a/isaaclab_arena_environments/galileo_g1_locomanip_pick_and_place_environment.py b/isaaclab_arena_environments/galileo_g1_locomanip_pick_and_place_environment.py index f61b08c6a..c802cda20 100644 --- a/isaaclab_arena_environments/galileo_g1_locomanip_pick_and_place_environment.py +++ b/isaaclab_arena_environments/galileo_g1_locomanip_pick_and_place_environment.py @@ -42,10 +42,10 @@ def get_env(self, args_cli: argparse.Namespace): blue_sorting_bin.set_initial_pose( Pose( position_xyz=(-0.2450, -1.6272, -0.2641), - rotation_wxyz=(0.0, 0.0, 0.0, 1.0), + rotation_xyzw=(0.0, 0.0, 1.0, 0.0), ) ) - embodiment.set_initial_pose(Pose(position_xyz=(0.0, 0.18, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + embodiment.set_initial_pose(Pose(position_xyz=(0.0, 0.18, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) if ( args_cli.embodiment == "g1_wbc_pink" diff --git a/isaaclab_arena_environments/galileo_pick_and_place_environment.py b/isaaclab_arena_environments/galileo_pick_and_place_environment.py index 379cb704e..53a6d3150 100644 --- a/isaaclab_arena_environments/galileo_pick_and_place_environment.py +++ b/isaaclab_arena_environments/galileo_pick_and_place_environment.py @@ -37,7 +37,7 @@ def get_env(self, args_cli: argparse.Namespace): # -> IsaacLabArenaEnvironment: pick_up_object.set_initial_pose( Pose( position_xyz=(0.55, 0.0, 0.33), - rotation_wxyz=(0.0, 0.0, 0.7071068, -0.7071068), + rotation_xyzw=(0.0, 0.7071068, -0.7071068, 0.0), ) ) diff --git a/isaaclab_arena_environments/gr1_open_microwave_environment.py b/isaaclab_arena_environments/gr1_open_microwave_environment.py index 382432c35..f7e3ee8c9 100644 --- a/isaaclab_arena_environments/gr1_open_microwave_environment.py +++ b/isaaclab_arena_environments/gr1_open_microwave_environment.py @@ -31,7 +31,7 @@ def get_env(self, args_cli: argparse.Namespace): # -> IsaacLabArenaEnvironment: args_cli.embodiment ) embodiment = self.asset_registry.get_asset_by_name(args_cli.embodiment)(enable_cameras=args_cli.enable_cameras) - embodiment.set_initial_pose(Pose(position_xyz=(-0.4, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + embodiment.set_initial_pose(Pose(position_xyz=(-0.4, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) if args_cli.teleop_device is not None: teleop_device = self.device_registry.get_device_by_name(args_cli.teleop_device)() @@ -41,7 +41,7 @@ def get_env(self, args_cli: argparse.Namespace): # -> IsaacLabArenaEnvironment: # Put the microwave on the packing table. microwave_pose = Pose( position_xyz=(0.4, -0.00586, 0.22773), - rotation_wxyz=(0.7071068, 0, 0, -0.7071068), + rotation_xyzw=(0, 0, -0.7071068, 0.7071068), ) microwave.set_initial_pose(microwave_pose) @@ -50,7 +50,7 @@ def get_env(self, args_cli: argparse.Namespace): # -> IsaacLabArenaEnvironment: object = self.asset_registry.get_asset_by_name(args_cli.object)() object_pose = Pose( position_xyz=(0.466, -0.437, 0.154), - rotation_wxyz=(0.5, -0.5, 0.5, -0.5), + rotation_xyzw=(-0.5, 0.5, -0.5, 0.5), ) object.set_initial_pose(object_pose) assets.append(object) diff --git a/isaaclab_arena_environments/gr1_put_and_close_door_environment.py b/isaaclab_arena_environments/gr1_put_and_close_door_environment.py index 7237b2477..bf6a83f84 100644 --- a/isaaclab_arena_environments/gr1_put_and_close_door_environment.py +++ b/isaaclab_arena_environments/gr1_put_and_close_door_environment.py @@ -99,7 +99,7 @@ def __post_init__(self): for key, value in MIMIC_DATAGEN_CONFIG_DEFAULTS.items(): setattr(self.datagen_config, key, value) - camera_offset = Pose(position_xyz=(0.12515, 0.0, 0.06776), rotation_wxyz=(0.57469, 0.11204, -0.17712, -0.79108)) + camera_offset = Pose(position_xyz=(0.12515, 0.0, 0.06776), rotation_xyzw=(0.11204, -0.17712, -0.79108, 0.57469)) embodiment = self.asset_registry.get_asset_by_name(args_cli.embodiment)( enable_cameras=args_cli.enable_cameras, camera_offset=camera_offset ) @@ -125,7 +125,7 @@ def __post_init__(self): embodiment.set_initial_pose( Pose( position_xyz=(3.943, -1.0, 0.995), - rotation_wxyz=(0.7071068, 0.0, 0.0, 0.7071068), + rotation_xyzw=(0.0, 0.0, 0.7071068, 0.7071068), ) ) @@ -161,6 +161,7 @@ def __post_init__(self): ) scene = Scene( assets=[kitchen_background, kitchen_counter_top, pickup_object, light, refrigerator, refrigerator_shelf] + # assets=[pickup_object] ) # Create pick and place task diff --git a/isaaclab_arena_environments/gr1_turn_stand_mixer_knob_environment.py b/isaaclab_arena_environments/gr1_turn_stand_mixer_knob_environment.py index b1a3af430..00d846c29 100644 --- a/isaaclab_arena_environments/gr1_turn_stand_mixer_knob_environment.py +++ b/isaaclab_arena_environments/gr1_turn_stand_mixer_knob_environment.py @@ -31,7 +31,7 @@ def get_env(self, args_cli: argparse.Namespace): # -> IsaacLabArenaEnvironment: args_cli.embodiment ) embodiment = self.asset_registry.get_asset_by_name(args_cli.embodiment)(enable_cameras=args_cli.enable_cameras) - embodiment.set_initial_pose(Pose(position_xyz=(-0.4, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + embodiment.set_initial_pose(Pose(position_xyz=(-0.4, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) if args_cli.teleop_device is not None: teleop_device = self.device_registry.get_device_by_name(args_cli.teleop_device)() @@ -41,7 +41,7 @@ def get_env(self, args_cli: argparse.Namespace): # -> IsaacLabArenaEnvironment: # Put the microwave on the packing table. stand_mixer_pose = Pose( position_xyz=(0.4, -0.00586, 0.22773), - rotation_wxyz=(0.7071068, 0, 0, -0.7071068), + rotation_xyzw=(0, 0, -0.7071068, 0.7071068), ) stand_mixer.set_initial_pose(stand_mixer_pose) @@ -49,7 +49,7 @@ def get_env(self, args_cli: argparse.Namespace): # -> IsaacLabArenaEnvironment: object = self.asset_registry.get_asset_by_name(args_cli.object)() object_pose = Pose( position_xyz=(0.466, -0.437, 0.154), - rotation_wxyz=(0.5, -0.5, 0.5, -0.5), + rotation_xyzw=(-0.5, 0.5, -0.5, 0.5), ) object.set_initial_pose(object_pose) assets.append(object) diff --git a/isaaclab_arena_environments/kitchen_pick_and_place_environment.py b/isaaclab_arena_environments/kitchen_pick_and_place_environment.py index 8fac9bbc4..b936ecd42 100644 --- a/isaaclab_arena_environments/kitchen_pick_and_place_environment.py +++ b/isaaclab_arena_environments/kitchen_pick_and_place_environment.py @@ -60,7 +60,6 @@ def get_env(self, args_cli: argparse.Namespace): # -> IsaacLabArenaEnvironment: name="destination_location", prim_path="{ENV_REGEX_NS}/kitchen/Cabinet_B_02", parent_asset=background, - object_type=ObjectType.RIGID, ) if args_cli.teleop_device is not None: teleop_device = self.device_registry.get_device_by_name(args_cli.teleop_device)() diff --git a/isaaclab_arena_environments/lift_object_environment.py b/isaaclab_arena_environments/lift_object_environment.py index 69e768a12..f1ae8d278 100644 --- a/isaaclab_arena_environments/lift_object_environment.py +++ b/isaaclab_arena_environments/lift_object_environment.py @@ -43,8 +43,8 @@ def get_env(self, args_cli: argparse.Namespace): # -> IsaacLabArenaEnvironment: teleop_device = None # Set all positions - background.set_initial_pose(Pose(position_xyz=(0.5, 0, 0), rotation_wxyz=(0.707, 0, 0, 0.707))) - pick_up_object.set_initial_pose(Pose(position_xyz=(0.5, 0, 0.055), rotation_wxyz=(1, 0, 0, 0))) + background.set_initial_pose(Pose(position_xyz=(0.5, 0, 0), rotation_xyzw=(0, 0, 0.707, 0.707))) + pick_up_object.set_initial_pose(Pose(position_xyz=(0.5, 0, 0.055), rotation_xyzw=(0, 0, 0, 1))) ground_plane.set_initial_pose(Pose(position_xyz=(0.0, 0.0, -1.05))) # Compose the scene diff --git a/isaaclab_arena_environments/mdp/env_callbacks.py b/isaaclab_arena_environments/mdp/env_callbacks.py index 697928b0b..746971cbf 100644 --- a/isaaclab_arena_environments/mdp/env_callbacks.py +++ b/isaaclab_arena_environments/mdp/env_callbacks.py @@ -41,14 +41,15 @@ def assembly_env_cfg_callback(env_cfg: IsaacLabArenaManagerBasedRLEnvCfg) -> Isa Returns: The modified environment configuration. """ - from isaaclab.sim import PhysxCfg, SimulationCfg + from isaaclab.sim import SimulationCfg + from isaaclab_physx.physics.physx_manager_cfg import PhysxCfg from isaaclab.sim.spawners.materials import RigidBodyMaterialCfg # Simulation settings optimized for assembly tasks env_cfg.sim = SimulationCfg( dt=1 / 60, # 60Hz - balance between speed and stability render_interval=2, - physx=PhysxCfg( + physics=PhysxCfg( solver_type=1, max_position_iteration_count=192, # Important to avoid interpenetration max_velocity_iteration_count=1, diff --git a/isaaclab_arena_environments/press_button_environment.py b/isaaclab_arena_environments/press_button_environment.py index 8b60e0e38..46aeb4ed8 100644 --- a/isaaclab_arena_environments/press_button_environment.py +++ b/isaaclab_arena_environments/press_button_environment.py @@ -37,7 +37,7 @@ def get_env(self, args_cli: argparse.Namespace): # -> IsaacLabArenaEnvironment: teleop_device = None # Put the coffee_machine on the packing table. - press_object_pose = Pose(position_xyz=(0.7, 0.4, 0.19), rotation_wxyz=(0.7071, 0.0, 0.0, -0.7071)) + press_object_pose = Pose(position_xyz=(0.7, 0.4, 0.19), rotation_xyzw=(0.0, 0.0, -0.7071, 0.7071)) press_object.set_initial_pose(press_object_pose) # Compose the scene diff --git a/isaaclab_arena_environments/sorting_environment.py b/isaaclab_arena_environments/sorting_environment.py index 579455153..4e2334ca2 100644 --- a/isaaclab_arena_environments/sorting_environment.py +++ b/isaaclab_arena_environments/sorting_environment.py @@ -32,7 +32,7 @@ def get_env(self, args_cli: argparse.Namespace): background.set_initial_pose( Pose( position_xyz=(0.3, 0.0, 0.0), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) @@ -44,7 +44,7 @@ def get_env(self, args_cli: argparse.Namespace): embodiment.set_initial_pose( Pose( position_xyz=(-0.4, 0.0, 0.0), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) @@ -68,7 +68,7 @@ def get_env(self, args_cli: argparse.Namespace): destination_location_1.set_initial_pose( Pose( position_xyz=(0.0, 0.1, 0.1), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) @@ -76,7 +76,7 @@ def get_env(self, args_cli: argparse.Namespace): destination_location_2.set_initial_pose( Pose( position_xyz=(0.0, -0.1, 0.1), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) @@ -84,7 +84,7 @@ def get_env(self, args_cli: argparse.Namespace): pick_up_object_1.set_initial_pose( Pose( position_xyz=(0.0, 0.3, 0.1), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) @@ -92,7 +92,7 @@ def get_env(self, args_cli: argparse.Namespace): pick_up_object_2.set_initial_pose( Pose( position_xyz=(0.0, -0.3, 0.1), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) diff --git a/isaaclab_arena_environments/tabletop_gearmesh_environment.py b/isaaclab_arena_environments/tabletop_gearmesh_environment.py index a44ebc497..5bd754d96 100644 --- a/isaaclab_arena_environments/tabletop_gearmesh_environment.py +++ b/isaaclab_arena_environments/tabletop_gearmesh_environment.py @@ -45,34 +45,34 @@ def get_env(self, args_cli: argparse.Namespace): # -> IsaacLabArenaEnvironment: else: teleop_device = None - background.set_initial_pose(Pose(position_xyz=(0.55, 0.0, 0.0), rotation_wxyz=(0.707, 0, 0, 0.707))) + background.set_initial_pose(Pose(position_xyz=(0.55, 0.0, 0.0), rotation_xyzw=(0, 0, 0.707, 0.707))) # Set initial poses for all 4 gears gear_base.set_initial_pose( Pose( position_xyz=(0.6, 0.0, 0.0), # Gear base position - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) medium_gear.set_initial_pose( Pose( position_xyz=(0.5, 0.2, 0.0), # Medium gear to be assembled - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) small_gear.set_initial_pose( Pose( position_xyz=(0.6, 0.0, 0.0), # Small reference gear - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) large_gear.set_initial_pose( Pose( position_xyz=(0.6, 0.0, 0.0), # Large reference gear - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) diff --git a/isaaclab_arena_environments/tabletop_peginsert_environment.py b/isaaclab_arena_environments/tabletop_peginsert_environment.py index bfefbf11e..72fd59e3b 100644 --- a/isaaclab_arena_environments/tabletop_peginsert_environment.py +++ b/isaaclab_arena_environments/tabletop_peginsert_environment.py @@ -35,19 +35,19 @@ def get_env(self, args_cli: argparse.Namespace): # -> IsaacLabArenaEnvironment: else: teleop_device = None - background.set_initial_pose(Pose(position_xyz=(0.55, 0.0, 0.0), rotation_wxyz=(0.707, 0, 0, 0.707))) + background.set_initial_pose(Pose(position_xyz=(0.55, 0.0, 0.0), rotation_xyzw=(0, 0, 0.707, 0.707))) pick_up_object.set_initial_pose( Pose( position_xyz=(0.45, 0.0, 0.0), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) destination_object.set_initial_pose( Pose( position_xyz=(0.45, 0.1, 0.0), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) diff --git a/isaaclab_arena_environments/tabletop_place_upright_environment.py b/isaaclab_arena_environments/tabletop_place_upright_environment.py index e6199e590..2c8cdd574 100644 --- a/isaaclab_arena_environments/tabletop_place_upright_environment.py +++ b/isaaclab_arena_environments/tabletop_place_upright_environment.py @@ -52,7 +52,7 @@ class EventCfgPlaceUprightMug: # Add the asset registry from the arena migration package background = self.asset_registry.get_asset_by_name(args_cli.background)() placeable_object = self.asset_registry.get_asset_by_name(args_cli.object)( - initial_pose=Pose(position_xyz=(0.05, 0.0, 0.75), rotation_wxyz=(0.0, 1.0, 0.0, 0.0)) + initial_pose=Pose(position_xyz=(0.05, 0.0, 0.75), rotation_xyzw=(1.0, 0.0, 0.0, 0.0)) ) if args_cli.embodiment == "agibot": embodiment = self.asset_registry.get_asset_by_name(args_cli.embodiment)( @@ -71,10 +71,10 @@ class EventCfgPlaceUprightMug: embodiment.set_initial_pose( Pose( position_xyz=(-0.60, 0.0, 0.0), - rotation_wxyz=(1.0, 0.0, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) - background.set_initial_pose(Pose(position_xyz=(0.50, 0.0, 0.625), rotation_wxyz=(0.7071, 0, 0, 0.7071))) + background.set_initial_pose(Pose(position_xyz=(0.50, 0.0, 0.625), rotation_xyzw=(0, 0, 0.7071, 0.7071))) background.object_cfg.spawn.scale = (1.0, 1.0, 0.60) ground_plane = self.asset_registry.get_asset_by_name("ground_plane")() diff --git a/isaaclab_arena_examples/relations/dummy_object_placer_notebook.py b/isaaclab_arena_examples/relations/dummy_object_placer_notebook.py index d9f56616e..34823d3c6 100644 --- a/isaaclab_arena_examples/relations/dummy_object_placer_notebook.py +++ b/isaaclab_arena_examples/relations/dummy_object_placer_notebook.py @@ -55,7 +55,7 @@ def run_dummy_object_placer_demo(): # Mark desk as the anchor for relation solving (not subject to optimization) desk.add_relation(IsAnchor()) - desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + desk.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) # Center box is on the desk center_box.add_relation(On(desk, clearance_m=0.01)) @@ -127,8 +127,8 @@ def run_dummy_multi_anchor_demo(): ) # Anchor objects (fixed positions) - table.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) - chair.set_initial_pose(Pose(position_xyz=(2.0, 0.0, 0.0), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + table.set_initial_pose(Pose(position_xyz=(0.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) + chair.set_initial_pose(Pose(position_xyz=(2.0, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) table.add_relation(IsAnchor()) chair.add_relation(IsAnchor()) @@ -242,3 +242,24 @@ def run_dummy_no_collision_demo(): run_dummy_no_collision_demo() # %% + + +dummy_object_1 = DummyObject() +dummy_object_2 = DummyObject() +dummy_object_3 = DummyObject() +table = DummyObject() + +dummy_object_1.add_relation(On(table)) +dummy_object_2.add_relation(On(table)) +dummy_object_3.add_relation(On(table)) +table.add_relation(IsAnchor()) + +dummy_object_1.add_relation(NoCollision(dummy_object_2)) +dummy_object_1.add_relation(NoCollision(dummy_object_3)) +dummy_object_2.add_relation(NoCollision(dummy_object_3)) + + + + + + diff --git a/isaaclab_arena_examples/relations/isaac_sim_object_placer_notebook.py b/isaaclab_arena_examples/relations/isaac_sim_object_placer_notebook.py index 54d5e8eb7..d61246549 100644 --- a/isaaclab_arena_examples/relations/isaac_sim_object_placer_notebook.py +++ b/isaaclab_arena_examples/relations/isaac_sim_object_placer_notebook.py @@ -11,7 +11,6 @@ """Example notebook demonstrating ObjectPlacer with real Isaac Sim objects.""" # NOTE: When running as a notebook, first run this cell to launch the simulation app: -import pinocchio # noqa: F401 from isaaclab.app import AppLauncher print("Launching simulation app once in notebook") diff --git a/isaaclab_arena_examples/relations/relation_solver_visualization_notebook.py b/isaaclab_arena_examples/relations/relation_solver_visualization_notebook.py index 8a791e75d..5fe1dcf6c 100644 --- a/isaaclab_arena_examples/relations/relation_solver_visualization_notebook.py +++ b/isaaclab_arena_examples/relations/relation_solver_visualization_notebook.py @@ -155,18 +155,18 @@ def run_visualization_demo(): # Create parent object parent = DummyObject(name="parent", bounding_box=parent_bbox) - parent.set_initial_pose(Pose(position_xyz=parent_pos, rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) + parent.set_initial_pose(Pose(position_xyz=parent_pos, rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) parent.add_relation(IsAnchor()) # Create first child - placed to the RIGHT of parent child1 = DummyObject(name="child1", bounding_box=child_bbox) child1.add_relation(NextTo(parent, side=Side.POSITIVE_X, distance_m=distance_m)) - child1.set_initial_pose(Pose(position_xyz=(0.5, 0.0, 0.05), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) # Initial guess + child1.set_initial_pose(Pose(position_xyz=(0.5, 0.0, 0.05), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) # Initial guess # Create second child - placed to the RIGHT of child1 (chained placement) child2 = DummyObject(name="child2", bounding_box=child_bbox) child2.add_relation(NextTo(child1, side=Side.POSITIVE_X, distance_m=distance_m)) - child2.set_initial_pose(Pose(position_xyz=(0.8, 0.0, 0.05), rotation_wxyz=(1.0, 0.0, 0.0, 0.0))) # Initial guess + child2.set_initial_pose(Pose(position_xyz=(0.8, 0.0, 0.05), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) # Initial guess # Create solver solver = RelationSolver(params=RelationSolverParams(verbose=False)) @@ -193,7 +193,7 @@ def run_visualization_demo(): # Note: child2's relation parent is child1, so we need child1 at a fixed position # and include parent in objects list since child1 has a relation to parent child1.set_initial_pose( - Pose(position_xyz=(0.45, 0.0, 0.05), rotation_wxyz=(1.0, 0.0, 0.0, 0.0)) + Pose(position_xyz=(0.45, 0.0, 0.05), rotation_xyzw=(0.0, 0.0, 0.0, 1.0)) ) # Ideal position for child1 X, Y, losses_child2 = create_loss_heatmap_2d( solver=solver, diff --git a/isaaclab_arena_g1/g1_env/mdp/actions/g1_decoupled_wbc_pink_action.py b/isaaclab_arena_g1/g1_env/mdp/actions/g1_decoupled_wbc_pink_action.py index a6dd0fdd0..9ab23525a 100644 --- a/isaaclab_arena_g1/g1_env/mdp/actions/g1_decoupled_wbc_pink_action.py +++ b/isaaclab_arena_g1/g1_env/mdp/actions/g1_decoupled_wbc_pink_action.py @@ -7,6 +7,7 @@ import numpy as np import torch +import warp as wp from collections.abc import Sequence from scipy.spatial.transform import Rotation as R from typing import TYPE_CHECKING @@ -45,7 +46,7 @@ if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv - from isaaclab_arena_g1.g1_env.mdp.actions.g1_decoupled_wbc_pink_action_cfg import G1DecoupledWBCPinkActionCfg + from isaaclab_arena.embodiments.g1.mdp.actions.g1_decoupled_wbc_pink_action_cfg import G1DecoupledWBCPinkActionCfg class G1DecoupledWBCPinkAction(G1DecoupledWBCJointAction): @@ -194,9 +195,9 @@ def process_actions(self, actions: torch.Tensor): action = [left_hand_state: dim=1, 0 for open, 1 for close, right_hand_state: dim=1, 0 for open, 1 for close, left_arm_pos: dim=3, xyz position, - left_arm_quat: dim=4, wxyz quaternion, + left_arm_quat: dim=4, xyzw quaternion, right_arm_pos: dim=3, xyz position, - right_arm_quat: dim=4, wxyz quaternion, + right_arm_quat: dim=4, xyzw quaternion, navigate_cmd: dim=3, xyz velocity, base_height_cmd: dim=1, height, torso_orientation_rpy_cmd: dim=3, rpy] @@ -215,22 +216,13 @@ def process_actions(self, actions: torch.Tensor): """ # Extract upper body left/right arm pos/quat from actions left_arm_pos = actions_clone[:, LEFT_WRIST_POS_START_IDX:LEFT_WRIST_POS_END_IDX].squeeze(0).cpu() - left_arm_quat = actions_clone[:, LEFT_WRIST_QUAT_START_IDX:LEFT_WRIST_QUAT_END_IDX].squeeze(0).cpu().numpy() + left_arm_quat = actions_clone[:, LEFT_WRIST_QUAT_START_IDX:LEFT_WRIST_QUAT_END_IDX].squeeze(0).cpu() right_arm_pos = actions_clone[:, RIGHT_WRIST_POS_START_IDX:RIGHT_WRIST_POS_END_IDX].squeeze(0).cpu() - right_arm_quat = actions_clone[:, RIGHT_WRIST_QUAT_START_IDX:RIGHT_WRIST_QUAT_END_IDX].squeeze(0).cpu().numpy() - - # Replace zero-norm quaternions with identity (e.g. zero actions from env.step(zeros)) - _IDENTITY_QUAT_WXYZ = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64) - for q in (left_arm_quat, right_arm_quat): - if np.linalg.norm(q) < 1e-8: - q[:] = _IDENTITY_QUAT_WXYZ + right_arm_quat = actions_clone[:, RIGHT_WRIST_QUAT_START_IDX:RIGHT_WRIST_QUAT_END_IDX].squeeze(0).cpu() # Convert from pos/quat to 4x4 transform matrix - # Scipy requires quat xyzw, IsaacLab uses wxyz so a conversion is needed - left_arm_quat_xyzw = np.roll(left_arm_quat, -1) - right_arm_quat_xyzw = np.roll(right_arm_quat, -1) - left_rotmat = R.from_quat(left_arm_quat_xyzw).as_matrix() - right_rotmat = R.from_quat(right_arm_quat_xyzw).as_matrix() + left_rotmat = R.from_quat(left_arm_quat).as_matrix() + right_rotmat = R.from_quat(right_arm_quat).as_matrix() left_arm_pose = np.eye(4) left_arm_pose[:3, :3] = left_rotmat @@ -294,8 +286,8 @@ def process_actions(self, actions: torch.Tensor): target_xy = torch.tensor(target_xy_heading[:2]) target_heading = torch.tensor(target_xy_heading[2]) - current_xy = self._asset.data.root_link_pos_w - current_heading = self._asset.data.heading_w + current_xy = wp.to_torch(self._asset.data.root_link_pos_w) + current_heading = wp.to_torch(self._asset.data.heading_w) check_xy_reached = self.navigation_p_controller.check_xy_within_threshold(target_xy, current_xy) check_heading_reached = self.navigation_p_controller.check_heading_within_threshold( @@ -367,35 +359,3 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: env_ids: A list of environment IDs to reset. If None, all environments are reset. """ self._raw_actions[env_ids] = torch.zeros(self.action_dim, device=self.device) - - def preprocess_actions(self, actions: torch.Tensor) -> torch.Tensor: - """Transform wrist positions and orientations from world frame to robot base frame. - - Args: - actions: The input actions tensor, shape (num_envs, action_dim). - - Returns: - The processed actions tensor (same shape as input). - """ - actions = actions.clone() - - robot_base_pos = self._asset.data.root_link_pos_w[:, :3] - robot_base_quat = self._asset.data.root_link_quat_w - - left_wrist_pos_world = actions[:, LEFT_WRIST_POS_START_IDX:LEFT_WRIST_POS_END_IDX] - right_wrist_pos_world = actions[:, RIGHT_WRIST_POS_START_IDX:RIGHT_WRIST_POS_END_IDX] - left_wrist_pos_base = math_utils.quat_apply_inverse(robot_base_quat, left_wrist_pos_world - robot_base_pos) - right_wrist_pos_base = math_utils.quat_apply_inverse(robot_base_quat, right_wrist_pos_world - robot_base_pos) - - left_wrist_quat_world = actions[:, LEFT_WRIST_QUAT_START_IDX:LEFT_WRIST_QUAT_END_IDX] - right_wrist_quat_world = actions[:, RIGHT_WRIST_QUAT_START_IDX:RIGHT_WRIST_QUAT_END_IDX] - robot_base_quat_inv = math_utils.quat_inv(robot_base_quat) - left_wrist_quat_base = math_utils.quat_mul(robot_base_quat_inv, left_wrist_quat_world) - right_wrist_quat_base = math_utils.quat_mul(robot_base_quat_inv, right_wrist_quat_world) - - actions[:, LEFT_WRIST_POS_START_IDX:LEFT_WRIST_POS_END_IDX] = left_wrist_pos_base - actions[:, LEFT_WRIST_QUAT_START_IDX:LEFT_WRIST_QUAT_END_IDX] = left_wrist_quat_base - actions[:, RIGHT_WRIST_POS_START_IDX:RIGHT_WRIST_POS_END_IDX] = right_wrist_pos_base - actions[:, RIGHT_WRIST_QUAT_START_IDX:RIGHT_WRIST_QUAT_END_IDX] = right_wrist_quat_base - - return actions diff --git a/isaaclab_arena_g1/g1_whole_body_controller/wbc_policy/run_policy.py b/isaaclab_arena_g1/g1_whole_body_controller/wbc_policy/run_policy.py index 0cb2659ec..d1f2db2fd 100644 --- a/isaaclab_arena_g1/g1_whole_body_controller/wbc_policy/run_policy.py +++ b/isaaclab_arena_g1/g1_whole_body_controller/wbc_policy/run_policy.py @@ -5,6 +5,7 @@ import numpy as np import torch +import warp as wp import isaaclab.utils.math as math_utils from isaaclab.assets import ArticulationData @@ -61,8 +62,8 @@ def prepare_observations( - torso_ang_vel: Torso angular velocity """ # Get robot joint observations - sim_joint_pos = robot_data.joint_pos.cpu().numpy() - sim_joint_vel = robot_data.joint_vel.cpu().numpy() + sim_joint_pos = wp.to_torch(robot_data.joint_pos).cpu().numpy() + sim_joint_vel = wp.to_torch(robot_data.joint_vel).cpu().numpy() num_joints = len(robot_data.joint_names) # Convert joints data from Lab's order to GR00T's order saved in config yaml @@ -75,19 +76,21 @@ def prepare_observations( # Prepare obs dict for WBC policy input to G1DecoupledWholeBodyPolicy class assert wbc_joint_pos.shape == wbc_joint_vel.shape == wbc_joint_acc.shape == (num_envs, num_joints) - root_link_pos_w = robot_data.root_link_pos_w.cpu().numpy() - root_link_quat_w = robot_data.root_link_quat_w.cpu().numpy() - base_pose_w = np.concatenate((root_link_pos_w, root_link_quat_w), axis=1) - base_lin_vel_b = robot_data.root_link_lin_vel_b.cpu().numpy() - base_ang_vel_b = robot_data.root_link_ang_vel_b.cpu().numpy() + root_link_pos_w = wp.to_torch(robot_data.root_link_pos_w).cpu().numpy() + root_link_quat_w_xyzw = wp.to_torch(robot_data.root_link_quat_w).cpu().numpy() + root_link_quat_w_wxyz = np.concatenate((root_link_quat_w_xyzw[:, 3:4], root_link_quat_w_xyzw[:, :3]), axis=1) + base_pose_w = np.concatenate((root_link_pos_w, root_link_quat_w_wxyz), axis=1) + base_lin_vel_b = wp.to_torch(robot_data.root_link_lin_vel_b).cpu().numpy() + base_ang_vel_b = wp.to_torch(robot_data.root_link_ang_vel_b).cpu().numpy() base_vel_b = np.concatenate((base_lin_vel_b, base_ang_vel_b), axis=1) # torso link in world frame - torso_link_pose_w = robot_data.body_link_state_w[:, robot_data.body_names.index("torso_link"), :] - torso_link_quat_w = torso_link_pose_w[:, 3:7] # w, x, y, z + torso_link_pose_w = wp.to_torch(robot_data.body_link_state_w)[:, robot_data.body_names.index("torso_link"), :] + torso_link_quat_w_xyzw = torso_link_pose_w[:, 3:7] + torso_link_quat_w_wxyz = torch.cat((torso_link_quat_w_xyzw[:, 3:4], torso_link_quat_w_xyzw[:, :3]), dim=1) torso_link_ang_vel_w = torso_link_pose_w[:, -3:] - torso_link_ang_vel_b = math_utils.quat_apply_inverse(torso_link_quat_w, torso_link_ang_vel_w) + torso_link_ang_vel_b = math_utils.quat_apply_inverse(torso_link_quat_w_xyzw, torso_link_ang_vel_w) # Prepare obs tmers wbc_obs = { @@ -98,7 +101,7 @@ def prepare_observations( "floating_base_pose": base_pose_w, # wrt world frame, used to project gravity vector to local frame "floating_base_vel": base_vel_b, # wrt body frame "floating_base_acc": np.zeros((num_envs, 6)), # Not used by Standing Waist Height Policy - "torso_quat": torso_link_quat_w.cpu().numpy(), + "torso_quat": torso_link_quat_w_wxyz.cpu().numpy(), "torso_ang_vel": torso_link_ang_vel_b.cpu().numpy(), } return wbc_obs diff --git a/isaaclab_arena_gr00t/lerobot/config/droid_manip_config.yaml b/isaaclab_arena_gr00t/lerobot/config/droid_manip_config.yaml new file mode 100644 index 000000000..6e7bec9d6 --- /dev/null +++ b/isaaclab_arena_gr00t/lerobot/config/droid_manip_config.yaml @@ -0,0 +1,59 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +# Gr00tDatasetConfig Example Configuration +# This file shows how to configure the Gr00tDatasetConfig dataclass using YAML +# Example for GR1 tabletop manipulation task + +# Root directory for all data storage +data_root: "/datasets/" + +# Instruction given to the policy in natural language +language_instruction: "Pick and Place" +task_index: 0 + +# Name of the HDF5 file to use for the dataset +hdf5_name: "curobo_v3_with_camera.hdf5" + +# Mimic-generated HDF5 datafield names +state_name_sim: "robot_joint_pos" +action_name_sim: "processed_actions" +pov_cam_name_sim: "robot_pov_cam_rgb" + +# Gr00t-LeRobot datafield names +state_name_lerobot: "observation.state" +action_name_lerobot: "action" +video_name_lerobot: "observation.images.ego_view" +task_description_lerobot: "annotation.human.action.task_description" + +# Parquet configuration +chunks_size: 1000 + +# Video configuration +fps: 50 + +# File path templates +data_path: "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet" +video_path: "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4" + +# Configuration file paths +modality_template_path: "isaaclab_arena_gr00t/embodiments/gr1/modality.json" +modality_fname: "modality.json" +episodes_fname: "episodes.jsonl" +tasks_fname: "tasks.jsonl" +info_template_path: "isaaclab_arena_gr00t/embodiments/gr1/info.json" +info_fname: "info.json" + +# policy specific parameters (gr00t demonstration data stored in lerobot format) +policy_joints_config_path: "isaaclab_arena_gr00t/embodiments/gr1/gr00t_26dof_joint_space.yaml" +robot_type: "gr1" + +# Robot simulation specific parameters +action_joints_config_path: "isaaclab_arena_gr00t/embodiments/gr1/36dof_joint_space.yaml" +state_joints_config_path: "isaaclab_arena_gr00t/embodiments/gr1/54dof_joint_space.yaml" + +# Image configuration (height, width, channels) +original_image_size: [512, 512, 3] +target_image_size: [512, 512, 3] diff --git a/isaaclab_arena_gr00t/policy/gr00t_closedloop_policy.py b/isaaclab_arena_gr00t/policy/gr00t_closedloop_policy.py index 8643accee..edd21939d 100644 --- a/isaaclab_arena_gr00t/policy/gr00t_closedloop_policy.py +++ b/isaaclab_arena_gr00t/policy/gr00t_closedloop_policy.py @@ -7,16 +7,31 @@ from __future__ import annotations import argparse -import gymnasium as gym -import torch +import os +import sys from dataclasses import dataclass, field from typing import Any +# Prepend GR00T deps when loaded without re-exec (e.g. eval_runner, tests before conftest re-exec). +_GROOT_DEPS_DIR = os.environ.get("GROOT_DEPS_DIR") +if _GROOT_DEPS_DIR and _GROOT_DEPS_DIR not in sys.path: + sys.path.insert(0, _GROOT_DEPS_DIR) + +import gymnasium as gym +import torch + +from gr00t.data.embodiment_tags import EmbodimentTag from gr00t.policy.gr00t_policy import Gr00tPolicy from isaaclab_arena.policy.action_chunking import ActionChunkingState from isaaclab_arena.policy.policy_base import PolicyBase from isaaclab_arena.utils.multiprocess import get_local_rank, get_world_size +from isaaclab_arena_gr00t.utils.eagle_config_compat import apply_eagle_config_compat +from isaaclab_arena_g1.g1_whole_body_controller.wbc_policy.policy.policy_constants import ( + NUM_BASE_HEIGHT_CMD, + NUM_NAVIGATE_CMD, + NUM_TORSO_ORIENTATION_RPY_CMD, +) from isaaclab_arena_gr00t.policy.config.gr00t_closedloop_policy_config import Gr00tClosedloopPolicyConfig, TaskMode from isaaclab_arena_gr00t.policy.gr00t_core import ( Gr00tBasePolicyArgs, @@ -137,9 +152,32 @@ def add_args_to_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentPars ) return parser - # ------------------------------------------------------------------ # - # Public API - # ------------------------------------------------------------------ # + def load_policy_joints_config(self, policy_config_path: Path) -> dict[str, Any]: + """Load the GR00T policy joint config from the data config.""" + return load_robot_joints_config_from_yaml(policy_config_path) + + def load_sim_state_joints_config(self, state_config_path: Path) -> dict[str, Any]: + """Load the simulation state joint config from the data config.""" + return load_robot_joints_config_from_yaml(state_config_path) + + def load_sim_action_joints_config(self, action_config_path: Path) -> dict[str, Any]: + """Load the simulation action joint config from the data config.""" + return load_robot_joints_config_from_yaml(action_config_path) + + def load_policy(self) -> Gr00tPolicy: + """Load the dataset, whose iterator will be used as the policy.""" + assert Path( + self.policy_config.model_path + ).exists(), f"Dataset path {self.policy_config.dataset_path} does not exist" + + apply_eagle_config_compat() + + return Gr00tPolicy( + model_path=self.policy_config.model_path, + embodiment_tag=EmbodimentTag[self.policy_config.embodiment_tag], + device=self.device, + strict=True, + ) def set_task_description(self, task_description: str | None) -> str: """Set the language instruction of the task being evaluated.""" diff --git a/isaaclab_arena_gr00t/tests/conftest.py b/isaaclab_arena_gr00t/tests/conftest.py index 185414432..e62abd64f 100644 --- a/isaaclab_arena_gr00t/tests/conftest.py +++ b/isaaclab_arena_gr00t/tests/conftest.py @@ -3,13 +3,17 @@ # # SPDX-License-Identifier: Apache-2.0 -# Isaac Sim makes testing complicated. During shutdown Isaac Sim will -# terminate the surrounding pytest process with exit code 0, regardless -# of whether the tests passed or failed. -# To work around this, we stash the session object and set a flag -# when a test fails. This flag is checked in isaaclab_arena.tests.utils.subprocess.py -# prior to closing the simulation app, in order to generate the correct exit code. +# Re-exec pytest with PYTHONPATH so GR00T deps load first (set GROOT_DEPS_DIR when running GR00T tests). +import os +import sys +_groot_deps = os.environ.get("GROOT_DEPS_DIR") +if _groot_deps and os.environ.get("_GROOT_PYTHONPATH_APPLIED") != "1": + os.environ["PYTHONPATH"] = _groot_deps + os.pathsep + os.environ.get("PYTHONPATH", "") + os.environ["_GROOT_PYTHONPATH_APPLIED"] = "1" + os.execv(sys.executable, [sys.executable, "-m", "pytest"] + sys.argv[1:]) + +# Isaac Sim exits with 0 on shutdown; stash session and set tests_failed so subprocess.py can report failure. import isaaclab_arena.tests.conftest as arena_conftest diff --git a/isaaclab_arena_gr00t/tests/test_gr00t_closedloop_policy.py b/isaaclab_arena_gr00t/tests/test_gr00t_closedloop_policy.py index 0b7283c21..7b5df9422 100644 --- a/isaaclab_arena_gr00t/tests/test_gr00t_closedloop_policy.py +++ b/isaaclab_arena_gr00t/tests/test_gr00t_closedloop_policy.py @@ -3,6 +3,9 @@ # # SPDX-License-Identifier: Apache-2.0 +import os +import sys + import yaml import pytest @@ -63,7 +66,10 @@ def gr00t_finetuned_model_path(tmp_path_factory): args.append("--color_jitter_params") # Tyro expects key-value pairs as separate arguments args.extend(["brightness", "0.3", "contrast", "0.4", "saturation", "0.5", "hue", "0.08"]) - run_subprocess(args) + env = os.environ.copy() + if os.environ.get("GROOT_DEPS_DIR"): + env["PYTHONPATH"] = os.environ["GROOT_DEPS_DIR"] + os.pathsep + env.get("PYTHONPATH", "") + run_subprocess(args, env=env) return model_dir / "checkpoint-10" diff --git a/isaaclab_arena_gr00t/utils/eagle_config_compat.py b/isaaclab_arena_gr00t/utils/eagle_config_compat.py new file mode 100644 index 000000000..bb7023e65 --- /dev/null +++ b/isaaclab_arena_gr00t/utils/eagle_config_compat.py @@ -0,0 +1,41 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers. +# SPDX-License-Identifier: Apache-2.0 + +"""Eagle config compat: alias _attn_implementation_autoset and set flash_attention_2 on loaded configs.""" + +_APPLIED = False + + +def apply_eagle_config_compat() -> None: + """Apply once per process. Idempotent.""" + global _APPLIED + if _APPLIED: + return + import transformers + import transformers.configuration_utils as configuration_utils + + _orig_getattribute = configuration_utils.PretrainedConfig.__getattribute__ + + def _compat_getattribute(self, name: str): + if name == "_attn_implementation_autoset": + try: + return _orig_getattribute(self, "_attn_implementation_internal") + except AttributeError: + pass + return _orig_getattribute(self, name) + + configuration_utils.PretrainedConfig.__getattribute__ = _compat_getattribute + + _orig_from_pretrained = transformers.AutoConfig.from_pretrained + + @classmethod + def _wrapped_from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): + config = _orig_from_pretrained(pretrained_model_name_or_path, *args, **kwargs) + for sub in ("text_config", "vision_config"): + sub_config = getattr(config, sub, None) + if sub_config is not None and getattr(sub_config, "_attn_implementation", None) != "flash_attention_2": + sub_config._attn_implementation = "flash_attention_2" + return config + + transformers.AutoConfig.from_pretrained = _wrapped_from_pretrained + _APPLIED = True diff --git a/submodules/IsaacLab b/submodules/IsaacLab index 5528d986d..28cf07780 160000 --- a/submodules/IsaacLab +++ b/submodules/IsaacLab @@ -1 +1 @@ -Subproject commit 5528d986d8909825a29f3c97656108abf054a261 +Subproject commit 28cf077800daff5eab76e02fe09ffc02f377aa1b