diff --git a/.github/actions/docker-build/action.yml b/.github/actions/docker-build/action.yml index 4595d47fd033..7f88241cfb8c 100644 --- a/.github/actions/docker-build/action.yml +++ b/.github/actions/docker-build/action.yml @@ -52,13 +52,18 @@ runs: local dockerfile_path="$4" local context_path="$5" + # Skip build if image already exists locally (e.g. built by a prior job on the same runner) + if docker image inspect "$image_tag" > /dev/null 2>&1; then + echo "Image $image_tag already exists locally, skipping build." + return 0 + fi + echo "Building Docker image: $image_tag" echo "Using Dockerfile: $dockerfile_path" echo "Build context: $context_path" # Build Docker image docker buildx build --progress=plain --platform linux/amd64 \ - -t isaac-lab-dev \ -t $image_tag \ --build-arg ISAACSIM_BASE_IMAGE_ARG="$isaacsim_base_image" \ --build-arg ISAACSIM_VERSION_ARG="$isaacsim_version" \ @@ -71,7 +76,8 @@ runs: --load $context_path echo "✅ Docker image built successfully: $image_tag" - docker images | grep isaac-lab-dev + echo "Current local Docker images:" + docker images } # Call the function with provided parameters diff --git a/.github/actions/ecr-build-push-pull/README.md b/.github/actions/ecr-build-push-pull/README.md new file mode 100644 index 000000000000..2e306c978b8e --- /dev/null +++ b/.github/actions/ecr-build-push-pull/README.md @@ -0,0 +1,24 @@ +# ecr-build-push-pull + +Builds a Docker image and pushes it to ECR, or pulls it if the tag already exists. +ECR is also used as the BuildKit layer cache. + +## Usage + +```yaml +- uses: ./.github/actions/ecr-build-push-pull + with: + image-tag: ${{ env.DOCKER_IMAGE_TAG }} + isaacsim-base-image: nvcr.io/nvidia/isaac-sim + isaacsim-version: 6.0.0 + dockerfile-path: docker/Dockerfile.base + cache-tag: cache-base + ecr-url: (optional, complete url for ECR storage) +``` + +## ECR URL resolution order + +1. `ecr-url` input +2. `ECR_CACHE_URL` environment variable on the runner +3. SSM parameter `/github-runner//ecr-cache-url` +4. If none resolve, ECR is skipped and the image is built locally diff --git a/.github/actions/ecr-build-push-pull/action.yml b/.github/actions/ecr-build-push-pull/action.yml new file mode 100644 index 000000000000..32985546e853 --- /dev/null +++ b/.github/actions/ecr-build-push-pull/action.yml @@ -0,0 +1,195 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +name: 'ECR Build-Push-Pull' +description: > + Builds a Docker image and pushes it to ECR, using ECR as the layer cache. + If the image already exists in ECR (same tag), pulls it instead of building. + Drop-in replacement for docker-build/action.yml with ECR-backed caching. + +inputs: + image-tag: + description: 'Tag for the Docker image (e.g. my-image:latest).' + required: true + isaacsim-base-image: + description: 'IsaacSim base image (passed as ISAACSIM_BASE_IMAGE_ARG build-arg).' + required: true + isaacsim-version: + description: 'IsaacSim version (passed as ISAACSIM_VERSION_ARG build-arg).' + required: true + dockerfile-path: + description: 'Path to Dockerfile, relative to the repository root' + default: 'docker/Dockerfile.base' + required: false + ecr-url: + description: > + ECR repository URL (e.g. "123456789.dkr.ecr.us-west-2.amazonaws.com/my-repo"). + Resolved in the following order: + 1. ecr-url input, if provided. + 2. ECR_CACHE_URL environment variable on the runner. + 3. SSM parameter /github-runner//ecr-cache-url. + 4. If still empty, ECR cache is skipped and the image is built locally. + required: false + default: '' + cache-tag: + description: Tag used for the ECR layer cache image (e.g. "cache-base", "cache-curobo"). + required: false + default: 'cache' + +runs: + using: composite + steps: + - name: ECR Build, Push, Pull + shell: bash + env: + INPUT_ECR_URL: ${{ inputs.ecr-url }} + run: | + IMAGE_TAG="${{ inputs.image-tag }}" + + # Prefer explicit input; fall back to ECR_CACHE_URL env var; then SSM + + ECR_URL="${INPUT_ECR_URL:-}" + + if [ -z "${ECR_URL}" ]; then + echo "🟢 ecr-url input not set, trying ECR_CACHE_URL env var..." + ECR_URL="${ECR_CACHE_URL:-}" + if [ -n "${ECR_URL}" ]; then + echo "🟢 Using ECR_CACHE_URL env var: ${ECR_URL}" + fi + fi + + if [ -z "${ECR_URL}" ]; then + echo "🟢 ECR_CACHE_URL env var not set, trying SSM..." + + IMDS_TOKEN=$(curl -sf -X PUT "http://169.254.169.254/latest/api/token" \ + -H "X-aws-ec2-metadata-token-ttl-seconds: 21600") || true + INSTANCE_ID=$(curl -sf -H "X-aws-ec2-metadata-token: ${IMDS_TOKEN}" \ + "http://169.254.169.254/latest/meta-data/instance-id") || true + INSTANCE_REGION=$(curl -sf -H "X-aws-ec2-metadata-token: ${IMDS_TOKEN}" \ + "http://169.254.169.254/latest/meta-data/placement/region") || true + + if [ -n "${INSTANCE_ID}" ]; then + ECR_URL=$(aws ssm get-parameter \ + --name "/github-runner/${INSTANCE_ID}/ecr-cache-url" \ + --region "${INSTANCE_REGION}" \ + --query 'Parameter.Value' --output text 2>/dev/null) || ECR_URL="" + if [ -n "${ECR_URL}" ]; then + echo "🟢 Resolved ECR URL from SSM (/github-runner/${INSTANCE_ID}/ecr-cache-url): ${ECR_URL}" + else + echo "🟢 SSM parameter not found for instance ${INSTANCE_ID}" + fi + else + echo "🟢 Not running on EC2 or Instance Metadata Service is unavailable, skipping SSM lookup" + fi + fi + + if [ -z "${ECR_URL}" ]; then + echo "🟠 ECR URL cannot be resolved. Building locally." + else + # Derive registry host and AWS region from the URL + REGISTRY=$(echo "${ECR_URL}" | cut -d'/' -f1) + AWS_REGION=$(echo "${REGISTRY}" | sed 's/.*\.dkr\.ecr\.\(.*\)\.amazonaws\.com/\1/') + + # Validate ECR URL + if [ "${AWS_REGION}" = "${REGISTRY}" ]; then + echo "🔴 Invalid ECR URL — cannot extract AWS region: ${ECR_URL}" + echo "🔴 Expected format: .dkr.ecr..amazonaws.com/" + exit 1 + fi + + # Derive ECR image tag by slugifying the local image-tag + # (colons/slashes not valid in ECR tags) + ECR_TAG=$(echo "${IMAGE_TAG}" | tr ':/' '--') + ECR_IMAGE="${ECR_URL}:${ECR_TAG}" + CACHE_IMAGE="${ECR_URL}:${{ inputs.cache-tag }}" + + # The runner's credential store is broken ("not implemented"), so use a temp config dir + # with credsStore disabled. Copy existing ~/.docker/config.json first to preserve other + # registry auths (e.g. nvcr.io), then strip credsStore/credHelpers. + # + DOCKER_CONFIG_DIR=$(mktemp -d) + trap 'rm -rf "${DOCKER_CONFIG_DIR}"' EXIT + if [ -f "${HOME}/.docker/config.json" ]; then + python3 -c "import json; cfg=json.load(open('${HOME}/.docker/config.json')); cfg['credsStore']=''; cfg.pop('credHelpers',None); json.dump(cfg,open('${DOCKER_CONFIG_DIR}/config.json','w'))" + else + echo '{"credsStore":""}' > "${DOCKER_CONFIG_DIR}/config.json" + fi + # + # Export so all subsequent docker calls in this job inherit it + echo "DOCKER_CONFIG=${DOCKER_CONFIG_DIR}" >> "$GITHUB_ENV" + export DOCKER_CONFIG="${DOCKER_CONFIG_DIR}" + + # Login to ECR registry with token provided by aws cli + echo "🟢 Logging into ECR registry..." + aws ecr get-login-password --region "${AWS_REGION}" | \ + docker login --username AWS --password-stdin "${REGISTRY}" + + # Try pulling the exact image from ECR to skip the build + echo "🟢 Checking for existing image ${ECR_IMAGE} in ECR..." + PULL_OUTPUT=$(docker pull "${ECR_IMAGE}" 2>&1) && { + echo "🟢 Image found in ECR, pulled: ${ECR_IMAGE}" + docker tag "${ECR_IMAGE}" "${IMAGE_TAG}" + echo "🟢 Tagged as: ${IMAGE_TAG} — skipping build." + exit 0 + } || { + echo "🟠 Image can't be pulled from ECR, will build locally:" + echo -e "\e[90m${PULL_OUTPUT}\e[0m" + } + fi + + # Log into nvcr.io so the Dockerfile can pull nvcr.io base images. + # Credential is written into the temp config dir when ECR is enabled, + # and into the default config when it is not. + if [ -n "${{ env.NGC_API_KEY }}" ]; then + echo "🟢 Logging into nvcr.io..." + docker login -u \$oauthtoken -p ${{ env.NGC_API_KEY }} nvcr.io + else + echo "🟠 NGC_API_KEY not set — skipping nvcr.io login (normal for fork PRs)" + fi + + # Build the image + + BUILD_ARGS=( + --progress=plain + --platform linux/amd64 + -f "${{ inputs.dockerfile-path }}" + --build-arg "ISAACSIM_BASE_IMAGE_ARG=${{ inputs.isaacsim-base-image }}" + --build-arg "ISAACSIM_VERSION_ARG=${{ inputs.isaacsim-version }}" + --build-arg "ISAACSIM_ROOT_PATH_ARG=/isaac-sim" + --build-arg "ISAACLAB_PATH_ARG=/workspace/isaaclab" + --build-arg "DOCKER_USER_HOME_ARG=/root" + -t "${IMAGE_TAG}" + ) + + # Add ECR layer caching + if [ -n "${ECR_URL}" ]; then + BUILD_ARGS+=( + --cache-from "type=registry,ref=${CACHE_IMAGE}" + --cache-to "type=registry,ref=${CACHE_IMAGE},mode=max" + -t "${ECR_IMAGE}" + ) + fi + + # Set up buildx builder with docker-container driver (required for registry cache export). + BUILDER_NAME="ci-builder-${{ github.run_id }}-${{ github.job }}" + + # Create builder (or fall back to reusing existing one on workflow re-runs) + docker buildx create --use --driver docker-container --name "${BUILDER_NAME}" \ + || docker buildx use "${BUILDER_NAME}" + + echo "🟢 Building ${IMAGE_TAG}..." + docker buildx build --load "${BUILD_ARGS[@]}" . + + # Clean up builder + docker buildx rm "${BUILDER_NAME}" || true + + echo "🟢 Build completed for ${IMAGE_TAG}, pushing to ECR..." + + if [ -n "${ECR_URL}" ]; then + docker push "${ECR_IMAGE}" + echo "🟢 Pushed ${ECR_IMAGE} to ECR" + fi + + echo "🟢 Image ${IMAGE_TAG} is ready!" diff --git a/.github/actions/run-tests/action.yml b/.github/actions/run-tests/action.yml index 9363ba9481b2..396b71b2c290 100644 --- a/.github/actions/run-tests/action.yml +++ b/.github/actions/run-tests/action.yml @@ -156,11 +156,11 @@ runs: # Copy test results with error handling echo "📋 Attempting to copy test results..." if docker cp $container_name:/workspace/isaaclab/tests/$result_file "$reports_dir/$result_file" 2>/dev/null; then - echo "✅ Test results copied successfully" + echo "✅ Test results copied successfully to $reports_dir/$result_file" else echo "❌ Failed to copy specific result file, trying to copy all test results..." if docker cp $container_name:/workspace/isaaclab/tests/ "$reports_dir/" 2>/dev/null; then - echo "✅ All test results copied successfully" + echo "✅ All test results copied successfully to $reports_dir/" # Look for any XML files and use the first one found if [ -f "$reports_dir/full_report.xml" ]; then mv "$reports_dir/full_report.xml" "$reports_dir/$result_file" diff --git a/.github/workflows/build.yml b/.github/workflows/build.yaml similarity index 63% rename from .github/workflows/build.yml rename to .github/workflows/build.yaml index 28df43417118..d4089454a2a9 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yaml @@ -3,15 +3,22 @@ # # SPDX-License-Identifier: BSD-3-Clause -name: Build and Test +name: Docker + Tests on: pull_request: + types: [opened, synchronize, reopened] + paths: + - 'source/**' + - 'docker/**' + - 'tools/**' + - 'apps/**' + - '.github/workflows/build.yaml' + - '.github/actions/**' branches: - - develop - main + - develop - 'release/**' - - feature/isaacsim-6-0 # Concurrency control to prevent parallel runs on the same PR concurrency: @@ -31,59 +38,103 @@ env: DOCKER_IMAGE_TAG: isaac-lab-dev:${{ github.event_name == 'pull_request' && format('pr-{0}', github.event.pull_request.number) || github.ref_name }}-${{ github.sha }} jobs: + build: + name: Build base Docker image + runs-on: [self-hosted, gpu] + steps: + - name: Checkout Code + uses: actions/checkout@v4 + with: + fetch-depth: 1 + lfs: true + + - name: Build and push to ECR + uses: ./.github/actions/ecr-build-push-pull + with: + image-tag: ${{ env.DOCKER_IMAGE_TAG }} + isaacsim-base-image: ${{ env.ISAACSIM_BASE_IMAGE }} + isaacsim-version: ${{ env.ISAACSIM_BASE_VERSION }} + dockerfile-path: docker/Dockerfile.base + cache-tag: cache-base + + # ECR images must be pruned with a lifecycle policy. + - name: Clean up stale Docker images + if: always() + run: docker image prune -f --filter "until=24h" || true + + build-curobo: + name: Build cuRobo Docker image + runs-on: [self-hosted, gpu] + steps: + - name: Checkout Code + uses: actions/checkout@v4 + with: + fetch-depth: 1 + lfs: true + + - name: Build and push to ECR + uses: ./.github/actions/ecr-build-push-pull + with: + image-tag: ${{ env.DOCKER_IMAGE_TAG }}-curobo + isaacsim-base-image: ${{ env.ISAACSIM_BASE_IMAGE }} + isaacsim-version: ${{ env.ISAACSIM_BASE_VERSION }} + dockerfile-path: docker/Dockerfile.curobo + cache-tag: cache-curobo + + # ECR images must be pruned with a lifecycle policy. + - name: Clean up stale Docker images + if: always() + run: docker image prune -f --filter "until=24h" || true + + test-isaaclab-tasks: runs-on: [self-hosted, gpu] timeout-minutes: 180 continue-on-error: true + needs: build steps: - name: Checkout Code uses: actions/checkout@v4 with: - fetch-depth: 0 + fetch-depth: 1 lfs: true - - name: Build Docker Image - uses: ./.github/actions/docker-build + - name: Pull image from ECR + uses: ./.github/actions/ecr-build-push-pull with: image-tag: ${{ env.DOCKER_IMAGE_TAG }} isaacsim-base-image: ${{ env.ISAACSIM_BASE_IMAGE }} isaacsim-version: ${{ env.ISAACSIM_BASE_VERSION }} + dockerfile-path: docker/Dockerfile.base + cache-tag: cache-base - name: Run IsaacLab Tasks Tests uses: ./.github/actions/run-tests with: test-path: "tools" - result-file: "isaaclab-tasks-report.xml" - container-name: "isaac-lab-tasks-test-$$" + result-file: "${{ github.job }}-report.xml" + container-name: "isaac-lab-tasks-test-${{ github.run_id }}-${{ github.run_attempt }}" image-tag: ${{ env.DOCKER_IMAGE_TAG }} pytest-options: "" filter-pattern: "isaaclab_tasks" include-files: "test_multi_agent_environments.py,test_pickplace_stack_environments.py,test_environments.py,test_factory_environments.py,test_environments_training.py,test_cartpole_showcase_environments.py,test_teleop_environments.py" - - name: Copy Test Results from IsaacLab Tasks Container - run: | - CONTAINER_NAME="isaac-lab-tasks-test-$$" - if docker ps -a | grep -q $CONTAINER_NAME; then - echo "Copying test results from IsaacLab Tasks container..." - docker cp $CONTAINER_NAME:/workspace/isaaclab/tests/isaaclab-tasks-report.xml reports/ 2>/dev/null || echo "No test results to copy from IsaacLab Tasks container" - fi - - name: Upload IsaacLab Tasks Test Results uses: actions/upload-artifact@v4 if: always() with: name: isaaclab-tasks-test-results - path: reports/isaaclab-tasks-report.xml + path: reports/${{ github.job }}-report.xml retention-days: 1 compression-level: 9 - name: Check Test Results for Fork PRs if: github.event.pull_request.head.repo.full_name != github.repository run: | - if [ -f "reports/isaaclab-tasks-report.xml" ]; then + if [ -f "reports/${{ github.job }}-report.xml" ]; then # Check if the test results contain any failures - if grep -q 'failures="[1-9]' reports/isaaclab-tasks-report.xml || grep -q 'errors="[1-9]' reports/isaaclab-tasks-report.xml; then + if grep -q 'failures="[1-9]' reports/${{ github.job }}-report.xml || grep -q 'errors="[1-9]' reports/${{ github.job }}-report.xml; then echo "Tests failed for PR from fork. The test report is in the logs. Failing the job." exit 1 fi @@ -92,59 +143,59 @@ jobs: exit 1 fi + # ECR images must be pruned with a lifecycle policy. + - name: Clean up stale Docker images + if: always() + run: docker image prune -f --filter "until=24h" || true + test-isaaclab-tasks-2: runs-on: [self-hosted, gpu] timeout-minutes: 180 continue-on-error: true + needs: build steps: - name: Checkout Code uses: actions/checkout@v4 with: - fetch-depth: 0 + fetch-depth: 1 lfs: true - - name: Build Docker Image - uses: ./.github/actions/docker-build + - name: Pull image from ECR + uses: ./.github/actions/ecr-build-push-pull with: image-tag: ${{ env.DOCKER_IMAGE_TAG }} isaacsim-base-image: ${{ env.ISAACSIM_BASE_IMAGE }} isaacsim-version: ${{ env.ISAACSIM_BASE_VERSION }} + dockerfile-path: docker/Dockerfile.base + cache-tag: cache-base - name: Run IsaacLab Tasks Tests 2 uses: ./.github/actions/run-tests with: test-path: "tools" - result-file: "isaaclab-tasks-2-report.xml" - container-name: "isaac-lab-tasks-2-test-$$" + result-file: "${{ github.job }}-report.xml" + container-name: "isaac-lab-tasks-2-test-${{ github.run_id }}-${{ github.run_attempt }}" image-tag: ${{ env.DOCKER_IMAGE_TAG }} pytest-options: "" filter-pattern: "isaaclab_tasks" include-files: "test_teleop_environments_with_stage_in_memory.py,test_lift_teddy_bear.py,test_environment_determinism.py,test_hydra.py,test_env_cfg_no_forbidden_imports.py,test_rl_device_separation.py,test_cartpole_showcase_environments_with_stage_in_memory.py,test_environments_with_stage_in_memory.py" - - name: Copy Test Results from IsaacLab Tasks 2 Container - run: | - CONTAINER_NAME="isaac-lab-tasks-2-test-$$" - if docker ps -a | grep -q $CONTAINER_NAME; then - echo "Copying test results from IsaacLab Tasks 2 container..." - docker cp $CONTAINER_NAME:/workspace/isaaclab/tests/isaaclab-tasks-2-report.xml reports/ 2>/dev/null || echo "No test results to copy from IsaacLab Tasks 2 container" - fi - - name: Upload IsaacLab Tasks 2 Test Results uses: actions/upload-artifact@v4 if: always() with: name: isaaclab-tasks-2-test-results - path: reports/isaaclab-tasks-2-report.xml + path: reports/${{ github.job }}-report.xml retention-days: 1 compression-level: 9 - name: Check Test Results for Fork PRs if: github.event.pull_request.head.repo.full_name != github.repository run: | - if [ -f "reports/isaaclab-tasks-2-report.xml" ]; then + if [ -f "reports/${{ github.job }}-report.xml" ]; then # Check if the test results contain any failures - if grep -q 'failures="[1-9]' reports/isaaclab-tasks-2-report.xml || grep -q 'errors="[1-9]' reports/isaaclab-tasks-2-report.xml; then + if grep -q 'failures="[1-9]' reports/${{ github.job }}-report.xml || grep -q 'errors="[1-9]' reports/${{ github.job }}-report.xml; then echo "Tests failed for PR from fork. The test report is in the logs. Failing the job." exit 1 fi @@ -153,58 +204,58 @@ jobs: exit 1 fi + # ECR images must be pruned with a lifecycle policy. + - name: Clean up stale Docker images + if: always() + run: docker image prune -f --filter "until=24h" || true + test-general: runs-on: [self-hosted, gpu] timeout-minutes: 180 + needs: build steps: - name: Checkout Code uses: actions/checkout@v4 with: - fetch-depth: 0 + fetch-depth: 1 lfs: true - - name: Build Docker Image - uses: ./.github/actions/docker-build + - name: Pull image from ECR + uses: ./.github/actions/ecr-build-push-pull with: image-tag: ${{ env.DOCKER_IMAGE_TAG }} isaacsim-base-image: ${{ env.ISAACSIM_BASE_IMAGE }} isaacsim-version: ${{ env.ISAACSIM_BASE_VERSION }} + dockerfile-path: docker/Dockerfile.base + cache-tag: cache-base - name: Run General Tests id: run-general-tests uses: ./.github/actions/run-tests with: test-path: "tools" - result-file: "general-tests-report.xml" - container-name: "isaac-lab-general-test-$$" + result-file: "${{ github.job }}-report.xml" + container-name: "isaac-lab-general-test-${{ github.run_id }}-${{ github.run_attempt }}" image-tag: ${{ env.DOCKER_IMAGE_TAG }} pytest-options: "" filter-pattern: "not isaaclab_tasks" - - name: Copy Test Results from General Tests Container - run: | - CONTAINER_NAME="isaac-lab-general-test-$$" - if docker ps -a | grep -q $CONTAINER_NAME; then - echo "Copying test results from General Tests container..." - docker cp $CONTAINER_NAME:/workspace/isaaclab/tests/general-tests-report.xml reports/ 2>/dev/null || echo "No test results to copy from General Tests container" - fi - - name: Upload General Test Results uses: actions/upload-artifact@v4 if: always() with: name: general-test-results - path: reports/general-tests-report.xml + path: reports/${{ github.job }}-report.xml retention-days: 1 compression-level: 9 - name: Check Test Results for Fork PRs if: github.event.pull_request.head.repo.full_name != github.repository run: | - if [ -f "reports/general-tests-report.xml" ]; then + if [ -f "reports/${{ github.job }}-report.xml" ]; then # Check if the test results contain any failures - if grep -q 'failures="[1-9]' reports/general-tests-report.xml || grep -q 'errors="[1-9]' reports/general-tests-report.xml; then + if grep -q 'failures="[1-9]' reports/${{ github.job }}-report.xml || grep -q 'errors="[1-9]' reports/${{ github.job }}-report.xml; then echo "Tests failed for PR from fork. The test report is in the logs. Failing the job." exit 1 fi @@ -213,60 +264,59 @@ jobs: exit 1 fi + # ECR images must be pruned with a lifecycle policy. + - name: Clean up stale Docker images + if: always() + run: docker image prune -f --filter "until=24h" || true + test-curobo: runs-on: [self-hosted, gpu] timeout-minutes: 120 continue-on-error: true + needs: build-curobo steps: - name: Checkout Code uses: actions/checkout@v4 with: - fetch-depth: 0 + fetch-depth: 1 lfs: true - - name: Build Docker Image (cuRobo) - uses: ./.github/actions/docker-build + - name: Pull cuRobo image from ECR + uses: ./.github/actions/ecr-build-push-pull with: image-tag: ${{ env.DOCKER_IMAGE_TAG }}-curobo isaacsim-base-image: ${{ env.ISAACSIM_BASE_IMAGE }} isaacsim-version: ${{ env.ISAACSIM_BASE_VERSION }} - dockerfile-path: 'docker/Dockerfile.curobo' + dockerfile-path: docker/Dockerfile.curobo + cache-tag: cache-curobo - name: Run cuRobo and SkillGen Tests uses: ./.github/actions/run-tests with: test-path: "tools" - result-file: "curobo-tests-report.xml" - container-name: "isaac-lab-curobo-test-$$" + result-file: "${{ github.job }}-report.xml" + container-name: "isaac-lab-curobo-test-${{ github.run_id }}-${{ github.run_attempt }}" image-tag: ${{ env.DOCKER_IMAGE_TAG }}-curobo pytest-options: "" filter-pattern: "" curobo-only: "true" - - name: Copy Test Results from cuRobo Container - run: | - CONTAINER_NAME="isaac-lab-curobo-test-$$" - if docker ps -a | grep -q $CONTAINER_NAME; then - echo "Copying test results from cuRobo container..." - docker cp $CONTAINER_NAME:/workspace/isaaclab/tests/curobo-tests-report.xml reports/ 2>/dev/null || echo "No test results to copy from cuRobo container" - fi - - name: Upload cuRobo Test Results uses: actions/upload-artifact@v4 if: always() with: name: curobo-test-results - path: reports/curobo-tests-report.xml + path: reports/${{ github.job }}-report.xml retention-days: 1 compression-level: 9 - name: Check Test Results for Fork PRs if: github.event.pull_request.head.repo.full_name != github.repository run: | - if [ -f "reports/curobo-tests-report.xml" ]; then + if [ -f "reports/${{ github.job }}-report.xml" ]; then # Check if the test results contain any failures - if grep -q 'failures="[1-9]' reports/curobo-tests-report.xml || grep -q 'errors="[1-9]' reports/curobo-tests-report.xml; then + if grep -q 'failures="[1-9]' reports/${{ github.job }}-report.xml || grep -q 'errors="[1-9]' reports/${{ github.job }}-report.xml; then echo "Tests failed for PR from fork. The test report is in the logs. Failing the job." exit 1 fi @@ -275,21 +325,28 @@ jobs: exit 1 fi + # ECR images must be pruned with a lifecycle policy. + - name: Clean up stale Docker images + if: always() + run: docker image prune -f --filter "until=24h" || true + combine-results: - needs: [test-isaaclab-tasks, test-isaaclab-tasks-2, test-general, test-curobo] - runs-on: [self-hosted, gpu] + needs: [build, build-curobo, test-isaaclab-tasks, test-isaaclab-tasks-2, test-general, test-curobo] + runs-on: ubuntu-latest if: always() steps: - name: Checkout Code uses: actions/checkout@v4 with: - fetch-depth: 0 + fetch-depth: 1 lfs: false + sparse-checkout: .github/actions + sparse-checkout-cone-mode: true - name: Create Reports Directory run: | - mkdir -p reports + mkdir -pv reports - name: Download Test Results uses: actions/download-artifact@v4 @@ -306,10 +363,12 @@ jobs: continue-on-error: true - name: Download General Test Results + if: always() uses: actions/download-artifact@v4 with: name: general-test-results path: reports/ + continue-on-error: true - name: Download cuRobo Test Results uses: actions/download-artifact@v4 @@ -319,6 +378,7 @@ jobs: continue-on-error: true - name: Combine All Test Results + if: always() uses: ./.github/actions/combine-results with: tests-dir: "reports" diff --git a/.github/workflows/check-links.yml b/.github/workflows/check-links.yml index ac142552172d..d77832d3f481 100644 --- a/.github/workflows/check-links.yml +++ b/.github/workflows/check-links.yml @@ -3,15 +3,11 @@ # # SPDX-License-Identifier: BSD-3-Clause -name: Check Documentation Links +name: Documentation Links on: - # Run on pull requests that modify documentation + # Run on all pull requests pull_request: - paths: - - 'docs/**' - - '**.md' - - '.github/workflows/check-links.yml' # Run on pushes to main branches push: branches: @@ -19,10 +15,6 @@ on: - develop - 'release/**' - 'feature/isaacsim-6-0' - paths: - - 'docs/**' - - '**.md' - - '.github/workflows/check-links.yml' # Allow manual trigger workflow_dispatch: # Run weekly to catch external links that break over time @@ -42,9 +34,32 @@ jobs: - name: Checkout code uses: actions/checkout@v4 with: - fetch-depth: 0 + fetch-depth: 2 + filter: tree:0 + + - name: Check if latest commit touched documentation + id: check_docs + run: | + # Get the files changed by this PR/commit. + # For PRs: HEAD is the merge commit, HEAD^1 is the base branch tip. + # Diffing HEAD^1..HEAD shows exactly what the PR introduces — files only + # changed on the base branch (not by the PR) are identical in both and + # won't appear. Works for push events too (last commit vs its parent). + CHANGED=$(git diff --name-only HEAD^1 HEAD || true) + echo "Files changed:" + echo "$CHANGED" + + # Check against doc file types and doc paths + if echo "$CHANGED" | grep -qE '\.md$|\.rst$|^docs/|\.github/workflows/check-links\.yml'; then + echo "Documentation or workflow files changed. Proceeding..." + echo "run_job=true" >> "$GITHUB_OUTPUT" + else + echo "No documentation files changed in this specific commit. Skipping." + echo "run_job=false" >> "$GITHUB_OUTPUT" + fi - name: Restore lychee cache + if: steps.check_docs.outputs.run_job == 'true' || github.event_name == 'workflow_dispatch' || github.event_name == 'schedule' uses: actions/cache@v4 with: path: .lycheecache @@ -52,9 +67,11 @@ jobs: restore-keys: cache-lychee- - name: Run Link Checker + if: steps.check_docs.outputs.run_job == 'true' || github.event_name == 'workflow_dispatch' || github.event_name == 'schedule' uses: lycheeverse/lychee-action@v2 with: # Check all markdown files and documentation + # Excluded are known crawl blockers args: >- --verbose --no-progress @@ -76,9 +93,17 @@ jobs: --exclude 'user@' --exclude 'helm\.ngc\.nvidia\.com' --exclude 'slurm\.schedmd\.com' - --max-retries 3 - --retry-wait-time 5 - --timeout 30 + --exclude 'graphics\.pixar\.com' + --exclude 'openpbs\.org' + --exclude 'docutils\.sourceforge\.io' + --exclude 'huggingface\.co/datasets/nvidia/PhysicalAI-Robotics-NuRec' + --exclude 'huggingface\.co/nvidia/COMPASS' + --exclude 'huggingface\.co/nvidia/X-Mobility' + --exclude 'openusd\.org' + --max-retries 5 + --retry-wait-time 10 + --timeout 20 + --user-agent 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36' --accept 200,201,202,203,204,206,301,302,303,307,308,429 --scheme https --scheme http @@ -94,7 +119,7 @@ jobs: token: ${{ secrets.GITHUB_TOKEN }} - name: Print results to logs - if: always() + if: always() && (steps.check_docs.outputs.run_job == 'true' || github.event_name == 'workflow_dispatch' || github.event_name == 'schedule') run: | echo "========================================" echo "Link Checker Results:" diff --git a/.github/workflows/daily-compatibility.yml b/.github/workflows/daily-compatibility.yml index bbf59e45160d..bb2845f18252 100644 --- a/.github/workflows/daily-compatibility.yml +++ b/.github/workflows/daily-compatibility.yml @@ -70,9 +70,9 @@ jobs: steps: - name: Checkout Code - uses: actions/checkout@v3 + uses: actions/checkout@v4 with: - fetch-depth: 0 + fetch-depth: 1 lfs: true - name: Build Docker Image @@ -127,9 +127,9 @@ jobs: steps: - name: Checkout Code - uses: actions/checkout@v3 + uses: actions/checkout@v4 with: - fetch-depth: 0 + fetch-depth: 1 lfs: true - name: Build Docker Image @@ -168,15 +168,17 @@ jobs: combine-compat-results: needs: [test-isaaclab-tasks-compat, test-general-compat] - runs-on: [self-hosted, gpu] + runs-on: ubuntu-latest if: always() steps: - name: Checkout Code - uses: actions/checkout@v3 + uses: actions/checkout@v4 with: - fetch-depth: 0 + fetch-depth: 1 lfs: false + sparse-checkout: .github/actions + sparse-checkout-cone-mode: true - name: Create Reports Directory run: | @@ -217,15 +219,17 @@ jobs: notify-compatibility-status: needs: [setup-versions, combine-compat-results] - runs-on: [self-hosted, gpu] + runs-on: ubuntu-latest if: always() steps: - name: Checkout Code - uses: actions/checkout@v3 + uses: actions/checkout@v4 with: - fetch-depth: 0 + fetch-depth: 1 lfs: false + sparse-checkout: .github/actions + sparse-checkout-cone-mode: true - name: Create Compatibility Report run: | diff --git a/.github/workflows/docs.yaml b/.github/workflows/docs.yaml index 9b27d63c1562..de58be47af5a 100644 --- a/.github/workflows/docs.yaml +++ b/.github/workflows/docs.yaml @@ -3,7 +3,7 @@ # # SPDX-License-Identifier: BSD-3-Clause -name: Build & deploy docs +name: Docs on: push: @@ -13,6 +13,8 @@ on: - 'release/**' - 'feature/isaacsim-6-0' pull_request: + # we're skipping the branches and paths filter to allow docs to be built on any PR because heredoc is used + # additionally, we have a check that determines what version of docs will be built types: [opened, synchronize, reopened] concurrency: @@ -20,8 +22,8 @@ concurrency: cancel-in-progress: true jobs: - check-secrets: - name: Check secrets + doc-build-type: + name: Detect Doc Build Type runs-on: ubuntu-latest outputs: trigger-deploy: ${{ steps.trigger-deploy.outputs.defined }} @@ -30,19 +32,21 @@ jobs: env: REPO_NAME: ${{ secrets.REPO_NAME }} if: "${{ github.repository == env.REPO_NAME && (github.ref == 'refs/heads/main' || github.ref == 'refs/heads/develop' || startsWith(github.ref, 'refs/heads/release/')) }}" - run: echo "defined=true" >> "$GITHUB_OUTPUT" + run: echo "defined=true" >> "$GITHUB_OUTPUT"; echo "Docs will be built multi-version and deployed" - build-docs: - name: Build Docs + build-latest-docs: + name: Build Latest Docs runs-on: ubuntu-latest - needs: [check-secrets] + needs: [doc-build-type] + # run on non-deploy branches to build current version docs only + if: needs.doc-build-type.outputs.trigger-deploy != 'true' steps: - name: Checkout code - uses: actions/checkout@v2 + uses: actions/checkout@v4 - name: Setup python - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: "3.12" architecture: x64 @@ -51,15 +55,43 @@ jobs: working-directory: ./docs run: pip install -r requirements.txt - - name: Check branch docs building + - name: Build current version docs working-directory: ./docs - if: needs.check-secrets.outputs.trigger-deploy != 'true' run: make current-docs + - name: Upload docs artifact + uses: actions/upload-artifact@v4 + with: + name: docs-html + path: ./docs/_build + + build-multi-docs: + name: Build Multi-Version Docs + runs-on: ubuntu-latest + needs: [doc-build-type] + # run on deploy branches to create multi-version docs + if: needs.doc-build-type.outputs.trigger-deploy == 'true' + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Setup python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + architecture: x64 + + - name: Install dev requirements + working-directory: ./docs + run: pip install -r requirements.txt + - name: Generate multi-version docs working-directory: ./docs run: | git fetch --prune --unshallow --tags + git checkout --detach HEAD + git for-each-ref --format="%(refname:short)" refs/heads/ | xargs -r git branch -D make multi-docs - name: Upload docs artifact @@ -71,8 +103,9 @@ jobs: deploy-docs: name: Deploy Docs runs-on: ubuntu-latest - needs: [check-secrets, build-docs] - if: needs.check-secrets.outputs.trigger-deploy == 'true' + needs: [doc-build-type, build-multi-docs] + # deploy only on "deploy" branches + if: needs.doc-build-type.outputs.trigger-deploy == 'true' steps: - name: Download docs artifact diff --git a/.github/workflows/license-check.yaml b/.github/workflows/license-check.yaml index ca7ead9aa345..9e1914121779 100644 --- a/.github/workflows/license-check.yaml +++ b/.github/workflows/license-check.yaml @@ -3,7 +3,7 @@ # # SPDX-License-Identifier: BSD-3-Clause -name: Check Python Dependency Licenses +name: Python Dependency Licenses Check on: pull_request: @@ -19,7 +19,9 @@ jobs: steps: - name: Checkout code - uses: actions/checkout@v3 + uses: actions/checkout@v4 + with: + filter: tree:0 # - name: Install jq # run: sudo apt-get update && sudo apt-get install -y jq @@ -48,7 +50,7 @@ jobs: - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: '3.11' # Adjust as needed @@ -63,10 +65,6 @@ jobs: pip install --upgrade pip pip install 'isaacsim[all,extscache]==${{ vars.ISAACSIM_BASE_VERSION || '5.1.0' }}' --extra-index-url https://pypi.nvidia.com chmod +x ./isaaclab.sh # Make sure the script is executable - # Install torch - pip install -U torch==2.9.0 torchvision==0.24.0 --index-url https://download.pytorch.org/whl/cu128 - # clean up pip cache - pip cache purge # install all lab dependencies ./isaaclab.sh -i diff --git a/.github/workflows/license-exceptions.json b/.github/workflows/license-exceptions.json index fff08cf14207..e17b2fadd1ff 100644 --- a/.github/workflows/license-exceptions.json +++ b/.github/workflows/license-exceptions.json @@ -464,6 +464,16 @@ "license": "PSF-2.0", "comment": "PSF-2.0 / OSRB" }, + { + "package": "standard-distutils", + "license": "Python Software Foundation License", + "comment": "PSFL / OSRB" + }, + { + "package": "vhacdx", + "license": "UNKNOWN", + "comment": "BSD/MIT" + }, { "package": "cuda-bindings", "license": "LicenseRef-NVIDIA-SOFTWARE-LICENSE", diff --git a/.github/workflows/postmerge-ci.yml b/.github/workflows/postmerge-ci.yml index e19613cd9dd3..c77aada7a74b 100644 --- a/.github/workflows/postmerge-ci.yml +++ b/.github/workflows/postmerge-ci.yml @@ -42,7 +42,7 @@ jobs: - name: Checkout Code uses: actions/checkout@v4 with: - fetch-depth: 0 + fetch-depth: 1 lfs: true - name: Set up QEMU diff --git a/.github/workflows/pre-commit.yaml b/.github/workflows/pre-commit.yaml index f71f1f373899..983e2429b310 100644 --- a/.github/workflows/pre-commit.yaml +++ b/.github/workflows/pre-commit.yaml @@ -3,7 +3,7 @@ # # SPDX-License-Identifier: BSD-3-Clause -name: Run linters using pre-commit +name: Run `pre-commit` on: pull_request: @@ -13,8 +13,8 @@ jobs: pre-commit: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v3 - - uses: actions/setup-python@v3 + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 with: python-version: "3.12" - uses: pre-commit/action@v3.0.0 diff --git a/.vscode/tools/launch.template.json b/.vscode/tools/launch.template.json index a44f114c822b..9cc6c7051871 100644 --- a/.vscode/tools/launch.template.json +++ b/.vscode/tools/launch.template.json @@ -11,8 +11,14 @@ "program": "${file}", "console": "integratedTerminal" }, + // Attach to listening debugpy + // Usage (VSCode but should be similar for other IDEs): + // 1. Set your breakpoints + // 2. Run your code under debugpy like so: + // ./isaaclab.sh -p -m debugpy --listen 3000 --wait-for-client -c "from isaaclab.cli import cli; cli()" [cli_args] + // 3. Select this "Python: Debugger Attach" configuration and press green play button or F5 { - "name": "Python: Attach (windows-x86_64/linux-x86_64)", + "name": "Python: Debugger Attach", "type": "python", "request": "attach", "port": 3000, diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 10d84f479390..201c70bf727e 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -90,6 +90,7 @@ Guidelines for modifications: * Jan Kerner * Jean Tampon * Jeonghwan Kim +* Ji Yuan Feng * Jia Lin Yuan * Jiakai Zhang * Jinghuan Shang diff --git a/apps/isaaclab.python.headless.kit b/apps/isaaclab.python.headless.kit index 51175a8643ca..28e30f64c4c0 100644 --- a/apps/isaaclab.python.headless.kit +++ b/apps/isaaclab.python.headless.kit @@ -109,13 +109,6 @@ app.extensions.excluded = ["omni.kit.pip_archive"] # Extensions ############################### -[settings.exts."omni.kit.registry.nucleus"] -registries = [ - { name = "kit/default", url = "https://ovextensionsprod.blob.core.windows.net/exts/kit/prod/${kit_version_short}/shared" }, - { name = "kit/sdk", url = "https://ovextensionsprod.blob.core.windows.net/exts/kit/prod/sdk/${kit_version_short}/${kit_git_hash}" }, - { name = "kit/community", url = "https://dw290v42wisod.cloudfront.net/exts/kit/community" }, -] - [settings.app.extensions] skipPublishVerification = false registryEnabled = true diff --git a/apps/isaaclab.python.headless.rendering.kit b/apps/isaaclab.python.headless.rendering.kit index 4ff987a93276..12229aca0eca 100644 --- a/apps/isaaclab.python.headless.rendering.kit +++ b/apps/isaaclab.python.headless.rendering.kit @@ -109,13 +109,6 @@ metricsAssembler.changeListenerEnabled = false # explicitly disable omni.kit.pip_archive to prevent conflicting dependencies app.extensions.excluded = ["omni.kit.pip_archive"] -[settings.exts."omni.kit.registry.nucleus"] -registries = [ - { name = "kit/default", url = "https://ovextensionsprod.blob.core.windows.net/exts/kit/prod/${kit_version_short}/shared" }, - { name = "kit/sdk", url = "https://ovextensionsprod.blob.core.windows.net/exts/kit/prod/sdk/${kit_version_short}/${kit_git_hash}" }, - { name = "kit/community", url = "https://dw290v42wisod.cloudfront.net/exts/kit/community" }, -] - [settings.app.python] # These disable the kit app from also printing out python output, which gets confusing interceptSysStdOutput = false diff --git a/apps/isaaclab.python.kit b/apps/isaaclab.python.kit index f3a7d3083dae..8fa08e6ac3db 100644 --- a/apps/isaaclab.python.kit +++ b/apps/isaaclab.python.kit @@ -262,13 +262,6 @@ outDirectory = "${data}" # Extensions ############################### -[settings.exts."omni.kit.registry.nucleus"] -registries = [ - { name = "kit/default", url = "https://ovextensionsprod.blob.core.windows.net/exts/kit/prod/${kit_version_short}/shared" }, - { name = "kit/sdk", url = "https://ovextensionsprod.blob.core.windows.net/exts/kit/prod/sdk/${kit_version_short}/${kit_git_hash}" }, - { name = "kit/community", url = "https://dw290v42wisod.cloudfront.net/exts/kit/community" }, -] - [settings.app.extensions] skipPublishVerification = false registryEnabled = true diff --git a/apps/isaaclab.python.rendering.kit b/apps/isaaclab.python.rendering.kit index e4993f93d5a3..ad58b159186a 100644 --- a/apps/isaaclab.python.rendering.kit +++ b/apps/isaaclab.python.rendering.kit @@ -110,13 +110,6 @@ fabricUpdateJointStates = false ### When Direct GPU mode is enabled (suppressReadback=true) use direct interop between PhysX GPU and Fabric fabricUseGPUInterop = true -[settings.exts."omni.kit.registry.nucleus"] -registries = [ - { name = "kit/default", url = "https://ovextensionsprod.blob.core.windows.net/exts/kit/prod/${kit_version_short}/shared" }, - { name = "kit/sdk", url = "https://ovextensionsprod.blob.core.windows.net/exts/kit/prod/sdk/${kit_version_short}/${kit_git_hash}" }, - { name = "kit/community", url = "https://dw290v42wisod.cloudfront.net/exts/kit/community" }, -] - [settings.app.python] # These disable the kit app from also printing out python output, which gets confusing interceptSysStdOutput = false diff --git a/docker/Dockerfile.base b/docker/Dockerfile.base index 85c57547312e..f13d7918c9a6 100644 --- a/docker/Dockerfile.base +++ b/docker/Dockerfile.base @@ -40,10 +40,10 @@ USER root RUN apt-get update && \ apt-get install -y --no-install-recommends \ build-essential \ - cmake \ - git \ libglib2.0-0 \ ncurses-term \ + cmake \ + git \ wget # Install packages needed to build imgui-bundle on arm64 as no prebuilt pip wheel is provided @@ -54,8 +54,11 @@ RUN if [ "$(dpkg --print-architecture)" = "arm64" ]; then \ libx11-dev libxcursor-dev libxi-dev libxinerama-dev libxrandr-dev; \ fi -# Copy the Isaac Lab directory (files to exclude are defined in .dockerignore) -COPY ../ ${ISAACLAB_PATH} +# copy files necessary for installing Isaac Lab and its dependencies +# this way we minimize the chance of cache invalidation on file edits +COPY ../isaaclab.* ../environment.yml ../pyproject.toml ${ISAACLAB_PATH}/ +COPY ../tools/ ${ISAACLAB_PATH}/tools/ +COPY ../source/ ${ISAACLAB_PATH}/source/ # Fix the line endings for the shell scripts (Windows git may add \r) RUN find ${ISAACLAB_PATH} -type f -name "*.sh" -exec sed -i 's/\r$//' {} + @@ -69,11 +72,9 @@ RUN ln -sf ${ISAACSIM_ROOT_PATH} ${ISAACLAB_PATH}/_isaac_sim # Install toml dependency RUN ${ISAACLAB_PATH}/isaaclab.sh -p -m pip install toml -# Install apt dependencies for extensions that declare them in their extension.toml -RUN ${ISAACLAB_PATH}/isaaclab.sh -p ${ISAACLAB_PATH}/tools/install_deps.py apt ${ISAACLAB_PATH}/source - -# Apt Cleanup -RUN apt-get -y autoremove && apt-get clean && \ +# Install apt dependencies for extensions that declare them in their extension.toml, +RUN ${ISAACLAB_PATH}/isaaclab.sh -p ${ISAACLAB_PATH}/tools/install_deps.py apt ${ISAACLAB_PATH}/source && \ + apt-get -y autoremove && apt-get clean && \ rm -rf /var/lib/apt/lists/* # for singularity usage, have to create the directories that will binded @@ -116,6 +117,9 @@ RUN echo "export ISAACLAB_PATH=${ISAACLAB_PATH}" >> ${HOME}/.bashrc && \ echo "shopt -s histappend" >> /root/.bashrc && \ echo "PROMPT_COMMAND='history -a'" >> /root/.bashrc +# copy the rest of the files +COPY ../ ${ISAACLAB_PATH}/ + # make working directory as the Isaac Lab directory # this is the default directory when the container is run WORKDIR ${ISAACLAB_PATH} diff --git a/docker/Dockerfile.curobo b/docker/Dockerfile.curobo index 8882ff024081..36072612f576 100644 --- a/docker/Dockerfile.curobo +++ b/docker/Dockerfile.curobo @@ -36,17 +36,23 @@ ENV DEBIAN_FRONTEND=noninteractive USER root -# Install dependencies and remove cache -RUN --mount=type=cache,target=/var/cache/apt \ - apt-get update && apt-get install -y --no-install-recommends \ +# Install dependencies +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ build-essential \ - cmake \ - git \ libglib2.0-0 \ ncurses-term \ - wget && \ - apt -y autoremove && apt clean autoclean && \ - rm -rf /var/lib/apt/lists/* + cmake \ + git \ + wget + +# Install packages needed to build imgui-bundle on arm64 as no prebuilt pip wheel is provided +RUN if [ "$(dpkg --print-architecture)" = "arm64" ]; then \ + apt-get update && \ + apt-get install -y --no-install-recommends \ + libgl1-mesa-dev libopengl-dev libglx-dev \ + libx11-dev libxcursor-dev libxi-dev libxinerama-dev libxrandr-dev; \ + fi # Detect Ubuntu version and install CUDA 12.8 via NVIDIA network repo (cuda-keyring) RUN set -euo pipefail && \ @@ -83,8 +89,11 @@ ENV PATH=${CUDA_HOME}/bin:${PATH} ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH} ENV TORCH_CUDA_ARCH_LIST=8.0+PTX -# Copy the Isaac Lab directory (files to exclude are defined in .dockerignore) -COPY ../ ${ISAACLAB_PATH} +# copy files necessary for installing Isaac Lab and its dependencies +# this way we minimize the chance of cache invalidation on file edits +COPY ../isaaclab.* ../environment.yml ../pyproject.toml ${ISAACLAB_PATH}/ +COPY ../tools/ ${ISAACLAB_PATH}/tools/ +COPY ../source/ ${ISAACLAB_PATH}/source/ # Fix the line endings for the shell scripts (Windows git may add \r) RUN find ${ISAACLAB_PATH} -type f -name "*.sh" -exec sed -i 's/\r$//' {} + @@ -98,10 +107,9 @@ RUN ln -sf ${ISAACSIM_ROOT_PATH} ${ISAACLAB_PATH}/_isaac_sim # Install toml dependency RUN ${ISAACLAB_PATH}/isaaclab.sh -p -m pip install toml -# Install apt dependencies for extensions that declare them in their extension.toml -RUN --mount=type=cache,target=/var/cache/apt \ - ${ISAACLAB_PATH}/isaaclab.sh -p ${ISAACLAB_PATH}/tools/install_deps.py apt ${ISAACLAB_PATH}/source && \ - apt -y autoremove && apt clean autoclean && \ +# Install apt dependencies for extensions that declare them in their extension.toml, +RUN ${ISAACLAB_PATH}/isaaclab.sh -p ${ISAACLAB_PATH}/tools/install_deps.py apt ${ISAACLAB_PATH}/source && \ + apt-get -y autoremove && apt-get clean && \ rm -rf /var/lib/apt/lists/* # for singularity usage, have to create the directories that will binded @@ -164,6 +172,9 @@ RUN echo "export ISAACLAB_PATH=${ISAACLAB_PATH}" >> ${HOME}/.bashrc && \ echo "shopt -s histappend" >> /root/.bashrc && \ echo "PROMPT_COMMAND='history -a'" >> /root/.bashrc +# copy the rest of the files +COPY ../ ${ISAACLAB_PATH}/ + # make working directory as the Isaac Lab directory # this is the default directory when the container is run WORKDIR ${ISAACLAB_PATH} diff --git a/docs/source/api/lab/isaaclab.app.rst b/docs/source/api/lab/isaaclab.app.rst index b170fa8fa8ff..c1440dd631a0 100644 --- a/docs/source/api/lab/isaaclab.app.rst +++ b/docs/source/api/lab/isaaclab.app.rst @@ -111,5 +111,5 @@ Simulation App Launcher :members: -.. _livestream: https://docs.omniverse.nvidia.com/app_isaacsim/app_isaacsim/manual_livestream_clients.html +.. _livestream: https://docs.isaacsim.omniverse.nvidia.com/latest/installation/manual_livestream_clients.html .. _`WebRTC Livestream`: https://docs.isaacsim.omniverse.nvidia.com/latest/installation/manual_livestream_clients.html#isaac-sim-short-webrtc-streaming-client diff --git a/docs/source/experimental-features/newton-physics-integration/visualization.rst b/docs/source/experimental-features/newton-physics-integration/visualization.rst index f54435733934..cbda8eb88ef7 100644 --- a/docs/source/experimental-features/newton-physics-integration/visualization.rst +++ b/docs/source/experimental-features/newton-physics-integration/visualization.rst @@ -59,8 +59,8 @@ Launch visualizers from the command line with ``--visualizer``: .. code-block:: bash - # Launch all visualizers - python scripts/reinforcement_learning/rsl_rl/train.py --task Isaac-Cartpole-v0 --visualizer omniverse newton rerun + # Launch all visualizers (comma-delimited list, no spaces) + python scripts/reinforcement_learning/rsl_rl/train.py --task Isaac-Cartpole-v0 --visualizer kit,newton,rerun # Launch just newton visualizer python scripts/reinforcement_learning/rsl_rl/train.py --task Isaac-Cartpole-v0 --visualizer newton @@ -77,18 +77,20 @@ If ``--headless`` is given, no visualizers will be launched. Configuration ~~~~~~~~~~~~~ -Launching visualizers with the command line will use default visualizer configurations. Default configs can be found and edited in ``source/isaaclab/isaaclab/visualizers``. +Launching visualizers with the command line will use default visualizer configurations. Visualizer backends live in the ``isaaclab_visualizers`` package (e.g. ``source/isaaclab_visualizers/isaaclab_visualizers/kit``, ``newton``, ``rerun``). -You can also configure custom visualizers in the code by defining new ``VisualizerCfg`` instances for the ``SimulationCfg``, for example: +You can also configure custom visualizers in the code by defining ``VisualizerCfg`` instances for the ``SimulationCfg``, for example: .. code-block:: python from isaaclab.sim import SimulationCfg - from isaaclab.visualizers import NewtonVisualizerCfg, OVVisualizerCfg, RerunVisualizerCfg + from isaaclab_visualizers.kit import KitVisualizerCfg + from isaaclab_visualizers.newton import NewtonVisualizerCfg + from isaaclab_visualizers.rerun import RerunVisualizerCfg sim_cfg = SimulationCfg( visualizer_cfgs=[ - OVVisualizerCfg( + KitVisualizerCfg( viewport_name="Visualizer Viewport", create_viewport=True, dock_position="SAME", @@ -128,9 +130,9 @@ Omniverse Visualizer .. code-block:: python - from isaaclab.visualizers import OVVisualizerCfg + from isaaclab_visualizers.kit import KitVisualizerCfg - visualizer_cfg = OVVisualizerCfg( + visualizer_cfg = KitVisualizerCfg( # Viewport settings viewport_name="Visualizer Viewport", # Viewport window name create_viewport=True, # Create new viewport vs. use existing @@ -176,8 +178,6 @@ Newton Visualizer - Look around * - **Mouse Scroll** - Zoom in/out - * - **Space** - - Pause/resume rendering (physics continues) * - **H** - Toggle UI sidebar * - **ESC** @@ -187,7 +187,7 @@ Newton Visualizer .. code-block:: python - from isaaclab.visualizers import NewtonVisualizerCfg + from isaaclab_visualizers.newton import NewtonVisualizerCfg visualizer_cfg = NewtonVisualizerCfg( # Window settings @@ -233,12 +233,14 @@ Rerun Visualizer .. code-block:: python - from isaaclab.visualizers import RerunVisualizerCfg + from isaaclab_visualizers.rerun import RerunVisualizerCfg visualizer_cfg = RerunVisualizerCfg( # Server settings app_id="isaaclab-simulation", # Application identifier for viewer + grpc_port=9876, # gRPC endpoint for logging SDK connection web_port=9090, # Port for local web viewer (launched in browser) + bind_address="0.0.0.0", # Endpoint host formatting/reuse checks # Camera settings camera_position=(8.0, 8.0, 3.0), # Initial camera position (x, y, z) @@ -252,6 +254,10 @@ Rerun Visualizer record_to_rrd="recording.rrd", # Path to save .rrd file (None = no recording) ) +Rerun startup uses the Python SDK through ``newton.viewer.ViewerRerun`` (no external ``rerun`` CLI process +management). If ``grpc_port`` is already active, Isaac Lab reuses that server. If ``web_port`` is occupied while +starting a new server, initialization fails with a clear port-conflict error. + Performance Note ---------------- diff --git a/docs/source/features/hydra.rst b/docs/source/features/hydra.rst index 40f52c895cd9..65b11072e95c 100644 --- a/docs/source/features/hydra.rst +++ b/docs/source/features/hydra.rst @@ -237,6 +237,32 @@ discovers ``PresetCfg`` subclasses and converts their fields into a presets dict python train.py --task=Isaac-Reach-Franka-v0 \ env.sim.physics=newton +The ``default`` field can be set to ``None`` to make an optional feature that is disabled unless +explicitly selected on the command line: + +.. code-block:: python + + @configclass + class CameraPresetCfg(PresetCfg): + default = None + small: CameraCfg = CameraCfg(width=64, height=64) + large: CameraCfg = CameraCfg(width=256, height=256) + + @configclass + class SceneCfg: + camera: CameraPresetCfg = CameraPresetCfg() + +When no CLI argument is given, ``camera`` resolves to ``None`` (no camera): + +.. code-block:: bash + + # camera is None — no camera overhead + python train.py --task=Isaac-Reach-Franka-v0 + + # activate camera with the "large" preset + python train.py --task=Isaac-Reach-Franka-v0 \ + env.scene.camera=large + Using Presets ^^^^^^^^^^^^^ diff --git a/docs/source/features/isaac_teleop.rst b/docs/source/features/isaac_teleop.rst index 7a06f7bdf77e..2d0d8423e17b 100644 --- a/docs/source/features/isaac_teleop.rst +++ b/docs/source/features/isaac_teleop.rst @@ -446,7 +446,7 @@ If the built-in retargeters do not cover your use case, you can implement a cust #. Connect to existing source nodes (``HandsSource``, ``ControllersSource``) or create a new ``IDeviceIOSource`` subclass for custom input devices. -See the `Retargeting Engine README `_ +See the `Isaac Teleop repository `_ and `Contributing Guide `_ for details. diff --git a/docs/source/how-to/optimize_stage_creation.rst b/docs/source/how-to/optimize_stage_creation.rst index 980f0a1388b2..be7d4320401c 100644 --- a/docs/source/how-to/optimize_stage_creation.rst +++ b/docs/source/how-to/optimize_stage_creation.rst @@ -9,7 +9,7 @@ What These Features Do **Fabric Cloning** -- Clones environments using Fabric library (see `USD Fabric USDRT Documentation `_) +- Clones environments using Fabric library (see `USD Fabric USDRT Documentation `_) - Partially supported and enabled by default on some environments (see `Limitations`_ section for a list) **Stage in Memory** diff --git a/docs/source/how-to/record_animation.rst b/docs/source/how-to/record_animation.rst index 2d675684fbf4..61b4ca823d25 100644 --- a/docs/source/how-to/record_animation.rst +++ b/docs/source/how-to/record_animation.rst @@ -122,6 +122,6 @@ After the stop time is reached, a file will be saved to: .. _Stage Recorder: https://docs.omniverse.nvidia.com/extensions/latest/ext_animation_stage-recorder.html -.. _Fabric: https://docs.omniverse.nvidia.com/kit/docs/usdrt/latest/docs/usd_fabric_usdrt.html +.. _Fabric: https://docs.omniverse.nvidia.com/kit/docs/usdrt.scenegraph/latest/usd_fabric_usdrt.html .. _Omniverse Launcher: https://docs.omniverse.nvidia.com/launcher/latest/index.html .. _tutorial on layering in Omniverse: https://www.youtube.com/watch?v=LTwmNkSDh-c&ab_channel=NVIDIAOmniverse diff --git a/docs/source/overview/core-concepts/renderers.rst b/docs/source/overview/core-concepts/renderers.rst index e2b9f5375684..7f1dae6791ca 100644 --- a/docs/source/overview/core-concepts/renderers.rst +++ b/docs/source/overview/core-concepts/renderers.rst @@ -46,6 +46,26 @@ Core concepts # Lazily loads ovrtx when instantiated; may fail if isaaclab_ov / ovrtx is not installed renderer: BaseRenderer = Renderer(OVRTXRendererCfg()) +Installing the OVRTX renderer +------------------------------ + +The OVRTX renderer is provided by the ``isaaclab_ov`` extension and requires the +`ovrtx `_ package (hosted on +``pypi.nvidia.com``). + +Install via the Isaac Lab CLI: + +.. code-block:: bash + + # Install isaaclab_ov (and its ovrtx dependency) alongside the core package + ./isaaclab.sh -i ov + +Or install manually with pip: + +.. code-block:: bash + + pip install --extra-index-url https://pypi.nvidia.com -e source/isaaclab_ov + - **Opaque render data**: The render data object returned by :meth:`~isaaclab.renderers.BaseRenderer.create_render_data` is passed to subsequent renderer methods. It should be completely opaque to the caller: inspecting or modifying it via get/set attributes is an anti-pattern and breaks the API contract. diff --git a/docs/source/overview/developer-guide/development.rst b/docs/source/overview/developer-guide/development.rst index 48a5019609fa..3d9ddbf014a3 100644 --- a/docs/source/overview/developer-guide/development.rst +++ b/docs/source/overview/developer-guide/development.rst @@ -61,19 +61,19 @@ More specifically, when an extension is enabled, the python module specified in While loading extensions into Omniverse happens automatically, using the python package in standalone applications requires additional steps. To simplify the build process and avoid the need to understand the `premake `__ -build system used by Omniverse, we directly use the `setuptools `__ +build system used by Omniverse, we directly use the `setuptools `__ python package to build the python module provided by the extensions. This is done by the ``setup.py`` file in the extension directory. .. note:: The ``setup.py`` file is not required for extensions that are only loaded into Omniverse - using the `Extension Manager `__. + using the `Extension Manager `__. Lastly, the ``tests`` directory contains the unit tests for the extension. These are written using the `unittest `__ framework. It is important to note that Omniverse also provides a similar -`testing framework `__. +`testing framework `__. However, it requires going through the build process and does not support testing of the python module in standalone applications. @@ -81,7 +81,7 @@ Custom Extension Dependency Management ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Certain extensions may have dependencies which require the installation of additional packages before the extension -can be used. While Python dependencies are handled by the `setuptools `__ +can be used. While Python dependencies are handled by the `setuptools `__ package and specified in the ``setup.py`` file, non-Python dependencies such as `ROS `__ packages or `apt `__ packages are not handled by setuptools. Handling these kinds of dependencies requires an additional procedure. diff --git a/docs/source/overview/developer-guide/vs_code.rst b/docs/source/overview/developer-guide/vs_code.rst index a19889d1bdb8..11946a71c769 100644 --- a/docs/source/overview/developer-guide/vs_code.rst +++ b/docs/source/overview/developer-guide/vs_code.rst @@ -55,6 +55,25 @@ following links: * `Isaac Sim VSCode support `__ +Attach to a Running ``debugpy`` Session +--------------------------------------- + +The generated ``.vscode/launch.json`` includes a ``Python: Debugger Attach`` +configuration that starts to listen on port ``localhost:3000`` for the debugpy session. + +To use it: + +1. Set your breakpoints. +2. Run your code under debugpy like so: + + .. code-block:: bash + + ./isaaclab.sh -p -m debugpy --listen 3000 --wait-for-client -c "from isaaclab.cli import cli; cli()" [cli_args] + +3. In VS Code, select the ``Python: Debugger Attach`` configuration from the Run and Debug panel + and press the green play button or ``F5``. VS Code will connect to the debugpy server + running on ``localhost:3000``. + Configuring the python interpreter ---------------------------------- diff --git a/docs/source/overview/imitation-learning/humanoids_imitation.rst b/docs/source/overview/imitation-learning/humanoids_imitation.rst index f1d6a1d2decd..07f77faf7338 100644 --- a/docs/source/overview/imitation-learning/humanoids_imitation.rst +++ b/docs/source/overview/imitation-learning/humanoids_imitation.rst @@ -432,6 +432,7 @@ Follow the same data collection, annotation, and generation process as demonstra ./isaaclab.sh -p scripts/tools/record_demos.py \ --device cpu \ + --visualizer kit \ --task Isaac-PickPlace-Locomanipulation-G1-Abs-v0 \ --dataset_file ./datasets/dataset_g1_locomanip.hdf5 \ --num_demos 5 @@ -446,6 +447,7 @@ Follow the same data collection, annotation, and generation process as demonstra ./isaaclab.sh -p scripts/tools/replay_demos.py \ --device cpu \ + --visualizer kit \ --task Isaac-PickPlace-Locomanipulation-G1-Abs-v0 \ --dataset_file ./datasets/dataset_g1_locomanip.hdf5 @@ -455,6 +457,7 @@ Follow the same data collection, annotation, and generation process as demonstra ./isaaclab.sh -p scripts/imitation_learning/isaaclab_mimic/annotate_demos.py \ --device cpu \ + --visualizer kit \ --task Isaac-Locomanipulation-G1-Abs-Mimic-v0 \ --input_file ./datasets/dataset_g1_locomanip.hdf5 \ --output_file ./datasets/dataset_annotated_g1_locomanip.hdf5 @@ -492,6 +495,7 @@ Visualize the trained policy performance: ./isaaclab.sh -p scripts/imitation_learning/robomimic/play.py \ --device cpu \ + --visualizer kit \ --task Isaac-PickPlace-Locomanipulation-G1-Abs-v0 \ --num_rollouts 50 \ --horizon 400 \ @@ -707,8 +711,8 @@ dataset. Some of them are captured from a humanoid-viewpoint to match the camera For example, when using the asset ``hand_hold-voyager-babyboom``, the relevant files are: -- ``stage.usdz``: a USDZ archive that bundles 3D Gaussian splatting (``volume.nurec``), a collision mesh (``mesh.usd``), etc. -- ``occupancy_map.yaml`` and ``occupancy_map.png``: occupancy map for path planning and navigation. +- `stage.usdz `__: a USDZ archive that bundles 3D Gaussian splatting (``volume.nurec``), a collision mesh (``mesh.usd``), etc. +- `occupancy_map.yaml `__ and `occupancy_map.png `__: occupancy map for path planning and navigation. Download the files and place them under ````. @@ -734,18 +738,18 @@ Then run the following command: --background_usd_path /stage.usdz \ --background_occupancy_yaml_file /occupancy_map.yaml \ --init_camera_view \ - --randomize_placement + --randomize_placement \ + --high_res_video The key parameters are: - ``--background_usd_path``: Path to the NuRec USD asset. - ``--background_occupancy_yaml_file``: Path to the occupancy map file. - ``--init_camera_view``: Set the viewport camera behind the robot at the start of episode. -- ``--high_res_video``: An optional argument to generate a higher resolution video (540x960) for the ego-centric camera view. Default resolution is 160x256. +- ``--high_res_video``: Generate a higher resolution video (540x960) for the ego-centric camera view. On successful task completion, an HDF5 dataset is generated containing camera observations. You can convert -the ego-centric camera view to MP4. If you use ``--high_res_video`` during data generation, match the -dimension(540x960) in the command below. +the ego-centric camera view to MP4. .. code:: bash @@ -753,5 +757,5 @@ dimension(540x960) in the command below. --input_file /generated_dataset_g1_locomanipulation_sdg_with_background.hdf5 \ --output_dir / \ --input_keys robot_pov_cam \ - --video_width 256 \ - --video_height 160 + --video_width 960 \ + --video_height 540 diff --git a/docs/source/policy_deployment/03_compass_with_NuRec/compass_navigation_policy_with_NuRec.rst b/docs/source/policy_deployment/03_compass_with_NuRec/compass_navigation_policy_with_NuRec.rst new file mode 100644 index 000000000000..54c92465585c --- /dev/null +++ b/docs/source/policy_deployment/03_compass_with_NuRec/compass_navigation_policy_with_NuRec.rst @@ -0,0 +1,398 @@ +Training & Deploying COMPASS Navigation Policy with Real2Sim NuRec +==================================================================== + +This tutorial shows you how to train and deploy COMPASS navigation policies using NuRec Real2Sim assets in the Isaac Lab simulation environment. +COMPASS (Cross-embodiment Mobility Policy via Residual RL and Skill Synthesis) is a novel framework for cross-embodiment mobility. +For more details about the project, please visit the `COMPASS Repository`_. + +Setup +----- + +.. note:: + + This tutorial is for **Ubuntu 22.04** and the **OVX with RTX platform**. It is intended for Isaac Lab 3.0.0 and Isaac Sim 6.0.0. + +Create a Workspace +~~~~~~~~~~~~~~~~~~ + +Start by creating a dedicated workspace directory for this project: + +.. code-block:: bash + + mkdir compass-nurec + cd compass-nurec + +Terminal 1 — Isaac Lab & Isaac Sim Installation +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Follow these steps in your first terminal to install Isaac Sim and Isaac Lab. + +1. Install Isaac Sim 6.0. The supported installation methods are: + + - **pip install** (recommended): Follow the `Isaac Sim pip Installation Guide`_. + - **Binary download**: Download the pre-built binary from the `Isaac Sim Installation Guide`_. + +2. Clone the Isaac Lab repository and check out the ``v3.0`` tag (or the ``develop`` branch for testing): + +.. code-block:: bash + + git clone https://github.com/isaac-sim/IsaacLab.git + cd IsaacLab + git fetch origin + git checkout v3.0 + +3. Set the required environment variables: + +.. note:: + + This step is only required if Isaac Sim was installed via the **binary download**. + If you installed Isaac Sim via ``pip``, these variables are not needed. + +.. code-block:: bash + + # Isaac Sim root directory + export ISAACSIM_PATH="${HOME}/isaacsim" + + # Isaac Sim python executable + export ISAACSIM_PYTHON_EXE="${ISAACSIM_PATH}/python.sh" + +4. Verify that Isaac Sim starts correctly: + +.. code-block:: bash + + ${ISAACSIM_PATH}/isaac-sim.sh + +5. Create a symbolic link to Isaac Sim inside the Isaac Lab directory: + +.. code-block:: bash + + ln -s ${ISAACSIM_PATH} _isaac_sim + +6. Create a dedicated conda environment and install Isaac Lab: + +.. code-block:: bash + + # Create the conda environment + ./isaaclab.sh --conda env_isaaclab_3.0_compass + + # Activate the environment + conda activate env_isaaclab_3.0_compass + + # Install Isaac Lab + ./isaaclab.sh --install + +7. Verify the Isaac Lab installation: + +.. code-block:: bash + + ./isaaclab.sh -p scripts/tutorials/00_sim/create_empty.py --visualizer kit + +Terminal 2 — COMPASS Installation +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Open a second terminal and follow these steps to install the COMPASS repository. + +1. Activate the conda environment created in Terminal 1: + +.. code-block:: bash + + conda deactivate + conda activate env_isaaclab_3.0_compass + +2. Clone the COMPASS repository and check out the NuRec-compatible branch: + +.. code-block:: bash + + git clone https://github.com/NVlabs/COMPASS.git + cd COMPASS + git fetch + git checkout samc/support_nurec_assets_isaaclab_3.0 + +3. Set the path to your Isaac Lab installation: + +.. code-block:: bash + + export ISAACLAB_PATH= + +4. Install the Python dependencies: + +.. code-block:: bash + + ${ISAACLAB_PATH}/isaaclab.sh -p -m pip install -r requirements.txt + +5. Install the X-Mobility base policy package: + +.. code-block:: bash + + ${ISAACLAB_PATH}/isaaclab.sh -p -m pip install x_mobility/x_mobility-0.1.0-py3-none-any.whl + +6. Install the ``mobility_es`` Isaac Lab extension: + +.. code-block:: bash + + cd compass/rl_env + ${ISAACLAB_PATH}/isaaclab.sh -p -m pip install -e exts/mobility_es + cd - + +Testing the Setup +~~~~~~~~~~~~~~~~~ + +Run the following command from the ``COMPASS`` directory to verify the setup: + +.. code-block:: bash + + cd compass/rl_env + ${ISAACLAB_PATH}/isaaclab.sh -p scripts/play.py --enable_cameras --visualizer kit + cd - + +Downloading Assets & Checkpoints +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +**1. Pre-trained X-Mobility checkpoint** + +Download the checkpoint from: +https://huggingface.co/nvidia/X-Mobility/blob/main/x_mobility-nav2-semantic_action_path.ckpt + +**2. COMPASS USD Assets** + +Download the pre-packaged COMPASS USD assets: +https://huggingface.co/nvidia/COMPASS/blob/main/compass_usds.zip + +**3. NuRec Real2Sim Assets** + +Download the NuRec Real2Sim assets from the `PhysicalAI-Robotics-NuRec dataset`_ on Hugging Face: + +.. note:: + + You need to agree to the dataset terms on Hugging Face before accessing the files. + +The dataset provides several environments. For COMPASS, download the environment(s) you need +(e.g., ``nova_carter-galileo``) and place the extracted files under the COMPASS extension directory: + +.. code-block:: bash + + .. Ensure that you are in COMPASS root directory + compass/rl_env/exts/mobility_es/mobility_es/usd// + +For example, for the Galileo environment: + +.. code-block:: bash + + compass/rl_env/exts/mobility_es/mobility_es/usd/nova_carter-galileo/ + ├── stage.usdz + ├── occupancy_map.yaml + └── occupancy_map.png + +.. note:: + + The following environments are available in the dataset: + + .. list-table:: + :header-rows: 1 + :widths: 30 50 20 + + * - Environment Name + - Description + - Contains Mesh + * - ``nova_carter-galileo`` + - Galileo lab in NVIDIA — aisles, shelves, and boxes + - Yes + * - ``nova_carter-cafe`` + - NVIDIA cafe — open area with natural lighting + - Yes + * - ``nova_carter-wormhole`` + - Conference room some chairs and tables + - Yes + * - ``hand_hold-endeavor-livingroom`` + - Living room in NVIDIA Endeavor building + - Yes + +Training the Policy +------------------- + +Training Configuration +~~~~~~~~~~~~~~~~~~~~~~ + +The training configuration for NuRec Real2Sim environments is specified in ``configs/train_config_real2sim.gin``. +This configuration includes optimized settings for Real2Sim environments: + +- **Collision checking distances**: Tuned for Real2Sim environments (0.5m for both goal and start poses) +- **Precomputed valid poses**: Enabled for faster pose sampling in constrained environments +- **compute_orientations**: Set to True to compute orientations for the start and goal poses. +- **Environment spacing**: Set to 500m to accommodate larger Real2Sim scenes + +Training the Residual RL Specialist +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Execute the following command from the ``COMPASS`` directory (Terminal 2) to train a residual RL specialist policy for NuRec Real2Sim environments. + +.. code-block:: bash + + ${ISAACLAB_PATH:?}/isaaclab.sh -p run.py \ + -c configs/train_config_real2sim.gin \ + -o \ + -b \ + --embodiment \ + --environment nova_carter-galileo \ + --num_envs 64 \ + --video \ + --video_interval 1 \ + --visualizer kit \ + --enable_cameras \ + --headless \ + --precompute_valid_poses + +Where: +- ````: Directory where training outputs and checkpoints will be saved +- ````: Path to the downloaded X-Mobility checkpoint +- ````: One of ``h1``, ``spot``, ``carter``, ``g1``, or ``digit`` +- ``--environment nova_carter-galileo``: Specifies the NuRec Real2Sim Galileo environment + +The training will run for the number of iterations specified in the config file (default: 1000 iterations). +The resulting checkpoint will be stored in ``/checkpoints/`` with the filename ``model_.pt``. +Videos will be saved in ``/videos/``. + +.. note:: + + The GPU memory usage is proportional to the number of environments. For example, 64 environments will use around 30-40GB memory. + Reduce ``--num_envs`` if you have limited GPU memory. + + For Real2Sim environments, it's recommended to use ``--precompute_valid_poses`` flag to precompute valid pose locations, + which significantly speeds up pose sampling in constrained environments. + For very tight spaces, it's recommended to use ``--compute_orientations`` flag to compute orientations for the start and goal poses. + Orientation computation is slow so use it only if necessary. + +Advanced Training Options +~~~~~~~~~~~~~~~~~~~~~~~~~ + +You can customize the training by modifying the gin config file or using command-line arguments: + +- **Collision distances**: Adjust ``goal_pose_collision_distance`` and ``start_pose_collision_distance`` in the config +- **Precompute valid poses**: Set ``precompute_valid_poses = True`` in the config or use ``--precompute_valid_poses`` flag +- **Compute orientations**: Set ``compute_orientations = True`` in the config or use ``--compute_orientations`` flag +- **Number of iterations**: Modify ``num_iterations`` in the config +- **Number of environments**: Use ``--num_envs`` command-line argument + +Testing the Trained Policy +-------------------------- + +Evaluate the trained policy +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Execute the following command from the ``COMPASS`` directory to evaluate the trained policy checkpoint. + +.. code-block:: bash + + ${ISAACLAB_PATH:?}/isaaclab.sh -p run.py \ + -c configs/eval_config_real2sim.gin \ + -o \ + -b \ + -p \ + --embodiment \ + --environment nova_carter-galileo \ + --num_envs \ + --video \ + --video_interval 1 + --enable_cameras \ + --visualizer kit \ + --headless + +Where: +- ````: Path to the trained residual policy checkpoint (e.g., ``/checkpoints/model_1000.pt``) +- ``--video``: Enable video recording during evaluation +- ``--video_interval``: Record video every N iterations + +The evaluation will run for the number of iterations specified in the eval config file. +Videos will be saved in ``/videos/``. + +Model Export +------------ + +Export to ONNX or JIT +~~~~~~~~~~~~~~~~~~~~~ + +Export the trained residual RL specialist policy to ONNX or JIT formats for deployment: + +.. code-block:: bash + + python3 onnx_conversion.py \ + -b \ + -r \ + -e \ + -o \ + -j + +Convert ONNX to TensorRT +~~~~~~~~~~~~~~~~~~~~~~~~~ + +For optimized inference, convert the ONNX model to TensorRT: + +.. code-block:: bash + + python3 trt_conversion.py \ + -o \ + -t + +Deployment +---------- + +ROS2 Deployment +~~~~~~~~~~~~~~~ + +The trained COMPASS policy can be deployed using the ROS2 deployment framework. +Refer to the `COMPASS ROS2 Deployment Guide`_ for detailed instructions on deploying the policy in simulation or on real robots. + +The ROS2 deployment supports: +- Isaac Sim integration for simulation testing +- Zero-shot sim-to-real transfer for real robot deployment +- Object navigation integration with object localization modules + +Sim-to-Real Deployment +~~~~~~~~~~~~~~~~~~~~~~ + +COMPASS policies trained on NuRec Real2Sim environments are designed for zero-shot sim-to-real transfer. +The Real2Sim assets provide a bridge between simulation and reality, enabling policies trained in simulation to work directly on real robots. + +For sim-to-real deployment: +1. Export the trained policy to ONNX or TensorRT format (see Model Export section) +2. Use the ROS2 deployment framework to run inference on the real robot +3. Integrate with visual SLAM (e.g., cuVSLAM) for robot state estimation +4. The policy will output velocity commands based on camera observations and goal poses + +Troubleshooting +--------------- + +Common Issues +~~~~~~~~~~~~~ + +**Issue: "Failed to sample collision free poses"** + +This error occurs when the collision checking is too strict for the environment. Solutions: +- Reduce ``goal_pose_collision_distance`` and ``start_pose_collision_distance`` in the config +- Enable ``precompute_valid_poses`` to precompute valid pose locations +- Check that the occupancy map files are correctly placed in the ``omap/`` directory + +**Issue: High GPU memory usage** + +- Reduce the number of environments using ``--num_envs`` +- For Real2Sim environments, start with 32-64 environments + +**Issue: Slow pose sampling** + +- Enable ``--precompute_valid_poses`` flag to precompute valid poses +- This is especially important for Real2Sim environments with constrained spaces + +Configuration Tips +~~~~~~~~~~~~~~~~~~ + +For NuRec Real2Sim environments: +- Use collision distances of 0.5m or less for more constrained environments +- Enable precomputed valid poses for constrained environments +- Use environment spacing of 500m to accommodate larger scenes + +.. _Isaac Lab Installation Guide: https://isaac-sim.github.io/IsaacLab/v2.0.0/source/setup/installation/index.html +.. _Isaac Sim Installation Guide: https://docs.isaacsim.omniverse.nvidia.com/latest/installation/install_workstation.html +.. _Isaac Sim pip Installation Guide: https://docs.isaacsim.omniverse.nvidia.com/latest/installation/install_python.html +.. _COMPASS Repository: https://github.com/NVlabs/COMPASS +.. _COMPASS ROS2 Deployment Guide: https://github.com/NVlabs/COMPASS/tree/main/ros2_deployment +.. _PhysicalAI-Robotics-NuRec dataset: https://huggingface.co/datasets/nvidia/PhysicalAI-Robotics-NuRec diff --git a/docs/source/policy_deployment/index.rst b/docs/source/policy_deployment/index.rst index 3ee100f22174..32a10a42034f 100644 --- a/docs/source/policy_deployment/index.rst +++ b/docs/source/policy_deployment/index.rst @@ -11,3 +11,4 @@ Below, you’ll find detailed examples of various policies for training and depl 00_hover/hover_policy 01_io_descriptors/io_descriptors_101 02_gear_assembly/gear_assembly_policy + 03_compass_with_NuRec/compass_navigation_policy_with_NuRec diff --git a/docs/source/refs/contributing.rst b/docs/source/refs/contributing.rst index e21dd9dc6d07..21a2b8d07342 100644 --- a/docs/source/refs/contributing.rst +++ b/docs/source/refs/contributing.rst @@ -364,23 +364,59 @@ Key rules for ``.pyi`` stubs: `__). * Group imports from the same submodule on one line. Use parenthesized multi-line imports if the line exceeds 100 characters. -* Only use **relative imports** (``from .something import ...``) — ``lazy_loader`` - does not support absolute imports in stubs. +* Use **relative imports** (``from .something import ...``) for local submodule + symbols. Absolute wildcard imports (``from pkg import *``) are only used for + cross-package fallbacks (see below). * Include the standard Isaac Lab license header. **Cross-package fallback** — for modules that re-export names from another package -(e.g. task MDP modules that delegate to ``isaaclab.envs.mdp``): +(e.g. task MDP modules that delegate to ``isaaclab.envs.mdp``), add a wildcard +import for the external package in the ``.pyi`` stub: + +.. code:: python + + # isaaclab_tasks/.../mdp/__init__.pyi + __all__ = ["MyReward", "MyObservation"] + + from .rewards import MyReward + from .observations import MyObservation + + from isaaclab.envs.mdp import * + +The ``__init__.py`` stays the same as the standard pattern — just ``lazy_export()`` +with no arguments: .. code:: python # isaaclab_tasks/.../mdp/__init__.py from isaaclab.utils.module import lazy_export - lazy_export(packages=["isaaclab.envs.mdp"]) + lazy_export() + +At runtime, ``lazy_export`` parses the ``.pyi`` stub and uses the absolute wildcard +import (``from isaaclab.envs.mdp import *``) as a fallback: any name not found in +the local submodules is looked up in the specified package. This also gives type +checkers and IDEs full visibility into the re-exported symbols. + +**Relative wildcard re-exports** — the stub can also use ``from .submodule import *`` +to eagerly export all public names from a local submodule. This is resolved at +import time (not lazily) and is useful when a submodule's public API is large or +changes frequently. + +.. note:: + + Relative wildcard re-exports bypass lazy loading and eagerly import every public + name from the submodule at package init time. In general, we advise against using + them unless absolutely necessary. Prefer listing explicit named imports in the stub + so that the public API surface is clear, reviewable, and remains lazily loaded. + +.. code:: python + + # isaaclab_tasks/.../mdp/__init__.pyi + from .rewards import * + from .observations import * -The ``lazy_export`` helper in :mod:`isaaclab.utils.module` wraps ``lazy_loader.attach_stub`` -and, when ``packages`` is provided, adds a runtime fallback that scans the specified -packages for names not found in the local stub. + from isaaclab.envs.mdp import * **Ensuring .pyi stubs are distributed** diff --git a/docs/source/refs/issues.rst b/docs/source/refs/issues.rst index c4bb56182e51..42a78ec27ffb 100644 --- a/docs/source/refs/issues.rst +++ b/docs/source/refs/issues.rst @@ -74,7 +74,7 @@ If you use an instanceable assets for markers, the marker class removes all the This is then replicated across other references of the same asset since physics properties of instanceable assets are stored in the instanceable asset's USD file and not in its stage reference's USD file. -.. _instanceable assets: https://docs.omniverse.nvidia.com/app_isaacsim/app_isaacsim/tutorial_gym_instanceable_assets.html +.. _instanceable assets: https://docs.isaacsim.omniverse.nvidia.com/latest/isaac_lab_tutorials/tutorial_instanceable_assets.html .. _Omniverse Isaac Sim documentation: https://docs.isaacsim.omniverse.nvidia.com/latest/overview/known_issues.html# diff --git a/docs/source/tutorials/00_sim/create_empty.rst b/docs/source/tutorials/00_sim/create_empty.rst index 89f28201c410..1a9740ab93af 100644 --- a/docs/source/tutorials/00_sim/create_empty.rst +++ b/docs/source/tutorials/00_sim/create_empty.rst @@ -69,7 +69,7 @@ Configuring the simulation context When launching the simulator from a standalone script, the user has complete control over playing, pausing and stepping the simulator. All these operations are handled through the **simulation -context**. It takes care of various timeline events and also configures the `physics scene`_ for +context**. It takes care of various timeline events and also configures the physics scene for simulation. In Isaac Lab, the :class:`sim.SimulationContext` class inherits from Isaac Sim's @@ -164,4 +164,3 @@ following tutorial where we will learn how to add assets to the stage. .. _`Isaac Sim Workflows`: https://docs.isaacsim.omniverse.nvidia.com/latest/introduction/workflows.html .. _carb: https://docs.omniverse.nvidia.com/kit/docs/carbonite/latest/index.html -.. _`physics scene`: https://docs.omniverse.nvidia.com/prod_extensions/prod_extensions/ext_physics.html#physics-scene diff --git a/docs/source/tutorials/00_sim/spawn_prims.rst b/docs/source/tutorials/00_sim/spawn_prims.rst index 66941ddb1c7f..299f15adbe56 100644 --- a/docs/source/tutorials/00_sim/spawn_prims.rst +++ b/docs/source/tutorials/00_sim/spawn_prims.rst @@ -188,5 +188,5 @@ demonstrates the basic concepts of scene designing in Isaac Lab and how to use t we will now look at how to interact with the scene and the simulation. -.. _`USD documentation`: https://graphics.pixar.com/usd/docs/index.html +.. _`USD documentation`: https://openusd.org/docs/index.html .. _`different light prims`: https://youtu.be/c7qyI8pZvF4?feature=shared diff --git a/docs/source/tutorials/03_envs/register_rl_env_gym.rst b/docs/source/tutorials/03_envs/register_rl_env_gym.rst index 53e653a42755..c1b232a82eb1 100644 --- a/docs/source/tutorials/03_envs/register_rl_env_gym.rst +++ b/docs/source/tutorials/03_envs/register_rl_env_gym.rst @@ -136,7 +136,7 @@ and whether to render, are used to override the default configuration. .. literalinclude:: ../../../../scripts/environments/random_agent.py :language: python - :start-at: # create environment configuration + :start-at: # parse configuration via Hydra :end-at: env = gym.make(args_cli.task, cfg=env_cfg) Once creating the environment, the rest of the execution follows the standard resetting and stepping. diff --git a/scripts/benchmarks/benchmark_lazy_export.py b/scripts/benchmarks/benchmark_lazy_export.py new file mode 100644 index 000000000000..999dc7b5e1f4 --- /dev/null +++ b/scripts/benchmarks/benchmark_lazy_export.py @@ -0,0 +1,177 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Benchmark lazy_export import performance via parse_env_cfg. + +Measures the cold end-to-end time to construct environment configs — the real +user-facing cost that includes gym registry walking, lazy_export stub parsing, +fallback resolution, module imports, and config class instantiation. + +Each iteration fully purges task modules and the gym registration guard so +that the next ``import isaaclab_tasks`` re-walks every task package and +re-registers every gym environment, matching a fresh-process cold start. + +The report separates **package loading** (``import isaaclab_tasks`` — registry +walk + gym registrations) from **config construction** +(``load_cfg_from_registry`` — module import + class instantiation). + +This script does NOT require Isaac Sim or a GPU. + +Usage:: + + ./isaaclab.sh -p scripts/benchmarks/benchmark_lazy_export.py + ./isaaclab.sh -p scripts/benchmarks/benchmark_lazy_export.py --iterations 20 + ./isaaclab.sh -p scripts/benchmarks/benchmark_lazy_export.py --tasks Isaac-Velocity-Flat-Anymal-D-v0 +""" + +from __future__ import annotations + +import argparse +import builtins +import importlib +import io +import statistics +import sys +import time +import warnings + +import gymnasium + +with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + import isaaclab_tasks # noqa: F401 + +from isaaclab_tasks.utils.parse_cfg import load_cfg_from_registry + +_REPRESENTATIVE_TASKS = [ + "Isaac-Cartpole-v0", + "Isaac-Humanoid-v0", + "Isaac-Velocity-Flat-Anymal-D-v0", + "Isaac-Reach-Franka-v0", + "Isaac-Lift-Cube-Franka-v0", + "Isaac-Dexsuite-Kuka-Allegro-Reorient-v0", + "Isaac-Navigation-Flat-Anymal-C-v0", + "Isaac-Stack-Cube-Franka-v0", +] + + +def _purge_all() -> None: + """Fully purge task modules, gym registrations, and the registration guard.""" + to_delete = [k for k in sys.modules if k.startswith("isaaclab_tasks.")] + for k in to_delete: + del sys.modules[k] + del sys.modules["isaaclab_tasks"] + + builtins._isaaclab_tasks_registered = False + + isaac_ids = [name for name in gymnasium.registry if name.startswith("Isaac-")] + for name in isaac_ids: + del gymnasium.registry[name] + + +def _short_name(task: str) -> str: + return task.replace("Isaac-", "").replace("-v0", "") + + +def benchmark( + tasks: list[str], iterations: int +) -> tuple[dict[str, list[float]], dict[str, list[float]], dict[str, list[float]]]: + """Benchmark with package loading and config construction timed separately. + + Returns: + Three dicts keyed by task name, each mapping to a list of per-iteration + times in ms: (pkg_load_times, cfg_construct_times, total_times). + """ + pkg_results: dict[str, list[float]] = {t: [] for t in tasks} + cfg_results: dict[str, list[float]] = {t: [] for t in tasks} + total_results: dict[str, list[float]] = {t: [] for t in tasks} + + for _ in range(iterations): + for task in tasks: + _purge_all() + + old_stdout, old_stderr = sys.stdout, sys.stderr + sys.stdout = io.StringIO() + sys.stderr = io.StringIO() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + try: + t0 = time.perf_counter_ns() + importlib.import_module("isaaclab_tasks") + t1 = time.perf_counter_ns() + load_cfg_from_registry(task, "env_cfg_entry_point") + t2 = time.perf_counter_ns() + finally: + sys.stdout = old_stdout + sys.stderr = old_stderr + + pkg_ms = (t1 - t0) / 1_000_000 + cfg_ms = (t2 - t1) / 1_000_000 + pkg_results[task].append(pkg_ms) + cfg_results[task].append(cfg_ms) + total_results[task].append(pkg_ms + cfg_ms) + + return pkg_results, cfg_results, total_results + + +def _print_table(title: str, results: dict[str, list[float]], unit: str = "ms") -> None: + print(f"\n{'=' * 80}") + print(f" {title}") + print(f"{'=' * 80}") + print(f" {'Task':<40} {'median':>10} {'mean':>10} {'stdev':>10} ({unit})") + print(f" {'-' * 40} {'-' * 10} {'-' * 10} {'-' * 10}") + + all_medians: list[float] = [] + for task, times in results.items(): + short = _short_name(task) + med = statistics.median(times) + avg = statistics.mean(times) + std = statistics.stdev(times) if len(times) > 1 else 0.0 + all_medians.append(med) + print(f" {short:<40} {med:>10.2f} {avg:>10.2f} {std:>10.2f}") + + total_med = sum(all_medians) + avg_med = total_med / len(all_medians) if all_medians else 0 + print(f" {'-' * 40} {'-' * 10} {'-' * 10} {'-' * 10}") + print(f" {'TOTAL (sum of medians)':<40} {total_med:>10.2f}") + print(f" {'AVERAGE (per task)':<40} {avg_med:>10.2f}") + print() + + +def main(): + parser = argparse.ArgumentParser(description="Benchmark lazy_export via cold parse_env_cfg.") + parser.add_argument("--iterations", type=int, default=20, help="Iterations per benchmark.") + parser.add_argument( + "--tasks", + nargs="*", + default=None, + help="Task IDs to benchmark. Defaults to a representative set.", + ) + args = parser.parse_args() + n = args.iterations + + tasks = args.tasks or _REPRESENTATIVE_TASKS + valid = [t for t in tasks if t in gymnasium.registry] + skipped = [t for t in tasks if t not in gymnasium.registry] + if skipped: + print(f"[WARN] Skipping unregistered tasks: {skipped}") + tasks = valid + + if not tasks: + print("[ERROR] No valid tasks found.") + return + + print(f"Benchmarking cold parse_env_cfg with {n} iterations") + print(f"Tasks ({len(tasks)}): {[_short_name(t) for t in tasks]}") + + pkg, cfg, total = benchmark(tasks, n) + + _print_table("Package loading (import isaaclab_tasks — registry walk)", pkg) + _print_table("Config construction (load_cfg_from_registry)", cfg) + _print_table("Total (package loading + config construction)", total) + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmarks/benchmark_rsl_rl.py b/scripts/benchmarks/benchmark_rsl_rl.py index c256905998c9..e0c6eb68b5d5 100644 --- a/scripts/benchmarks/benchmark_rsl_rl.py +++ b/scripts/benchmarks/benchmark_rsl_rl.py @@ -63,6 +63,7 @@ imports_time_begin = time.perf_counter_ns() +import importlib.metadata as metadata from datetime import datetime import gymnasium as gym @@ -74,7 +75,7 @@ from isaaclab.utils.dict import print_dict from isaaclab.utils.io import dump_yaml -from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper +from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper, handle_deprecated_rsl_rl_cfg import isaaclab_tasks # noqa: F401 from isaaclab_tasks.utils import get_checkpoint_path, launch_simulation, resolve_task_config @@ -199,6 +200,9 @@ def main( task_startup_time_end = time.perf_counter_ns() + # handle deprecated configurations (e.g. legacy policy -> actor/critic migration) + agent_cfg = handle_deprecated_rsl_rl_cfg(agent_cfg, metadata.version("rsl-rl-lib")) + # create runner from rsl-rl runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device) # write git state to logs @@ -232,13 +236,13 @@ def main( # prepare RL timing dict collection_fps = ( 1 - / (np.array(log_data["Perf/collection time"])) + / (np.array(log_data["Perf/collection_time"])) * env.unwrapped.num_envs * agent_cfg.num_steps_per_env * world_size ) rl_training_times = { - "Collection Time": (np.array(log_data["Perf/collection time"]) / 1000).tolist(), + "Collection Time": (np.array(log_data["Perf/collection_time"]) / 1000).tolist(), "Learning Time": (np.array(log_data["Perf/learning_time"]) / 1000).tolist(), "Collection FPS": collection_fps.tolist(), "Total FPS": log_data["Perf/total_fps"] * world_size, diff --git a/scripts/environments/random_agent.py b/scripts/environments/random_agent.py index 6a40060d64b1..e98bd17f951e 100644 --- a/scripts/environments/random_agent.py +++ b/scripts/environments/random_agent.py @@ -5,11 +5,14 @@ """Script to an environment with random action agent.""" -"""Launch Isaac Sim Simulator first.""" - import argparse +import sys + +import gymnasium as gym +import torch -from isaaclab.app import AppLauncher +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.utils import add_launcher_args, launch_simulation, resolve_task_config # add argparse arguments parser = argparse.ArgumentParser(description="Random agent for Isaac Lab environments.") @@ -19,54 +22,57 @@ parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") parser.add_argument("--task", type=str, default=None, help="Name of the task.") # append AppLauncher cli args -AppLauncher.add_app_launcher_args(parser) +add_launcher_args(parser) # parse the arguments -args_cli = parser.parse_args() - -# launch omniverse app -app_launcher = AppLauncher(args_cli) -simulation_app = app_launcher.app +args_cli, hydra_args = parser.parse_known_args() -"""Rest everything follows.""" - -import gymnasium as gym -import torch - -import isaaclab_tasks # noqa: F401 -from isaaclab_tasks.utils import parse_env_cfg +# pass remaining args to Hydra +sys.argv = [sys.argv[0]] + hydra_args # PLACEHOLDER: Extension template (do not remove this comment) def main(): """Random actions agent with Isaac Lab environment.""" - # create environment configuration - env_cfg = parse_env_cfg( - args_cli.task, device=args_cli.device, num_envs=args_cli.num_envs, use_fabric=not args_cli.disable_fabric - ) - # create environment - env = gym.make(args_cli.task, cfg=env_cfg) - - # print info (this is vectorized environment) - print(f"[INFO]: Gym observation space: {env.observation_space}") - print(f"[INFO]: Gym action space: {env.action_space}") - # reset environment - env.reset() - # simulate environment - while simulation_app.is_running(): - # run everything in inference mode - with torch.inference_mode(): - # sample actions from -1 to 1 - actions = 2 * torch.rand(env.action_space.shape, device=env.unwrapped.device) - 1 - # apply actions - env.step(actions) - - # close the simulator - env.close() + + torch.manual_seed(42) + + # parse configuration via Hydra (supports preset selection, e.g. env.sim.physics=newton) + env_cfg, _ = resolve_task_config(args_cli.task, "") + + with launch_simulation(env_cfg, args_cli): + # override with CLI arguments + env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs + env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device + if args_cli.disable_fabric: + env_cfg.sim.use_fabric = False + + # create environment + env = gym.make(args_cli.task, cfg=env_cfg) + + # print info (this is vectorized environment) + print(f"[INFO]: Gym observation space: {env.observation_space}") + print(f"[INFO]: Gym action space: {env.action_space}") + # reset environment + env.reset() + # simulate environment + sim = env.unwrapped.sim + while True: + if sim.visualizers: + # visualizer mode: run until the visualizer window is closed + if not any(v.is_running() and not v.is_closed for v in sim.visualizers): + break + # run everything in inference mode + with torch.inference_mode(): + # sample actions from -1 to 1 + actions = 2 * torch.rand(env.action_space.shape, device=env.unwrapped.device) - 1 + # apply actions + env.step(actions) + + # close the simulator + env.close() if __name__ == "__main__": # run the main function main() - # close sim app - simulation_app.close() diff --git a/scripts/environments/zero_agent.py b/scripts/environments/zero_agent.py index edd9317a6287..f2a64bd29d9a 100644 --- a/scripts/environments/zero_agent.py +++ b/scripts/environments/zero_agent.py @@ -5,11 +5,14 @@ """Script to run an environment with zero action agent.""" -"""Launch Isaac Sim Simulator first.""" - import argparse +import sys + +import gymnasium as gym +import torch -from isaaclab.app import AppLauncher +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.utils import add_launcher_args, launch_simulation, resolve_task_config # add argparse arguments parser = argparse.ArgumentParser(description="Zero agent for Isaac Lab environments.") @@ -19,54 +22,57 @@ parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") parser.add_argument("--task", type=str, default=None, help="Name of the task.") # append AppLauncher cli args -AppLauncher.add_app_launcher_args(parser) +add_launcher_args(parser) # parse the arguments -args_cli = parser.parse_args() - -# launch omniverse app -app_launcher = AppLauncher(args_cli) -simulation_app = app_launcher.app +args_cli, hydra_args = parser.parse_known_args() -"""Rest everything follows.""" - -import gymnasium as gym -import torch - -import isaaclab_tasks # noqa: F401 -from isaaclab_tasks.utils import parse_env_cfg +# pass remaining args to Hydra +sys.argv = [sys.argv[0]] + hydra_args # PLACEHOLDER: Extension template (do not remove this comment) +MAX_STEPS = 100 def main(): """Zero actions agent with Isaac Lab environment.""" - # parse configuration - env_cfg = parse_env_cfg( - args_cli.task, device=args_cli.device, num_envs=args_cli.num_envs, use_fabric=not args_cli.disable_fabric - ) - # create environment - env = gym.make(args_cli.task, cfg=env_cfg) - - # print info (this is vectorized environment) - print(f"[INFO]: Gym observation space: {env.observation_space}") - print(f"[INFO]: Gym action space: {env.action_space}") - # reset environment - env.reset() - # simulate environment - while simulation_app.is_running(): - # run everything in inference mode - with torch.inference_mode(): - # compute zero actions - actions = torch.zeros(env.action_space.shape, device=env.unwrapped.device) - # apply actions - env.step(actions) - - # close the simulator - env.close() + + torch.manual_seed(42) + + # parse configuration via Hydra (supports preset selection, e.g. env.sim.physics=newton) + env_cfg, _ = resolve_task_config(args_cli.task, "") + + with launch_simulation(env_cfg, args_cli): + # override with CLI arguments + env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs + env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device + if args_cli.disable_fabric: + env_cfg.sim.use_fabric = False + + # create environment + env = gym.make(args_cli.task, cfg=env_cfg) + + # print info (this is vectorized environment) + print(f"[INFO]: Gym observation space: {env.observation_space}") + print(f"[INFO]: Gym action space: {env.action_space}") + # reset environment + env.reset() + # simulate environment + # keep running while any visualizer is open, otherwise fall back to MAX_STEPS + sim = env.unwrapped.sim + actions = torch.zeros(env.action_space.shape, device=env.unwrapped.device) + while True: + if sim.visualizers: + # visualizer mode: run until the visualizer window is closed + if not any(v.is_running() and not v.is_closed for v in sim.visualizers): + break + # run everything in inference mode + with torch.inference_mode(): + # apply actions + env.step(actions) + # close the simulator + env.close() if __name__ == "__main__": # run the main function main() - # close sim app - simulation_app.close() diff --git a/scripts/imitation_learning/locomanipulation_sdg/generate_data.py b/scripts/imitation_learning/locomanipulation_sdg/generate_data.py index 9343e4e5e5fb..b26de0c7b632 100644 --- a/scripts/imitation_learning/locomanipulation_sdg/generate_data.py +++ b/scripts/imitation_learning/locomanipulation_sdg/generate_data.py @@ -470,7 +470,7 @@ def setup_navigation_scene( print(f"Failed to randomize fixture placement for {fixture.entity_name}", flush=True) return None - sync_simulation_state(env) + sync_simulation_state(env) occupancy_map = merge_occupancy_maps([occupancy_map, fixture.get_occupancy_map()]) initial_state = env.load_input_data(input_episode_data, 0) diff --git a/scripts/reinforcement_learning/rl_games/play.py b/scripts/reinforcement_learning/rl_games/play.py index e4e81ea7d65f..eb2390af90d6 100644 --- a/scripts/reinforcement_learning/rl_games/play.py +++ b/scripts/reinforcement_learning/rl_games/play.py @@ -169,28 +169,31 @@ def main(): if agent.is_rnn: agent.init_rnn() # simulate environment - while True: - start_time = time.time() - with torch.inference_mode(): - obs = agent.obs_to_torch(obs) - actions = agent.get_action(obs, is_deterministic=agent.is_deterministic) - obs, _, dones, _ = env.step(actions) - - if len(dones) > 0: - if agent.is_rnn and agent.states is not None: - for s in agent.states: - s[:, dones, :] = 0.0 - if args_cli.video: - timestep += 1 - if timestep == args_cli.video_length: - break - - 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() + try: + while True: + start_time = time.time() + with torch.inference_mode(): + obs = agent.obs_to_torch(obs) + actions = agent.get_action(obs, is_deterministic=agent.is_deterministic) + obs, _, dones, _ = env.step(actions) + + if len(dones) > 0: + if agent.is_rnn and agent.states is not None: + for s in agent.states: + s[:, dones, :] = 0.0 + if args_cli.video: + timestep += 1 + if timestep == args_cli.video_length: + break + + 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() + except KeyboardInterrupt: + pass if __name__ == "__main__": diff --git a/scripts/reinforcement_learning/rl_games/train.py b/scripts/reinforcement_learning/rl_games/train.py index f37f910fe292..5ad13b401bb3 100644 --- a/scripts/reinforcement_learning/rl_games/train.py +++ b/scripts/reinforcement_learning/rl_games/train.py @@ -221,15 +221,16 @@ def main(): wandb.config.update({"env_cfg": env_cfg.to_dict()}) wandb.config.update({"agent_cfg": agent_cfg}) - if args_cli.checkpoint is not None: - runner.run({"train": True, "play": False, "sigma": train_sigma, "checkpoint": resume_path}) - else: - runner.run({"train": True, "play": False, "sigma": train_sigma}) - - print(f"Training time: {round(time.time() - start_time, 2)} seconds") - - # close the simulator - env.close() + try: + if args_cli.checkpoint is not None: + runner.run({"train": True, "play": False, "sigma": train_sigma, "checkpoint": resume_path}) + else: + runner.run({"train": True, "play": False, "sigma": train_sigma}) + print(f"Training time: {round(time.time() - start_time, 2)} seconds") + # close the simulator + env.close() + except KeyboardInterrupt: + pass if __name__ == "__main__": diff --git a/scripts/reinforcement_learning/rsl_rl/play.py b/scripts/reinforcement_learning/rsl_rl/play.py index e4828152b15e..f790f627a22a 100644 --- a/scripts/reinforcement_learning/rsl_rl/play.py +++ b/scripts/reinforcement_learning/rsl_rl/play.py @@ -6,23 +6,32 @@ """Script to play a checkpoint if an RL agent from RSL-RL.""" import argparse +import importlib.metadata as metadata import os import sys import time import gymnasium as gym import torch +from packaging import version from rsl_rl.runners import DistillationRunner, OnPolicyRunner -from isaaclab.envs import DirectMARLEnvCfg +from isaaclab.envs import DirectMARLEnvCfg, DirectRLEnvCfg, ManagerBasedRLEnvCfg from isaaclab.utils.assets import retrieve_file_path from isaaclab.utils.dict import print_dict -from isaaclab_rl.rsl_rl import RslRlVecEnvWrapper, export_policy_as_jit, export_policy_as_onnx +from isaaclab_rl.rsl_rl import ( + RslRlBaseRunnerCfg, + RslRlVecEnvWrapper, + export_policy_as_jit, + export_policy_as_onnx, + handle_deprecated_rsl_rl_cfg, +) from isaaclab_rl.utils.pretrained_checkpoint import get_published_pretrained_checkpoint import isaaclab_tasks # noqa: F401 -from isaaclab_tasks.utils import add_launcher_args, get_checkpoint_path, launch_simulation, resolve_task_config +from isaaclab_tasks.utils import add_launcher_args, get_checkpoint_path, launch_simulation +from isaaclab_tasks.utils.hydra import hydra_task_config # local imports import cli_args # isort: skip @@ -57,10 +66,13 @@ sys.argv = [sys.argv[0]] + hydra_args +# Check for installed RSL-RL version +installed_version = metadata.version("rsl-rl-lib") + -def main(): +@hydra_task_config(args_cli.task, args_cli.agent) +def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlBaseRunnerCfg): """Play with RSL-RL agent.""" - env_cfg, agent_cfg = resolve_task_config(args_cli.task, args_cli.agent) with launch_simulation(env_cfg, args_cli): # grab task name for checkpoint path task_name = args_cli.task.split(":")[-1] @@ -70,7 +82,11 @@ def main(): agent_cfg = 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 + # handle deprecated configurations + agent_cfg = handle_deprecated_rsl_rl_cfg(agent_cfg, installed_version) + # 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 @@ -118,6 +134,7 @@ def main(): 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": @@ -129,22 +146,32 @@ def main(): # obtain the trained policy for inference policy = runner.get_inference_policy(device=env.unwrapped.device) - try: - policy_nn = runner.alg.policy - except AttributeError: - policy_nn = runner.alg.actor_critic - - 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 the trained policy to JIT and ONNX formats 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") + + if version.parse(installed_version) >= version.parse("4.0.0"): + # use the new export functions for rsl-rl >= 4.0.0 + runner.export_policy_to_jit(path=export_model_dir, filename="policy.pt") + runner.export_policy_to_onnx(path=export_model_dir, filename="policy.onnx") + policy_nn = None # Not needed for rsl-rl >= 4.0.0 + else: + # extract the neural network for rsl-rl < 4.0.0 + if version.parse(installed_version) >= version.parse("2.3.0"): + policy_nn = runner.alg.policy + else: + 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 to JIT and ONNX + 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 @@ -152,23 +179,33 @@ def main(): obs = env.get_observations() timestep = 0 # simulate environment - while True: - start_time = time.time() - with torch.inference_mode(): - actions = policy(obs) - obs, _, dones, _ = env.step(actions) - policy_nn.reset(dones) - if args_cli.video: - timestep += 1 - if timestep == args_cli.video_length: - break - - 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() + try: + while True: + start_time = time.time() + # run everything in inference mode + with torch.inference_mode(): + # agent stepping + actions = policy(obs) + # env stepping + obs, _, dones, _ = env.step(actions) + # reset recurrent states for episodes that have terminated + if version.parse(installed_version) >= version.parse("4.0.0"): + policy.reset(dones) + else: + policy_nn.reset(dones) + if args_cli.video: + timestep += 1 + if timestep == args_cli.video_length: + break + + 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() + except KeyboardInterrupt: + pass if __name__ == "__main__": diff --git a/scripts/reinforcement_learning/rsl_rl/train.py b/scripts/reinforcement_learning/rsl_rl/train.py index 5d7c90b95447..7ca2d3156da5 100644 --- a/scripts/reinforcement_learning/rsl_rl/train.py +++ b/scripts/reinforcement_learning/rsl_rl/train.py @@ -19,14 +19,15 @@ from packaging import version from rsl_rl.runners import DistillationRunner, OnPolicyRunner -from isaaclab.envs import DirectMARLEnvCfg, ManagerBasedRLEnvCfg +from isaaclab.envs import DirectMARLEnvCfg, DirectRLEnvCfg, ManagerBasedRLEnvCfg from isaaclab.utils.dict import print_dict from isaaclab.utils.io import dump_yaml -from isaaclab_rl.rsl_rl import RslRlVecEnvWrapper +from isaaclab_rl.rsl_rl import RslRlBaseRunnerCfg, RslRlVecEnvWrapper, handle_deprecated_rsl_rl_cfg import isaaclab_tasks # noqa: F401 -from isaaclab_tasks.utils import add_launcher_args, get_checkpoint_path, launch_simulation, resolve_task_config +from isaaclab_tasks.utils import add_launcher_args, get_checkpoint_path, launch_simulation +from isaaclab_tasks.utils.hydra import hydra_task_config # local imports import cli_args # isort: skip @@ -85,9 +86,9 @@ exit(1) -def main(): +@hydra_task_config(args_cli.task, args_cli.agent) +def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlBaseRunnerCfg): """Train with RSL-RL agent.""" - env_cfg, agent_cfg = resolve_task_config(args_cli.task, args_cli.agent) with launch_simulation(env_cfg, args_cli): # override configurations with non-hydra CLI arguments agent_cfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli) @@ -96,6 +97,9 @@ def main(): args_cli.max_iterations if args_cli.max_iterations is not None else agent_cfg.max_iterations ) + # handle deprecated configurations + agent_cfg = handle_deprecated_rsl_rl_cfg(agent_cfg, installed_version) + # set the environment seed # note: certain randomizations occur in the environment initialization so we set the seed here env_cfg.seed = agent_cfg.seed @@ -193,12 +197,13 @@ def main(): 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) - - print(f"Training time: {round(time.time() - start_time, 2)} seconds") - - # close the simulator - env.close() + try: + runner.learn(num_learning_iterations=agent_cfg.max_iterations, init_at_random_ep_len=True) + print(f"Training time: {round(time.time() - start_time, 2)} seconds") + # close the simulator + env.close() + except KeyboardInterrupt: + pass if __name__ == "__main__": diff --git a/scripts/reinforcement_learning/sb3/play.py b/scripts/reinforcement_learning/sb3/play.py index 1ce78c3566c4..a6f222d346c1 100644 --- a/scripts/reinforcement_learning/sb3/play.py +++ b/scripts/reinforcement_learning/sb3/play.py @@ -160,22 +160,25 @@ def main(): obs = env.reset() timestep = 0 # simulate environment - while True: - start_time = time.time() - with torch.inference_mode(): - actions, _ = agent.predict(obs, deterministic=True) - obs, _, _, _ = env.step(actions) - if args_cli.video: - timestep += 1 - if timestep == args_cli.video_length: - break - - 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() + try: + while True: + start_time = time.time() + with torch.inference_mode(): + actions, _ = agent.predict(obs, deterministic=True) + obs, _, _, _ = env.step(actions) + if args_cli.video: + timestep += 1 + if timestep == args_cli.video_length: + break + + 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() + except KeyboardInterrupt: + pass if __name__ == "__main__": diff --git a/scripts/reinforcement_learning/skrl/play.py b/scripts/reinforcement_learning/skrl/play.py index f6572123f789..0349d405967e 100644 --- a/scripts/reinforcement_learning/skrl/play.py +++ b/scripts/reinforcement_learning/skrl/play.py @@ -194,27 +194,30 @@ def main(): obs, _ = env.reset() timestep = 0 # simulate environment - while True: - start_time = time.time() - - with torch.inference_mode(): - outputs = runner.agent.act(obs, timestep=0, timesteps=0) - if hasattr(env, "possible_agents"): - actions = {a: outputs[-1][a].get("mean_actions", outputs[0][a]) for a in env.possible_agents} - else: - actions = outputs[-1].get("mean_actions", outputs[0]) - obs, _, _, _, _ = env.step(actions) - if args_cli.video: - timestep += 1 - if timestep == args_cli.video_length: - break - - 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() + try: + while True: + start_time = time.time() + + with torch.inference_mode(): + outputs = runner.agent.act(obs, timestep=0, timesteps=0) + if hasattr(env, "possible_agents"): + actions = {a: outputs[-1][a].get("mean_actions", outputs[0][a]) for a in env.possible_agents} + else: + actions = outputs[-1].get("mean_actions", outputs[0]) + obs, _, _, _, _ = env.step(actions) + if args_cli.video: + timestep += 1 + if timestep == args_cli.video_length: + break + + 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() + except KeyboardInterrupt: + pass if __name__ == "__main__": diff --git a/scripts/reinforcement_learning/skrl/train.py b/scripts/reinforcement_learning/skrl/train.py index 31a26967683c..750ebeb8798d 100644 --- a/scripts/reinforcement_learning/skrl/train.py +++ b/scripts/reinforcement_learning/skrl/train.py @@ -208,12 +208,13 @@ def main(): runner.agent.load(resume_path) # run training - runner.run() - - print(f"Training time: {round(time.time() - start_time, 2)} seconds") - - # close the simulator - env.close() + try: + runner.run() + print(f"Training time: {round(time.time() - start_time, 2)} seconds") + # close the simulator + env.close() + except KeyboardInterrupt: + pass if __name__ == "__main__": diff --git a/scripts/tools/convert_mjcf.py b/scripts/tools/convert_mjcf.py index 9679bea2d6c6..d693a493bd6e 100644 --- a/scripts/tools/convert_mjcf.py +++ b/scripts/tools/convert_mjcf.py @@ -24,8 +24,8 @@ --merge-mesh Merge meshes where possible to optimize the model. (default: False) --collision-from-visuals Generate collision geometry from visual geometries. (default: False) --collision-type Type of collision geometry to use. (default: "default") - --self-collision Activate self-collisions between links. (default: False) - + --self-collision Activate self-collisions between links. (default: False) + --import-physics-scene Import the physics scene from the MJCF file. (default: False) """ """Launch Isaac Sim Simulator first.""" @@ -62,7 +62,12 @@ default=False, help="Activate self-collisions between links of the articulation.", ) - +parser.add_argument( + "--import-physics-scene", + action="store_true", + default=False, + help="Import the physics scene (worldbody, defaults) from the MJCF file. Use --no-import-physics-scene to disable.", +) # append AppLauncher cli args AppLauncher.add_app_launcher_args(parser) # parse the arguments @@ -105,6 +110,7 @@ def main(): collision_from_visuals=args_cli.collision_from_visuals, collision_type=args_cli.collision_type, self_collision=args_cli.self_collision, + import_physics_scene=args_cli.import_physics_scene, ) # Print info diff --git a/source/isaaclab/config/extension.toml b/source/isaaclab/config/extension.toml index d6cad837b5f9..4acb8613d6d9 100644 --- a/source/isaaclab/config/extension.toml +++ b/source/isaaclab/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "4.5.5" +version = "4.5.11" # Description title = "Isaac Lab framework for Robot Learning" diff --git a/source/isaaclab/docs/CHANGELOG.rst b/source/isaaclab/docs/CHANGELOG.rst index d06330d160b3..a448076ec2c2 100644 --- a/source/isaaclab/docs/CHANGELOG.rst +++ b/source/isaaclab/docs/CHANGELOG.rst @@ -1,6 +1,91 @@ Changelog --------- +4.5.12 (2026-03-10) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Unified :class:`~isaaclab.envs.mdp.randomize_physics_scene_gravity` into a single + class-based term that auto-detects the active physics backend (PhysX or Newton). + +Removed +^^^^^^^ + +* Removed ``randomize_newton_physics_scene_gravity``. Use + :class:`~isaaclab.envs.mdp.randomize_physics_scene_gravity` instead, which + handles both PhysX and Newton backends automatically. + + +4.5.11 (2026-02-28) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added ``visualizers`` parameter to ``build_simulation_context()``. + + +4.5.10 (2026-03-09) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Refactored :func:`~isaaclab.utils.module.lazy_export` to infer fallback packages + and relative wildcard re-exports from the ``.pyi`` stub, making the stub the + single source of truth. The ``packages`` argument is deprecated. + + +4.5.9 (2026-03-08) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* :mod:`isaaclab_ov` is now always installed when using ``./isaaclab.sh -i`` (or + ``--install all``), but with ``--no-deps`` so it is importable (e.g. by + :mod:`isaaclab_tasks` presets) without pulling in the optional ``ovrtx`` dependency. + To install the ovrtx dependency for OVRTX rendering, run ``./isaaclab.sh -i ovrtx``. + +4.5.8 (2026-03-06) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added ``debugpy`` to :mod:`isaaclab` package dependencies to support debugging out of the box. + +4.5.7 (2026-03-06) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Extended ``test_articulation_iface.py`` with Newton backend mock tests — added + Newton-specific mock view setup, sim config, and test parametrization alongside + existing PhysX tests. + +* Extended ``test_rigid_object_iface.py`` with Newton backend mock tests — added + Newton-specific mock view setup and test parametrization. + +* Fixed mask type handling in ``test_rigid_object_collection_iface.py`` to use + consistent mask types across backends. + + +4.5.6 (2026-03-06) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Reorganized Visualizers and SDP packages +* Added Visualizer unit tests +* Improved PhysX Scene Data Provider perf +* Tweaked default Visualizer Configs + + 4.5.5 (2026-03-05) ~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab/isaaclab/app/app_launcher.py b/source/isaaclab/isaaclab/app/app_launcher.py index 1e04efd3270c..f4ba77f1d2db 100644 --- a/source/isaaclab/isaaclab/app/app_launcher.py +++ b/source/isaaclab/isaaclab/app/app_launcher.py @@ -25,14 +25,16 @@ with contextlib.suppress(ModuleNotFoundError): import isaacsim # noqa: F401 - -from isaacsim import SimulationApp + from isaacsim import SimulationApp from isaaclab.app.settings_manager import get_settings_manager, initialize_carb_settings # import logger logger = logging.getLogger(__name__) +# Suppress noisy debug-level websocket frame logs from the Kit LiveSync server +logging.getLogger("websockets").setLevel(logging.WARNING) + class ExplicitAction(argparse.Action): """Custom action to track if an argument was explicitly passed by the user.""" @@ -44,6 +46,35 @@ def __call__(self, parser, namespace, values, option_string=None): setattr(namespace, f"{self.dest}_explicit", True) +def _parse_visualizer_csv(value: str) -> list[str]: + """Parse visualizer list from a single comma-delimited CLI token.""" + valid = {"kit", "newton", "rerun", "viser"} + token = (value or "").strip() + if not token: + raise argparse.ArgumentTypeError( + "Invalid --visualizer value: empty string. Use a comma-separated list, e.g. --visualizer kit,newton." + ) + if " " in token: + raise argparse.ArgumentTypeError( + "Invalid --visualizer value: spaces are not allowed. " + "Use a comma-separated list without spaces, e.g. --visualizer kit,newton,rerun,viser." + ) + + names = [item.strip().lower() for item in token.split(",")] + if any(not name for name in names): + raise argparse.ArgumentTypeError( + "Invalid --visualizer value: empty visualizer entry detected. " + "Use a comma-separated list without empty items." + ) + invalid = [name for name in names if name not in valid] + if invalid: + raise argparse.ArgumentTypeError( + f"Invalid --visualizer value(s): {', '.join(invalid)}. Valid options: {', '.join(sorted(valid))}." + ) + # De-duplicate while preserving order. + return list(dict.fromkeys(names)) + + class AppLauncher: """A utility class to launch Isaac Sim application based on command-line arguments and environment variables. @@ -241,7 +272,8 @@ def add_app_launcher_args(parser: argparse.ArgumentParser) -> None: - ``newton``: Use Newton visualizer. - ``viser``: Use Viser visualizer. - ``kit``: Use Omniverse Kit visualizer. - - Multiple visualizers can be specified: ``--visualizer rerun newton viser`` + - Multiple visualizers can be specified as a comma-delimited list: + ``--visualizer rerun,newton,viser`` - If not specified (default), NO visualizers will be initialized and headless mode is auto-enabled. Note: If visualizer configs are not defined in the simulation config, default configs will be @@ -252,7 +284,7 @@ def add_app_launcher_args(parser: argparse.ArgumentParser) -> None: * ``visualizer_max_worlds`` (int | None): Optional global override for the maximum number of worlds rendered in Newton-based visualizers (newton, rerun, viser). If omitted, each visualizer uses its - config default. Use ``0`` to render all worlds. + config default. Args: parser: An argument parser instance to be extended with the AppLauncher specific options. @@ -327,10 +359,9 @@ def add_app_launcher_args(parser: argparse.ArgumentParser) -> None: ) arg_group.add_argument( "--visualizer", - type=str, - nargs="+", + type=_parse_visualizer_csv, default=None, - help="Visualizer backends to enable (e.g., kit, newton, rerun, viser).", + help="Visualizer backends to enable as CSV (e.g., kit,newton,rerun,viser).", ) # Add the deprecated cpu flag to raise an error if it is used arg_group.add_argument("--cpu", action="store_true", help=argparse.SUPPRESS) @@ -403,8 +434,7 @@ def add_app_launcher_args(parser: argparse.ArgumentParser) -> None: default=AppLauncher._APPLAUNCHER_CFG_INFO["visualizer_max_worlds"][1], help=( "Optional global max worlds override for Newton-based visualizers (newton/rerun/viser). " - "If omitted, visualizer config defaults are used. " - "Set to 0 to render all worlds." + "If omitted, visualizer config defaults are used." ), ) # special flag for backwards compatibility @@ -979,7 +1009,7 @@ def _set_visualizer_settings(self, launcher_args: dict) -> None: with contextlib.suppress(Exception): visualizer_str = " ".join(visualizers) if visualizers else "" - get_settings_manager().set_string("/isaaclab/visualizer", visualizer_str) + get_settings_manager().set_string("/isaaclab/visualizer/types", visualizer_str) # Store as int setting where -1 means "use per-visualizer defaults". if visualizer_max_worlds is None: get_settings_manager().set_int("/isaaclab/visualizer/max_worlds", -1) diff --git a/source/isaaclab/isaaclab/cli/__init__.py b/source/isaaclab/isaaclab/cli/__init__.py index a10a9d812b5d..f7488b1824e9 100644 --- a/source/isaaclab/isaaclab/cli/__init__.py +++ b/source/isaaclab/isaaclab/cli/__init__.py @@ -39,7 +39,9 @@ def cli() -> None: "Install Isaac Lab sub-packages and RL frameworks.\n" "Accepts a comma-separated list of sub-package names, one of the RL frameworks, or a special value.\n" "\n" - "Sub-packages: assets, physx, contrib, mimic, newton, rl, tasks, teleop.\n" + "Sub-packages: assets, physx, contrib, mimic, newton, rl, tasks, teleop, visualizers.\n" + "Use -i ovrtx to install the ovrtx dependency for isaaclab_ov.\n" + "Visualizer selectors: visualizers[all|kit|newton|rerun|viser].\n" "RL frameworks: rl_games, rsl_rl, sb3, skrl, robomimic.\n" "\n" "Passing an RL framework name installs all sub-packages + that framework.\n" @@ -48,6 +50,7 @@ def cli() -> None: "- all - Install all sub-packages + all RL frameworks (default).\n" "- none - Install only the core 'isaaclab' package.\n" "- (-i or --install without value) - Install all sub-packages + all RL frameworks.\n" + "- quote visualizer selectors in bash, e.g. --install 'visualizers[rerun]'.\n" ), ) parser.add_argument( diff --git a/source/isaaclab/isaaclab/cli/commands/install.py b/source/isaaclab/isaaclab/cli/commands/install.py index 80aef5bda643..975083542e61 100644 --- a/source/isaaclab/isaaclab/cli/commands/install.py +++ b/source/isaaclab/isaaclab/cli/commands/install.py @@ -127,22 +127,100 @@ def _ensure_cuda_torch() -> None: # Valid sub-package names that can be passed to --install. # Each sub-package maps to a source directory named "isaaclab_" under source/. -VALID_ISAACLAB_SUBPACKAGES: set[str] = {"assets", "physx", "contrib", "mimic", "newton", "rl", "tasks", "teleop"} +VALID_ISAACLAB_SUBPACKAGES: set[str] = { + "assets", + "ovrtx", + "physx", + "contrib", + "mimic", + "newton", + "rl", + "tasks", + "teleop", + "visualizers", +} +# Sub-packages that are always installed but with --no-deps when install_type is "all", +# so they are importable (e.g. for config types) without pulling in optional heavy deps. +INSTALL_NO_DEPS_SUBPACKAGES: set[str] = {"ov"} + +# -i ovrtx installs this dependency only (isaaclab_ov is already installed with --no-deps). +# Keep in sync with isaaclab_ov/setup.py INSTALL_REQUIRES. +OVRTX_PIP_SPEC: str = "ovrtx>=0.2.0,<0.3.0" + +VALID_VISUALIZER_EXTRAS: set[str] = {"all", "kit", "newton", "rerun", "viser"} # RL framework names accepted. # Passing one of these installs all extensions + that framework. VALID_RL_FRAMEWORKS: set[str] = {"rl_games", "rsl_rl", "sb3", "skrl", "robomimic"} -def _install_isaaclab_extensions(extensions: list[str] | None = None) -> None: +def _split_install_items(install_type: str) -> list[str]: + """Split comma-separated install items, ignoring commas inside brackets.""" + parts: list[str] = [] + buf: list[str] = [] + bracket_depth = 0 + for ch in install_type: + if ch == "[": + bracket_depth += 1 + elif ch == "]": + bracket_depth = max(0, bracket_depth - 1) + if ch == "," and bracket_depth == 0: + token = "".join(buf).strip() + if token: + parts.append(token) + buf = [] + else: + buf.append(ch) + token = "".join(buf).strip() + if token: + parts.append(token) + return parts + + +def _parse_visualizer_selector(token: str) -> str | None: + """Parse visualizer selector token like 'visualizers[rerun]' into '[rerun]'.""" + if token == "visualizers": + return "[all]" + prefix = "visualizers[" + if not (token.startswith(prefix) and token.endswith("]")): + return None + + extras_raw = token[len(prefix) : -1].strip() + if not extras_raw: + return "[all]" + + extras = [x.strip() for x in extras_raw.split(",") if x.strip()] + invalid = [x for x in extras if x not in VALID_VISUALIZER_EXTRAS] + if invalid: + valid = ", ".join(sorted(VALID_VISUALIZER_EXTRAS)) + print_warning( + f"Unknown visualizer extra(s) in '{token}': {', '.join(invalid)}. " + f"Valid visualizer extras: {valid}. Skipping visualizers selector." + ) + return None + + return f"[{','.join(extras)}]" + + +def _install_isaaclab_extensions( + extensions: list[str] | None = None, + extension_extras: dict[str, str] | None = None, + exclude: set[str] | None = None, +) -> None: """Install Isaac Lab extensions from the source directory. Scans ``source/`` for sub-directories that contain a ``setup.py`` and installs each one as an editable pip package. Args: - extensions: Optional, list of source directory names to install - If ``None`` is provided, every extension found under ``source/`` is installed. + extensions: Optional, list of source directory names to install. + If ``None`` is provided, every extension found under ``source/`` + is installed (subject to *exclude*). + extension_extras: Optional mapping from extension source directory + name to pip extras selector (e.g. + ``{"isaaclab_visualizers": "[rerun]"}``). + exclude: Optional set of source directory names to skip even when + *extensions* is ``None``. """ python_exe = extract_python_exe() source_dir = ISAACLAB_ROOT / "source" @@ -151,26 +229,65 @@ def _install_isaaclab_extensions(extensions: list[str] | None = None) -> None: print_warning(f"Source directory not found: {source_dir}") return - # recursively look into directories and install them - # this does not check dependencies between extensions - # source directory + # Collect installable extensions from source/. + install_items = [] for item in source_dir.iterdir(): - if item.is_dir() and (item / "setup.py").exists(): - # Skip extensions not in the requested list. - if extensions is not None and item.name not in extensions: - continue - print_info(f"Installing extension: {item.name}") - # If the directory contains setup.py then install the python module. - run_command( - [ - python_exe, - "-m", - "pip", - "install", - "--editable", - str(item), - ] - ) + if not (item.is_dir() and (item / "setup.py").exists()): + continue + if extensions is not None and item.name not in extensions: + continue + if exclude and item.name in exclude: + continue + install_items.append(item) + + # Install order matters for local editable deps: + # packages like isaaclab_visualizers depend on the local isaaclab package. + install_items.sort(key=lambda item: (item.name != "isaaclab", item.name)) + + for item in install_items: + print_info(f"Installing extension: {item.name}") + extras_suffix = (extension_extras or {}).get(item.name, "") + install_target = f"{item}{extras_suffix}" + run_command( + [ + python_exe, + "-m", + "pip", + "install", + "--editable", + install_target, + ] + ) + + +def _install_ovrtx_dependency() -> None: + """Install the ovrtx dependency (for use with isaaclab_ov).""" + python_exe = extract_python_exe() + print_info("Installing ovrtx dependency for isaaclab_ov...") + run_command([python_exe, "-m", "pip", "install", OVRTX_PIP_SPEC]) + + +def _install_no_deps_extensions() -> None: + """Install extensions listed in INSTALL_NO_DEPS_SUBPACKAGES with --no-deps.""" + python_exe = extract_python_exe() + source_dir = ISAACLAB_ROOT / "source" + for short_name in INSTALL_NO_DEPS_SUBPACKAGES: + pkg_name = f"isaaclab_{short_name}" + pkg_path = source_dir / pkg_name + if not (pkg_path.is_dir() and (pkg_path / "setup.py").exists()): + continue + print_info(f"Installing {pkg_name} (no dependencies) for importability...") + run_command( + [ + python_exe, + "-m", + "pip", + "install", + "--editable", + str(pkg_path), + "--no-deps", + ] + ) def _install_extra_frameworks(framework_name: str = "all") -> None: @@ -246,26 +363,50 @@ def command_install(install_type: str = "all") -> None: print_info(f"Python executable: {python_exe}") # Decide which source directories (source/isaaclab/*) to install. - # "all" : install everything + all RL frameworks + # "all" : install everything + all RL frameworks (no-deps extensions installed separately with --no-deps) # "none" : core isaaclab only, no RL frameworks # RL framework : install everything + only that RL framework (e.g. "skrl") # "a,b" : core + selected sub-package directories, no RL frameworks + # Extensions in INSTALL_NO_DEPS_SUBPACKAGES are excluded from the main loop and installed with --no-deps. + no_deps_dirs = {f"isaaclab_{name}" for name in INSTALL_NO_DEPS_SUBPACKAGES} + install_ovrtx = False + if install_type == "all": extensions = None + exclude = no_deps_dirs + extension_extras = {"isaaclab_visualizers": "[all]"} framework_type = "all" elif install_type == "none": extensions = ["isaaclab"] + exclude = None + extension_extras = {} framework_type = "none" elif install_type in VALID_RL_FRAMEWORKS: - # Single RL framework name: install all extensions + only that framework. extensions = None + exclude = no_deps_dirs + extension_extras = {"isaaclab_visualizers": "[all]"} framework_type = install_type else: # Parse comma-separated sub-package names into source directory names. extensions = ["isaaclab"] # core is always required - for name in (s.strip() for s in install_type.split(",") if s.strip()): + exclude = None # explicit selection — no exclusions + extension_extras = {} + for name in _split_install_items(install_type): + visualizer_extras = _parse_visualizer_selector(name) + if visualizer_extras is not None: + if "isaaclab_visualizers" not in extensions: + extensions.append("isaaclab_visualizers") + extension_extras["isaaclab_visualizers"] = visualizer_extras + continue if name in VALID_ISAACLAB_SUBPACKAGES: - extensions.append(f"isaaclab_{name}") + if name == "ovrtx": + install_ovrtx = True # install ovrtx dependency only; isaaclab_ov already present + elif name == "visualizers": + if "isaaclab_visualizers" not in extensions: + extensions.append("isaaclab_visualizers") + extension_extras["isaaclab_visualizers"] = "[all]" + else: + extensions.append(f"isaaclab_{name}") else: valid = sorted(VALID_ISAACLAB_SUBPACKAGES) + sorted(VALID_RL_FRAMEWORKS) print_warning(f"Unknown sub-package '{name}'. Valid values: {', '.join(valid)}. Skipping.") @@ -316,7 +457,16 @@ def command_install(install_type: str = "all") -> None: _ensure_cuda_torch() # Install the python modules for the extensions in Isaac Lab. - _install_isaaclab_extensions(extensions) + _install_isaaclab_extensions(extensions, extension_extras, exclude) + + # Install no-deps extensions (e.g. isaaclab_ov) with --no-deps so they are + # importable without pulling in optional deps like ovrtx. + if install_type == "all" or install_type in VALID_RL_FRAMEWORKS: + _install_no_deps_extensions() + + # Install ovrtx when user requested -i ovrtx (the specific dependency for isaaclab_ov). + if install_ovrtx: + _install_ovrtx_dependency() # Install the python packages for supported reinforcement learning frameworks. print_info("Installing extra requirements such as learning frameworks...") diff --git a/source/isaaclab/isaaclab/cli/utils.py b/source/isaaclab/isaaclab/cli/utils.py index 669b9574fe0f..a5f16b82e734 100644 --- a/source/isaaclab/isaaclab/cli/utils.py +++ b/source/isaaclab/isaaclab/cli/utils.py @@ -193,6 +193,8 @@ def run_command( except subprocess.CalledProcessError as e: print_error(f'Command failed with code {e.returncode}: "{command_str}"') sys.exit(e.returncode) + except KeyboardInterrupt: + sys.exit(130) def extract_python_exe(allow_isaacsim_python: bool = True) -> str: diff --git a/source/isaaclab/isaaclab/cloner/__init__.pyi b/source/isaaclab/isaaclab/cloner/__init__.pyi index ff0055fdf17c..69adc7cd7378 100644 --- a/source/isaaclab/isaaclab/cloner/__init__.pyi +++ b/source/isaaclab/isaaclab/cloner/__init__.pyi @@ -11,6 +11,7 @@ __all__ = [ "filter_collisions", "grid_transforms", "make_clone_plan", + "resolve_visualizer_clone_fn", "usd_replicate", ] @@ -21,5 +22,6 @@ from .cloner_utils import ( filter_collisions, grid_transforms, make_clone_plan, + resolve_visualizer_clone_fn, usd_replicate, ) diff --git a/source/isaaclab/isaaclab/cloner/cloner_cfg.py b/source/isaaclab/isaaclab/cloner/cloner_cfg.py index d5e21e0dd85e..72887017cfe6 100644 --- a/source/isaaclab/isaaclab/cloner/cloner_cfg.py +++ b/source/isaaclab/isaaclab/cloner/cloner_cfg.py @@ -73,6 +73,9 @@ class TemplateCloneCfg: physics_clone_fn: callable | None = None """Function used to perform physics replication.""" + visualizer_clone_fn: callable | None = None + """Optional function used to build precomputed visualizer artifacts from the clone plan.""" + clone_strategy: callable = random """Function used to build prototype-to-environment mapping. Default is :func:`random`.""" diff --git a/source/isaaclab/isaaclab/cloner/cloner_utils.py b/source/isaaclab/isaaclab/cloner/cloner_utils.py index 3bc9ea70dc57..3413cf1ef8fa 100644 --- a/source/isaaclab/isaaclab/cloner/cloner_utils.py +++ b/source/isaaclab/isaaclab/cloner/cloner_utils.py @@ -6,7 +6,9 @@ from __future__ import annotations import itertools +import logging import math +from collections.abc import Callable from typing import TYPE_CHECKING import torch @@ -14,10 +16,13 @@ from pxr import Gf, Sdf, Usd, UsdGeom, Vt import isaaclab.sim as sim_utils +from isaaclab.physics.scene_data_requirements import SceneDataRequirement, VisualizerPrebuiltArtifacts if TYPE_CHECKING: from .cloner_cfg import TemplateCloneCfg +logger = logging.getLogger(__name__) + def clone_from_template(stage: Usd.Stage, num_clones: int, template_clone_cfg: TemplateCloneCfg) -> None: """Clone assets from a template root into per-environment destinations. @@ -56,7 +61,7 @@ def clone_from_template(stage: Usd.Stage, num_clones: int, template_clone_cfg: T src_paths, dest_paths, clone_masking = make_clone_plan(src, dest, num_clones, cfg.clone_strategy, cfg.device) # Spawn the first instance of clones from prototypes, then deactivate the prototypes, those first instances - # will be served as sources for usd and physx replication. + # will be served as sources for usd and physics replication. proto_idx = clone_masking.to(torch.int32).argmax(dim=1) proto_mask = torch.zeros_like(clone_masking) proto_mask.scatter_(1, proto_idx.view(-1, 1).to(torch.long), clone_masking.any(dim=1, keepdim=True)) @@ -70,6 +75,8 @@ def clone_from_template(stage: Usd.Stage, num_clones: int, template_clone_cfg: T replicate_args = [clone_path_fmt.format(0)], [clone_path_fmt], world_indices, mapping if cfg.clone_physics and cfg.physics_clone_fn is not None: cfg.physics_clone_fn(stage, *replicate_args, positions=positions, device=cfg.device) + if cfg.visualizer_clone_fn is not None: + cfg.visualizer_clone_fn(stage, *replicate_args, positions=positions, device=cfg.device) if cfg.clone_usd: # parse env_origins directly from clone_path usd_replicate(stage, *replicate_args, positions=positions) @@ -79,6 +86,8 @@ def clone_from_template(stage: Usd.Stage, num_clones: int, template_clone_cfg: T replicate_args = selected_src, dest_paths, world_indices, clone_masking if cfg.clone_physics and cfg.physics_clone_fn is not None: cfg.physics_clone_fn(stage, *replicate_args, positions=positions, device=cfg.device) + if cfg.visualizer_clone_fn is not None: + cfg.visualizer_clone_fn(stage, *replicate_args, positions=positions, device=cfg.device) if cfg.clone_usd: usd_replicate(stage, *replicate_args) @@ -156,14 +165,13 @@ def usd_replicate( positions: Optional positions (``[E, 3]``) -> ``xformOp:translate``. quaternions: Optional orientations (``[E, 4]``) in ``xyzw`` -> ``xformOp:orient``. - Returns: - None """ rl = stage.GetRootLayer() # Group replication by destination path depth so ancestors land before deeper paths. # This avoids composition issues for nested or interdependent specs. def dp_depth(template: str) -> int: + """Return destination prim path depth for stable parent-first replication.""" dp = template.format(0) return Sdf.Path(dp).pathElementCount @@ -240,8 +248,6 @@ def filter_collisions( prim_paths: Per-clone prim paths. global_paths: Optional global-collider paths. - Returns: - None """ scene_prim = stage.GetPrimAtPath(physicsscene_path) @@ -375,3 +381,37 @@ def grid_transforms(N: int, spacing: float = 1.0, up_axis: str = "z", device="cp ori = torch.zeros((N, 4), device=device) ori[:, 3] = 1.0 # w=1 for identity quaternion return pos, ori + + +def resolve_visualizer_clone_fn( + physics_backend: str, + requirements: SceneDataRequirement, + stage, + set_visualizer_artifact: Callable[[VisualizerPrebuiltArtifacts | None], None], +): + """Return an optional visualizer prebuild hook for clone workflows. + + Args: + physics_backend: Active physics backend name. + requirements: Aggregated scene-data requirements. + stage: USD stage used by the clone callback. + set_visualizer_artifact: Callback for storing prebuilt visualizer artifacts. + + Returns: + Clone callback when the prebuild path is supported; otherwise ``None``. + """ + if "physx" not in physics_backend or not requirements.requires_newton_model: + return None + try: + from isaaclab_newton.cloner.newton_replicate import ( + create_newton_visualizer_prebuild_clone_fn, + ) + except (ImportError, ModuleNotFoundError) as exc: + logger.warning("Visualizer prebuild hook unavailable: failed to import backend helper.") + logger.debug("Visualizer prebuild import failure details: %s", exc) + return None + + return create_newton_visualizer_prebuild_clone_fn( + stage=stage, + set_visualizer_artifact=set_visualizer_artifact, + ) diff --git a/source/isaaclab/isaaclab/controllers/rmp_flow.py b/source/isaaclab/isaaclab/controllers/rmp_flow.py index e0fc1cd46ea7..e8f456eb10e1 100644 --- a/source/isaaclab/isaaclab/controllers/rmp_flow.py +++ b/source/isaaclab/isaaclab/controllers/rmp_flow.py @@ -3,19 +3,19 @@ # # SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations + import os +import numpy as np import torch -from isaacsim.core.prims import SingleArticulation - # enable motion generation extensions from isaacsim.core.utils.extensions import enable_extension, get_extension_path_from_name enable_extension("isaacsim.robot_motion.lula") enable_extension("isaacsim.robot_motion.motion_generation") -from isaacsim.robot_motion.motion_generation import ArticulationMotionPolicy from isaacsim.robot_motion.motion_generation.lula.motion_policies import RmpFlow, RmpFlowSmoothed import isaaclab.sim as sim_utils @@ -52,10 +52,8 @@ def __init__(self, cfg: RmpFlowControllerCfg, device: str): cfg: The configuration for the controller. device: The device to use for computation. """ - # store input self.cfg = cfg self._device = device - # display info print(f"[INFO]: Loading RMPFlow controller URDF from: {self.cfg.urdf_file}") """ @@ -71,39 +69,37 @@ def num_actions(self) -> int: Operations. """ - def initialize(self, prim_paths_expr: str): + def initialize(self, num_robots: int, joint_names: list[str]): """Initialize the controller. Args: - prim_paths_expr: The expression to find the articulation prim paths. + num_robots: Number of robot instances (environments). + joint_names: Ordered list of all joint names from the articulation. """ - # obtain the simulation time physics_dt = sim_utils.SimulationContext.instance().get_physics_dt() - # find all prims - self._prim_paths = sim_utils.find_matching_prim_paths(prim_paths_expr) - self.num_robots = len(self._prim_paths) - # resolve controller + self.num_robots = num_robots + self._physics_dt = physics_dt + if self.cfg.name == "rmp_flow": controller_cls = RmpFlow elif self.cfg.name == "rmp_flow_smoothed": controller_cls = RmpFlowSmoothed else: raise ValueError(f"Unsupported controller in Lula library: {self.cfg.name}") - # create all franka robots references and their controllers - self.articulation_policies = list() - for prim_path in self._prim_paths: - # add robot reference - robot = SingleArticulation(prim_path) - robot.initialize() - # download files if they are not local + name_to_idx = {name: i for i, name in enumerate(joint_names)} + + self._rmpflow_policies: list[RmpFlow | RmpFlowSmoothed] = [] + self._active_indices: list[np.ndarray] = [] + self._watched_indices: list[np.ndarray] = [] + + for _ in range(num_robots): local_urdf_file = retrieve_file_path(_resolve_rmpflow_path(self.cfg.urdf_file), force_download=True) local_collision_file = retrieve_file_path( _resolve_rmpflow_path(self.cfg.collision_file), force_download=True ) local_config_file = retrieve_file_path(_resolve_rmpflow_path(self.cfg.config_file), force_download=True) - # add controller rmpflow = controller_cls( robot_description_path=local_collision_file, urdf_path=local_urdf_file, @@ -112,53 +108,59 @@ def initialize(self, prim_paths_expr: str): maximum_substep_size=physics_dt / self.cfg.evaluations_per_frame, ignore_robot_state_updates=self.cfg.ignore_robot_state_updates, ) - # wrap rmpflow to connect to the Franka robot articulation - articulation_policy = ArticulationMotionPolicy(robot, rmpflow, physics_dt) - self.articulation_policies.append(articulation_policy) - # get number of active joints - self.active_dof_names = self.articulation_policies[0].get_motion_policy().get_active_joints() + + active_indices = np.array([name_to_idx[n] for n in rmpflow.get_active_joints()], dtype=np.intp) + watched_indices = np.array([name_to_idx[n] for n in rmpflow.get_watched_joints()], dtype=np.intp) + + self._rmpflow_policies.append(rmpflow) + self._active_indices.append(active_indices) + self._watched_indices.append(watched_indices) + + self.active_dof_names = self._rmpflow_policies[0].get_active_joints() self.num_dof = len(self.active_dof_names) - # create buffers - # -- for storing command + self._command = torch.zeros(self.num_robots, self.num_actions, device=self._device) - # -- for policy output self.dof_pos_target = torch.zeros((self.num_robots, self.num_dof), device=self._device) self.dof_vel_target = torch.zeros((self.num_robots, self.num_dof), device=self._device) - def reset_idx(self, robot_ids: torch.Tensor = None): + def reset_idx(self, robot_ids: torch.Tensor | None = None): """Reset the internals.""" - # if no robot ids are provided, then reset all robots if robot_ids is None: robot_ids = torch.arange(self.num_robots, device=self._device) - # reset policies for specified robots for index in robot_ids: - self.articulation_policies[index].motion_policy.reset() + self._rmpflow_policies[index].reset() def set_command(self, command: torch.Tensor): """Set target end-effector pose command.""" - # store command self._command[:] = command - def compute(self) -> tuple[torch.Tensor, torch.Tensor]: + def compute( + self, joint_positions: torch.Tensor, joint_velocities: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: """Performs inference with the controller. + Args: + joint_positions: Current joint positions, shape ``[num_robots, num_joints]``. + joint_velocities: Current joint velocities, shape ``[num_robots, num_joints]``. + Returns: The target joint positions and velocity commands. """ - # convert command to numpy command = self._command.cpu().numpy() - # compute control actions - for i, policy in enumerate(self.articulation_policies): - # enable type-hinting - policy: ArticulationMotionPolicy - # set rmpflow target to be the current position of the target cube. - policy.get_motion_policy().set_end_effector_target( - target_position=command[i, 0:3], target_orientation=command[i, 3:7] + all_pos = joint_positions.cpu().numpy() + all_vel = joint_velocities.cpu().numpy() + + for i, rmpflow in enumerate(self._rmpflow_policies): + rmpflow.set_end_effector_target(target_position=command[i, 0:3], target_orientation=command[i, 3:7]) + active_pos = all_pos[i][self._active_indices[i]] + active_vel = all_vel[i][self._active_indices[i]] + watched_pos = all_pos[i][self._watched_indices[i]] + watched_vel = all_vel[i][self._watched_indices[i]] + + pos_targets, vel_targets = rmpflow.compute_joint_targets( + active_pos, active_vel, watched_pos, watched_vel, self._physics_dt ) - # apply action on the robot - action = policy.get_next_articulation_action() - # copy actions into buffer - self.dof_pos_target[i, :] = torch.from_numpy(action.joint_positions[:]).to(self.dof_pos_target) - self.dof_vel_target[i, :] = torch.from_numpy(action.joint_velocities[:]).to(self.dof_vel_target) + self.dof_pos_target[i, :] = torch.from_numpy(pos_targets[:]).to(self.dof_pos_target) + self.dof_vel_target[i, :] = torch.from_numpy(vel_targets[:]).to(self.dof_vel_target) return self.dof_pos_target, self.dof_vel_target diff --git a/source/isaaclab/isaaclab/devices/openxr/retargeters/__init__.py b/source/isaaclab/isaaclab/devices/openxr/retargeters/__init__.py index cd062fe92613..85e47be523a6 100644 --- a/source/isaaclab/isaaclab/devices/openxr/retargeters/__init__.py +++ b/source/isaaclab/isaaclab/devices/openxr/retargeters/__init__.py @@ -7,4 +7,4 @@ from isaaclab.utils.module import lazy_export -lazy_export(packages=["isaaclab_teleop.deprecated.openxr.retargeters"]) +lazy_export() diff --git a/source/isaaclab/isaaclab/devices/openxr/retargeters/__init__.pyi b/source/isaaclab/isaaclab/devices/openxr/retargeters/__init__.pyi new file mode 100644 index 000000000000..c094ea29539b --- /dev/null +++ b/source/isaaclab/isaaclab/devices/openxr/retargeters/__init__.pyi @@ -0,0 +1,6 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab_teleop.deprecated.openxr.retargeters import * diff --git a/source/isaaclab/isaaclab/devices/openxr/retargeters/manipulator/__init__.py b/source/isaaclab/isaaclab/devices/openxr/retargeters/manipulator/__init__.py index 1ceac8277fdb..cc18f59367d3 100644 --- a/source/isaaclab/isaaclab/devices/openxr/retargeters/manipulator/__init__.py +++ b/source/isaaclab/isaaclab/devices/openxr/retargeters/manipulator/__init__.py @@ -7,4 +7,4 @@ from isaaclab.utils.module import lazy_export -lazy_export(packages=["isaaclab_teleop.deprecated.openxr.retargeters.manipulator"]) +lazy_export() diff --git a/source/isaaclab/isaaclab/devices/openxr/retargeters/manipulator/__init__.pyi b/source/isaaclab/isaaclab/devices/openxr/retargeters/manipulator/__init__.pyi new file mode 100644 index 000000000000..5895081bca64 --- /dev/null +++ b/source/isaaclab/isaaclab/devices/openxr/retargeters/manipulator/__init__.pyi @@ -0,0 +1,6 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab_teleop.deprecated.openxr.retargeters.manipulator import * diff --git a/source/isaaclab/isaaclab/envs/direct_marl_env.py b/source/isaaclab/isaaclab/envs/direct_marl_env.py index f022046b182f..0ee0bc643fac 100644 --- a/source/isaaclab/isaaclab/envs/direct_marl_env.py +++ b/source/isaaclab/isaaclab/envs/direct_marl_env.py @@ -136,7 +136,10 @@ def __init__(self, cfg: DirectMARLEnvCfg, render_mode: str | None = None, **kwar # viewport is not available in other rendering modes so the function will throw a warning # FIXME: This needs to be fixed in the future when we unify the UI functionalities even for # non-rendering modes. - if self.sim.has_gui: + # Initialize when GUI is available OR when visualizers are active (headless rendering) + # Visualizers support camera updates via sim.set_camera_view() which forwards to all active visualizers + has_visualizers = bool(self.sim.get_setting("/isaaclab/visualizer")) + if self.sim.has_gui or has_visualizers: self.viewport_camera_controller = ViewportCameraController(self, self.cfg.viewer) else: self.viewport_camera_controller = None @@ -222,7 +225,10 @@ def __init__(self, cfg: DirectMARLEnvCfg, render_mode: str | None = None, **kwar def __del__(self): """Cleanup for the environment.""" - self.close() + import sys + + if not sys.is_finalizing(): + self.close() """ Properties. diff --git a/source/isaaclab/isaaclab/envs/direct_rl_env.py b/source/isaaclab/isaaclab/envs/direct_rl_env.py index dd259ed6602f..9ef788bb159d 100644 --- a/source/isaaclab/isaaclab/envs/direct_rl_env.py +++ b/source/isaaclab/isaaclab/envs/direct_rl_env.py @@ -141,7 +141,10 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs # viewport is not available in other rendering modes so the function will throw a warning # FIXME: This needs to be fixed in the future when we unify the UI functionalities even for # non-rendering modes. - if self.sim.has_gui: + # Initialize when GUI is available OR when visualizers are active (headless rendering) + # Visualizers support camera updates via sim.set_camera_view() which forwards to all active visualizers + has_visualizers = bool(self.sim.get_setting("/isaaclab/visualizer")) + if self.sim.has_gui or has_visualizers: self.viewport_camera_controller = ViewportCameraController(self, self.cfg.viewer) else: self.viewport_camera_controller = None @@ -243,7 +246,10 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs def __del__(self): """Cleanup for the environment.""" - self.close() + import sys + + if not sys.is_finalizing(): + self.close() """ Properties. diff --git a/source/isaaclab/isaaclab/envs/manager_based_env.py b/source/isaaclab/isaaclab/envs/manager_based_env.py index 91725cf78c33..be37f21d390e 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_env.py @@ -148,7 +148,10 @@ def __init__(self, cfg: ManagerBasedEnvCfg): # viewport is not available in other rendering modes so the function will throw a warning # FIXME: This needs to be fixed in the future when we unify the UI functionalities even for # non-rendering modes. - if self.sim.has_gui: + # Initialize when GUI is available OR when visualizers are active (headless rendering) + # Visualizers support camera updates via sim.set_camera_view() which forwards to all active visualizers + has_visualizers = bool(self.sim.get_setting("/isaaclab/visualizer")) + if self.sim.has_gui or has_visualizers: self.viewport_camera_controller = ViewportCameraController(self, self.cfg.viewer) else: self.viewport_camera_controller = None @@ -212,7 +215,10 @@ def __init__(self, cfg: ManagerBasedEnvCfg): def __del__(self): """Cleanup for the environment.""" - self.close() + import sys + + if not sys.is_finalizing(): + self.close() """ Properties. diff --git a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py index eb80c75d80ba..d08b7e3be3a7 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py @@ -271,7 +271,9 @@ def render(self, recompute: bool = False) -> np.ndarray | None: return None elif self.render_mode == "rgb_array": # check that if any render could have happened - if not self.sim.has_gui and not self.sim.has_offscreen_render: + # Check for GUI, offscreen rendering, or visualizers + has_visualizers = bool(self.sim.get_setting("/isaaclab/visualizer")) + if not (self.sim.has_gui or self.sim.has_offscreen_render or has_visualizers): raise RuntimeError( f"Cannot render '{self.render_mode}' - no GUI and offscreen rendering not enabled." " If running headless, make sure --enable_cameras is set." diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_actions_cfg.py b/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_actions_cfg.py index 5c5a7dac602c..a988785bddb3 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_actions_cfg.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_actions_cfg.py @@ -45,9 +45,6 @@ class OffsetCfg: controller: RmpFlowControllerCfg = MISSING - articulation_prim_expr: str = MISSING # The expression to find the articulation prim paths. - """The configuration for the RMPFlow controller.""" - use_relative_mode: bool = False """ Defaults to False. diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py index be5d2787fa79..741de96808c9 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py @@ -181,16 +181,19 @@ def apply_actions(self): joint_pos = wp.to_torch(self._asset.data.joint_pos)[:, self._joint_ids] # compute the delta in joint-space if ee_quat_curr.norm() != 0: - joint_pos_des, joint_vel_des = self._rmpflow_controller.compute() + joint_pos_des, joint_vel_des = self._rmpflow_controller.compute( + wp.to_torch(self._asset.data.joint_pos), wp.to_torch(self._asset.data.joint_vel) + ) else: joint_pos_des = joint_pos.clone() + joint_vel_des = torch.zeros_like(joint_pos_des) # set the joint position command self._asset.set_joint_position_target_index(target=joint_pos_des, joint_ids=self._joint_ids) self._asset.set_joint_velocity_target_index(target=joint_vel_des, joint_ids=self._joint_ids) def reset(self, env_ids: Sequence[int] | None = None) -> None: self._raw_actions[env_ids] = 0.0 - self._rmpflow_controller.initialize(self.cfg.articulation_prim_expr) + self._rmpflow_controller.initialize(self.num_envs, list(self._asset.joint_names)) """ Helper functions. diff --git a/source/isaaclab/isaaclab/envs/mdp/events.py b/source/isaaclab/isaaclab/envs/mdp/events.py index 57111684a3ec..5c93a0c5eda0 100644 --- a/source/isaaclab/isaaclab/envs/mdp/events.py +++ b/source/isaaclab/isaaclab/envs/mdp/events.py @@ -501,50 +501,153 @@ def randomize_rigid_body_collider_offsets( asset.root_view.set_contact_offsets(contact_offset, env_ids.cpu()) -def randomize_physics_scene_gravity( - env: ManagerBasedEnv, - env_ids: torch.Tensor | None, - gravity_distribution_params: tuple[list[float], list[float]], - operation: Literal["add", "scale", "abs"], - distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform", -): +class randomize_physics_scene_gravity(ManagerTermBase): """Randomize gravity by adding, scaling, or setting random values. - This function allows randomizing gravity of the physics scene. The function samples random values from the - given distribution parameters and adds, scales, or sets the values into the physics simulation based on the - operation. + Automatically detects the active physics backend (PhysX or Newton) and applies + the appropriate gravity randomization strategy: - The distribution parameters are lists of two elements each, representing the lower and upper bounds of the - distribution for the x, y, and z components of the gravity vector. The function samples random values for each - component independently. + - **PhysX**: samples a single gravity vector and sets it scene-wide via the PhysX + simulation view. All environments share the same gravity. + - **Newton**: samples per-environment gravity vectors and writes them in-place to + the Newton model's per-world gravity array on GPU. - .. attention:: - This function applied the same gravity for all the environments. + The distribution parameters are tuples of two lists with three floats each, + representing the lower and upper bounds for the x, y, and z gravity components [m/s^2]. - .. tip:: - This function uses CPU tensors to assign gravity. + Args: + cfg: The configuration of the event term. + env: The environment instance. """ - # get the current gravity - gravity = torch.tensor(env.sim.cfg.gravity, device="cpu").unsqueeze(0) - dist_param_0 = torch.tensor(gravity_distribution_params[0], device="cpu") - dist_param_1 = torch.tensor(gravity_distribution_params[1], device="cpu") - gravity = _randomize_prop_by_op( - gravity, - (dist_param_0, dist_param_1), - None, - slice(None), - operation=operation, - distribution=distribution, - ) - # unbatch the gravity tensor into a list - gravity = gravity[0].tolist() - # set the gravity into the physics simulation (local: carb/physx only available with Kit) - import carb # noqa: PLC0415 - import omni.physics.tensors.impl.api as physx # noqa: PLC0415 + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): + super().__init__(cfg, env) - physics_sim_view: physx.SimulationView = sim_utils.SimulationContext.instance().physics_sim_view - physics_sim_view.set_gravity(carb.Float3(*gravity)) + manager_name = env.sim.physics_manager.__name__.lower() + if "newton" in manager_name: + self._backend = "newton" + self._init_newton(cfg, env) + else: + self._backend = "physx" + self._init_physx(env) + + distribution = cfg.params.get("distribution", "uniform") + if distribution == "uniform": + self._dist_fn = math_utils.sample_uniform + elif distribution == "log_uniform": + self._dist_fn = math_utils.sample_log_uniform + elif distribution == "gaussian": + self._dist_fn = math_utils.sample_gaussian + else: + raise NotImplementedError( + f"Unknown distribution: '{distribution}' for gravity randomization." + " Please use 'uniform', 'log_uniform', or 'gaussian'." + ) + + operation = cfg.params["operation"] + if operation not in ("add", "scale", "abs"): + raise NotImplementedError( + f"Unknown operation: '{operation}' for gravity randomization. Please use 'add', 'scale', or 'abs'." + ) + + gravity_distribution_params = cfg.params["gravity_distribution_params"] + self._dist_param_0 = torch.tensor(gravity_distribution_params[0], device=env.device, dtype=torch.float32) + self._dist_param_1 = torch.tensor(gravity_distribution_params[1], device=env.device, dtype=torch.float32) + + def __call__( + self, + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + gravity_distribution_params: tuple[list[float], list[float]], + operation: Literal["add", "scale", "abs"], + distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform", + ): + """Randomize gravity for the specified environments. + + Args: + env: The environment instance. + env_ids: The environment IDs to randomize. If None, all environments are randomized. + gravity_distribution_params: Distribution parameters as a tuple of two lists, each + with 3 floats corresponding to (x, y, z) gravity components [m/s^2]. Updated + into pre-allocated tensors each call to support curriculum-driven range changes. + operation: The operation to apply ('add', 'scale', or 'abs'). + distribution: The distribution type (cached at init, param ignored at runtime). + """ + self._dist_param_0[0] = gravity_distribution_params[0][0] + self._dist_param_1[0] = gravity_distribution_params[1][0] + self._dist_param_0[1] = gravity_distribution_params[0][1] + self._dist_param_1[1] = gravity_distribution_params[1][1] + self._dist_param_0[2] = gravity_distribution_params[0][2] + self._dist_param_1[2] = gravity_distribution_params[1][2] + + if self._backend == "newton": + self._call_newton(env, env_ids, operation) + else: + self._call_physx(env, operation) + + def _init_newton(self, cfg: EventTermCfg, env: ManagerBasedEnv): + """Cache Newton manager reference and solver notification flag.""" + import isaaclab_newton.physics.newton_manager as newton_manager_module # noqa: PLC0415 + from newton.solvers import SolverNotifyFlags # noqa: PLC0415 + + self._newton_manager = newton_manager_module.NewtonManager + self._notify_model_properties = SolverNotifyFlags.MODEL_PROPERTIES + + def _call_newton( + self, + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + operation: str, + ): + """Apply per-environment gravity via Newton's per-world gravity array on GPU.""" + model = self._newton_manager.get_model() + if model is None or model.gravity is None: + raise RuntimeError("Newton model is not initialized. Cannot randomize gravity.") + + gravity = wp.to_torch(model.gravity) + + if env_ids is None: + env_ids = env.scene._ALL_INDICES + if len(env_ids) == 0: + return + + num = len(env_ids) + random_values = self._dist_fn( + self._dist_param_0.unsqueeze(0).expand(num, -1), + self._dist_param_1.unsqueeze(0).expand(num, -1), + (num, 3), + device=env.device, + ) + + if operation == "abs": + gravity[env_ids] = random_values + elif operation == "add": + gravity[env_ids] += random_values + elif operation == "scale": + gravity[env_ids] *= random_values + + self._newton_manager.add_model_change(self._notify_model_properties) + + def _init_physx(self, env: ManagerBasedEnv): + """Cache the ``carb`` module and PhysX simulation view for scene-wide gravity updates.""" + import carb # noqa: PLC0415 + + self._carb = carb + self._physics_sim_view = sim_utils.SimulationContext.instance().physics_sim_view + + def _call_physx(self, env: ManagerBasedEnv, operation: str): + """Sample a single gravity vector and apply it scene-wide via the PhysX simulation view.""" + gravity = torch.tensor(env.sim.cfg.gravity, device="cpu").unsqueeze(0) + gravity = _randomize_prop_by_op( + gravity, + (self._dist_param_0.cpu(), self._dist_param_1.cpu()), + None, + slice(None), + operation=operation, + distribution="uniform", + ) + gravity = gravity[0].tolist() + self._physics_sim_view.set_gravity(self._carb.Float3(*gravity)) class randomize_actuator_gains(ManagerTermBase): diff --git a/source/isaaclab/isaaclab/envs/ui/viewport_camera_controller.py b/source/isaaclab/isaaclab/envs/ui/viewport_camera_controller.py index ac535e2ee3e5..b28b26e86063 100644 --- a/source/isaaclab/isaaclab/envs/ui/viewport_camera_controller.py +++ b/source/isaaclab/isaaclab/envs/ui/viewport_camera_controller.py @@ -12,8 +12,7 @@ import numpy as np import torch - -import omni.kit.app +import warp as wp from isaaclab.assets.articulation.articulation import Articulation @@ -75,6 +74,8 @@ def __init__(self, env: ManagerBasedEnv | DirectRLEnv, cfg: ViewerCfg): self.update_view_to_world() # subscribe to post update event so that camera view can be updated at each rendering step + import omni.kit.app + app_interface = omni.kit.app.get_app_interface() app_event_stream = app_interface.get_post_update_event_stream() self._viewport_camera_update_handle = app_event_stream.create_subscription_to_pop( @@ -160,8 +161,9 @@ def update_view_to_asset_root(self, asset_name: str): self.cfg.asset_name = asset_name # set origin type to asset_root self.cfg.origin_type = "asset_root" - # update the camera origins - self.viewer_origin = self._env.scene[self.cfg.asset_name].data.root_pos_w[self.cfg.env_index] + # update the camera origins (convert Warp array to torch tensor first, then index) + root_pos = wp.to_torch(self._env.scene[self.cfg.asset_name].data.root_pos_w) + self.viewer_origin = root_pos[self.cfg.env_index] # update the camera view self.update_view_location() @@ -193,8 +195,9 @@ def update_view_to_asset_body(self, asset_name: str, body_name: str): self.cfg.asset_name = asset_name # set origin type to asset_body self.cfg.origin_type = "asset_body" - # update the camera origins - self.viewer_origin = self._env.scene[self.cfg.asset_name].data.body_pos_w[self.cfg.env_index, body_id].view(3) + # update the camera origins (convert Warp array to torch tensor first, then index) + body_pos = wp.to_torch(self._env.scene[self.cfg.asset_name].data.body_pos_w) + self.viewer_origin = body_pos[self.cfg.env_index, body_id] # update the camera view self.update_view_location() diff --git a/source/isaaclab/isaaclab/physics/__init__.pyi b/source/isaaclab/isaaclab/physics/__init__.pyi index b4f910866210..14d0d4d61739 100644 --- a/source/isaaclab/isaaclab/physics/__init__.pyi +++ b/source/isaaclab/isaaclab/physics/__init__.pyi @@ -8,7 +8,11 @@ __all__ = [ "PhysicsEvent", "PhysicsManager", "PhysicsCfg", + "BaseSceneDataProvider", + "SceneDataProvider", ] +from .base_scene_data_provider import BaseSceneDataProvider from .physics_manager import CallbackHandle, PhysicsEvent, PhysicsManager from .physics_manager_cfg import PhysicsCfg +from .scene_data_provider import SceneDataProvider diff --git a/source/isaaclab/isaaclab/sim/scene_data_providers/scene_data_provider.py b/source/isaaclab/isaaclab/physics/base_scene_data_provider.py similarity index 98% rename from source/isaaclab/isaaclab/sim/scene_data_providers/scene_data_provider.py rename to source/isaaclab/isaaclab/physics/base_scene_data_provider.py index 6c5de77e85eb..e5b709da0ce7 100644 --- a/source/isaaclab/isaaclab/sim/scene_data_providers/scene_data_provider.py +++ b/source/isaaclab/isaaclab/physics/base_scene_data_provider.py @@ -11,7 +11,7 @@ from typing import Any -class SceneDataProvider(ABC): +class BaseSceneDataProvider(ABC): """Backend-agnostic scene data provider interface.""" @abstractmethod diff --git a/source/isaaclab/isaaclab/physics/physics_manager_cfg.py b/source/isaaclab/isaaclab/physics/physics_manager_cfg.py index d4cda2cc5e1d..0ff2348b591a 100644 --- a/source/isaaclab/isaaclab/physics/physics_manager_cfg.py +++ b/source/isaaclab/isaaclab/physics/physics_manager_cfg.py @@ -3,11 +3,6 @@ # # SPDX-License-Identifier: BSD-3-Clause -# Copyright (c) 2022-2026, The Isaac Lab Project Developers. -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - """Base configuration for physics managers.""" from __future__ import annotations diff --git a/source/isaaclab/isaaclab/physics/scene_data_provider.py b/source/isaaclab/isaaclab/physics/scene_data_provider.py new file mode 100644 index 000000000000..0095a413158d --- /dev/null +++ b/source/isaaclab/isaaclab/physics/scene_data_provider.py @@ -0,0 +1,44 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Factory for creating scene data provider instances.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from isaaclab.utils.backend_utils import FactoryBase + +from .base_scene_data_provider import BaseSceneDataProvider + +if TYPE_CHECKING: + from isaaclab.sim import SimulationContext + + +class SceneDataProvider(FactoryBase, BaseSceneDataProvider): + """Factory for creating scene data provider instances.""" + + _backend_class_names = {"physx": "PhysxSceneDataProvider", "newton": "NewtonSceneDataProvider"} + + @classmethod + def _get_backend(cls, stage, simulation_context: SimulationContext, *args, **kwargs) -> str: + manager_name = simulation_context.physics_manager.__name__.lower() + if "newton" in manager_name: + return "newton" + if "physx" in manager_name: + return "physx" + raise ValueError(f"Unknown physics manager: {manager_name}") + + @classmethod + def _get_module_name(cls, backend: str) -> str: + return f"isaaclab_{backend}.scene_data_providers" + + def __new__(cls, stage, simulation_context: SimulationContext, *args, **kwargs) -> BaseSceneDataProvider: + """Create a new scene data provider based on the active physics backend.""" + result = super().__new__(cls, stage, simulation_context, *args, **kwargs) + if not isinstance(result, BaseSceneDataProvider): + name = type(result).__name__ + raise TypeError(f"Backend scene data provider {name!r} must inherit from BaseSceneDataProvider.") + return result diff --git a/source/isaaclab/isaaclab/physics/scene_data_requirements.py b/source/isaaclab/isaaclab/physics/scene_data_requirements.py new file mode 100644 index 000000000000..616592d9b1e0 --- /dev/null +++ b/source/isaaclab/isaaclab/physics/scene_data_requirements.py @@ -0,0 +1,144 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Central requirement resolution for scene-data consumers. + +This module is intentionally type-based (not config-import based) so requirement +checks stay robust even when optional backend packages are not installed. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class SceneDataRequirement: + """Capabilities required from a scene data provider.""" + + requires_newton_model: bool = False + requires_usd_stage: bool = False + + +@dataclass(frozen=True) +class VisualizerPrebuiltArtifacts: + """Prebuilt model/state payload shared from scene setup to providers. + + This gets produced during clone-time visualizer prebuild and then read by + scene data providers as a fast path (instead of rebuilding from USD). + """ + + model: Any + state: Any + rigid_body_paths: list[str] + articulation_paths: list[str] + num_envs: int + + +_VISUALIZER_REQUIREMENTS: dict[str, SceneDataRequirement] = { + "kit": SceneDataRequirement(requires_usd_stage=True), + "newton": SceneDataRequirement(requires_newton_model=True), + "rerun": SceneDataRequirement(requires_newton_model=True), + "viser": SceneDataRequirement(requires_newton_model=True), +} + +_RENDERER_REQUIREMENTS: dict[str, SceneDataRequirement] = { + "isaac_rtx": SceneDataRequirement(requires_usd_stage=True), + "newton_warp": SceneDataRequirement(requires_newton_model=True), + "ovrtx": SceneDataRequirement(requires_newton_model=True, requires_usd_stage=True), +} + + +def supported_visualizer_types() -> tuple[str, ...]: + """Return supported visualizer type names in sorted order. + + Returns: + Sorted tuple of supported visualizer type names. + """ + return tuple(sorted(_VISUALIZER_REQUIREMENTS)) + + +def supported_renderer_types() -> tuple[str, ...]: + """Return supported renderer type names in sorted order. + + Returns: + Sorted tuple of supported renderer type names. + """ + return tuple(sorted(_RENDERER_REQUIREMENTS)) + + +def requirement_for_visualizer_type(visualizer_type: str) -> SceneDataRequirement: + """Resolve scene-data requirements for one visualizer type. + + Args: + visualizer_type: Visualizer type name. + + Returns: + Requirement object for the given visualizer type. + + Raises: + ValueError: If ``visualizer_type`` is unknown. + """ + requirement = _VISUALIZER_REQUIREMENTS.get(visualizer_type) + if requirement is None: + supported = ", ".join(repr(v) for v in supported_visualizer_types()) + raise ValueError(f"Unknown visualizer type {visualizer_type!r}. Supported types: {supported}.") + return requirement + + +def requirement_for_renderer_type(renderer_type: str) -> SceneDataRequirement: + """Resolve scene-data requirements for one renderer type. + + Args: + renderer_type: Renderer type name. + + Returns: + Requirement object for the given renderer type. + + Raises: + ValueError: If ``renderer_type`` is unknown. + """ + requirement = _RENDERER_REQUIREMENTS.get(renderer_type) + if requirement is None: + supported = ", ".join(repr(v) for v in supported_renderer_types()) + raise ValueError(f"Unknown renderer type {renderer_type!r}. Supported types: {supported}.") + return requirement + + +def aggregate_requirements(requirements: Iterable[SceneDataRequirement]) -> SceneDataRequirement: + """Combine a sequence of requirements using logical OR. + + Args: + requirements: Requirement objects to combine. + + Returns: + Combined requirement object. + """ + requires_newton_model = False + requires_usd_stage = False + for requirement in requirements: + requires_newton_model |= requirement.requires_newton_model + requires_usd_stage |= requirement.requires_usd_stage + return SceneDataRequirement(requires_newton_model=requires_newton_model, requires_usd_stage=requires_usd_stage) + + +def resolve_scene_data_requirements( + visualizer_types: Iterable[str], + renderer_types: Iterable[str] = (), +) -> SceneDataRequirement: + """Resolve combined scene-data requirements from visualizer and renderer types. + + Args: + visualizer_types: Visualizer type names to resolve. + renderer_types: Renderer type names to resolve. + + Returns: + Combined requirement object. + """ + requirements = [requirement_for_visualizer_type(viz_type) for viz_type in visualizer_types] + requirements.extend(requirement_for_renderer_type(renderer_type) for renderer_type in renderer_types) + return aggregate_requirements(requirements) diff --git a/source/isaaclab/isaaclab/scene/interactive_scene.py b/source/isaaclab/isaaclab/scene/interactive_scene.py index 3c8428bff9eb..306c0dd390b2 100644 --- a/source/isaaclab/isaaclab/scene/interactive_scene.py +++ b/source/isaaclab/isaaclab/scene/interactive_scene.py @@ -28,6 +28,7 @@ RigidObjectCollection, RigidObjectCollectionCfg, ) +from isaaclab.physics.scene_data_requirements import resolve_scene_data_requirements from isaaclab.sensors import ContactSensorCfg, FrameTransformerCfg, SensorBase, SensorBaseCfg from isaaclab.sim import SimulationContext from isaaclab.sim.utils.stage import get_current_stage, get_current_stage_id @@ -137,15 +138,18 @@ def __init__(self, cfg: InteractiveSceneCfg): self.sim = SimulationContext.instance() self.stage = get_current_stage() self.stage_id = get_current_stage_id() + self.sim.clear_scene_data_visualizer_prebuilt_artifact() self.physics_backend = self.sim.physics_manager.__name__.lower() + visualizer_clone_fn = None + requested_viz_types = set(self.sim.resolve_visualizer_types()) if "physx" in self.physics_backend: from isaaclab_physx.cloner import physx_replicate physics_clone_fn = physx_replicate elif "newton" in self.physics_backend: - from isaaclab_newton.cloner import newton_replicate + from isaaclab_newton.cloner import newton_physics_replicate - physics_clone_fn = newton_replicate + physics_clone_fn = newton_physics_replicate else: raise ValueError(f"Unsupported physics backend: {self.physics_backend}") # physics scene path @@ -158,6 +162,7 @@ def __init__(self, cfg: InteractiveSceneCfg): clone_in_fabric=self.cfg.clone_in_fabric, device=self.device, physics_clone_fn=physics_clone_fn, + visualizer_clone_fn=None, ) # create source prim @@ -173,8 +178,32 @@ def __init__(self, cfg: InteractiveSceneCfg): ) self._global_prim_paths = list() - if self._is_scene_setup_from_cfg(): + has_scene_cfg_entities = self._is_scene_setup_from_cfg() + if has_scene_cfg_entities: self._add_entities_from_cfg() + + requirements = resolve_scene_data_requirements( + visualizer_types=requested_viz_types, + renderer_types=self._sensor_renderer_types(), + ) + self.sim.update_scene_data_requirements(requirements) + visualizer_clone_fn = cloner.resolve_visualizer_clone_fn( + physics_backend=self.physics_backend, + requirements=requirements, + stage=self.stage, + set_visualizer_artifact=self.sim.set_scene_data_visualizer_prebuilt_artifact, + ) + if visualizer_clone_fn is not None: + logger.debug( + "Enabling visualizer artifact prebuild for clone path " + "(backend=%s, requires_newton_model=%s, requires_usd_stage=%s).", + self.physics_backend, + requirements.requires_newton_model, + requirements.requires_usd_stage, + ) + self.cloner_cfg.visualizer_clone_fn = visualizer_clone_fn + + if has_scene_cfg_entities: self.clone_environments(copy_from_source=(not self.cfg.replicate_physics)) # Collision filtering is PhysX-specific (PhysxSchema.PhysxSceneAPI) if self.cfg.filter_collisions and "physx" in self.physics_backend: @@ -209,8 +238,22 @@ def clone_environments(self, copy_from_source: bool = False): if not copy_from_source: # skip physx cloning, this means physx will walk and parse the stage one by one faithfully self.cloner_cfg.physics_clone_fn(self.stage, *replicate_args, device=self.cloner_cfg.device) + if self.cloner_cfg.visualizer_clone_fn is not None: + self.cloner_cfg.visualizer_clone_fn(self.stage, *replicate_args, device=self.cloner_cfg.device) cloner.usd_replicate(self.stage, *replicate_args) + def _sensor_renderer_types(self) -> list[str]: + """Return renderer type names used by scene sensors.""" + renderer_types: list[str] = [] + for sensor in self._sensors.values(): + sensor_cfg = getattr(sensor, "cfg", None) + renderer_cfg = getattr(sensor_cfg, "renderer_cfg", None) + if renderer_cfg is None: + continue + renderer_type = getattr(renderer_cfg, "renderer_type", "default") + renderer_types.append(renderer_type) + return renderer_types + def filter_collisions(self, global_prim_paths: list[str] | None = None): """Filter environments collisions. diff --git a/source/isaaclab/isaaclab/sim/converters/mjcf_converter.py b/source/isaaclab/isaaclab/sim/converters/mjcf_converter.py index fd1eaf1a6d57..2c8fe992a35a 100644 --- a/source/isaaclab/isaaclab/sim/converters/mjcf_converter.py +++ b/source/isaaclab/isaaclab/sim/converters/mjcf_converter.py @@ -79,6 +79,7 @@ def _convert_asset(self, cfg: MjcfConverterCfg): collision_from_visuals=cfg.collision_from_visuals, collision_type=cfg.collision_type, allow_self_collision=cfg.self_collision, + import_scene=cfg.import_physics_scene, ) importer = MJCFImporter(import_config) diff --git a/source/isaaclab/isaaclab/sim/converters/mjcf_converter_cfg.py b/source/isaaclab/isaaclab/sim/converters/mjcf_converter_cfg.py index fc0574c77f3a..a9dbe620c7da 100644 --- a/source/isaaclab/isaaclab/sim/converters/mjcf_converter_cfg.py +++ b/source/isaaclab/isaaclab/sim/converters/mjcf_converter_cfg.py @@ -28,11 +28,14 @@ class MjcfConverterCfg(AssetConverterBaseCfg): collision_from_visuals: bool = False """Generate collision geometry from visual geometries. Defaults to False.""" - collision_type: str = "default" - """Type of collision geometry to use. Defaults to ``"default"``. + collision_type: str = "Convex Hull" + """Type of collision geometry to use. Defaults to ``"Convex Hull"``. - Supported values are ``"default"``, ``"Convex Hull"``, and ``"Convex Decomposition"``. + Supported values are ``"Convex Hull"``, and ``"Convex Decomposition"``. """ self_collision: bool = False """Activate self-collisions between links of the articulation. Defaults to False.""" + + import_physics_scene: bool = False + """Import the physics scene (time step per second, gravity, etc.) from the MJCF file. Defaults to False.""" diff --git a/source/isaaclab/isaaclab/sim/converters/urdf_converter.py b/source/isaaclab/isaaclab/sim/converters/urdf_converter.py index 96c1ff169bfa..a2c41483a46f 100644 --- a/source/isaaclab/isaaclab/sim/converters/urdf_converter.py +++ b/source/isaaclab/isaaclab/sim/converters/urdf_converter.py @@ -141,12 +141,7 @@ def _convert_asset(self, cfg: UrdfConverterCfg): # step 3: apply optional importer features if cfg.collision_from_visuals: - collision_type_map = { - "convex_hull": "Convex Hull", - "convex_decomposition": "Convex Decomposition", - } - collision_type = collision_type_map.get(cfg.collider_type, "Convex Hull") - importer_utils.collision_from_visuals(stage, collision_type) + importer_utils.collision_from_visuals(stage, cfg.collision_type) importer_utils.enable_self_collision(stage, cfg.self_collision) diff --git a/source/isaaclab/isaaclab/sim/converters/urdf_converter_cfg.py b/source/isaaclab/isaaclab/sim/converters/urdf_converter_cfg.py index 888ef79f68bb..feb13f61ff20 100644 --- a/source/isaaclab/isaaclab/sim/converters/urdf_converter_cfg.py +++ b/source/isaaclab/isaaclab/sim/converters/urdf_converter_cfg.py @@ -134,13 +134,13 @@ class NaturalFrequencyGainsCfg: collision_from_visuals = False """Whether to create collision geometry from visual geometry. Defaults to False.""" - collider_type: Literal["convex_hull", "convex_decomposition"] = "convex_hull" + collision_type: Literal["Convex Hull", "Convex Decomposition"] = "Convex Hull" """The collision shape simplification. Defaults to "convex_hull". Supported values are: - * ``"convex_hull"``: The collision shape is simplified to a convex hull. - * ``"convex_decomposition"``: The collision shape is decomposed into smaller convex shapes for a closer fit. + * ``"Convex Hull"``: The collision shape is simplified to a convex hull. + * ``"Convex Decomposition"``: The collision shape is decomposed into smaller convex shapes for a closer fit. """ self_collision: bool = False @@ -152,3 +152,6 @@ class NaturalFrequencyGainsCfg: .. deprecated:: This option is no longer supported by the URDF importer 3.0. A warning is logged if enabled. """ + + merge_mesh: bool = False + """Merge meshes where possible to optimize the model. Defaults to False.""" diff --git a/source/isaaclab/isaaclab/sim/simulation_context.py b/source/isaaclab/isaaclab/sim/simulation_context.py index 89266cc079e4..efca55208ba6 100644 --- a/source/isaaclab/isaaclab/sim/simulation_context.py +++ b/source/isaaclab/isaaclab/sim/simulation_context.py @@ -21,24 +21,22 @@ import isaaclab.sim as sim_utils import isaaclab.sim.utils.stage as stage_utils from isaaclab.app.settings_manager import SettingsManager -from isaaclab.physics import PhysicsManager +from isaaclab.physics import BaseSceneDataProvider, PhysicsManager, SceneDataProvider +from isaaclab.physics.scene_data_requirements import ( + SceneDataRequirement, + VisualizerPrebuiltArtifacts, + resolve_scene_data_requirements, +) from isaaclab.sim.utils import create_new_stage from isaaclab.utils.version import has_kit -from isaaclab.visualizers import ( - KitVisualizerCfg, - NewtonVisualizerCfg, - RerunVisualizerCfg, - ViserVisualizerCfg, - Visualizer, -) +from isaaclab.visualizers.base_visualizer import BaseVisualizer -from .scene_data_providers import SceneDataProvider from .simulation_cfg import SimulationCfg from .spawners import DomeLightCfg, GroundPlaneCfg logger = logging.getLogger(__name__) -# Visualizer type names (CLI and config). App launcher stores --visualizer a b c as space-separated. +# Visualizer type names (CLI and config). App launcher parses CSV and stores as a space-separated setting. _VISUALIZER_TYPES = ("newton", "rerun", "viser", "kit") @@ -160,8 +158,10 @@ def __init__(self, cfg: SimulationCfg | None = None): self._apply_render_cfg_settings() # Initialize visualizer state (provider/visualizers are created lazily during initialize_visualizers()). - self._scene_data_provider: SceneDataProvider | None = None - self._visualizers: list[Visualizer] = [] + self._scene_data_provider: BaseSceneDataProvider | None = None + self._visualizers: list[BaseVisualizer] = [] + self._scene_data_requirements = SceneDataRequirement() + self._visualizer_prebuilt_artifact: VisualizerPrebuiltArtifacts | None = None self._visualizer_step_counter = 0 # Default visualization dt used before/without visualizer initialization. physics_dt = getattr(self.cfg.physics, "dt", None) @@ -334,7 +334,7 @@ def is_rendering(self) -> bool: self._has_gui or self._has_offscreen_render or self.get_setting("/isaaclab/render/rtx_sensors") - or bool(self.get_setting("/isaaclab/visualizer")) + or bool(self.get_setting("/isaaclab/visualizer/types")) ) def get_physics_dt(self) -> float: @@ -342,40 +342,63 @@ def get_physics_dt(self) -> float: return self.physics_manager.get_physics_dt() def _create_default_visualizer_configs(self, requested_visualizers: list[str]) -> list: - """Create default visualizer configs for requested types.""" + """Create default visualizer configs for requested types. + + Loads only the requested visualizer submodule (e.g. isaaclab_visualizers.rerun) + so dependencies for other backends are not imported. + """ + import importlib + default_configs = [] + cfg_class_names = { + "kit": "KitVisualizerCfg", + "newton": "NewtonVisualizerCfg", + "rerun": "RerunVisualizerCfg", + "viser": "ViserVisualizerCfg", + } for viz_type in requested_visualizers: try: - if viz_type == "newton": - default_configs.append(NewtonVisualizerCfg()) - elif viz_type == "rerun": - default_configs.append(RerunVisualizerCfg()) - elif viz_type == "viser": - default_configs.append(ViserVisualizerCfg()) - elif viz_type == "kit": - default_configs.append(KitVisualizerCfg()) - else: + if viz_type not in _VISUALIZER_TYPES: logger.warning( f"[SimulationContext] Unknown visualizer type '{viz_type}' requested. " f"Valid types: {', '.join(repr(t) for t in _VISUALIZER_TYPES)}. Skipping." ) + continue + mod = importlib.import_module(f"isaaclab_visualizers.{viz_type}") + cfg_cls = getattr(mod, cfg_class_names[viz_type]) + default_configs.append(cfg_cls()) + except (ImportError, ModuleNotFoundError) as exc: + # isaaclab_visualizers is optional; log once at warning level + if "isaaclab_visualizers" in str(exc): + logger.warning( + "[SimulationContext] Visualizer '%s' skipped: isaaclab_visualizers is not installed. " + "Install with: pip install isaaclab_visualizers[%s]", + viz_type, + viz_type, + ) + else: + logger.error( + "[SimulationContext] Failed to create default config for visualizer '%s': %s", + viz_type, + exc, + ) except Exception as exc: logger.error(f"[SimulationContext] Failed to create default config for visualizer '{viz_type}': {exc}") return default_configs def _get_cli_visualizer_types(self) -> list[str]: """Return list of visualizer types requested via CLI (setting).""" - requested = self.get_setting("/isaaclab/visualizer") - if not requested: + requested = self.get_setting("/isaaclab/visualizer/types") + if not isinstance(requested, str) or not requested.strip(): return [] - parts = [p.strip() for p in requested.split(",") if p.strip()] - return [v for part in parts for v in part.split() if v] + # App launcher writes this as a single string; accept comma and/or whitespace separators. + return [value for chunk in requested.split(",") for value in chunk.split() if value] def _get_cli_visualizer_max_worlds_override(self) -> tuple[bool, int | None]: """Return CLI override for visualizer max worlds. Returns: - Tuple of (has_override, value), where value=None means render all worlds. + Tuple of (has_override, value), where value=None means no override. """ value = self.get_setting("/isaaclab/visualizer/max_worlds") if value is None: @@ -386,11 +409,9 @@ def _get_cli_visualizer_max_worlds_override(self) -> tuple[bool, int | None]: logger.warning("[SimulationContext] Invalid /isaaclab/visualizer/max_worlds setting: %r", value) return False, None - # -1 means no CLI override; 0 means explicit "all worlds". + # -1 means no CLI override. if max_worlds < 0: return False, None - if max_worlds == 0: - return True, None return True, max_worlds def _apply_visualizer_cli_overrides(self, visualizer_cfgs: list[Any]) -> None: @@ -408,13 +429,8 @@ def _apply_visualizer_cli_overrides(self, visualizer_cfgs: list[Any]) -> None: cfg.max_worlds = max_worlds_override def resolve_visualizer_types(self) -> list[str]: - """Resolve visualizer types from config or CLI settings.""" - visualizer_cfgs = self.cfg.visualizer_cfgs - if visualizer_cfgs is None: - return self._get_cli_visualizer_types() - - if not isinstance(visualizer_cfgs, list): - visualizer_cfgs = [visualizer_cfgs] + """Resolve effective visualizer types with CLI overrides applied.""" + visualizer_cfgs = self._resolve_visualizer_cfgs() return [cfg.visualizer_type for cfg in visualizer_cfgs if getattr(cfg, "visualizer_type", None)] def _resolve_visualizer_cfgs(self) -> list[Any]: @@ -458,7 +474,13 @@ def initialize_visualizers(self) -> None: if not visualizer_cfgs: return - self.initialize_scene_data_provider(visualizer_cfgs) + # Resolve visualizer-driven requirements once and keep optional artifact payload untouched. + visualizer_types = [ + cfg.visualizer_type for cfg in visualizer_cfgs if getattr(cfg, "visualizer_type", None) is not None + ] + requirements = resolve_scene_data_requirements(visualizer_types=visualizer_types) + self._scene_data_requirements = requirements + self.initialize_scene_data_provider() self._visualizers = [] for cfg in visualizer_cfgs: @@ -466,9 +488,13 @@ def initialize_visualizers(self) -> None: visualizer = cfg.create_visualizer() visualizer.initialize(self._scene_data_provider) self._visualizers.append(visualizer) - logger.info(f"Initialized visualizer: {type(visualizer).__name__} (type: {cfg.visualizer_type})") except Exception as exc: - logger.error(f"Failed to initialize visualizer '{cfg.visualizer_type}' ({type(cfg).__name__}): {exc}") + logger.exception( + "Failed to initialize visualizer '%s' (%s): %s", + cfg.visualizer_type, + type(cfg).__name__, + exc, + ) if not self._visualizers and self._scene_data_provider is not None: close_provider = getattr(self._scene_data_provider, "close", None) @@ -476,20 +502,37 @@ def initialize_visualizers(self) -> None: close_provider() self._scene_data_provider = None - def initialize_scene_data_provider(self, visualizer_cfgs: list[Any]) -> SceneDataProvider: + def initialize_scene_data_provider(self) -> BaseSceneDataProvider: if self._scene_data_provider is None: - if "newton" in self.physics_manager.__name__.lower(): - from .scene_data_providers import NewtonSceneDataProvider + self._scene_data_provider = SceneDataProvider(self.stage, self) + return self._scene_data_provider - self._scene_data_provider = NewtonSceneDataProvider(visualizer_cfgs, self.stage, self) - else: - from .scene_data_providers import PhysxSceneDataProvider + def get_scene_data_requirements(self) -> SceneDataRequirement: + """Return scene-data requirements resolved from visualizers/renderers.""" + return self._scene_data_requirements - self._scene_data_provider = PhysxSceneDataProvider(visualizer_cfgs, self.stage, self) - return self._scene_data_provider + def update_scene_data_requirements(self, requirements: SceneDataRequirement) -> None: + """Update scene-data requirements.""" + self._scene_data_requirements = requirements + + def get_scene_data_visualizer_prebuilt_artifact(self) -> VisualizerPrebuiltArtifacts | None: + """Return optional prebuilt visualizer artifact.""" + return self._visualizer_prebuilt_artifact + + def set_scene_data_visualizer_prebuilt_artifact(self, artifact: VisualizerPrebuiltArtifacts | None) -> None: + """Set or clear the optional visualizer prebuilt artifact. + + The scene (clone flow) writes this once, and providers can read it + during initialization as a fast path. + """ + self._visualizer_prebuilt_artifact = artifact + + def clear_scene_data_visualizer_prebuilt_artifact(self) -> None: + """Clear optional prebuilt artifact in provider context.""" + self.set_scene_data_visualizer_prebuilt_artifact(None) @property - def visualizers(self) -> list[Visualizer]: + def visualizers(self) -> list[BaseVisualizer]: """Returns the list of active visualizers.""" return self._visualizers @@ -562,15 +605,14 @@ def update_visualizers(self, dt: float) -> None: visualizers_to_remove = [] for viz in self._visualizers: try: - if viz.is_rendering_paused(): - continue - if viz.is_closed: - logger.info("Visualizer closed: %s", type(viz).__name__) + if viz.is_closed or not viz.is_running(): + if viz.is_closed: + logger.info("Visualizer closed: %s", type(viz).__name__) + else: + logger.info("Visualizer not running: %s", type(viz).__name__) visualizers_to_remove.append(viz) continue - if not viz.is_running(): - logger.info("Visualizer not running: %s", type(viz).__name__) - visualizers_to_remove.append(viz) + if viz.is_rendering_paused(): continue while viz.is_training_paused() and viz.is_running(): viz.step(0.0) @@ -706,6 +748,7 @@ def build_simulation_context( add_ground_plane: bool = False, add_lighting: bool = False, auto_add_lighting: bool = False, + visualizers: list[str] | None = None, ) -> Iterator[SimulationContext]: """Context manager to build a simulation context with the provided settings. @@ -718,6 +761,10 @@ def build_simulation_context( add_ground_plane: Whether to add a ground plane. Defaults to False. add_lighting: Whether to add a dome light. Defaults to False. auto_add_lighting: Whether to auto-add lighting if GUI present. Defaults to False. + visualizers: List of visualizer backend keys to enable (e.g. ``["kit", "newton", "rerun"]``). + Valid types: ``"kit"``, ``"newton"``, ``"rerun"``, ``"viser"``. + When provided, sets the ``/isaaclab/visualizer/types`` setting so the + existing visualizer resolution machinery picks them up. Defaults to None. Yields: The simulation context to use for the simulation. @@ -733,11 +780,14 @@ def build_simulation_context( sim = SimulationContext(sim_cfg) + if visualizers: + sim.set_setting("/isaaclab/visualizer/types", " ".join(visualizers)) + if add_ground_plane: cfg = GroundPlaneCfg() cfg.func("/World/defaultGroundPlane", cfg) - if add_lighting or (auto_add_lighting and sim.get_setting("/isaaclab/has_gui")): + if add_lighting or (auto_add_lighting and (sim.get_setting("/isaaclab/has_gui") or visualizers)): cfg = DomeLightCfg( color=(0.1, 0.1, 0.1), enable_color_temperature=True, color_temperature=5500, intensity=10000 ) diff --git a/source/isaaclab/isaaclab/utils/backend_utils.py b/source/isaaclab/isaaclab/utils/backend_utils.py index 30d2517d1564..2e4ee55c62df 100644 --- a/source/isaaclab/isaaclab/utils/backend_utils.py +++ b/source/isaaclab/isaaclab/utils/backend_utils.py @@ -6,8 +6,6 @@ import importlib import logging -from isaaclab.sim.simulation_context import SimulationContext - logger = logging.getLogger(__name__) @@ -43,6 +41,8 @@ def _get_backend(cls, *args, **kwargs) -> str: context is initialized yet. """ # Import lazily to avoid import cycles at module load time. + from isaaclab.sim.simulation_context import SimulationContext + manager_name = SimulationContext.instance().physics_manager.__name__.lower() if "newton" in manager_name: return "newton" @@ -51,6 +51,11 @@ def _get_backend(cls, *args, **kwargs) -> str: else: raise ValueError(f"Unknown physics manager: {manager_name}") + @classmethod + def _get_module_name(cls, backend: str) -> str: + """Return module path that hosts backend implementation for a given backend key.""" + return f"isaaclab_{backend}.{cls._module_subpath}" + def __new__(cls, *args, **kwargs): """Create a new instance of an implementation based on the backend.""" backend = cls._get_backend(*args, **kwargs) @@ -62,7 +67,7 @@ def __new__(cls, *args, **kwargs): # This is done to only import the module once. if backend not in cls._registry: # Construct the module name from the backend and the determined subpath. - module_name = f"isaaclab_{backend}.{cls._module_subpath}" + module_name = cls._get_module_name(backend) try: module = importlib.import_module(module_name) class_name = getattr(cls, "_backend_class_names", {}).get(backend, cls.__name__) diff --git a/source/isaaclab/isaaclab/utils/module.py b/source/isaaclab/isaaclab/utils/module.py index 580450d952f4..a1138bb6f875 100644 --- a/source/isaaclab/isaaclab/utils/module.py +++ b/source/isaaclab/isaaclab/utils/module.py @@ -12,145 +12,205 @@ import os import sys import tempfile +import warnings from collections.abc import Callable import lazy_loader as lazy -def _filter_stub(stub_file: str) -> str | None: - """Return a path to a filtered copy of *stub_file* that ``lazy_loader`` can parse. +def _parse_stub(stub_file: str) -> tuple[str | None, list[str], list[str]]: + """Parse a ``.pyi`` stub in a single AST pass. - ``lazy_loader.attach_stub`` only supports relative (``from .x import y``) - imports and rejects absolute imports and star (``*``) imports. This helper - strips those unsupported nodes from the AST so the remaining (local) - relative imports can still be resolved through ``attach_stub``. + Returns: + A 3-tuple of ``(filtered_path, fallback_packages, relative_wildcards)``. - Returns the path to a temporary filtered ``.pyi`` file, or *None* if no - filtering was needed (i.e. the original stub is already compatible). + *filtered_path* is a temporary ``.pyi`` containing only explicit + relative imports (what ``lazy_loader`` can handle), or ``None`` when + no filtering was needed. The temporary file may be empty when the + stub contained only wildcard imports; this is intentional — passing + the original stub to ``lazy_loader`` would raise ``ValueError`` + because it does not support absolute or wildcard imports. + + *fallback_packages* lists fully-qualified package names extracted from + absolute wildcard imports (``from pkg import *``). + + *relative_wildcards* lists submodule names extracted from relative + wildcard imports (``from .mod import *``). """ with open(stub_file) as f: source = f.read() tree = ast.parse(source) + fallback_packages: list[str] = [] + relative_wildcards: list[str] = [] + filtered_body: list[ast.stmt] = [] needs_filter = False - for node in ast.iter_child_nodes(tree): - if isinstance(node, ast.ImportFrom): - if node.level != 1 or any(alias.name == "*" for alias in node.names): - needs_filter = True - break - - if not needs_filter: - return None - filtered = ast.Module(body=[], type_ignores=[]) for node in tree.body: if isinstance(node, ast.ImportFrom): - if node.level != 1: - continue - if any(alias.name == "*" for alias in node.names): + is_star = any(alias.name == "*" for alias in node.names) + if node.level == 1 and not is_star: + filtered_body.append(node) continue - filtered.body.append(node) + if node.level == 0 and is_star and node.module: + fallback_packages.append(node.module) + elif node.level == 1 and is_star and node.module: + relative_wildcards.append(node.module) + needs_filter = True + else: + filtered_body.append(node) + + if not needs_filter: + return None, fallback_packages, relative_wildcards + filtered = ast.Module(body=filtered_body, type_ignores=[]) with tempfile.NamedTemporaryFile(mode="w", suffix=".pyi", delete=False) as tmp: tmp.write(ast.unparse(filtered)) - return tmp.name + + return tmp.name, fallback_packages, relative_wildcards def lazy_export( *, packages: list[str] | tuple[str, ...] | None = None, ) -> tuple[Callable[[str], object], Callable[[], list[str]], list[str]]: - """Lazy-load names from a ``.pyi`` stub, with optional cross-package fallback. + """Lazy-load names from a ``.pyi`` stub. - Call with no arguments to lazily export everything declared in the - adjacent ``.pyi`` stub:: - - from isaaclab.utils.module import lazy_export + The ``.pyi`` stub is the single source of truth for what a module exports. + ``lazy_export()`` reads the stub and derives everything from it: - lazy_export() + * ``from .rewards import foo, bar`` — lazy-loads specific names from a + local submodule (existing ``lazy_loader`` behaviour). + * ``from .rewards import *`` — eagerly imports the submodule and + re-exports all of its public names at ``lazy_export()`` time. + * ``from isaaclab.envs.mdp import *`` — sets up a lazy fallback so that + any name not found locally is resolved from the specified package. - When a module re-exports names from another package at runtime (e.g. - task MDP modules that fall back to ``isaaclab.envs.mdp``), pass the - package names to scan as a fallback:: + Basic usage (no wildcards):: from isaaclab.utils.module import lazy_export - lazy_export(packages=["isaaclab.envs.mdp"]) + lazy_export() + + With a ``.pyi`` stub that contains ``from isaaclab.envs.mdp import *`` + and/or ``from .rewards import *``, no extra arguments are needed — + ``lazy_export()`` infers the behaviour from the stub. Args: - packages: Fully-qualified package names to fall back to when a - name is not found in the local ``.pyi`` stub. When *None* - (the default), only the stub is used. + packages: **Deprecated.** Fallback packages are now inferred from + absolute wildcard imports in the ``.pyi`` stub. Passing this + argument still works but emits a :class:`DeprecationWarning`. + + Raises: + ImportError: If the ``.pyi`` stub declares ``from pkg import *`` but + *pkg* is not installed. """ caller_globals = sys._getframe(1).f_globals package_name: str = caller_globals["__name__"] caller_file: str = caller_globals["__file__"] - if packages is None: - __getattr__, __dir__, __all__ = lazy.attach_stub(package_name, caller_file) - mod = sys.modules[package_name] - setattr(mod, "__getattr__", __getattr__) - setattr(mod, "__dir__", __dir__) - setattr(mod, "__all__", __all__) - return __getattr__, __dir__, __all__ + if packages is not None: + warnings.warn( + "The 'packages' argument to lazy_export() is deprecated. " + "Add 'from import *' to your .pyi stub instead.", + DeprecationWarning, + stacklevel=2, + ) stub_file = f"{os.path.splitext(caller_file)[0]}.pyi" has_stub = os.path.exists(stub_file) - if has_stub: - filtered_stub = _filter_stub(stub_file) - if filtered_stub is not None: - stub_getattr, stub_dir, __all__ = lazy.attach_stub(package_name, filtered_stub) - os.unlink(filtered_stub) - else: - stub_getattr, stub_dir, __all__ = lazy.attach_stub(package_name, caller_file) + fallback_packages: list[str] = list(packages) if packages else [] + relative_wildcards: list[str] = [] - if not has_stub: + if has_stub: + filtered_path, stub_fallbacks, relative_wildcards = _parse_stub(stub_file) + if stub_fallbacks: + fallback_packages = list(dict.fromkeys(fallback_packages + stub_fallbacks)) + + stub_path = filtered_path if filtered_path is not None else caller_file + stub_getattr, stub_dir, __all__ = lazy.attach_stub(package_name, stub_path) + if filtered_path is not None: + os.unlink(filtered_path) + else: __all__: list[str] = [] - def _pkg_getattr(name: str): - for pkg in packages: - try: - mod = importlib.import_module(pkg) - if hasattr(mod, name): - val = getattr(mod, name) - sys.modules[package_name].__dict__[name] = val - return val - except (ImportError, ModuleNotFoundError): - continue - raise AttributeError(f"module {package_name!r} has no attribute {name!r}") + mod = sys.modules[package_name] - def _pkg_dir(): - names: list[str] = [] - for pkg in packages: + # -- Eagerly resolve relative wildcard imports (from .X import *) ------ + for rel_mod_name in relative_wildcards: + fq_name = f"{package_name}.{rel_mod_name}" + sub = importlib.import_module(fq_name) + exported = getattr(sub, "__all__", [n for n in dir(sub) if not n.startswith("_")]) + for name in exported: + mod.__dict__[name] = getattr(sub, name) + if name not in __all__: + __all__.append(name) + + # -- Build lazy fallback for absolute wildcard imports ----------------- + _sentinel = object() + + if fallback_packages: + _resolved_pkgs: list = [] + for _pkg_name in fallback_packages: try: - mod = importlib.import_module(pkg) - names.extend(n for n in dir(mod) if not n.startswith("_")) - except (ImportError, ModuleNotFoundError): - continue - return sorted(set(names)) + _resolved_pkgs.append(importlib.import_module(_pkg_name)) + except (ImportError, ModuleNotFoundError) as e: + raise ImportError( + f"lazy_export() in {package_name!r}: .pyi stub declares " + f"'from {_pkg_name} import *' but the package is not installed." + ) from e + + def _pkg_getattr(name: str): + for pkg_mod in _resolved_pkgs: + val = getattr(pkg_mod, name, _sentinel) + if val is not _sentinel: + mod.__dict__[name] = val + return val + raise AttributeError(f"module {package_name!r} has no attribute {name!r}") + + if has_stub: + _stub_getattr = stub_getattr + _stub_dir = stub_dir + _dir_cache: list[str] | None = None + + def __getattr__(name: str): + try: + return _stub_getattr(name) + except AttributeError: + return _pkg_getattr(name) + + def __dir__(): + nonlocal _dir_cache + if _dir_cache is None: + pkg_names = {n for p in _resolved_pkgs for n in dir(p) if not n.startswith("_")} + _dir_cache = sorted(set(_stub_dir()) | pkg_names) + return _dir_cache - if has_stub: + else: + _dir_cache: list[str] | None = None - def __getattr__(name: str): - try: - return stub_getattr(name) - except AttributeError: + def __getattr__(name: str): return _pkg_getattr(name) - def __dir__(): - return sorted(set(stub_dir()) | set(_pkg_dir())) + def __dir__(): + nonlocal _dir_cache + if _dir_cache is None: + _dir_cache = sorted(n for p in _resolved_pkgs for n in dir(p) if not n.startswith("_")) + return _dir_cache + elif has_stub: + __getattr__ = stub_getattr + __dir__ = stub_dir else: def __getattr__(name: str): - return _pkg_getattr(name) + raise AttributeError(f"module {package_name!r} has no attribute {name!r}") def __dir__(): - return _pkg_dir() + return [] - mod = sys.modules[package_name] setattr(mod, "__getattr__", __getattr__) setattr(mod, "__dir__", __dir__) setattr(mod, "__all__", __all__) diff --git a/source/isaaclab/isaaclab/visualizers/__init__.py b/source/isaaclab/isaaclab/visualizers/__init__.py index d2835fe21b61..c5bcae09f363 100644 --- a/source/isaaclab/isaaclab/visualizers/__init__.py +++ b/source/isaaclab/isaaclab/visualizers/__init__.py @@ -3,52 +3,12 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Visualizer Registry. +"""Visualizer base and factory entrypoints.""" -This module uses a registry pattern to decouple visualizer instantiation -from specific types. Configs can create visualizers via the -`create_visualizer()` factory method. -""" +from __future__ import annotations +from .base_visualizer import BaseVisualizer from .visualizer import Visualizer -from typing import Any +from .visualizer_cfg import VisualizerCfg -from isaaclab.utils.module import lazy_export - -lazy_export() - -_VISUALIZER_REGISTRY: dict[str, Any] = {} - - -def get_visualizer_class(name: str) -> type[Visualizer] | None: - """Get a visualizer class by name (lazy-loaded).""" - if name in _VISUALIZER_REGISTRY: - return _VISUALIZER_REGISTRY[name] - - try: - if name == "newton": - from .newton_visualizer import NewtonVisualizer - - _VISUALIZER_REGISTRY["newton"] = NewtonVisualizer - return NewtonVisualizer - if name == "viser": - from .viser_visualizer import ViserVisualizer - - _VISUALIZER_REGISTRY["viser"] = ViserVisualizer - return ViserVisualizer - if name == "kit": - from .kit_visualizer import KitVisualizer - - _VISUALIZER_REGISTRY["kit"] = KitVisualizer - return KitVisualizer - if name == "rerun": - from .rerun_visualizer import RerunVisualizer - - _VISUALIZER_REGISTRY["rerun"] = RerunVisualizer - return RerunVisualizer - return None - except ImportError as exc: - import warnings - - warnings.warn(f"Failed to load visualizer '{name}': {exc}", ImportWarning) - return None +__all__ = ["BaseVisualizer", "Visualizer", "VisualizerCfg"] diff --git a/source/isaaclab/isaaclab/visualizers/__init__.pyi b/source/isaaclab/isaaclab/visualizers/__init__.pyi index 7125b7f70373..0fcbfd637dc9 100644 --- a/source/isaaclab/isaaclab/visualizers/__init__.pyi +++ b/source/isaaclab/isaaclab/visualizers/__init__.pyi @@ -4,15 +4,11 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ - "KitVisualizerCfg", - "NewtonVisualizerCfg", - "RerunVisualizerCfg", - "ViserVisualizerCfg", + "BaseVisualizer", + "Visualizer", "VisualizerCfg", ] -from .kit_visualizer_cfg import KitVisualizerCfg -from .newton_visualizer_cfg import NewtonVisualizerCfg -from .rerun_visualizer_cfg import RerunVisualizerCfg -from .viser_visualizer_cfg import ViserVisualizerCfg +from .base_visualizer import BaseVisualizer +from .visualizer import Visualizer from .visualizer_cfg import VisualizerCfg diff --git a/source/isaaclab/isaaclab/visualizers/base_visualizer.py b/source/isaaclab/isaaclab/visualizers/base_visualizer.py new file mode 100644 index 000000000000..2480bc89bbaa --- /dev/null +++ b/source/isaaclab/isaaclab/visualizers/base_visualizer.py @@ -0,0 +1,322 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Base class for visualizers.""" + +from __future__ import annotations + +import logging +import random +import re +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from isaaclab.physics import SceneDataProvider + + from .visualizer_cfg import VisualizerCfg + + +logger = logging.getLogger(__name__) + + +class BaseVisualizer(ABC): + """Base class for all visualizer backends. + + Lifecycle: __init__() -> initialize() -> step() (repeated) -> close() + """ + + def __init__(self, cfg: VisualizerCfg): + """Initialize visualizer with config. + + Args: + cfg: Visualizer configuration. + """ + self.cfg = cfg + self._scene_data_provider = None + self._is_initialized = False + self._is_closed = False + + @abstractmethod + def initialize(self, scene_data_provider: SceneDataProvider) -> None: + """Initialize visualizer resources. + + Args: + scene_data_provider: Scene data provider used by the visualizer. + """ + raise NotImplementedError + + @abstractmethod + def step(self, dt: float) -> None: + """Update visualization for one step. + + Args: + dt: Time step in seconds. + """ + raise NotImplementedError + + @abstractmethod + def close(self) -> None: + """Clean up resources.""" + raise NotImplementedError + + @abstractmethod + def is_running(self) -> bool: + """Check if visualizer is still running (e.g., window not closed). + + Returns: + ``True`` if the visualizer is running, otherwise ``False``. + """ + raise NotImplementedError + + def is_training_paused(self) -> bool: + """Check if training is paused by visualizer controls. + + Returns: + ``True`` if training is paused, otherwise ``False``. + """ + return False + + def is_rendering_paused(self) -> bool: + """Check if rendering is paused by visualizer controls. + + Returns: + ``True`` if rendering is paused, otherwise ``False``. + """ + return False + + @property + def is_initialized(self) -> bool: + """Check if initialize() has been called.""" + return self._is_initialized + + @property + def is_closed(self) -> bool: + """Check if close() has been called.""" + return self._is_closed + + def supports_markers(self) -> bool: + """Check if visualizer supports VisualizationMarkers. + + Returns: + ``True`` if marker rendering is supported, otherwise ``False``. + """ + return False + + def supports_live_plots(self) -> bool: + """Check if visualizer supports LivePlots. + + Returns: + ``True`` if live plots are supported, otherwise ``False``. + """ + return False + + def requires_forward_before_step(self) -> bool: + """Whether simulation should run forward() before step(). + + Returns: + ``True`` when forward kinematics should run before stepping. + """ + return False + + def pumps_app_update(self) -> bool: + """Whether this visualizer calls omni.kit.app.get_app().update() in step(). + + Returns True for visualizers (e.g. KitVisualizer) that already pump the Kit + app loop, so SimulationContext.render() can skip its own app.update() call + and avoid double-rendering. + """ + return False + + def get_visualized_env_ids(self) -> list[int] | None: + """Return env IDs this visualizer is displaying, if any. + + Returns: + Visualized environment ids, or ``None`` for all environments. + """ + return getattr(self, "_env_ids", None) + + def _compute_visualized_env_ids(self) -> list[int] | None: + """Compute which environment indices to visualize from config. + + Returns: + Selected environment ids, or ``None`` to visualize all environments. + """ + if self._scene_data_provider is None: + return None + filter_mode = getattr(self.cfg, "env_filter_mode", "none") + if filter_mode == "none": + return None + + num_envs = self._scene_data_provider.get_metadata().get("num_envs", 0) + if num_envs <= 0: + logger.debug("[Visualizer] num_envs is 0 or missing from provider metadata; env filtering disabled.") + return None + if filter_mode == "env_ids": + env_ids_cfg = getattr(self.cfg, "env_filter_ids", None) + if env_ids_cfg is not None and len(env_ids_cfg) > 0: + return [i for i in env_ids_cfg if 0 <= i < num_envs] + return None + if filter_mode == "random_n": + count = int(getattr(self.cfg, "env_filter_random_n", 0)) + if count <= 0: + return None + count = min(count, num_envs) + seed = int(getattr(self.cfg, "env_filter_seed", 0)) + rng = random.Random(seed) + return sorted(rng.sample(range(num_envs), count)) + logger.warning("[Visualizer] Unknown env_filter_mode='%s'; defaulting to all envs.", filter_mode) + return None + + def get_rendering_dt(self) -> float | None: + """Get rendering time step. + + Returns: + Rendering time step override, or ``None`` to use interface default. + """ + return None + + def set_camera_view(self, eye: tuple, target: tuple) -> None: + """Set camera view position. + + Args: + eye: Camera eye position. + target: Camera target position. + """ + pass + + def _resolve_camera_pose_from_usd_path( + self, usd_path: str + ) -> tuple[tuple[float, float, float], tuple[float, float, float]] | None: + """Resolve camera pose/target from provider camera transforms. + + Args: + usd_path: Concrete USD camera path. + + Returns: + Eye/target tuple when available, otherwise ``None``. + """ + if self._scene_data_provider is None: + return None + transforms = self._scene_data_provider.get_camera_transforms() + if not transforms: + return None + + env_id, template_path = self._resolve_template_camera_path(usd_path) + camera_transform = self._lookup_camera_transform(transforms, template_path, env_id) + if camera_transform is None: + return None + pos, ori = camera_transform + + pos_t = (float(pos[0]), float(pos[1]), float(pos[2])) + ori_t = (float(ori[0]), float(ori[1]), float(ori[2]), float(ori[3])) + forward = self._quat_rotate_vec(ori_t, (0.0, 0.0, -1.0)) + target = (pos_t[0] + forward[0], pos_t[1] + forward[1], pos_t[2] + forward[2]) + return pos_t, target + + def _resolve_template_camera_path(self, usd_path: str) -> tuple[int, str]: + """Normalize concrete env camera path to templated camera path. + + Args: + usd_path: Concrete USD camera path. + + Returns: + Tuple of environment id and templated camera path. + """ + env_pattern = re.compile(r"(?P/World/envs/env_)(?P\d+)(?P/.*)") + if match := env_pattern.match(usd_path): + return int(match.group("id")), match.group("root") + "%d" + match.group("path") + return 0, usd_path + + def _lookup_camera_transform( + self, transforms: dict[str, Any], template_path: str, env_id: int + ) -> tuple[list[float], list[float]] | None: + """Fetch camera position/orientation for a templated path and environment. + + Args: + transforms: Camera transform dictionary from provider. + template_path: Templated camera path. + env_id: Environment id to query. + + Returns: + Position/orientation tuple when available, otherwise ``None``. + """ + order = transforms.get("order", []) + positions = transforms.get("positions", []) + orientations = transforms.get("orientations", []) + + if template_path not in order: + return None + idx = order.index(template_path) + if idx >= len(positions) or idx >= len(orientations): + return None + if env_id < 0 or env_id >= len(positions[idx]): + return None + pos = positions[idx][env_id] + ori = orientations[idx][env_id] + if pos is None or ori is None: + return None + return pos, ori + + @staticmethod + def _quat_rotate_vec( + quat_xyzw: tuple[float, float, float, float], vec: tuple[float, float, float] + ) -> tuple[float, float, float]: + """Rotate a vector by a quaternion. + + Args: + quat_xyzw: Quaternion in xyzw order. + vec: Input vector. + + Returns: + Rotated vector. + """ + import torch + + from isaaclab.utils.math import quat_apply + + quat = torch.tensor(quat_xyzw, dtype=torch.float32).unsqueeze(0) + vector = torch.tensor(vec, dtype=torch.float32).unsqueeze(0) + rotated = quat_apply(quat, vector)[0] + return (float(rotated[0]), float(rotated[1]), float(rotated[2])) + + def reset(self, soft: bool = False) -> None: + """Reset visualizer state. + + Args: + soft: Whether to perform a soft reset. + """ + pass + + def _log_initialization_table(self, logger: logging.Logger, title: str, rows: list[tuple[str, Any]]) -> None: + """Log a compact initialization table for a visualizer. + + Args: + logger: Logger used to emit the table. + title: Table title. + rows: Table row key/value pairs. + """ + from prettytable import PrettyTable + + table = PrettyTable() + table.title = title + table.field_names = ["Field", "Value"] + table.align["Field"] = "l" + table.align["Value"] = "l" + for key, value in rows: + table.add_row([key, value]) + logger.info("Visualizer initialization:\n%s", table.get_string()) + + def play(self) -> None: + """Handle simulation play/start. No-op by default.""" + pass + + def pause(self) -> None: + """Handle simulation pause. No-op by default.""" + pass + + def stop(self) -> None: + """Handle simulation stop. No-op by default.""" + pass diff --git a/source/isaaclab/isaaclab/visualizers/rerun_visualizer.py b/source/isaaclab/isaaclab/visualizers/rerun_visualizer.py deleted file mode 100644 index 627cca3744d5..000000000000 --- a/source/isaaclab/isaaclab/visualizers/rerun_visualizer.py +++ /dev/null @@ -1,266 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Rerun-based visualizer using rerun-sdk.""" - -from __future__ import annotations - -import contextlib -import logging -import os -from typing import TYPE_CHECKING - -import rerun as rr -import rerun.blueprint as rrb -from newton.viewer import ViewerRerun - -from .rerun_visualizer_cfg import RerunVisualizerCfg -from .visualizer import Visualizer - -logger = logging.getLogger(__name__) - -if TYPE_CHECKING: - from isaaclab.sim.scene_data_providers import SceneDataProvider - - -class NewtonViewerRerun(ViewerRerun): - """Isaac Lab wrapper for Newton's ViewerRerun.""" - - def __init__( - self, - app_id: str | None = None, - address: str | None = None, - web_port: int | None = None, - grpc_port: int | None = None, - keep_historical_data: bool = False, - keep_scalar_history: bool = True, - record_to_rrd: str | None = None, - metadata: dict | None = None, - ): - super().__init__( - app_id=app_id, - address=address, - web_port=web_port, - grpc_port=grpc_port or 9876, - serve_web_viewer=True, - keep_historical_data=keep_historical_data, - keep_scalar_history=keep_scalar_history, - record_to_rrd=record_to_rrd, - ) - - self._metadata = metadata or {} - self._log_metadata() - - def _log_metadata(self) -> None: - metadata_text = "# Isaac Lab Scene Metadata\n\n" - physics_backend = self._metadata.get("physics_backend", "unknown") - metadata_text += f"**Physics Backend:** {physics_backend}\n" - num_envs = self._metadata.get("num_envs", 0) - metadata_text += f"**Total Environments:** {num_envs}\n" - - for key, value in self._metadata.items(): - if key not in ["physics_backend", "num_envs"]: - metadata_text += f"**{key}:** {value}\n" - - rr.log("metadata", rr.TextDocument(metadata_text, media_type=rr.MediaType.MARKDOWN)) - - -class RerunVisualizer(Visualizer): - """Rerun web-based visualizer with time scrubbing and data inspection.""" - - def __init__(self, cfg: RerunVisualizerCfg): - super().__init__(cfg) - self.cfg: RerunVisualizerCfg = cfg - self._viewer: NewtonViewerRerun | None = None - self._model = None - self._state = None - self._is_initialized = False - self._sim_time = 0.0 - self._scene_data_provider = None - self._rerun_server_process = None - self._rerun_address: str | None = None - self._active_record_path: str | None = None - self._last_camera_pose: tuple[tuple[float, float, float], tuple[float, float, float]] | None = None - - def initialize(self, scene_data_provider: SceneDataProvider) -> None: - if self._is_initialized: - logger.debug("[RerunVisualizer] initialize() called while already initialized.") - return - if scene_data_provider is None: - raise RuntimeError("Rerun visualizer requires a scene_data_provider.") - - self._scene_data_provider = scene_data_provider - metadata = scene_data_provider.get_metadata() - self._env_ids = self._compute_visualized_env_ids() - if self._env_ids: - get_filtered_model = getattr(scene_data_provider, "get_newton_model_for_env_ids", None) - if callable(get_filtered_model): - self._model = get_filtered_model(self._env_ids) - else: - self._model = scene_data_provider.get_newton_model() - else: - self._model = scene_data_provider.get_newton_model() - self._state = scene_data_provider.get_newton_state(self._env_ids) - - try: - self._setup_rerun_server() - self._active_record_path = self.cfg.record_to_rrd - self._create_viewer(record_to_rrd=self.cfg.record_to_rrd, metadata=metadata) - logger.info( - "[RerunVisualizer] initialized | camera_pos=%s camera_target=%s", - self.cfg.camera_position, - self.cfg.camera_target, - ) - self._is_initialized = True - except Exception as exc: - logger.error(f"[RerunVisualizer] Failed to initialize viewer: {exc}") - raise - - def step(self, dt: float) -> None: - if not self._is_initialized or self._viewer is None or self._scene_data_provider is None: - return - - if self.cfg.camera_source == "usd_path": - self._update_camera_from_usd_path() - - self._state = self._scene_data_provider.get_newton_state(self._env_ids) - self._sim_time += dt - - self._viewer.begin_frame(self._sim_time) - self._viewer.log_state(self._state) - self._viewer.end_frame() - - def close(self) -> None: - if not self._is_initialized: - return - - try: - self._close_viewer(finalize_rrd=bool(self._active_record_path)) - except Exception as exc: - logger.warning(f"[RerunVisualizer] Error during close: {exc}") - - self._viewer = None - self._is_initialized = False - self._is_closed = True - self._active_record_path = None - if self._rerun_server_process is not None: - with contextlib.suppress(Exception): - self._rerun_server_process.terminate() - self._rerun_server_process = None - - def is_running(self) -> bool: - if self._viewer is None: - return False - return self._viewer.is_running() - - def is_training_paused(self) -> bool: - return False - - def supports_markers(self) -> bool: - return False - - def supports_live_plots(self) -> bool: - return False - - def _setup_rerun_server(self) -> None: - if not self.cfg.bind_address or self._rerun_server_process is not None: - return - import shutil - import subprocess - - rerun_bin = shutil.which("rerun") - if rerun_bin is None: - logger.warning("[RerunVisualizer] 'rerun' binary not found in PATH. Skipping external bind.") - return - - cmd = [ - rerun_bin, - "--serve-web", - "--bind", - self.cfg.bind_address, - "--port", - str(self.cfg.grpc_port), - "--web-viewer-port", - str(self.cfg.web_port), - ] - if self.cfg.open_browser: - cmd.append("--web-viewer") - self._rerun_server_process = subprocess.Popen(cmd) - logger.info( - "[RerunVisualizer] Server bind %s:%s, web %s", - self.cfg.bind_address, - self.cfg.grpc_port, - self.cfg.web_port, - ) - self._rerun_address = f"rerun+http://127.0.0.1:{self.cfg.grpc_port}/proxy" - - def _create_viewer(self, record_to_rrd: str | None, metadata: dict | None = None, reset_time: bool = True) -> None: - self._viewer = NewtonViewerRerun( - app_id=self.cfg.app_id, - address=self._rerun_address, - web_port=self.cfg.web_port, - grpc_port=self.cfg.grpc_port, - keep_historical_data=self.cfg.keep_historical_data, - keep_scalar_history=self.cfg.keep_scalar_history, - record_to_rrd=record_to_rrd, - metadata=metadata or {}, - ) - max_worlds = self.cfg.max_worlds - self._viewer.set_model(self._model, max_worlds=None if max_worlds in (None, 0) else max_worlds) - self._set_rerun_camera_view(self._resolve_initial_camera_pose()) - - if reset_time: - self._sim_time = 0.0 - - def _close_viewer(self, finalize_rrd: bool = False, record_path: str | None = None) -> None: - if self._viewer is None: - return - self._viewer.close() - finalized_path = record_path if record_path is not None else self._active_record_path - if finalize_rrd and finalized_path: - if os.path.exists(finalized_path): - size = os.path.getsize(finalized_path) - logger.info("[RerunVisualizer] Recording saved: %s (%s bytes)", finalized_path, size) - else: - logger.warning("[RerunVisualizer] Recording file not found: %s", finalized_path) - self._viewer = None - - def _resolve_initial_camera_pose(self) -> tuple[tuple[float, float, float], tuple[float, float, float]]: - if self.cfg.camera_source == "usd_path": - pose = self._resolve_camera_pose_from_usd_path(self.cfg.camera_usd_path) - if pose is not None: - return pose - logger.warning( - "[RerunVisualizer] camera_usd_path '%s' not found; using configured camera.", - self.cfg.camera_usd_path, - ) - return self.cfg.camera_position, self.cfg.camera_target - - def _set_rerun_camera_view(self, pose: tuple[tuple[float, float, float], tuple[float, float, float]]) -> None: - cam_pos, cam_target = pose - try: - blueprint = rrb.Blueprint( - rrb.Spatial3DView( - name="3D View", - origin="/", - eye_controls=rrb.EyeControls3D( - position=cam_pos, - look_target=cam_target, - ), - ), - collapse_panels=True, - ) - rr.send_blueprint(blueprint) - self._last_camera_pose = (cam_pos, cam_target) - except Exception as exc: - logger.warning(f"[RerunVisualizer] Could not set camera view: {exc}") - - def _update_camera_from_usd_path(self) -> None: - pose = self._resolve_camera_pose_from_usd_path(self.cfg.camera_usd_path) - if pose is None: - return - if self._last_camera_pose == pose: - return - self._set_rerun_camera_view(pose) diff --git a/source/isaaclab/isaaclab/visualizers/visualizer.py b/source/isaaclab/isaaclab/visualizers/visualizer.py index 87d74974f89d..b84617ed870e 100644 --- a/source/isaaclab/isaaclab/visualizers/visualizer.py +++ b/source/isaaclab/isaaclab/visualizers/visualizer.py @@ -3,216 +3,79 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Base class for visualizers.""" +"""Factory for creating visualizer instances.""" from __future__ import annotations -import logging -import random -import re -from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any +from isaaclab.utils.backend_utils import FactoryBase -if TYPE_CHECKING: - from isaaclab.sim.scene_data_providers import SceneDataProvider +from .base_visualizer import BaseVisualizer - from .visualizer_cfg import VisualizerCfg +# Visualizer types; each loads from isaaclab_visualizers. for minimal deps. +_VISUALIZER_TYPES = ("kit", "newton", "rerun", "viser") -logger = logging.getLogger(__name__) +class Visualizer(FactoryBase, BaseVisualizer): + """Factory for creating visualizer instances.""" + _backend_class_names = { + "kit": "KitVisualizer", + "newton": "NewtonVisualizer", + "rerun": "RerunVisualizer", + "viser": "ViserVisualizer", + } -class Visualizer(ABC): - """Base class for all visualizer backends. + @classmethod + def _get_backend(cls, cfg, *args, **kwargs) -> str: + """Resolve backend key from visualizer config. - Lifecycle: __init__() -> initialize() -> step() (repeated) -> close() - """ - - def __init__(self, cfg: VisualizerCfg): - """Initialize visualizer with config.""" - self.cfg = cfg - self._scene_data_provider = None - self._is_initialized = False - self._is_closed = False + Args: + cfg: Visualizer configuration instance. + *args: Unused positional arguments. + **kwargs: Unused keyword arguments. - @abstractmethod - def initialize(self, scene_data_provider: SceneDataProvider) -> None: - """Initialize visualizer resources.""" - raise NotImplementedError + Returns: + Backend key used by the factory. - @abstractmethod - def step(self, dt: float) -> None: - """Update visualization for one step. + Raises: + ValueError: If visualizer type is not registered. + """ + visualizer_type = getattr(cfg, "visualizer_type", None) + if visualizer_type not in _VISUALIZER_TYPES: + raise ValueError( + f"Visualizer type '{visualizer_type}' is not registered. Valid types: " + f"{', '.join(repr(k) for k in _VISUALIZER_TYPES)}." + ) + return visualizer_type + + @classmethod + def _get_module_name(cls, backend: str) -> str: + """Return module path for a visualizer backend. Args: - dt: Time step in seconds. + backend: Backend key. + + Returns: + Module import path for the backend. """ - raise NotImplementedError - - @abstractmethod - def close(self) -> None: - """Clean up resources.""" - raise NotImplementedError - - @abstractmethod - def is_running(self) -> bool: - """Check if visualizer is still running (e.g., window not closed).""" - raise NotImplementedError - - def is_training_paused(self) -> bool: - """Check if training is paused by visualizer controls.""" - return False - - def is_rendering_paused(self) -> bool: - """Check if rendering is paused by visualizer controls.""" - return False - - @property - def is_initialized(self) -> bool: - """Check if initialize() has been called.""" - return self._is_initialized - - @property - def is_closed(self) -> bool: - """Check if close() has been called.""" - return self._is_closed - - def supports_markers(self) -> bool: - """Check if visualizer supports VisualizationMarkers.""" - return False - - def supports_live_plots(self) -> bool: - """Check if visualizer supports LivePlots.""" - return False - - def requires_forward_before_step(self) -> bool: - """Whether simulation should run forward() before step().""" - return False - - def pumps_app_update(self) -> bool: - """Whether this visualizer calls omni.kit.app.get_app().update() in step(). - - Returns True for visualizers (e.g. KitVisualizer) that already pump the Kit - app loop, so SimulationContext.render() can skip its own app.update() call - and avoid double-rendering. + return f"isaaclab_visualizers.{backend}" + + def __new__(cls, cfg, *args, **kwargs) -> BaseVisualizer: + """Create a new visualizer instance based on the visualizer type. + + Args: + cfg: Visualizer configuration instance. + *args: Additional constructor positional arguments. + **kwargs: Additional constructor keyword arguments. + + Returns: + Instantiated backend visualizer. + + Raises: + TypeError: If backend class does not inherit from ``BaseVisualizer``. """ - return False - - def get_visualized_env_ids(self) -> list[int] | None: - """Return env IDs this visualizer is displaying, if any.""" - return getattr(self, "_env_ids", None) - - def _compute_visualized_env_ids(self) -> list[int] | None: - """Compute which env indices to show from config.""" - if self._scene_data_provider is None: - return None - filter_mode = getattr(self.cfg, "env_filter_mode", "none") - if filter_mode == "none": - return None - - num_envs = self._scene_data_provider.get_metadata().get("num_envs", 0) - if num_envs <= 0: - logger.debug("[Visualizer] num_envs is 0 or missing from provider metadata; env filtering disabled.") - return None - if filter_mode == "env_ids": - env_ids_cfg = getattr(self.cfg, "env_filter_ids", None) - if env_ids_cfg is not None and len(env_ids_cfg) > 0: - return [i for i in env_ids_cfg if 0 <= i < num_envs] - return None - if filter_mode == "random_n": - count = int(getattr(self.cfg, "env_filter_random_n", 0)) - if count <= 0: - return None - count = min(count, num_envs) - seed = int(getattr(self.cfg, "env_filter_seed", 0)) - rng = random.Random(seed) - return sorted(rng.sample(range(num_envs), count)) - logger.warning("[Visualizer] Unknown env_filter_mode='%s'; defaulting to all envs.", filter_mode) - return None - - def get_rendering_dt(self) -> float | None: - """Get rendering time step. Returns None to use interface default.""" - return None - - def set_camera_view(self, eye: tuple, target: tuple) -> None: - """Set camera view position. No-op by default.""" - pass - - def _resolve_camera_pose_from_usd_path( - self, usd_path: str - ) -> tuple[tuple[float, float, float], tuple[float, float, float]] | None: - """Resolve camera pose/target from provider camera transforms.""" - if self._scene_data_provider is None: - return None - transforms = self._scene_data_provider.get_camera_transforms() - if not transforms: - return None - - env_id, template_path = self._resolve_template_camera_path(usd_path) - camera_transform = self._lookup_camera_transform(transforms, template_path, env_id) - if camera_transform is None: - return None - pos, ori = camera_transform - - pos_t = (float(pos[0]), float(pos[1]), float(pos[2])) - ori_t = (float(ori[0]), float(ori[1]), float(ori[2]), float(ori[3])) - forward = self._quat_rotate_vec(ori_t, (0.0, 0.0, -1.0)) - target = (pos_t[0] + forward[0], pos_t[1] + forward[1], pos_t[2] + forward[2]) - return pos_t, target - - def _resolve_template_camera_path(self, usd_path: str) -> tuple[int, str]: - """Normalize concrete env camera path to templated camera path.""" - env_pattern = re.compile(r"(?P/World/envs/env_)(?P\d+)(?P/.*)") - if match := env_pattern.match(usd_path): - return int(match.group("id")), match.group("root") + "%d" + match.group("path") - return 0, usd_path - - def _lookup_camera_transform( - self, transforms: dict[str, Any], template_path: str, env_id: int - ) -> tuple[list[float], list[float]] | None: - """Fetch (position, orientation) for a templated camera path and env index.""" - order = transforms.get("order", []) - positions = transforms.get("positions", []) - orientations = transforms.get("orientations", []) - - if template_path not in order: - return None - idx = order.index(template_path) - if idx >= len(positions) or idx >= len(orientations): - return None - if env_id < 0 or env_id >= len(positions[idx]): - return None - pos = positions[idx][env_id] - ori = orientations[idx][env_id] - if pos is None or ori is None: - return None - return pos, ori - - @staticmethod - def _quat_rotate_vec( - quat_xyzw: tuple[float, float, float, float], vec: tuple[float, float, float] - ) -> tuple[float, float, float]: - import torch - - from isaaclab.utils.math import quat_apply - - quat = torch.tensor(quat_xyzw, dtype=torch.float32).unsqueeze(0) - vector = torch.tensor(vec, dtype=torch.float32).unsqueeze(0) - rotated = quat_apply(quat, vector)[0] - return (float(rotated[0]), float(rotated[1]), float(rotated[2])) - - def reset(self, soft: bool = False) -> None: - """Reset visualizer state. No-op by default.""" - pass - - def play(self) -> None: - """Handle simulation play/start. No-op by default.""" - pass - - def pause(self) -> None: - """Handle simulation pause. No-op by default.""" - pass - - def stop(self) -> None: - """Handle simulation stop. No-op by default.""" - pass + result = super().__new__(cls, cfg, *args, **kwargs) + if not isinstance(result, BaseVisualizer): + name = type(result).__name__ + raise TypeError(f"Backend visualizer {name!r} must inherit from BaseVisualizer.") + return result diff --git a/source/isaaclab/isaaclab/visualizers/visualizer_cfg.py b/source/isaaclab/isaaclab/visualizers/visualizer_cfg.py index 64e079129618..a96e3c04d2b5 100644 --- a/source/isaaclab/isaaclab/visualizers/visualizer_cfg.py +++ b/source/isaaclab/isaaclab/visualizers/visualizer_cfg.py @@ -12,7 +12,7 @@ from isaaclab.utils import configclass if TYPE_CHECKING: - from .visualizer import Visualizer + from .base_visualizer import BaseVisualizer @configclass @@ -21,8 +21,8 @@ class VisualizerCfg: Note: This is an abstract base class and should not be instantiated directly. - Use specific visualizer configs like NewtonVisualizerCfg, RerunVisualizerCfg, ViserVisualizerCfg, - or KitVisualizerCfg. + Use specific configs from isaaclab_visualizers: KitVisualizerCfg, NewtonVisualizerCfg, + RerunVisualizerCfg, or ViserVisualizerCfg (from isaaclab_visualizers.kit/.newton/.rerun/.viser). """ visualizer_type: str | None = None @@ -70,31 +70,30 @@ def get_visualizer_type(self) -> str | None: """ return self.visualizer_type - def create_visualizer(self) -> Visualizer: + def create_visualizer(self) -> BaseVisualizer: """Create visualizer instance from this config using factory pattern. + Loads the matching backend from isaaclab_visualizers (e.g. isaaclab_visualizers.rerun). + Raises: ValueError: If visualizer_type is None (base class used directly) or not registered. + ImportError: If isaaclab_visualizers or the requested backend extra is not installed. """ - from . import get_visualizer_class + from .visualizer import Visualizer if self.visualizer_type is None: raise ValueError( "Cannot create visualizer from base VisualizerCfg class. " - "Use a specific visualizer config: NewtonVisualizerCfg, RerunVisualizerCfg, " - "ViserVisualizerCfg, or KitVisualizerCfg." + "Use a specific config from isaaclab_visualizers " + "(e.g. KitVisualizerCfg, NewtonVisualizerCfg, RerunVisualizerCfg, ViserVisualizerCfg)." ) - visualizer_class = get_visualizer_class(self.visualizer_type) - if visualizer_class is None: - if self.visualizer_type in ("newton", "rerun", "viser"): + try: + return Visualizer(self) + except (ValueError, ImportError, ModuleNotFoundError) as exc: + if self.visualizer_type in ("newton", "rerun", "viser", "kit"): raise ImportError( - f"Visualizer '{self.visualizer_type}' requires the Newton Python module and its dependencies. " - "Install the Newton backend (e.g., newton package/isaaclab_newton) and retry." - ) - raise ValueError( - f"Visualizer type '{self.visualizer_type}' is not registered. " - "Valid types: 'newton', 'rerun', 'viser', 'kit'." - ) - - return visualizer_class(self) + f"Visualizer '{self.visualizer_type}' requires the isaaclab_visualizers package. " + f"Install with: pip install isaaclab_visualizers[{self.visualizer_type}]" + ) from exc + raise diff --git a/source/isaaclab/setup.py b/source/isaaclab/setup.py index aa4944d393cb..2e9cdfd721f5 100644 --- a/source/isaaclab/setup.py +++ b/source/isaaclab/setup.py @@ -47,14 +47,10 @@ "pytest-mock", "junitparser", "coverage==7.6.1", + "debugpy>=1.8.20", "flatdict==4.0.0", "flaky", "packaging", - # visualizers - "newton==1.0.0rc1", - "imgui-bundle>=1.92.5", - "rerun-sdk>=0.29.0", - "viser>=1.0.16", # Required by pydantic-core/imgui_bundle on Python 3.12 (Sentinel symbol). "typing_extensions>=4.14.0", "lazy_loader>=0.4", @@ -90,6 +86,11 @@ "rl": ["isaaclab_rl"], "tasks": ["isaaclab_tasks"], "teleop": ["isaaclab_teleop"], + "visualizers": ["isaaclab_visualizers[all]"], + "visualizers-kit": ["isaaclab_visualizers[kit]"], + "visualizers-newton": ["isaaclab_visualizers[newton]"], + "visualizers-rerun": ["isaaclab_visualizers[rerun]"], + "visualizers-viser": ["isaaclab_visualizers[viser]"], # Convenience: all sub-packages (does not include isaacsim) "all": [ "isaaclab_assets", @@ -100,6 +101,7 @@ "isaaclab_rl", "isaaclab_tasks", "isaaclab_teleop", + "isaaclab_visualizers[all]", ], } diff --git a/source/isaaclab/test/app/test_kwarg_launch.py b/source/isaaclab/test/app/test_kwarg_launch.py index df687b65d1dc..faaf5175c772 100644 --- a/source/isaaclab/test/app/test_kwarg_launch.py +++ b/source/isaaclab/test/app/test_kwarg_launch.py @@ -3,6 +3,8 @@ # # SPDX-License-Identifier: BSD-3-Clause +import argparse + import pytest import isaaclab.app.app_launcher as app_launcher_module @@ -44,7 +46,7 @@ def test_set_visualizer_settings_stores_values(monkeypatch: pytest.MonkeyPatch): launcher._set_visualizer_settings({"visualizer": ["viser", "rerun"], "visualizer_max_worlds": 0}) assert settings.values == { - "/isaaclab/visualizer": "viser rerun", + "/isaaclab/visualizer/types": "viser rerun", "/isaaclab/visualizer/max_worlds": 0, } @@ -68,3 +70,25 @@ def _raise_settings_error(): launcher = AppLauncher.__new__(AppLauncher) launcher._set_visualizer_settings({"visualizer": ["viser"], "visualizer_max_worlds": 3}) + + +def test_parse_visualizer_csv_accepts_comma_delimited_values(): + parsed = app_launcher_module._parse_visualizer_csv("kit,newton,rerun,viser") + assert parsed == ["kit", "newton", "rerun", "viser"] + + +def test_parse_visualizer_csv_rejects_spaces_between_entries(): + with pytest.raises(argparse.ArgumentTypeError, match="spaces are not allowed"): + app_launcher_module._parse_visualizer_csv("kit, newton") + + +def test_visualizer_csv_does_not_swallow_hydra_overrides(): + parser = argparse.ArgumentParser(add_help=False) + app_launcher_module.AppLauncher.add_app_launcher_args(parser) + + args, hydra_args = parser.parse_known_args( + ["--visualizer", "kit,newton,rerun", "presets=newton", "env.episode_length=10"] + ) + + assert args.visualizer == ["kit", "newton", "rerun"] + assert hydra_args == ["presets=newton", "env.episode_length=10"] diff --git a/source/isaaclab/test/assets/test_articulation_iface.py b/source/isaaclab/test/assets/test_articulation_iface.py index 5753e2cd75dd..64842ed7273f 100644 --- a/source/isaaclab/test/assets/test_articulation_iface.py +++ b/source/isaaclab/test/assets/test_articulation_iface.py @@ -57,6 +57,15 @@ except ImportError: pass +try: + from isaaclab_newton.assets.articulation.articulation import Articulation as NewtonArticulation + from isaaclab_newton.assets.articulation.articulation_data import ArticulationData as NewtonArticulationData + from isaaclab_newton.test.mock_interfaces.views import MockNewtonArticulationView as NewtonMockArticulationView + + BACKENDS.append("newton") +except ImportError: + pass + def create_physx_articulation( num_instances: int = 2, @@ -166,7 +175,121 @@ def create_newton_articulation( num_bodies: int = 7, device: str = "cuda:0", ): - raise NotImplementedError("Newton articulation is not supported yet") + """Create a test Newton Articulation instance with mocked dependencies.""" + import isaaclab_newton.assets.articulation.articulation_data as newton_data_module + + joint_names = [f"joint_{i}" for i in range(num_joints)] + body_names = [f"body_{i}" for i in range(num_bodies)] + + # Create Newton mock view + mock_view = NewtonMockArticulationView( + num_instances=num_instances, + num_bodies=num_bodies, + num_joints=num_joints, + device=device, + is_fixed_base=False, + joint_names=joint_names, + body_names=body_names, + ) + mock_view.set_random_mock_data() + mock_view._noop_setters = True + + # Mock NewtonManager (aliased as SimulationManager in Newton modules) + mock_model = MagicMock() + mock_model.gravity = wp.array(np.array([[0.0, 0.0, -9.81]], dtype=np.float32), dtype=wp.vec3f, device=device) + mock_state = MagicMock() + mock_control = MagicMock() + + mock_manager = MagicMock() + mock_manager.get_model.return_value = mock_model + mock_manager.get_state_0.return_value = mock_state + mock_manager.get_state_1.return_value = mock_state + mock_manager.get_control.return_value = mock_control + + # Patch SimulationManager in the Newton data module + original_sim_manager = newton_data_module.SimulationManager + newton_data_module.SimulationManager = mock_manager + + try: + data = NewtonArticulationData(mock_view, device) + finally: + newton_data_module.SimulationManager = original_sim_manager + + # Create Articulation shell (bypass __init__) + articulation = object.__new__(NewtonArticulation) + + articulation.cfg = ArticulationCfg( + prim_path="/World/Robot", + soft_joint_pos_limit_factor=1.0, + actuators={}, + ) + + object.__setattr__(articulation, "_root_view", mock_view) + object.__setattr__(articulation, "_device", device) + object.__setattr__(articulation, "_data", data) + + # Tendon names (Newton doesn't support tendons) + object.__setattr__(articulation, "_fixed_tendon_names", []) + object.__setattr__(articulation, "_spatial_tendon_names", []) + data.fixed_tendon_names = [] + data.spatial_tendon_names = [] + + # Mock wrench composers + mock_inst_wrench = MockWrenchComposer(articulation) + mock_perm_wrench = MockWrenchComposer(articulation) + object.__setattr__(articulation, "_instantaneous_wrench_composer", mock_inst_wrench) + object.__setattr__(articulation, "_permanent_wrench_composer", mock_perm_wrench) + + # Prevent __del__ / _clear_callbacks from raising AttributeError + object.__setattr__(articulation, "_initialize_handle", None) + object.__setattr__(articulation, "_invalidate_initialize_handle", None) + object.__setattr__(articulation, "_prim_deletion_handle", None) + object.__setattr__(articulation, "_debug_vis_handle", None) + + # Other required attributes + object.__setattr__(articulation, "actuators", {}) + object.__setattr__(articulation, "_has_implicit_actuators", False) + + # Newton uses wp.array for indices (not torch) + object.__setattr__(articulation, "_ALL_INDICES", wp.array(np.arange(num_instances, dtype=np.int32), device=device)) + object.__setattr__( + articulation, "_ALL_BODY_INDICES", wp.array(np.arange(num_bodies, dtype=np.int32), device=device) + ) + object.__setattr__( + articulation, "_ALL_JOINT_INDICES", wp.array(np.arange(num_joints, dtype=np.int32), device=device) + ) + + # Newton uses wp.bool masks + object.__setattr__(articulation, "_ALL_ENV_MASK", wp.ones((num_instances,), dtype=wp.bool, device=device)) + object.__setattr__(articulation, "_ALL_BODY_MASK", wp.ones((num_bodies,), dtype=wp.bool, device=device)) + object.__setattr__(articulation, "_ALL_JOINT_MASK", wp.ones((num_joints,), dtype=wp.bool, device=device)) + + # Tendon arrays (empty) + object.__setattr__(articulation, "_ALL_FIXED_TENDON_INDICES", wp.array(np.array([], dtype=np.int32), device=device)) + object.__setattr__(articulation, "_ALL_FIXED_TENDON_MASK", wp.ones((0,), dtype=wp.bool, device=device)) + object.__setattr__( + articulation, "_ALL_SPATIAL_TENDON_INDICES", wp.array(np.array([], dtype=np.int32), device=device) + ) + object.__setattr__(articulation, "_ALL_SPATIAL_TENDON_MASK", wp.ones((0,), dtype=wp.bool, device=device)) + + # Joint targets (Newton uses warp, not torch) + object.__setattr__( + articulation, + "_joint_pos_target_sim", + wp.zeros((num_instances, num_joints), dtype=wp.float32, device=device), + ) + object.__setattr__( + articulation, + "_joint_vel_target_sim", + wp.zeros((num_instances, num_joints), dtype=wp.float32, device=device), + ) + object.__setattr__( + articulation, + "_joint_effort_target_sim", + wp.zeros((num_instances, num_joints), dtype=wp.float32, device=device), + ) + + return articulation, mock_view def create_mock_articulation( @@ -689,6 +812,8 @@ def test_body_com_acc_w(self, backend, num_instances, num_joints, num_bodies, de @_default_dims @_default_devices def test_body_com_pose_b(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): + if backend == "newton": + pytest.xfail("Newton only stores CoM position, not orientation") art, _ = articulation_iface art.data.update(dt=0.01) _check_wp_array( @@ -727,6 +852,8 @@ def test_body_inertia(self, backend, num_instances, num_joints, num_bodies, devi def test_body_incoming_joint_wrench_b( self, backend, num_instances, num_joints, num_bodies, device, articulation_iface ): + if backend == "newton": + pytest.xfail("Newton does not support joint wrench reporting") art, _ = articulation_iface art.data.update(dt=0.01) _check_wp_array( @@ -831,6 +958,8 @@ def test_body_com_pos_b(self, backend, num_instances, num_joints, num_bodies, de @_default_dims @_default_devices def test_body_com_quat_b(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): + if backend == "newton": + pytest.xfail("Newton only stores CoM position, not orientation") art, _ = articulation_iface art.data.update(dt=0.01) _check_wp_array( @@ -1150,12 +1279,12 @@ def _make_bad_data_warp(shape: tuple, device: str, wp_dtype=wp.float32) -> wp.ar def _make_env_mask(num_instances: int, device: str, partial: bool) -> wp.array | None: - """Create an env_mask: None for all envs, or a partial int32 mask.""" + """Create an env_mask: None for all envs, or a partial bool mask.""" if not partial: return None - mask_np = np.zeros(num_instances, dtype=np.int32) - mask_np[0] = 1 - return wp.array(mask_np, dtype=wp.int32, device=device) + mask_np = np.zeros(num_instances, dtype=bool) + mask_np[0] = True + return wp.array(mask_np, dtype=wp.bool, device=device) def _make_env_ids(device: str, subset: bool) -> torch.Tensor | None: @@ -1166,11 +1295,11 @@ def _make_env_ids(device: str, subset: bool) -> torch.Tensor | None: def _make_item_mask(total: int, selected: list[int], device: str) -> wp.array: - """Create an int32 warp mask with 1s at `selected` indices, 0s elsewhere.""" - mask_np = np.zeros(total, dtype=np.int32) + """Create a bool warp mask with True at `selected` indices, False elsewhere.""" + mask_np = np.zeros(total, dtype=bool) for i in selected: - mask_np[i] = 1 - return wp.array(mask_np, dtype=wp.int32, device=device) + mask_np[i] = True + return wp.array(mask_np, dtype=wp.bool, device=device) # --------------------------------------------------------------------------- @@ -1483,6 +1612,8 @@ def test_body_writer_index( wp_dtype, trailing, ): + if backend == "newton" and method_base == "set_coms": + pytest.xfail("Newton only stores CoM position, not orientation") art, _ = articulation_iface art.data.update(dt=0.01) method = getattr(art, f"{method_base}_index") @@ -1565,6 +1696,8 @@ def test_body_writer_mask( wp_dtype, trailing, ): + if backend == "newton" and method_base == "set_coms": + pytest.xfail("Newton only stores CoM position, not orientation") art, _ = articulation_iface art.data.update(dt=0.01) method = getattr(art, f"{method_base}_mask") @@ -1677,6 +1810,9 @@ def test_joint_aliases(self, backend, num_instances, num_joints, num_bodies, dev # Tendon tests — parametrize, properties, finders, data, writers # --------------------------------------------------------------------------- +# Newton does not support tendons (always 0), so exclude it from tendon tests. +_tendon_backends = pytest.mark.parametrize("backend", [b for b in BACKENDS if b != "newton"], indirect=False) + _tendon_dims = pytest.mark.parametrize( "num_instances, num_joints, num_bodies, num_fixed_tendons, num_spatial_tendons", [ @@ -1690,7 +1826,7 @@ def test_joint_aliases(self, backend, num_instances, num_joints, num_bodies, dev class TestArticulationTendonProperties: """Test that tendon-related articulation properties return the correct types/values.""" - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_num_fixed_tendons( @@ -1707,7 +1843,7 @@ def test_num_fixed_tendons( art, _ = articulation_iface assert art.num_fixed_tendons == num_fixed_tendons - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_num_spatial_tendons( @@ -1724,7 +1860,7 @@ def test_num_spatial_tendons( art, _ = articulation_iface assert art.num_spatial_tendons == num_spatial_tendons - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_fixed_tendon_names( @@ -1744,7 +1880,7 @@ def test_fixed_tendon_names( assert len(names) == num_fixed_tendons assert all(isinstance(n, str) for n in names) - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_spatial_tendon_names( @@ -1768,7 +1904,7 @@ def test_spatial_tendon_names( class TestArticulationTendonFinders: """Test that tendon finder methods return (list[int], list[str]) tuples.""" - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_find_fixed_tendons_all( @@ -1792,7 +1928,7 @@ def test_find_fixed_tendons_all( assert all(isinstance(i, int) for i in indices) assert all(isinstance(n, str) for n in names) - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_find_fixed_tendons_single( @@ -1814,7 +1950,7 @@ def test_find_fixed_tendons_single( assert indices == [0] assert names == [first] - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_find_spatial_tendons_all( @@ -1836,7 +1972,7 @@ def test_find_spatial_tendons_all( assert len(indices) == num_spatial_tendons assert len(names) == num_spatial_tendons - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_find_spatial_tendons_single( @@ -1864,7 +2000,7 @@ class TestArticulationDataTendonState: # -- Fixed tendon data properties -- - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_fixed_tendon_stiffness( @@ -1887,7 +2023,7 @@ def test_fixed_tendon_stiffness( name="fixed_tendon_stiffness", ) - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_fixed_tendon_damping( @@ -1910,7 +2046,7 @@ def test_fixed_tendon_damping( name="fixed_tendon_damping", ) - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_fixed_tendon_limit_stiffness( @@ -1933,7 +2069,7 @@ def test_fixed_tendon_limit_stiffness( name="fixed_tendon_limit_stiffness", ) - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_fixed_tendon_rest_length( @@ -1956,7 +2092,7 @@ def test_fixed_tendon_rest_length( name="fixed_tendon_rest_length", ) - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_fixed_tendon_offset( @@ -1979,7 +2115,7 @@ def test_fixed_tendon_offset( name="fixed_tendon_offset", ) - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_fixed_tendon_pos_limits( @@ -2008,7 +2144,7 @@ def test_fixed_tendon_pos_limits( # -- Spatial tendon data properties -- - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_spatial_tendon_stiffness( @@ -2033,7 +2169,7 @@ def test_spatial_tendon_stiffness( name="spatial_tendon_stiffness", ) - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_spatial_tendon_damping( @@ -2058,7 +2194,7 @@ def test_spatial_tendon_damping( name="spatial_tendon_damping", ) - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_spatial_tendon_limit_stiffness( @@ -2083,7 +2219,7 @@ def test_spatial_tendon_limit_stiffness( name="spatial_tendon_limit_stiffness", ) - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_spatial_tendon_offset( @@ -2129,7 +2265,7 @@ def test_spatial_tendon_offset( class TestArticulationWritersFixedTendon: """Test fixed tendon writers/setters with all input combinations.""" - @_backends + @_tendon_backends @_tendon_dims @_default_devices @pytest.mark.parametrize( @@ -2194,7 +2330,7 @@ def test_fixed_tendon_writer_index( with pytest.raises((AssertionError, RuntimeError)): method(**{kwarg: _make_bad_data_warp((num_instances, num_fixed_tendons), device, wp_dtype)}) - @_backends + @_tendon_backends @_tendon_dims @_default_devices @pytest.mark.parametrize( @@ -2270,7 +2406,7 @@ def test_fixed_tendon_writer_mask( class TestArticulationWritersSpatialTendon: """Test spatial tendon writers/setters with all input combinations.""" - @_backends + @_tendon_backends @_tendon_dims @_default_devices @pytest.mark.parametrize( @@ -2331,7 +2467,7 @@ def test_spatial_tendon_writer_index( with pytest.raises((AssertionError, RuntimeError)): method(**{kwarg: _make_bad_data_warp((num_instances, num_spatial_tendons), device, wp_dtype)}) - @_backends + @_tendon_backends @_tendon_dims @_default_devices @pytest.mark.parametrize( @@ -2400,7 +2536,7 @@ def test_spatial_tendon_writer_mask( class TestArticulationWritersTendonToSim: """Smoke test write_*_tendon_properties_to_sim_index/mask methods.""" - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_write_fixed_tendon_properties_to_sim_index( @@ -2423,7 +2559,7 @@ def test_write_fixed_tendon_properties_to_sim_index( # subset envs art.write_fixed_tendon_properties_to_sim_index(env_ids=_make_env_ids(device, True)) - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_write_fixed_tendon_properties_to_sim_mask( @@ -2446,7 +2582,7 @@ def test_write_fixed_tendon_properties_to_sim_mask( # partial env mask art.write_fixed_tendon_properties_to_sim_mask(env_mask=_make_env_mask(num_instances, device, True)) - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_write_spatial_tendon_properties_to_sim_index( @@ -2469,7 +2605,7 @@ def test_write_spatial_tendon_properties_to_sim_index( # subset envs art.write_spatial_tendon_properties_to_sim_index(env_ids=_make_env_ids(device, True)) - @_backends + @_tendon_backends @_tendon_dims @_default_devices def test_write_spatial_tendon_properties_to_sim_mask( diff --git a/source/isaaclab/test/assets/test_rigid_object_collection_iface.py b/source/isaaclab/test/assets/test_rigid_object_collection_iface.py index f1b3fce07bee..143d32a7f3ba 100644 --- a/source/isaaclab/test/assets/test_rigid_object_collection_iface.py +++ b/source/isaaclab/test/assets/test_rigid_object_collection_iface.py @@ -61,6 +61,20 @@ except ImportError: pass +try: + from isaaclab_newton.assets.rigid_object_collection.rigid_object_collection import ( + RigidObjectCollection as NewtonRigidObjectCollection, + ) + from isaaclab_newton.assets.rigid_object_collection.rigid_object_collection_data import ( + RigidObjectCollectionData as NewtonRigidObjectCollectionData, + ) + from isaaclab_newton.test.mock_interfaces.mock_newton import MockWrenchComposer as NewtonMockWrenchComposer + from isaaclab_newton.test.mock_interfaces.views import MockNewtonCollectionView as NewtonMockCollectionView + + BACKENDS.append("newton") +except ImportError: + pass + def create_physx_rigid_object_collection( num_instances: int = 2, @@ -113,6 +127,88 @@ def create_physx_rigid_object_collection( return collection, mock_view +def create_newton_rigid_object_collection( + num_instances: int = 2, + num_bodies: int = 3, + device: str = "cuda:0", +): + """Create a test Newton RigidObjectCollection instance with mocked dependencies.""" + import isaaclab_newton.assets.rigid_object_collection.rigid_object_collection as newton_coll_module + import isaaclab_newton.assets.rigid_object_collection.rigid_object_collection_data as newton_data_module + + body_names = [f"object_{i}" for i in range(num_bodies)] + + # Create collection-specific mock view with (N, B) root shapes + mock_view = NewtonMockCollectionView( + num_envs=num_instances, + num_bodies=num_bodies, + device=device, + body_names=body_names, + ) + mock_view.set_random_mock_data() + mock_view._noop_setters = True + + # Mock NewtonManager (aliased as SimulationManager in Newton modules) + mock_model = MagicMock() + mock_model.gravity = wp.array(np.array([[0.0, 0.0, -9.81]], dtype=np.float32), dtype=wp.vec3f, device=device) + mock_state = MagicMock() + mock_control = MagicMock() + + mock_manager = MagicMock() + mock_manager.get_model.return_value = mock_model + mock_manager.get_state_0.return_value = mock_state + mock_manager.get_state_1.return_value = mock_state + mock_manager.get_control.return_value = mock_control + + # Patch SimulationManager in both data and collection modules + original_data_manager = newton_data_module.SimulationManager + original_coll_manager = newton_coll_module.SimulationManager + newton_data_module.SimulationManager = mock_manager + newton_coll_module.SimulationManager = mock_manager + + try: + data = NewtonRigidObjectCollectionData(mock_view, num_bodies, device) + finally: + newton_data_module.SimulationManager = original_data_manager + newton_coll_module.SimulationManager = original_coll_manager + + # Create collection shell (bypass __init__) + collection = object.__new__(NewtonRigidObjectCollection) + + rigid_objects = {f"object_{i}": RigidObjectCfg(prim_path=f"/World/Object_{i}") for i in range(num_bodies)} + collection.cfg = RigidObjectCollectionCfg(rigid_objects=rigid_objects) + + object.__setattr__(collection, "_root_view", mock_view) + object.__setattr__(collection, "_device", device) + object.__setattr__(collection, "_num_bodies", num_bodies) + object.__setattr__(collection, "_num_instances", num_instances) + object.__setattr__(collection, "_body_names_list", body_names) + object.__setattr__(collection, "_data", data) + data.body_names = body_names + + # Mock wrench composers (Newton-specific) + mock_inst_wrench = NewtonMockWrenchComposer(collection) + mock_perm_wrench = NewtonMockWrenchComposer(collection) + object.__setattr__(collection, "_instantaneous_wrench_composer", mock_inst_wrench) + object.__setattr__(collection, "_permanent_wrench_composer", mock_perm_wrench) + + # Prevent __del__ / _clear_callbacks from raising AttributeError + object.__setattr__(collection, "_initialize_handle", None) + object.__setattr__(collection, "_invalidate_initialize_handle", None) + object.__setattr__(collection, "_prim_deletion_handle", None) + object.__setattr__(collection, "_debug_vis_handle", None) + + # Index arrays (warp) + object.__setattr__( + collection, "_ALL_ENV_INDICES", wp.array(np.arange(num_instances, dtype=np.int32), device=device) + ) + object.__setattr__(collection, "_ALL_BODY_INDICES", wp.array(np.arange(num_bodies, dtype=np.int32), device=device)) + object.__setattr__(collection, "_ALL_ENV_MASK", wp.ones((num_instances,), dtype=wp.bool, device=device)) + object.__setattr__(collection, "_ALL_BODY_MASK", wp.ones((num_bodies,), dtype=wp.bool, device=device)) + + return collection, mock_view + + def create_mock_rigid_object_collection( num_instances: int = 2, num_bodies: int = 3, @@ -136,6 +232,8 @@ def get_rigid_object_collection( ): if backend == "physx": return create_physx_rigid_object_collection(num_instances, num_bodies, device) + elif backend == "newton": + return create_newton_rigid_object_collection(num_instances, num_bodies, device) elif backend.lower() == "mock": return create_mock_rigid_object_collection(num_instances, num_bodies, device) else: @@ -221,12 +319,12 @@ def _make_bad_data_warp(shape: tuple, device: str, wp_dtype=wp.float32) -> wp.ar def _make_env_mask(num_instances: int, device: str, partial: bool) -> wp.array | None: - """Create an env_mask: None for all envs, or a partial int32 mask.""" + """Create an env_mask: None for all envs, or a partial bool mask.""" if not partial: return None - mask_np = np.zeros(num_instances, dtype=np.int32) - mask_np[0] = 1 - return wp.array(mask_np, dtype=wp.int32, device=device) + mask_np = np.zeros(num_instances, dtype=bool) + mask_np[0] = True + return wp.array(mask_np, dtype=wp.bool, device=device) def _make_env_ids(device: str, subset: bool) -> torch.Tensor | None: @@ -244,11 +342,11 @@ def _make_body_ids(device: str, subset_ids: list[int] | None) -> torch.Tensor | def _make_item_mask(total: int, selected: list[int], device: str) -> wp.array: - """Create an int32 warp mask with 1s at `selected` indices, 0s elsewhere.""" - mask_np = np.zeros(total, dtype=np.int32) + """Create a bool warp mask with True at `selected` indices, False elsewhere.""" + mask_np = np.zeros(total, dtype=bool) for i in selected: - mask_np[i] = 1 - return wp.array(mask_np, dtype=wp.int32, device=device) + mask_np[i] = True + return wp.array(mask_np, dtype=wp.bool, device=device) # --------------------------------------------------------------------------- @@ -984,6 +1082,8 @@ class TestCollectionWritersBody: def test_body_writer_index( self, backend, num_instances, num_bodies, device, collection_iface, method_base, kwarg, wp_dtype, trailing ): + if backend == "newton" and method_base == "set_coms": + pytest.xfail("Newton set_coms expects vec3f (position only), not transformf (pose)") obj, _ = collection_iface obj.data.update(dt=0.01) method = getattr(obj, f"{method_base}_index") @@ -1049,6 +1149,8 @@ def _make_warp(n_envs, n_bods): def test_body_writer_mask( self, backend, num_instances, num_bodies, device, collection_iface, method_base, kwarg, wp_dtype, trailing ): + if backend == "newton" and method_base == "set_coms": + pytest.xfail("Newton set_coms expects vec3f (position only), not transformf (pose)") obj, _ = collection_iface obj.data.update(dt=0.01) method = getattr(obj, f"{method_base}_mask") diff --git a/source/isaaclab/test/assets/test_rigid_object_iface.py b/source/isaaclab/test/assets/test_rigid_object_iface.py index 99efb8514081..1dc89d2c4c54 100644 --- a/source/isaaclab/test/assets/test_rigid_object_iface.py +++ b/source/isaaclab/test/assets/test_rigid_object_iface.py @@ -57,6 +57,15 @@ except ImportError: pass +try: + from isaaclab_newton.assets.rigid_object.rigid_object import RigidObject as NewtonRigidObject + from isaaclab_newton.assets.rigid_object.rigid_object_data import RigidObjectData as NewtonRigidObjectData + from isaaclab_newton.test.mock_interfaces.views import MockNewtonArticulationView as NewtonMockArticulationView + + BACKENDS.append("newton") +except ImportError: + pass + def create_physx_rigid_object( num_instances: int = 2, @@ -106,6 +115,81 @@ def create_physx_rigid_object( return rigid_object, mock_view +def create_newton_rigid_object( + num_instances: int = 2, + device: str = "cuda:0", +): + """Create a test Newton RigidObject instance with mocked dependencies.""" + import isaaclab_newton.assets.rigid_object.rigid_object_data as newton_data_module + + body_names = ["body_0"] + + # Create Newton mock view (uses ArticulationView with num_bodies=1 for rigid objects) + mock_view = NewtonMockArticulationView( + num_instances=num_instances, + num_bodies=1, + num_joints=0, + device=device, + is_fixed_base=False, + joint_names=[], + body_names=body_names, + ) + mock_view.set_random_mock_data() + mock_view._noop_setters = True + + # Mock NewtonManager (aliased as SimulationManager in Newton modules) + mock_model = MagicMock() + mock_model.gravity = wp.array(np.array([[0.0, 0.0, -9.81]], dtype=np.float32), dtype=wp.vec3f, device=device) + mock_state = MagicMock() + mock_control = MagicMock() + + mock_manager = MagicMock() + mock_manager.get_model.return_value = mock_model + mock_manager.get_state_0.return_value = mock_state + mock_manager.get_state_1.return_value = mock_state + mock_manager.get_control.return_value = mock_control + + # Patch SimulationManager in the Newton data module + original_sim_manager = newton_data_module.SimulationManager + newton_data_module.SimulationManager = mock_manager + + try: + data = NewtonRigidObjectData(mock_view, device) + finally: + newton_data_module.SimulationManager = original_sim_manager + + # Create RigidObject shell (bypass __init__) + rigid_object = object.__new__(NewtonRigidObject) + + rigid_object.cfg = RigidObjectCfg(prim_path="/World/Object") + + object.__setattr__(rigid_object, "_root_view", mock_view) + object.__setattr__(rigid_object, "_device", device) + object.__setattr__(rigid_object, "_data", data) + + # Mock wrench composers + mock_inst_wrench = MockWrenchComposer(rigid_object) + mock_perm_wrench = MockWrenchComposer(rigid_object) + object.__setattr__(rigid_object, "_instantaneous_wrench_composer", mock_inst_wrench) + object.__setattr__(rigid_object, "_permanent_wrench_composer", mock_perm_wrench) + + # Prevent __del__ / _clear_callbacks from raising AttributeError + object.__setattr__(rigid_object, "_initialize_handle", None) + object.__setattr__(rigid_object, "_invalidate_initialize_handle", None) + object.__setattr__(rigid_object, "_prim_deletion_handle", None) + object.__setattr__(rigid_object, "_debug_vis_handle", None) + + # Newton uses wp.array for indices + object.__setattr__(rigid_object, "_ALL_INDICES", wp.array(np.arange(num_instances, dtype=np.int32), device=device)) + object.__setattr__(rigid_object, "_ALL_BODY_INDICES", wp.array(np.array([0], dtype=np.int32), device=device)) + + # Newton uses wp.bool masks + object.__setattr__(rigid_object, "_ALL_ENV_MASK", wp.ones((num_instances,), dtype=wp.bool, device=device)) + object.__setattr__(rigid_object, "_ALL_BODY_MASK", wp.ones((1,), dtype=wp.bool, device=device)) + + return rigid_object, mock_view + + def create_mock_rigid_object( num_instances: int = 2, device: str = "cuda:0", @@ -126,6 +210,8 @@ def get_rigid_object( ): if backend == "physx": return create_physx_rigid_object(num_instances, device) + elif backend == "newton": + return create_newton_rigid_object(num_instances, device) elif backend.lower() == "mock": return create_mock_rigid_object(num_instances, device) else: @@ -753,12 +839,12 @@ def _make_bad_data_warp(shape: tuple, device: str, wp_dtype=wp.float32) -> wp.ar def _make_env_mask(num_instances: int, device: str, partial: bool) -> wp.array | None: - """Create an env_mask: None for all envs, or a partial int32 mask.""" + """Create an env_mask: None for all envs, or a partial bool mask.""" if not partial: return None - mask_np = np.zeros(num_instances, dtype=np.int32) - mask_np[0] = 1 - return wp.array(mask_np, dtype=wp.int32, device=device) + mask_np = np.zeros(num_instances, dtype=bool) + mask_np[0] = True + return wp.array(mask_np, dtype=wp.bool, device=device) def _make_env_ids(device: str, subset: bool) -> torch.Tensor | None: @@ -769,11 +855,11 @@ def _make_env_ids(device: str, subset: bool) -> torch.Tensor | None: def _make_item_mask(total: int, selected: list[int], device: str) -> wp.array: - """Create an int32 warp mask with 1s at `selected` indices, 0s elsewhere.""" - mask_np = np.zeros(total, dtype=np.int32) + """Create a bool warp mask with True at `selected` indices, False elsewhere.""" + mask_np = np.zeros(total, dtype=bool) for i in selected: - mask_np[i] = 1 - return wp.array(mask_np, dtype=wp.int32, device=device) + mask_np[i] = True + return wp.array(mask_np, dtype=wp.bool, device=device) # --------------------------------------------------------------------------- @@ -926,6 +1012,8 @@ class TestRigidObjectWritersBody: def test_body_writer_index( self, backend, num_instances, device, rigid_object_iface, method_base, kwarg, wp_dtype, trailing ): + if backend == "newton" and method_base == "set_coms": + pytest.xfail("Newton set_coms expects vec3f (position only), not transformf (pose)") obj, _ = rigid_object_iface obj.data.update(dt=0.01) num_bodies = 1 @@ -992,6 +1080,8 @@ def _make_warp(n_envs, n_bods): def test_body_writer_mask( self, backend, num_instances, device, rigid_object_iface, method_base, kwarg, wp_dtype, trailing ): + if backend == "newton" and method_base == "set_coms": + pytest.xfail("Newton set_coms expects vec3f (position only), not transformf (pose)") obj, _ = rigid_object_iface obj.data.update(dt=0.01) num_bodies = 1 diff --git a/source/isaaclab/test/controllers/test_ik_configs/pink_ik_g1_test_configs.json b/source/isaaclab/test/controllers/test_ik_configs/pink_ik_g1_test_configs.json index ccd9f035ab5d..165d0cba8f4b 100644 --- a/source/isaaclab/test/controllers/test_ik_configs/pink_ik_g1_test_configs.json +++ b/source/isaaclab/test/controllers/test_ik_configs/pink_ik_g1_test_configs.json @@ -1,8 +1,8 @@ { "tolerances": { - "position": 0.003, + "position": 0.025, "pd_position": 0.002, - "rotation": 0.017, + "rotation": 0.030, "check_errors": true }, "allowed_steps_to_settle": 50, diff --git a/source/isaaclab/test/envs/test_env_rendering_logic.py b/source/isaaclab/test/envs/test_env_rendering_logic.py index 751717ac26ae..3a5531165109 100644 --- a/source/isaaclab/test/envs/test_env_rendering_logic.py +++ b/source/isaaclab/test/envs/test_env_rendering_logic.py @@ -16,6 +16,7 @@ import pytest import torch from isaaclab_physx.physics import IsaacEvents +from isaaclab_visualizers.kit import KitVisualizer, KitVisualizerCfg import isaaclab.sim as sim_utils from isaaclab.envs import ( @@ -29,8 +30,6 @@ from isaaclab.scene import InteractiveSceneCfg from isaaclab.sim import SimulationCfg, SimulationContext from isaaclab.utils import configclass -from isaaclab.visualizers.kit_visualizer import KitVisualizer -from isaaclab.visualizers.kit_visualizer_cfg import KitVisualizerCfg @configclass diff --git a/source/isaaclab/test/scene/test_interactive_scene.py b/source/isaaclab/test/scene/test_interactive_scene.py index 60e61790d416..087474baa59e 100644 --- a/source/isaaclab/test/scene/test_interactive_scene.py +++ b/source/isaaclab/test/scene/test_interactive_scene.py @@ -12,6 +12,8 @@ """Rest everything follows.""" +from types import SimpleNamespace + import pytest import torch import warp as wp @@ -127,6 +129,56 @@ def test_reset_to_env_ids_input_types(device, setup_scene): assert_state_equal(prev_state, scene.get_state()) +def test_clone_environments_non_cfg_invokes_visualizer_clone_fn(monkeypatch: pytest.MonkeyPatch): + """Non-cfg clone path should execute visualizer clone callback with replicate args.""" + scene = object.__new__(InteractiveScene) + scene.cfg = SimpleNamespace(replicate_physics=False, num_envs=3) + scene.stage = object() + scene.env_fmt = "/World/envs/env_{}" + scene._ALL_INDICES = torch.arange(3, dtype=torch.long) + scene._default_env_origins = torch.zeros((3, 3), dtype=torch.float32) + scene._is_scene_setup_from_cfg = lambda: False + + # Avoid binding this unit test to global SimulationContext singleton state. + monkeypatch.setattr(InteractiveScene, "device", property(lambda self: "cpu")) + + physics_calls = [] + visualizer_calls = [] + usd_calls = [] + + def _physics_clone_fn(stage, *args, **kwargs): + physics_calls.append((stage, args, kwargs)) + + def _visualizer_clone_fn(stage, *args, **kwargs): + visualizer_calls.append((stage, args, kwargs)) + + def _usd_replicate(stage, *args, **kwargs): + usd_calls.append((stage, args, kwargs)) + + scene.cloner_cfg = SimpleNamespace( + device="cpu", + physics_clone_fn=_physics_clone_fn, + visualizer_clone_fn=_visualizer_clone_fn, + ) + monkeypatch.setattr("isaaclab.scene.interactive_scene.cloner.usd_replicate", _usd_replicate) + + scene.clone_environments(copy_from_source=False) + assert len(physics_calls) == 1 + assert len(visualizer_calls) == 1 + assert len(usd_calls) == 1 + mapping = physics_calls[0][1][3] + assert mapping.dtype == torch.bool + assert mapping.shape == (1, scene.num_envs) + + physics_calls.clear() + visualizer_calls.clear() + usd_calls.clear() + scene.clone_environments(copy_from_source=True) + assert len(physics_calls) == 0 + assert len(visualizer_calls) == 1 + assert len(usd_calls) == 1 + + def assert_state_equal(s1: dict, s2: dict, path=""): """ Recursively assert that s1 and s2 have the same nested keys diff --git a/source/isaaclab/test/sim/test_cloner.py b/source/isaaclab/test/sim/test_cloner.py index 0c2b04161919..4784f59dcd0a 100644 --- a/source/isaaclab/test/sim/test_cloner.py +++ b/source/isaaclab/test/sim/test_cloner.py @@ -21,6 +21,8 @@ import isaaclab.sim as sim_utils from isaaclab.cloner import usd_replicate +from isaaclab.cloner.cloner_utils import resolve_visualizer_clone_fn +from isaaclab.physics.scene_data_requirements import SceneDataRequirement, VisualizerPrebuiltArtifacts from isaaclab.sim import build_simulation_context @@ -217,3 +219,81 @@ def test_clone_decorator_wildcard_patterns( f"Expected {len(parent_paths)} matching prims, got {len(all_matching)}. " "Spurious parent prims were likely created by the @clone decorator." ) + + +def test_resolve_visualizer_clone_fn_returns_none_when_not_physx_backend(): + """Resolver should ignore non-PhysX backends.""" + hook = resolve_visualizer_clone_fn( + physics_backend="newton", + requirements=SceneDataRequirement(requires_newton_model=True), + stage=object(), + set_visualizer_artifact=lambda artifact: artifact, + ) + assert hook is None + + +def test_resolve_visualizer_clone_fn_returns_none_when_newton_model_not_required(): + """Resolver should not load optional hook when requirement is not requested.""" + hook = resolve_visualizer_clone_fn( + physics_backend="physx", + requirements=SceneDataRequirement(requires_newton_model=False), + stage=object(), + set_visualizer_artifact=lambda artifact: artifact, + ) + assert hook is None + + +def test_resolve_visualizer_clone_fn_returns_callable_when_available(sim): + """Resolver should return a callable hook when backend helper is available.""" + pytest.importorskip("isaaclab_newton.cloner.newton_replicate") + hook = resolve_visualizer_clone_fn( + physics_backend="physx", + requirements=SceneDataRequirement(requires_newton_model=True), + stage=sim_utils.get_current_stage(), + set_visualizer_artifact=lambda artifact: artifact, + ) + assert callable(hook) + + +def test_physx_newton_requirement_hook_populates_prebuilt_artifact(sim, monkeypatch: pytest.MonkeyPatch): + """PhysX + Newton requirement path should populate prebuilt visualizer artifact.""" + newton_replicate = pytest.importorskip("isaaclab_newton.cloner.newton_replicate") + + class _FakeModel: + body_label = ["/World/envs/env_0/A", "/World/envs/env_1/A"] + articulation_label = ["/World/envs/env_0/Robot", "/World/envs/env_1/Robot"] + + fake_model = _FakeModel() + fake_state = object() + + def _fake_prebuild(*args, **kwargs): + return fake_model, fake_state + + monkeypatch.setattr(newton_replicate, "newton_visualizer_prebuild", _fake_prebuild) + + captured: list[VisualizerPrebuiltArtifacts] = [] + hook = resolve_visualizer_clone_fn( + physics_backend="physx", + requirements=SceneDataRequirement(requires_newton_model=True), + stage=sim_utils.get_current_stage(), + set_visualizer_artifact=lambda artifact: captured.append(artifact), + ) + + assert callable(hook) + hook( + stage=sim_utils.get_current_stage(), + sources=["/World/template/A"], + destinations=["/World/envs/env_{}/A"], + env_ids=torch.tensor([0, 1], dtype=torch.long), + mapping=torch.ones((1, 2), dtype=torch.bool), + device="cpu", + ) + + assert len(captured) == 1 + artifact = captured[0] + assert isinstance(artifact, VisualizerPrebuiltArtifacts) + assert artifact.model is fake_model + assert artifact.state is fake_state + assert artifact.rigid_body_paths == fake_model.body_label + assert artifact.articulation_paths == fake_model.articulation_label + assert artifact.num_envs == 2 diff --git a/source/isaaclab/test/sim/test_physx_scene_data_provider_visualizer_contract.py b/source/isaaclab/test/sim/test_physx_scene_data_provider_visualizer_contract.py new file mode 100644 index 000000000000..07766e5e484f --- /dev/null +++ b/source/isaaclab/test/sim/test_physx_scene_data_provider_visualizer_contract.py @@ -0,0 +1,124 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Unit tests for PhysxSceneDataProvider visualizer-facing contracts.""" + +from __future__ import annotations + +from types import SimpleNamespace + +from isaaclab_physx.scene_data_providers import PhysxSceneDataProvider + +from isaaclab.physics.scene_data_requirements import VisualizerPrebuiltArtifacts + + +def _make_provider(): + provider = object.__new__(PhysxSceneDataProvider) + return provider + + +def test_get_newton_model_for_env_ids_builds_and_caches_sorted_keys(): + provider = _make_provider() + provider._needs_newton_sync = True + provider._newton_model = "full-model" + provider._filtered_newton_model = None + provider._filtered_env_ids_key = None + + build_calls = [] + + def _fake_build(env_ids): + build_calls.append(env_ids) + provider._filtered_newton_model = f"filtered-{env_ids}" + + provider._build_filtered_newton_model = _fake_build + + # None asks for the full model. + assert provider.get_newton_model_for_env_ids(None) == "full-model" + + # First subset request builds using sorted env id key. + model_a = provider.get_newton_model_for_env_ids([3, 1]) + assert model_a == "filtered-[1, 3]" + assert build_calls == [[1, 3]] + + # Equivalent request should use cache and not rebuild. + model_b = provider.get_newton_model_for_env_ids([1, 3]) + assert model_b == "filtered-[1, 3]" + assert build_calls == [[1, 3]] + + # Different subset rebuilds. + model_c = provider.get_newton_model_for_env_ids([2]) + assert model_c == "filtered-[2]" + assert build_calls == [[1, 3], [2]] + + +def test_try_use_prebuilt_artifact_populates_provider_state(): + """Provider should consume scene-time prebuilt artifact as fast path.""" + provider = _make_provider() + artifact = VisualizerPrebuiltArtifacts( + model="prebuilt-model", + state="prebuilt-state", + rigid_body_paths=["/World/envs/env_0/A"], + articulation_paths=["/World/envs/env_0/Robot"], + num_envs=4, + ) + provider._simulation_context = SimpleNamespace(get_scene_data_visualizer_prebuilt_artifact=lambda: artifact) + + provider._xform_views = {"old": object()} + provider._view_body_index_map = {"old": [1]} + provider._view_order_tensors = {"old": object()} + provider._pose_buf_num_bodies = 7 + provider._positions_buf = object() + provider._orientations_buf = object() + provider._covered_buf = object() + provider._xform_mask_buf = object() + provider._env_id_to_body_indices = {0: [0]} + provider._filtered_newton_model = "old-filtered-model" + provider._filtered_newton_state = "old-filtered-state" + provider._filtered_env_ids_key = (0,) + provider._filtered_body_indices = [0] + + assert provider._try_use_prebuilt_newton_artifact() is True + assert provider._newton_model == "prebuilt-model" + assert provider._newton_state == "prebuilt-state" + assert provider._rigid_body_paths == ["/World/envs/env_0/A"] + assert provider._articulation_paths == ["/World/envs/env_0/Robot"] + assert provider._num_envs_at_last_newton_build == 4 + assert provider._xform_views == {} + assert provider._view_body_index_map == {} + assert provider._view_order_tensors == {} + assert provider._pose_buf_num_bodies == 0 + assert provider._positions_buf is None + assert provider._orientations_buf is None + assert provider._covered_buf is None + assert provider._xform_mask_buf is None + assert provider._env_id_to_body_indices == {} + assert provider._filtered_newton_model is None + assert provider._filtered_newton_state is None + assert provider._filtered_env_ids_key is None + assert provider._filtered_body_indices == [] + + +def test_try_use_prebuilt_artifact_respects_force_usd_fallback_flag(): + """Force flag should disable prebuilt fast path even when artifact is available.""" + provider = _make_provider() + provider._force_usd_fallback = True + artifact = VisualizerPrebuiltArtifacts( + model="prebuilt-model", + state="prebuilt-state", + rigid_body_paths=["/World/envs/env_0/A"], + articulation_paths=["/World/envs/env_0/Robot"], + num_envs=4, + ) + provider._simulation_context = SimpleNamespace(get_scene_data_visualizer_prebuilt_artifact=lambda: artifact) + + assert provider._try_use_prebuilt_newton_artifact() is False + + +def test_build_newton_model_from_usd_short_circuits_when_prebuilt_available(): + """If prebuilt artifact is available, USD fallback should not run.""" + provider = _make_provider() + provider._try_use_prebuilt_newton_artifact = lambda: True + provider._build_newton_model_from_usd() + assert provider._last_newton_model_build_source == "prebuilt" diff --git a/source/isaaclab/test/sim/test_simulation_context_visualizers.py b/source/isaaclab/test/sim/test_simulation_context_visualizers.py new file mode 100644 index 000000000000..3dc12911198b --- /dev/null +++ b/source/isaaclab/test/sim/test_simulation_context_visualizers.py @@ -0,0 +1,264 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Unit tests for SimulationContext visualizer orchestration.""" + +from __future__ import annotations + +from typing import Any, cast + +import isaaclab_visualizers.viser.viser_visualizer as viser_visualizer +import pytest +from isaaclab_visualizers.viser.viser_visualizer_cfg import ViserVisualizerCfg + +from isaaclab.sim.simulation_context import SimulationContext + + +class _FakePhysicsManager: + def __init__(self): + self.forward_calls = 0 + + def forward(self): + self.forward_calls += 1 + + +class _FakeProvider: + def __init__(self): + self.update_calls = [] + + def update(self, env_ids=None): + self.update_calls.append(env_ids) + + +class _FakeVisualizer: + def __init__( + self, + *, + env_ids=None, + running=True, + closed=False, + rendering_paused=False, + training_paused_steps=0, + raises_on_step=False, + requires_forward=False, + ): + self._env_ids = env_ids + self._running = running + self._closed = closed + self._rendering_paused = rendering_paused + self._training_paused_steps = training_paused_steps + self._raises_on_step = raises_on_step + self._requires_forward = requires_forward + self.step_calls = [] + self.close_calls = 0 + + @property + def is_closed(self): + return self._closed + + def is_running(self): + return self._running + + def is_rendering_paused(self): + return self._rendering_paused + + def is_training_paused(self): + if self._training_paused_steps > 0: + self._training_paused_steps -= 1 + return True + return False + + def step(self, dt): + self.step_calls.append(dt) + if self._raises_on_step: + raise RuntimeError("step failed") + + def close(self): + self.close_calls += 1 + self._closed = True + + def get_visualized_env_ids(self): + return self._env_ids + + def requires_forward_before_step(self): + return self._requires_forward + + +def _make_context(visualizers, provider=None): + ctx = object.__new__(SimulationContext) + ctx._visualizers = list(visualizers) + ctx._scene_data_provider = provider + ctx.physics_manager = _FakePhysicsManager() + ctx._visualizer_step_counter = 0 + return ctx + + +def test_update_scene_data_provider_unions_env_ids_and_forwards(): + provider = _FakeProvider() + viz_a = _FakeVisualizer(env_ids=[0, 2], requires_forward=True) + viz_b = _FakeVisualizer(env_ids=[2, 3]) + viz_c = _FakeVisualizer(env_ids=None) + ctx = _make_context([viz_a, viz_b, viz_c], provider=provider) + + ctx.update_scene_data_provider() + + assert ctx.physics_manager.forward_calls == 1 + assert provider.update_calls == [[0, 2, 3]] + assert ctx._visualizer_step_counter == 1 + + +def test_update_scene_data_provider_force_forward_with_no_visualizers(): + provider = _FakeProvider() + ctx = _make_context([], provider=provider) + ctx.update_scene_data_provider(force_require_forward=True) + assert ctx.physics_manager.forward_calls == 1 + assert provider.update_calls == [None] + + +def test_update_visualizers_removes_closed_nonrunning_and_failed(caplog): + provider = _FakeProvider() + closed_viz = _FakeVisualizer(closed=True) + stopped_viz = _FakeVisualizer(running=False) + failing_viz = _FakeVisualizer(raises_on_step=True) + paused_viz = _FakeVisualizer(rendering_paused=True) + healthy_viz = _FakeVisualizer(env_ids=[1]) + ctx = _make_context([closed_viz, stopped_viz, failing_viz, paused_viz, healthy_viz], provider=provider) + + with caplog.at_level("ERROR"): + ctx.update_visualizers(0.1) + + assert ctx._visualizers == [paused_viz, healthy_viz] + assert closed_viz.close_calls == 1 + assert stopped_viz.close_calls == 1 + assert failing_viz.close_calls == 1 + assert paused_viz.close_calls == 0 + assert healthy_viz.step_calls == [0.1] + assert any("Error stepping visualizer" in r.message for r in caplog.records) + + +def test_update_visualizers_handles_training_pause_loop(): + provider = _FakeProvider() + viz = _FakeVisualizer(training_paused_steps=1) + ctx = _make_context([viz], provider=provider) + + ctx.update_visualizers(0.2) + + assert viz.step_calls == [0.0, 0.2] + + +class _DummyViserSceneDataProvider: + def __init__(self): + self._metadata = {"num_envs": 4} + self.state_calls: list[list[int] | None] = [] + + def get_metadata(self) -> dict: + return self._metadata + + def get_newton_model(self): + return "dummy-model" + + def get_newton_state(self, env_ids: list[int] | None): + self.state_calls.append(env_ids) + return {"state_call": len(self.state_calls), "env_ids": env_ids} + + +class _DummyViserViewer: + def __init__(self): + self.calls = [] + + def begin_frame(self, sim_time: float) -> None: + self.calls.append(("begin_frame", sim_time)) + + def log_state(self, state) -> None: + self.calls.append(("log_state", state)) + + def end_frame(self) -> None: + self.calls.append(("end_frame",)) + + def is_running(self) -> bool: + return True + + +def test_viser_visualizer_initialize_and_step_uses_provider_state(monkeypatch: pytest.MonkeyPatch): + provider = _DummyViserSceneDataProvider() + viewer = _DummyViserViewer() + + def _fake_create_viewer(self, record_to_viser: str | None, metadata: dict | None = None): + assert record_to_viser is None + assert metadata == provider.get_metadata() + self._viewer = viewer + + monkeypatch.setattr(viser_visualizer.ViserVisualizer, "_create_viewer", _fake_create_viewer) + + visualizer = viser_visualizer.ViserVisualizer(ViserVisualizerCfg()) + visualizer.initialize(cast(Any, provider)) + visualizer.step(0.25) + + assert visualizer.is_initialized + assert provider.state_calls == [None, None] + assert visualizer._sim_time == pytest.approx(0.25) + assert viewer.calls[0][0] == "begin_frame" + assert viewer.calls[0][1] == pytest.approx(0.25) + assert viewer.calls[1] == ("log_state", {"state_call": 2, "env_ids": None}) + assert viewer.calls[2] == ("end_frame",) + + +@pytest.mark.parametrize( + ("cfg_max_worlds", "expected_max_worlds"), + [ + (None, None), + (0, 0), + (3, 3), + ], +) +def test_viser_visualizer_create_viewer_forwards_max_worlds( + monkeypatch: pytest.MonkeyPatch, cfg_max_worlds: int | None, expected_max_worlds: int | None +): + captured = {} + + class _FakeNewtonViewerViser: + def __init__( + self, + *, + port: int, + label: str | None, + verbose: bool, + share: bool, + record_to_viser: str | None, + metadata: dict | None = None, + ): + captured["init"] = { + "port": port, + "label": label, + "verbose": verbose, + "share": share, + "record_to_viser": record_to_viser, + "metadata": metadata, + } + + def set_model(self, model: Any, max_worlds: int | None) -> None: + captured["set_model"] = {"model": model, "max_worlds": max_worlds} + + monkeypatch.setattr(viser_visualizer, "NewtonViewerViser", _FakeNewtonViewerViser) + monkeypatch.setattr( + viser_visualizer.ViserVisualizer, + "_resolve_initial_camera_pose", + lambda self: ((1.0, 2.0, 3.0), (0.0, 0.0, 0.0)), + ) + monkeypatch.setattr(viser_visualizer.ViserVisualizer, "_set_viser_camera_view", lambda self, pose: None) + + cfg = ViserVisualizerCfg(max_worlds=cfg_max_worlds, open_browser=False) + visualizer = viser_visualizer.ViserVisualizer(cfg) + visualizer._model = "dummy-model" + visualizer._create_viewer(record_to_viser="record.viser", metadata={"num_envs": 8}) + + assert captured["set_model"] == {"model": "dummy-model", "max_worlds": expected_max_worlds} + + +def test_get_cli_visualizer_types_handles_non_string_setting_without_crashing(): + ctx = object.__new__(SimulationContext) + ctx.get_setting = lambda name: {"types": "newton,kit"} if name == "/isaaclab/visualizer/types" else None + + assert ctx._get_cli_visualizer_types() == [] diff --git a/source/isaaclab/test/sim/test_viser_visualizer.py b/source/isaaclab/test/sim/test_viser_visualizer.py deleted file mode 100644 index 0ad22ab1e76e..000000000000 --- a/source/isaaclab/test/sim/test_viser_visualizer.py +++ /dev/null @@ -1,190 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -# pyright: reportPrivateUsage=false - -from __future__ import annotations - -from typing import Any, cast - -import pytest - -pytest.importorskip("newton.viewer") - -import isaaclab.visualizers.viser_visualizer as viser_visualizer -from isaaclab.visualizers.viser_visualizer_cfg import ViserVisualizerCfg - - -class _DummySceneDataProvider: - def __init__(self): - self._metadata = {"num_envs": 4} - self.state_calls: list[list[int] | None] = [] - - def get_metadata(self) -> dict: - return self._metadata - - def get_newton_model(self): - return "dummy-model" - - def get_newton_state(self, env_ids: list[int] | None): - self.state_calls.append(env_ids) - return {"state_call": len(self.state_calls), "env_ids": env_ids} - - -class _DummyViewer: - def __init__(self): - self.calls = [] - - def begin_frame(self, sim_time: float) -> None: - self.calls.append(("begin_frame", sim_time)) - - def log_state(self, state) -> None: - self.calls.append(("log_state", state)) - - def end_frame(self) -> None: - self.calls.append(("end_frame",)) - - def close(self) -> None: - self.calls.append(("close",)) - - def is_running(self) -> bool: - return True - - -def test_newton_viewer_viser_forwards_init_args_and_metadata(monkeypatch: pytest.MonkeyPatch): - disable_calls = [] - monkeypatch.setattr( - viser_visualizer, "_disable_viser_runtime_client_rebuild_if_bundled", lambda: disable_calls.append(True) - ) - - captured = {} - - def _fake_base_init(self, *, port: int, label: str, verbose: bool, share: bool, record_to_viser: str | None): - captured.update( - {"port": port, "label": label, "verbose": verbose, "share": share, "record_to_viser": record_to_viser} - ) - - monkeypatch.setattr(viser_visualizer.ViewerViser, "__init__", _fake_base_init) - - metadata = {"num_envs": 16} - viewer = viser_visualizer.NewtonViewerViser( - port=9090, - label="unit-test", - verbose=False, - share=True, - record_to_viser="record.viser", - metadata=metadata, - ) - - assert disable_calls == [True] - assert captured == { - "port": 9090, - "label": "unit-test", - "verbose": False, - "share": True, - "record_to_viser": "record.viser", - } - assert viewer._metadata is metadata - - -def test_viser_visualizer_initialize_and_step_uses_provider_state(monkeypatch: pytest.MonkeyPatch): - provider = _DummySceneDataProvider() - viewer = _DummyViewer() - - def _fake_create_viewer(self, record_to_viser: str | None, metadata: dict | None = None): - assert record_to_viser is None - assert metadata == provider.get_metadata() - self._viewer = viewer - - monkeypatch.setattr(viser_visualizer.ViserVisualizer, "_create_viewer", _fake_create_viewer) - - visualizer = viser_visualizer.ViserVisualizer(ViserVisualizerCfg()) - visualizer.initialize(cast(Any, provider)) - visualizer.step(0.25) - - assert visualizer.is_initialized - assert provider.state_calls == [None, None] - assert visualizer._sim_time == pytest.approx(0.25) - assert viewer.calls[0][0] == "begin_frame" - assert viewer.calls[0][1] == pytest.approx(0.25) - assert viewer.calls[1] == ("log_state", {"state_call": 2, "env_ids": None}) - assert viewer.calls[2] == ("end_frame",) - - -@pytest.mark.parametrize( - ("cfg_max_worlds", "expected_max_worlds"), - [ - (None, None), - (0, None), - (3, 3), - ], -) -def test_viser_visualizer_create_viewer_forwards_max_worlds( - monkeypatch: pytest.MonkeyPatch, cfg_max_worlds: int | None, expected_max_worlds: int | None -): - captured = {} - - class _FakeNewtonViewerViser: - def __init__( - self, - *, - port: int, - label: str | None, - verbose: bool, - share: bool, - record_to_viser: str | None, - metadata: dict | None = None, - ): - captured["init"] = { - "port": port, - "label": label, - "verbose": verbose, - "share": share, - "record_to_viser": record_to_viser, - "metadata": metadata, - } - - def set_model(self, model: Any, max_worlds: int | None) -> None: - captured["set_model"] = {"model": model, "max_worlds": max_worlds} - - monkeypatch.setattr(viser_visualizer, "NewtonViewerViser", _FakeNewtonViewerViser) - monkeypatch.setattr( - viser_visualizer.ViserVisualizer, - "_resolve_initial_camera_pose", - lambda self: ((1.0, 2.0, 3.0), (0.0, 0.0, 0.0)), - ) - monkeypatch.setattr(viser_visualizer.ViserVisualizer, "_set_viser_camera_view", lambda self, pose: None) - - cfg = ViserVisualizerCfg(max_worlds=cfg_max_worlds) - visualizer = viser_visualizer.ViserVisualizer(cfg) - visualizer._model = "dummy-model" - visualizer._create_viewer(record_to_viser="record.viser", metadata={"num_envs": 8}) - - assert captured["set_model"] == {"model": "dummy-model", "max_worlds": expected_max_worlds} - - -def test_viser_visualizer_close_uses_recording_flag(monkeypatch: pytest.MonkeyPatch): - cfg = ViserVisualizerCfg(record_to_viser="record.viser") - visualizer = viser_visualizer.ViserVisualizer(cfg) - visualizer._is_initialized = True - visualizer._viewer = cast(Any, object()) - visualizer._active_record_path = "record.viser" - visualizer._pending_camera_pose = ((1.0, 1.0, 1.0), (0.0, 0.0, 0.0)) - - close_call = {} - - def _fake_close_viewer(self, finalize_viser: bool = False): - close_call["finalize_viser"] = finalize_viser - - monkeypatch.setattr(viser_visualizer.ViserVisualizer, "_close_viewer", _fake_close_viewer) - - visualizer.close() - - assert close_call["finalize_viser"] is True - assert visualizer._viewer is None - assert visualizer._is_initialized is False - assert visualizer._is_closed is True - assert visualizer._active_record_path is None - assert visualizer._pending_camera_pose is None diff --git a/source/isaaclab/test/visualizers/test_visualizer.py b/source/isaaclab/test/visualizers/test_visualizer.py new file mode 100644 index 000000000000..ed1baf9198c1 --- /dev/null +++ b/source/isaaclab/test/visualizers/test_visualizer.py @@ -0,0 +1,116 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Unit tests for visualizer config factory and base visualizer behavior.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from isaaclab.visualizers.base_visualizer import BaseVisualizer +from isaaclab.visualizers.visualizer import Visualizer +from isaaclab.visualizers.visualizer_cfg import VisualizerCfg + +# +# Config factory +# + + +def test_create_visualizer_raises_for_base_cfg(): + cfg = VisualizerCfg() + with pytest.raises(ValueError, match="Cannot create visualizer from base VisualizerCfg class"): + cfg.create_visualizer() + + +def test_create_visualizer_raises_for_unknown_type(): + cfg = VisualizerCfg(visualizer_type="unknown-backend") + with pytest.raises(ValueError, match="not registered"): + cfg.create_visualizer() + + +def test_create_visualizer_raises_import_error_when_backend_unavailable(monkeypatch): + monkeypatch.setattr(Visualizer, "_get_module_name", classmethod(lambda cls, backend: "does.not.exist")) + cfg = VisualizerCfg(visualizer_type="newton") + with pytest.raises(ImportError, match="isaaclab_visualizers"): + cfg.create_visualizer() + + +# +# Base visualizer (env filtering, camera pose) +# + + +class _DummyVisualizer(BaseVisualizer): + def initialize(self, scene_data_provider) -> None: + self._scene_data_provider = scene_data_provider + self._is_initialized = True + + def step(self, dt: float) -> None: + return + + def close(self) -> None: + self._is_closed = True + + def is_running(self) -> bool: + return True + + +def _make_cfg(**kwargs): + cfg = { + "env_filter_mode": "none", + "env_filter_ids": [0, 2, 4], + "env_filter_random_n": 2, + "env_filter_seed": 7, + } + cfg.update(kwargs) + return SimpleNamespace(**cfg) + + +class _FakeProvider: + def __init__(self, num_envs: int = 0, transforms: dict | None = None): + self._num_envs = num_envs + self._transforms = transforms + + def get_metadata(self) -> dict: + return {"num_envs": self._num_envs} + + def get_camera_transforms(self): + return self._transforms + + +def test_compute_visualized_env_ids_none_mode(): + viz = _DummyVisualizer(_make_cfg(env_filter_mode="none")) + viz._scene_data_provider = _FakeProvider(num_envs=8) + assert viz._compute_visualized_env_ids() is None + + +def test_compute_visualized_env_ids_from_ids_filters_out_of_range(): + viz = _DummyVisualizer(_make_cfg(env_filter_mode="env_ids", env_filter_ids=[-1, 0, 3, 99])) + viz._scene_data_provider = _FakeProvider(num_envs=4) + assert viz._compute_visualized_env_ids() == [0, 3] + + +def test_compute_visualized_env_ids_random_n_is_deterministic(): + cfg = _make_cfg(env_filter_mode="random_n", env_filter_random_n=3, env_filter_seed=123) + viz_a = _DummyVisualizer(cfg) + viz_b = _DummyVisualizer(cfg) + viz_a._scene_data_provider = _FakeProvider(num_envs=10) + viz_b._scene_data_provider = _FakeProvider(num_envs=10) + assert viz_a._compute_visualized_env_ids() == viz_b._compute_visualized_env_ids() + + +def test_resolve_camera_pose_from_usd_path_uses_provider_transforms(): + transforms = { + "order": ["/World/envs/env_%d/Camera"], + "positions": [[[1.0, 2.0, 3.0]]], + "orientations": [[[0.0, 0.0, 0.0, 1.0]]], + } + viz = _DummyVisualizer(_make_cfg()) + viz._scene_data_provider = _FakeProvider(num_envs=1, transforms=transforms) + pos, target = viz._resolve_camera_pose_from_usd_path("/World/envs/env_0/Camera") + assert pos == (1.0, 2.0, 3.0) + assert target == pytest.approx((1.0, 2.0, 2.0)) diff --git a/source/isaaclab_assets/isaaclab_assets/robots/kuka_allegro.py b/source/isaaclab_assets/isaaclab_assets/robots/kuka_allegro.py index ed5286d71f86..f521f45bff57 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/kuka_allegro.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/kuka_allegro.py @@ -108,6 +108,9 @@ "ring_joint_(0|1|2|3)": 0.01, "thumb_joint_(0|1|2|3)": 0.01, }, + armature={ + ".*": 0.01, + }, ), }, soft_joint_pos_limit_factor=1.0, diff --git a/source/isaaclab_newton/config/extension.toml b/source/isaaclab_newton/config/extension.toml index 072a0329036d..48f955abf7b0 100644 --- a/source/isaaclab_newton/config/extension.toml +++ b/source/isaaclab_newton/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "0.4.0" +version = "0.5.3" # Description title = "Newton simulation interfaces for IsaacLab core package" diff --git a/source/isaaclab_newton/docs/CHANGELOG.rst b/source/isaaclab_newton/docs/CHANGELOG.rst index 5ac9bd13c056..f9fdf7f2406c 100644 --- a/source/isaaclab_newton/docs/CHANGELOG.rst +++ b/source/isaaclab_newton/docs/CHANGELOG.rst @@ -1,6 +1,132 @@ Changelog --------- +0.5.4 (2026-02-28) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added contact sensor support via :class:`newton.sensors.SensorContact` with + Isaac Lab pattern conversion (``.*`` to fnmatch, USD path normalization) + inlined in :meth:`~isaaclab_newton.physics.NewtonManager.add_contact_sensor`. + +Changed +^^^^^^^ + +* Changed :class:`~isaaclab_newton.sensors.contact_sensor.ContactSensor` to + flatten Newton's per-world nested ``sensing_objs`` and ``counterparts`` + attributes. + +Fixed +^^^^^ + +* Fixed ``RigidObjectData.body_inertia`` shape from ``(N, B, 3, 3)`` to ``(N, B, 9)``. + + +0.5.3 (2026-03-09) +~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Fixed :attr:`~isaaclab_newton.assets.RigidObjectData.body_inertia` to return a + ``(num_instances, num_bodies, 9)`` float32 strided view, matching the articulation fix in 0.5.2. + +* Fixed non-contiguous array handling in ``RigidObjectData`` position, quaternion, and + spatial-vector extraction helpers. The ``source`` buffer shape and kernel dispatch ``dim`` + now use the input array's shape instead of the (possibly uninitialized) output shape. + +0.5.2 (2026-03-06) +~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Fixed :attr:`~isaaclab_newton.assets.ArticulationData.body_inertia` in + :class:`~isaaclab_newton.assets.ArticulationData` to return a ``(num_instances, num_bodies, 9)`` + float32 array as documented, instead of a ``(num_instances, num_bodies, 3, 3)`` array. The + ``(N, B, 3, 3)`` shape caused a broadcasting error in + :class:`~isaaclab.envs.mdp.events.randomize_rigid_body_mass` and a dimension mismatch when the + ``write_body_inertia_to_buffer_*`` kernels were called via + :meth:`~isaaclab_newton.assets.Articulation.set_inertias_index` and + :meth:`~isaaclab_newton.assets.Articulation.set_inertias_mask`. The fix creates a ``(N, B, 9)`` + view over the same memory using explicit strides, collapsing the two contiguous trailing + dimensions without copying data. + +* Fixed ``AttributeError: 'NoneType' object has no attribute 'device'`` in + :meth:`~isaaclab_newton.physics.NewtonManager.step` when ``use_cuda_graph=True`` but the CUDA + graph was not captured (e.g., when RTX/Fabric USD sync is active). The step condition now + checks ``cls._graph is not None`` directly instead of repeating the capture-time heuristic. + +0.5.1 (2026-03-06) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added :class:`~isaaclab_newton.assets.RigidObjectCollection` and + :class:`~isaaclab_newton.assets.RigidObjectCollectionData` for managing + collections of independent rigid bodies. Uses a single + ``ArticulationView`` with a combined fnmatch pattern to get direct + ``(num_envs, num_bodies)`` bindings into Newton's state, avoiding the + scatter/gather overhead needed by PhysX. + +* Added :class:`~isaaclab_newton.test.mock_interfaces.views.MockNewtonCollectionView` + for unit testing the collection data class without simulation. + +* Added Newton backend to the rigid object collection interface conformance + tests (``test_rigid_object_collection_iface.py``). + + +0.5.0 (2026-03-06) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added full Newton articulation test suite (``test_articulation.py``) — 194 passed, + 8 skipped, 4 xfailed — adapted from PhysX tests with Newton-specific imports, sim + config, and solver tolerance adjustments. + +* Added full Newton rigid body test suite (``test_rigid_object.py``) — 74 passed, + 25 skipped — adapted from PhysX tests with Newton-specific mass/COM APIs and + ``_newton_sim_context()`` helper for device/gravity/dt configuration. + +Fixed +^^^^^ + +* Fixed ``ArticulationData`` and ``RigidObjectData`` to rebind simulation pointers + on full sim reset via ``PHYSICS_READY`` callback, preventing stale warp array + references after ``sim.reset()`` recreates the Newton model. + +* Fixed ``ArticulationData`` to force ``eval_fk`` after joint state writes so that + link poses are consistent with joint positions before the next ``sim.step()``. + +* Fixed lazy initialization of ``TimestampedBuffer`` properties in + ``RigidObjectData`` (velocity-in-body-frame and deprecated state properties) + that were left as ``None`` and caused ``AttributeError`` on first access. + +* Fixed ``None`` guards for timestamp invalidation in ``RigidObject`` write methods + (``write_root_pose_to_sim``, ``write_root_velocity_to_sim``) to avoid + ``AttributeError`` when optional buffers have not been initialized. + +* Fixed ``is_contiguous`` usage in ``RigidObjectData`` — warp 1.12.0rc2 exposes it + as a property, not a method. + +* Fixed ``body_com_pose_b`` → ``body_com_pos_b`` kernel input naming in + ``RigidObjectData`` for ``root_com_pose_w`` and ``root_link_vel_w`` properties. + +* Fixed ``wp.from_torch()`` called on warp arrays in ``RigidObjectData`` body + inertia binding — replaced with direct ``.view()``/``.reshape()`` on warp arrays. + +* Improved CPU support in ``NewtonManager``: added device guards for CUDA graph + operations that are not available on CPU. + +* Fixed explicit mask resolution in asset write methods to correctly handle both + index-based and mask-based sparse writes. + + 0.4.1 (2026-03-03) ~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab_newton/isaaclab_newton/assets/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/assets/__init__.pyi index 52d1a435596a..b9359445960c 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/__init__.pyi +++ b/source/isaaclab_newton/isaaclab_newton/assets/__init__.pyi @@ -8,7 +8,10 @@ __all__ = [ "ArticulationData", "RigidObject", "RigidObjectData", + "RigidObjectCollection", + "RigidObjectCollectionData", ] from .articulation import Articulation, ArticulationData from .rigid_object import RigidObject, RigidObjectData +from .rigid_object_collection import RigidObjectCollection, RigidObjectCollectionData diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index fdacd5512552..5507581865d9 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -25,6 +25,7 @@ from isaaclab.actuators import ActuatorBase, ActuatorBaseCfg, ImplicitActuator from isaaclab.assets.articulation.base_articulation import BaseArticulation +from isaaclab.physics import PhysicsEvent from isaaclab.sim.utils.queries import find_first_matching_prim, get_all_matching_child_prims from isaaclab.utils.string import resolve_matching_names, resolve_matching_names_values from isaaclab.utils.types import ArticulationActions @@ -294,11 +295,11 @@ def write_data_to_sim(self): # apply actuator models self._apply_actuator_model() # write actions into simulation via Newton bindings - wp.copy(self.data._sim_bind_joint_effort, self._joint_effort_target_sim) + self.data._sim_bind_joint_effort.assign(self._joint_effort_target_sim) # position and velocity targets only for implicit actuators if self._has_implicit_actuators: - wp.copy(self.data._sim_bind_joint_position_target, self._joint_pos_target_sim) - wp.copy(self.data._sim_bind_joint_velocity_target, self._joint_vel_target_sim) + self.data._sim_bind_joint_position_target.assign(self._joint_pos_target_sim) + self.data._sim_bind_joint_velocity_target.assign(self._joint_vel_target_sim) def update(self, dt: float): """Updates the simulation data. @@ -445,7 +446,7 @@ def write_root_pose_to_sim_mask( or (num_instances,) with dtype wp.transformf. env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - self.write_root_link_pose_to_sim_mask(root_pose, env_mask=env_mask) + self.write_root_link_pose_to_sim_mask(root_pose=root_pose, env_mask=env_mask) def write_root_link_pose_to_sim_index( self, @@ -493,7 +494,7 @@ def write_root_link_pose_to_sim_index( self.data._root_link_state_w.timestamp = -1.0 if self.data._root_state_w is not None: self.data._root_state_w.timestamp = -1.0 - self.data._body_link_pose_w_timestamp = -1.0 # Forces a kinematic update to get the latest body link poses. + self.data._fk_timestamp = -1.0 # Forces a kinematic update to get the latest body link poses. if self.data._body_com_pose_w is not None: self.data._body_com_pose_w.timestamp = -1.0 if self.data._body_state_w is not None: @@ -525,8 +526,7 @@ def write_root_link_pose_to_sim_mask( or (num_instances,) with dtype wp.transformf. env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - if env_mask is None: - env_mask = self._ALL_ENV_MASK + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) self.assert_shape_and_dtype_mask(root_pose, (env_mask,), wp.transformf, "root_pose") wp.launch( @@ -549,7 +549,7 @@ def write_root_link_pose_to_sim_mask( self.data._root_link_state_w.timestamp = -1.0 if self.data._root_state_w is not None: self.data._root_state_w.timestamp = -1.0 - self.data._body_link_pose_w_timestamp = -1.0 # Forces a kinematic update to get the latest body link poses. + self.data._fk_timestamp = -1.0 # Forces a kinematic update to get the latest body link poses. if self.data._body_com_pose_w is not None: self.data._body_com_pose_w.timestamp = -1.0 if self.data._body_state_w is not None: @@ -593,7 +593,7 @@ def write_root_com_pose_to_sim_index( dim=env_ids.shape[0], inputs=[ root_pose, - self.data.body_com_pose_b, + self.data.body_com_pos_b, env_ids, ], outputs=[ @@ -613,7 +613,7 @@ def write_root_com_pose_to_sim_index( self.data._root_link_state_w.timestamp = -1.0 if self.data._root_state_w is not None: self.data._root_state_w.timestamp = -1.0 - self.data._body_link_pose_w_timestamp = -1.0 # Forces a kinematic update to get the latest body link poses. + self.data._fk_timestamp = -1.0 # Forces a kinematic update to get the latest body link poses. if self.data._body_com_pose_w is not None: self.data._body_com_pose_w.timestamp = -1.0 if self.data._body_state_w is not None: @@ -646,15 +646,14 @@ def write_root_com_pose_to_sim_mask( or (num_instances,) with dtype wp.transformf. env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - if env_mask is None: - env_mask = self._ALL_ENV_MASK + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) self.assert_shape_and_dtype_mask(root_pose, (env_mask,), wp.transformf, "root_pose") wp.launch( shared_kernels.set_root_com_pose_to_sim_mask, dim=root_pose.shape[0], inputs=[ root_pose, - self.data.body_com_pose_b, + self.data.body_com_pos_b, env_mask, ], outputs=[ @@ -674,7 +673,7 @@ def write_root_com_pose_to_sim_mask( self.data._root_link_state_w.timestamp = -1.0 if self.data._root_state_w is not None: self.data._root_state_w.timestamp = -1.0 - self.data._body_link_pose_w_timestamp = -1.0 # Forces a kinematic update to get the latest body link poses. + self.data._fk_timestamp = -1.0 # Forces a kinematic update to get the latest body link poses. if self.data._body_com_pose_w is not None: self.data._body_com_pose_w.timestamp = -1.0 if self.data._body_state_w is not None: @@ -810,8 +809,7 @@ def write_root_com_velocity_to_sim_mask( Shape is (num_instances, 6) or (num_instances,) with dtype wp.spatial_vectorf. env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - if env_mask is None: - env_mask = self._ALL_ENV_MASK + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) self.assert_shape_and_dtype_mask(root_velocity, (env_mask,), wp.spatial_vectorf, "root_velocity") wp.launch( shared_kernels.set_root_com_velocity_to_sim_mask, @@ -869,7 +867,7 @@ def write_root_link_velocity_to_sim_index( dim=env_ids.shape[0], inputs=[ root_velocity, - self.data.body_com_pose_b, + self.data.body_com_pos_b, self.data.root_link_pose_w, env_ids, self.data._num_bodies, @@ -916,15 +914,14 @@ def write_root_link_velocity_to_sim_mask( Shape is (num_instances, 6) or (num_instances,) with dtype wp.spatial_vectorf. env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - if env_mask is None: - env_mask = self._ALL_ENV_MASK + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) self.assert_shape_and_dtype_mask(root_velocity, (env_mask,), wp.spatial_vectorf, "root_velocity") wp.launch( shared_kernels.set_root_link_velocity_to_sim_mask, dim=root_velocity.shape[0], inputs=[ root_velocity, - self.data.body_com_pose_b, + self.data.body_com_pos_b, self.data.root_link_pose_w, env_mask, self.data._num_bodies, @@ -1013,6 +1010,8 @@ def write_joint_position_to_sim_index( ], device=self.device, ) + # Invalidate FK timestamp so body poses are recomputed on next access. + self.data._fk_timestamp = -1.0 # Need to invalidate the buffer to trigger the update with the new root pose. # Only invalidate if the buffer has been accessed (not None). if self.data._body_link_vel_w is not None: @@ -1049,10 +1048,8 @@ def write_joint_position_to_sim_mask( joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - if env_mask is None: - env_mask = self._ALL_ENV_MASK - if joint_mask is None: - joint_mask = self._ALL_JOINT_MASK + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + joint_mask = self._resolve_mask(joint_mask, self._ALL_JOINT_MASK) self.assert_shape_and_dtype_mask(position, (env_mask, joint_mask), wp.float32, "position") wp.launch( shared_kernels.write_2d_data_to_buffer_with_mask, @@ -1067,6 +1064,8 @@ def write_joint_position_to_sim_mask( ], device=self.device, ) + # Invalidate FK timestamp so body poses are recomputed on next access. + self.data._fk_timestamp = -1.0 # Need to invalidate the buffer to trigger the update with the new root pose. # Only invalidate if the buffer has been accessed (not None). if self.data._body_link_vel_w is not None: @@ -1145,10 +1144,8 @@ def write_joint_velocity_to_sim_mask( joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - if env_mask is None: - env_mask = self._ALL_ENV_MASK - if joint_mask is None: - joint_mask = self._ALL_JOINT_MASK + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + joint_mask = self._resolve_mask(joint_mask, self._ALL_JOINT_MASK) self.assert_shape_and_dtype_mask(velocity, (env_mask, joint_mask), wp.float32, "velocity") wp.launch( articulation_kernels.write_joint_vel_data_mask, @@ -1248,10 +1245,8 @@ def write_joint_stiffness_to_sim_mask( joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - if env_mask is None: - env_mask = self._ALL_ENV_MASK - if joint_mask is None: - joint_mask = self._ALL_JOINT_MASK + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + joint_mask = self._resolve_mask(joint_mask, self._ALL_JOINT_MASK) if isinstance(stiffness, float): wp.launch( articulation_kernels.float_data_to_buffer_with_mask, @@ -1363,10 +1358,8 @@ def write_joint_damping_to_sim_mask( joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - if env_mask is None: - env_mask = self._ALL_ENV_MASK - if joint_mask is None: - joint_mask = self._ALL_JOINT_MASK + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + joint_mask = self._resolve_mask(joint_mask, self._ALL_JOINT_MASK) if isinstance(damping, float): wp.launch( articulation_kernels.float_data_to_buffer_with_mask, @@ -1489,10 +1482,8 @@ def write_joint_position_limit_to_sim_mask( warn_limit_violation: Whether to use warning or info level logging when default joint positions exceed the new limits. Defaults to True. """ - if env_mask is None: - env_mask = self._ALL_ENV_MASK - if joint_mask is None: - joint_mask = self._ALL_JOINT_MASK + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + joint_mask = self._resolve_mask(joint_mask, self._ALL_JOINT_MASK) clamped_defaults = wp.zeros(1, dtype=wp.int32, device=self.device) if isinstance(limits, float): raise ValueError("Joint position limits must be a tensor or array, not a float.") @@ -1614,10 +1605,8 @@ def write_joint_velocity_limit_to_sim_mask( joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - if env_mask is None: - env_mask = self._ALL_ENV_MASK - if joint_mask is None: - joint_mask = self._ALL_JOINT_MASK + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + joint_mask = self._resolve_mask(joint_mask, self._ALL_JOINT_MASK) if isinstance(limits, float): wp.launch( articulation_kernels.float_data_to_buffer_with_mask, @@ -1735,10 +1724,8 @@ def write_joint_effort_limit_to_sim_mask( joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - if env_mask is None: - env_mask = self._ALL_ENV_MASK - if joint_mask is None: - joint_mask = self._ALL_JOINT_MASK + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + joint_mask = self._resolve_mask(joint_mask, self._ALL_JOINT_MASK) if isinstance(limits, float): wp.launch( articulation_kernels.float_data_to_buffer_with_mask, @@ -1856,10 +1843,8 @@ def write_joint_armature_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks - if env_mask is None: - env_mask = self._ALL_ENV_MASK - if joint_mask is None: - joint_mask = self._ALL_JOINT_MASK + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + joint_mask = self._resolve_mask(joint_mask, self._ALL_JOINT_MASK) if isinstance(armature, float): wp.launch( articulation_kernels.float_data_to_buffer_with_mask, @@ -1974,10 +1959,8 @@ def write_joint_friction_coefficient_to_sim_mask( joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - if env_mask is None: - env_mask = self._ALL_ENV_MASK - if joint_mask is None: - joint_mask = self._ALL_JOINT_MASK + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + joint_mask = self._resolve_mask(joint_mask, self._ALL_JOINT_MASK) if isinstance(joint_friction_coeff, float): wp.launch( articulation_kernels.float_data_to_buffer_with_mask, @@ -2080,10 +2063,8 @@ def set_masses_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks - if env_mask is None: - env_mask = self._ALL_ENV_MASK - if body_mask is None: - body_mask = self._ALL_BODY_MASK + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + body_mask = self._resolve_mask(body_mask, self._ALL_BODY_MASK) self.assert_shape_and_dtype_mask(masses, (env_mask, body_mask), wp.float32, "masses") wp.launch( shared_kernels.write_2d_data_to_buffer_with_mask, @@ -2171,16 +2152,15 @@ def set_coms_mask( aligned with the body frame. Args: - coms: Center of mass position of all bodies. Shape is (num_instances, num_bodies). In warp - the expected shape is (num_instances, num_bodies), with dtype wp.vec3f. + coms: Center of mass position of all bodies. Shape is (num_instances, num_bodies, 3) or + (num_instances, num_bodies, 7) (transformf convention — only position is used). In warp + the expected shape is (num_instances, num_bodies), with dtype wp.vec3f or wp.transformf. body_mask: Body mask. If None, then all bodies are used. Shape is (num_bodies,). env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks - if env_mask is None: - env_mask = self._ALL_ENV_MASK - if body_mask is None: - body_mask = self._ALL_BODY_MASK + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + body_mask = self._resolve_mask(body_mask, self._ALL_BODY_MASK) self.assert_shape_and_dtype_mask(coms, (env_mask, body_mask), wp.vec3f, "coms") wp.launch( shared_kernels.write_body_com_position_to_buffer_mask, @@ -2263,10 +2243,8 @@ def set_inertias_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks - if env_mask is None: - env_mask = self._ALL_ENV_MASK - if body_mask is None: - body_mask = self._ALL_BODY_MASK + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + body_mask = self._resolve_mask(body_mask, self._ALL_BODY_MASK) self.assert_shape_and_dtype_mask(inertias, (env_mask, body_mask), wp.float32, "inertias", trailing_dims=(9,)) wp.launch( shared_kernels.write_body_inertia_to_buffer_mask, @@ -2349,10 +2327,8 @@ def set_joint_position_target_mask( joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - if env_mask is None: - env_mask = self._ALL_ENV_MASK - if joint_mask is None: - joint_mask = self._ALL_JOINT_MASK + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + joint_mask = self._resolve_mask(joint_mask, self._ALL_JOINT_MASK) self.assert_shape_and_dtype_mask(target, (env_mask, joint_mask), wp.float32, "target") wp.launch( shared_kernels.write_2d_data_to_buffer_with_mask, @@ -2435,10 +2411,8 @@ def set_joint_velocity_target_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - if env_mask is None: - env_mask = self._ALL_ENV_MASK - if joint_mask is None: - joint_mask = self._ALL_JOINT_MASK + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + joint_mask = self._resolve_mask(joint_mask, self._ALL_JOINT_MASK) self.assert_shape_and_dtype_mask(target, (env_mask, joint_mask), wp.float32, "target") wp.launch( shared_kernels.write_2d_data_to_buffer_with_mask, @@ -2520,10 +2494,8 @@ def set_joint_effort_target_mask( joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - if env_mask is None: - env_mask = self._ALL_ENV_MASK - if joint_mask is None: - joint_mask = self._ALL_JOINT_MASK + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + joint_mask = self._resolve_mask(joint_mask, self._ALL_JOINT_MASK) self.assert_shape_and_dtype_mask(target, (env_mask, joint_mask), wp.float32, "target") wp.launch( shared_kernels.write_2d_data_to_buffer_with_mask, @@ -3229,6 +3201,13 @@ def _initialize_impl(self): # container for data access self._data = ArticulationData(self.root_view, self.device) + # Register callback to rebind simulation data after a full reset (model/state recreation). + self._physics_ready_handle = SimulationManager.register_callback( + lambda _: self._data._create_simulation_bindings(), + PhysicsEvent.PHYSICS_READY, + name=f"articulation_rebind_{self.cfg.prim_path}", + ) + # create buffers self._create_buffers() # process configuration @@ -3244,6 +3223,13 @@ def _initialize_impl(self): # Let the articulation data know that it is fully instantiated and ready to use. self.data.is_primed = True + def _clear_callbacks(self) -> None: + """Clears all registered callbacks, including the physics-ready rebind handle.""" + super()._clear_callbacks() + if hasattr(self, "_physics_ready_handle") and self._physics_ready_handle is not None: + self._physics_ready_handle.deregister() + self._physics_ready_handle = None + def _create_buffers(self): self._ALL_INDICES = wp.array(np.arange(self.num_instances, dtype=np.int32), device=self.device) self._ALL_ENV_MASK = wp.ones((self.num_instances,), dtype=wp.bool, device=self.device) @@ -3788,6 +3774,22 @@ def _resolve_spatial_tendon_ids( return self._ALL_SPATIAL_TENDON_INDICES return spatial_tendon_ids + def _resolve_mask(self, mask: wp.array | torch.Tensor | None, full_mask: wp.array) -> wp.array: + """Resolve a mask to a warp array. + + Args: + mask: Mask. If None, then all indices are used. + + Returns: + A warp array of mask. + """ + if mask is None: + return full_mask + + if isinstance(mask, torch.Tensor): + return wp.from_torch(mask, dtype=wp.bool) + return mask + """ Deprecated methods. """ @@ -3828,8 +3830,8 @@ def write_root_state_to_sim( ) if isinstance(root_state, wp.array): raise ValueError("The root state must be a torch tensor, not a warp array.") - self.write_root_link_pose_to_sim_index(root_state[:, :7], env_ids=env_ids) - self.write_root_com_velocity_to_sim_index(root_state[:, 7:], env_ids=env_ids) + self.write_root_link_pose_to_sim_index(root_pose=root_state[:, :7], env_ids=env_ids) + self.write_root_com_velocity_to_sim_index(root_velocity=root_state[:, 7:], env_ids=env_ids) def write_root_com_state_to_sim( self, @@ -3846,8 +3848,8 @@ def write_root_com_state_to_sim( ) if isinstance(root_state, wp.array): raise ValueError("The root state must be a torch tensor, not a warp array.") - self.write_root_com_pose_to_sim_index(root_state[:, :7], env_ids=env_ids) - self.write_root_com_velocity_to_sim_index(root_state[:, 7:], env_ids=env_ids) + self.write_root_com_pose_to_sim_index(root_pose=root_state[:, :7], env_ids=env_ids) + self.write_root_com_velocity_to_sim_index(root_velocity=root_state[:, 7:], env_ids=env_ids) def write_root_link_state_to_sim( self, @@ -3864,8 +3866,8 @@ def write_root_link_state_to_sim( ) if isinstance(root_state, wp.array): raise ValueError("The root state must be a torch tensor, not a warp array.") - self.write_root_link_pose_to_sim_index(root_state[:, :7], env_ids=env_ids) - self.write_root_link_velocity_to_sim_index(root_state[:, 7:], env_ids=env_ids) + self.write_root_link_pose_to_sim_index(root_pose=root_state[:, :7], env_ids=env_ids) + self.write_root_link_velocity_to_sim_index(root_velocity=root_state[:, 7:], env_ids=env_ids) def write_joint_state_to_sim( self, diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py index 5d650fed9d14..dce2c02bd2bb 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py @@ -64,6 +64,7 @@ def __init__(self, root_view: ArticulationView, device: str): # Set initial time stamp self._sim_timestamp = 0.0 self._is_primed = False + self._fk_timestamp = 0.0 # Convert to direction vector gravity = wp.to_torch(SimulationManager.get_model().gravity)[0] @@ -109,6 +110,9 @@ def update(self, dt: float) -> None: """ # update the simulation timestamp self._sim_timestamp += dt + # FK is current after a sim step — keep fk_timestamp in sync unless it was explicitly invalidated + if self._fk_timestamp >= 0.0: + self._fk_timestamp = self._sim_timestamp # Trigger an update of the joint and body com acceleration buffers at a higher frequency # since we do finite differencing. self.joint_acc @@ -563,7 +567,7 @@ def root_link_vel_w(self) -> wp.array: dim=self._num_instances, inputs=[ self.root_com_vel_w, - self.root_link_quat_w, + self.root_link_pose_w, self.body_com_pos_b, ], outputs=[ @@ -644,6 +648,9 @@ def body_link_pose_w(self) -> wp.array: This quantity is the pose of the articulation links' actor frame relative to the world. The orientation is provided in (x, y, z, w) format. """ + if self._fk_timestamp < self._sim_timestamp: + SimulationManager.forward() + self._fk_timestamp = self._sim_timestamp return self._sim_bind_body_link_pose_w @property @@ -798,7 +805,7 @@ def body_incoming_joint_wrench_b(self) -> wp.array: .. _PhysX documentation: https://nvidia-omniverse.github.io/PhysX/physx/5.5.1/docs/Articulations.html#link-incoming-joint-force .. _PhysX Tensor API: https://docs.omniverse.nvidia.com/kit/docs/omni_physics/latest/extensions/runtime/source/omni.physics.tensors/docs/api/python.html#omni.physics.tensors.impl.api.ArticulationView.get_link_incoming_joint_force """ - raise NotImplementedError + raise NotImplementedError("Not implemented for Newton") """ Joint state properties. @@ -1203,14 +1210,14 @@ def _create_simulation_bindings(self) -> None: indexed. Newton willing this is the case all the time, but we should pay attention to this if things look off. """ # Short-hand for the number of instances, number of links, and number of joints. - n_view = self._root_view.count - n_dof = self._root_view.joint_dof_count + self._num_instances = self._root_view.count + self._num_joints = self._root_view.joint_dof_count + self._num_bodies = self._root_view.link_count + self._num_fixed_tendons = 0 # self._root_view.max_fixed_tendons + self._num_spatial_tendons = 0 # self._root_view.max_spatial_tendons # -- root properties - if self._root_view.is_fixed_base: - self._sim_bind_root_link_pose_w = self._root_view.get_root_transforms(SimulationManager.get_state_0())[:, 0] - else: - self._sim_bind_root_link_pose_w = self._root_view.get_root_transforms(SimulationManager.get_state_0())[:, 0] + self._sim_bind_root_link_pose_w = self._root_view.get_root_transforms(SimulationManager.get_state_0())[:, 0] self._sim_bind_root_com_vel_w = self._root_view.get_root_velocities(SimulationManager.get_state_0()) if self._sim_bind_root_com_vel_w is not None: if self._root_view.is_fixed_base: @@ -1224,12 +1231,34 @@ def _create_simulation_bindings(self) -> None: if self._sim_bind_body_com_vel_w is not None: self._sim_bind_body_com_vel_w = self._sim_bind_body_com_vel_w[:, 0] self._sim_bind_body_mass = self._root_view.get_attribute("body_mass", SimulationManager.get_model())[:, 0] - self._sim_bind_body_inertia = self._root_view.get_attribute("body_inertia", SimulationManager.get_model())[:, 0] + # Newton stores body_inertia as (N, B, 3, 3) float32 (row-major 3x3 matrix per body). + # Reinterpret as (N, B, 9) float32 so the public interface matches the documented shape and + # so that the write_body_inertia_to_buffer_* kernels (which expect wp.array3d) receive a 3D array. + # The 9 elements of each 3x3 row-major matrix are contiguous in memory, so we keep the same + # pointer and outer strides but collapse the two trailing dimensions into one. + _body_inertia_raw = self._root_view.get_attribute("body_inertia", SimulationManager.get_model())[:, 0] + if _body_inertia_raw.ptr is not None and _body_inertia_raw.ndim == 4: + self._sim_bind_body_inertia = wp.array( + ptr=_body_inertia_raw.ptr, + dtype=wp.float32, + shape=(_body_inertia_raw.shape[0], _body_inertia_raw.shape[1], 9), + strides=(_body_inertia_raw.strides[0], _body_inertia_raw.strides[1], _body_inertia_raw.strides[3]), + device=_body_inertia_raw.device, + copy=False, + ) + else: + self._sim_bind_body_inertia = _body_inertia_raw self._sim_bind_body_external_wrench = self._root_view.get_attribute("body_f", SimulationManager.get_state_0())[ :, 0 ] + try: + self._sim_bind_body_parent_f = self._root_view.get_attribute( + "body_parent_f", SimulationManager.get_state_0() + )[:, 0] + except Exception: + self._sim_bind_body_parent_f = None # -- joint properties - if n_dof > 0: + if self._num_joints > 0: self._sim_bind_joint_pos_limits_lower = self._root_view.get_attribute( "joint_limit_lower", SimulationManager.get_model() )[:, 0] @@ -1269,31 +1298,40 @@ def _create_simulation_bindings(self) -> None: )[:, 0] else: # No joints (e.g., free-floating rigid body) - set bindings to empty arrays - self._sim_bind_joint_pos_limits_lower = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) - self._sim_bind_joint_pos_limits_upper = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) - self._sim_bind_joint_stiffness_sim = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) - self._sim_bind_joint_damping_sim = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) - self._sim_bind_joint_armature = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) - self._sim_bind_joint_friction_coeff = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) - self._sim_bind_joint_vel_limits_sim = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) - self._sim_bind_joint_effort_limits_sim = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) - self._sim_bind_joint_pos = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) - self._sim_bind_joint_vel = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) - self._sim_bind_joint_effort = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) - self._sim_bind_joint_position_target = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) - self._sim_bind_joint_velocity_target = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) + self._sim_bind_joint_pos_limits_lower = wp.zeros( + (self._num_instances, 0), dtype=wp.float32, device=self.device + ) + self._sim_bind_joint_pos_limits_upper = wp.zeros( + (self._num_instances, 0), dtype=wp.float32, device=self.device + ) + self._sim_bind_joint_stiffness_sim = wp.zeros( + (self._num_instances, 0), dtype=wp.float32, device=self.device + ) + self._sim_bind_joint_damping_sim = wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) + self._sim_bind_joint_armature = wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) + self._sim_bind_joint_friction_coeff = wp.zeros( + (self._num_instances, 0), dtype=wp.float32, device=self.device + ) + self._sim_bind_joint_vel_limits_sim = wp.zeros( + (self._num_instances, 0), dtype=wp.float32, device=self.device + ) + self._sim_bind_joint_effort_limits_sim = wp.zeros( + (self._num_instances, 0), dtype=wp.float32, device=self.device + ) + self._sim_bind_joint_pos = wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) + self._sim_bind_joint_vel = wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) + self._sim_bind_joint_effort = wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) + self._sim_bind_joint_position_target = wp.zeros( + (self._num_instances, 0), dtype=wp.float32, device=self.device + ) + self._sim_bind_joint_velocity_target = wp.zeros( + (self._num_instances, 0), dtype=wp.float32, device=self.device + ) def _create_buffers(self) -> None: """Create buffers for the root data.""" super()._create_buffers() - # Short-hand for the number of instances, number of links, and number of joints. - self._num_instances = self._root_view.count - self._num_joints = self._root_view.joint_dof_count - self._num_bodies = self._root_view.link_count - self._num_fixed_tendons = 0 # self._root_view.max_fixed_tendons - self._num_spatial_tendons = 0 # self._root_view.max_spatial_tendons - # Initialize history for finite differencing. If the articulation is fixed, the root com velocity is not # available, so we use zeros. if self._root_view.get_root_velocities(SimulationManager.get_state_0()) is None: @@ -1394,10 +1432,8 @@ def _create_buffers(self) -> None: self._joint_acc = TimestampedBuffer( shape=(self._num_instances, self._num_joints), dtype=wp.float32, device=self.device ) - # self._body_incoming_joint_wrench_b = TimestampedWarpBuffer( - # shape=(self._num_instances, self._num_joints), dtype=wp.spatial_vectorf, device=self.device - # ) # Empty memory pre-allocations + self._body_incoming_joint_wrench_b = None self._root_link_lin_vel_b = None self._root_link_ang_vel_b = None self._root_com_lin_vel_b = None diff --git a/source/isaaclab_newton/isaaclab_newton/assets/kernels.py b/source/isaaclab_newton/isaaclab_newton/assets/kernels.py index 436bd4cf1f6a..9a9aaf402d87 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/kernels.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/kernels.py @@ -161,7 +161,7 @@ def set_state_velocities_func( def get_link_velocity_in_com_frame_func( link_velocity_w: wp.spatial_vectorf, link_pose_w: wp.transformf, - body_com_pose_b: wp.transformf, + body_com_pos_b: wp.vec3f, ): """Compute COM velocity from link velocity by accounting for the COM offset. @@ -171,7 +171,7 @@ def get_link_velocity_in_com_frame_func( Args: link_velocity_w: Link spatial velocity in world frame (angular, linear). link_pose_w: Link pose in world frame. - body_com_pose_b: COM pose in body (link) frame. + body_com_pos_b: COM position in body (link) frame. Returns: COM spatial velocity in world frame (angular, linear). @@ -180,7 +180,7 @@ def get_link_velocity_in_com_frame_func( wp.spatial_top(link_velocity_w) + wp.cross( wp.spatial_bottom(link_velocity_w), - wp.quat_rotate(wp.transform_get_rotation(link_pose_w), wp.transform_get_translation(body_com_pose_b)), + wp.quat_rotate(wp.transform_get_rotation(link_pose_w), body_com_pos_b), ), wp.spatial_bottom(link_velocity_w), ) @@ -189,28 +189,24 @@ def get_link_velocity_in_com_frame_func( @wp.func def get_com_pose_in_link_frame_func( com_pose_w: wp.transformf, - com_pose_b: wp.transformf, + com_pos_b: wp.vec3f, ): """Compute link pose in world frame from COM pose by inverting the body-frame COM offset. This is the inverse of ``get_com_pose_from_link_pose_func``. Given the COM pose in - world frame and the COM offset in body frame, it recovers the link pose in world frame. + world frame and the COM position offset in body frame, it recovers the link pose in + world frame. Newton COM always has identity orientation, so only position is needed. Args: com_pose_w: COM pose in world frame. - com_pose_b: COM pose in body (link) frame. + com_pos_b: COM position in body (link) frame. Returns: Link pose in world frame. """ - T2 = wp.transform( - wp.quat_rotate( - wp.quat_inverse(wp.transform_get_rotation(com_pose_b)), -wp.transform_get_translation(com_pose_b) - ), - wp.quat_inverse(wp.transform_get_rotation(com_pose_b)), - ) - link_pose_w = com_pose_w * T2 - return link_pose_w + com_rot_w = wp.transform_get_rotation(com_pose_w) + link_pos_w = wp.transform_get_translation(com_pose_w) - wp.quat_rotate(com_rot_w, com_pos_b) + return wp.transform(link_pos_w, com_rot_w) """ @@ -535,7 +531,7 @@ def set_root_link_pose_to_sim_mask( @wp.kernel def set_root_com_pose_to_sim_index( data: wp.array(dtype=wp.transformf), - body_com_pose_b: wp.array2d(dtype=wp.transformf), + body_com_pos_b: wp.array2d(dtype=wp.vec3f), env_ids: wp.array(dtype=wp.int32), root_com_pose_w: wp.array(dtype=wp.transformf), root_link_pose_w: wp.array(dtype=wp.transformf), @@ -551,7 +547,7 @@ def set_root_com_pose_to_sim_index( Args: data: Input array of root COM poses. Shape is (num_selected_envs,). - body_com_pose_b: Input array of body COM poses in body frame. Shape is + body_com_pos_b: Input array of body COM positions in body frame. Shape is (num_envs, num_bodies). Only the first body (index 0) is used for the root. env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). root_com_pose_w: Output array where root COM poses are written. Shape is (num_envs,). @@ -570,7 +566,7 @@ def set_root_com_pose_to_sim_index( root_com_state_w[env_ids[i]] = set_state_transforms_func(root_com_state_w[env_ids[i]], data[i]) # Get the com pose in the link frame root_link_pose_w[env_ids[i]] = get_com_pose_in_link_frame_func( - root_com_pose_w[env_ids[i]], body_com_pose_b[env_ids[i], 0] + root_com_pose_w[env_ids[i]], body_com_pos_b[env_ids[i], 0] ) if root_link_state_w: root_link_state_w[env_ids[i]] = set_state_transforms_func( @@ -583,7 +579,7 @@ def set_root_com_pose_to_sim_index( @wp.kernel def set_root_com_pose_to_sim_mask( data: wp.array(dtype=wp.transformf), - body_com_pose_b: wp.array2d(dtype=wp.transformf), + body_com_pos_b: wp.array2d(dtype=wp.vec3f), env_mask: wp.array(dtype=wp.bool), root_com_pose_w: wp.array(dtype=wp.transformf), root_link_pose_w: wp.array(dtype=wp.transformf), @@ -599,7 +595,7 @@ def set_root_com_pose_to_sim_mask( Args: data: Input array of root COM poses. Shape is (num_instances,). - body_com_pose_b: Input array of body COM poses in body frame. Shape is + body_com_pos_b: Input array of body COM positions in body frame. Shape is (num_envs, num_bodies). Only the first body (index 0) is used for the root. env_mask: Input array of environment mask. Shape is (num_instances,). root_com_pose_w: Output array where root COM poses are written. Shape is (num_envs,). @@ -618,7 +614,7 @@ def set_root_com_pose_to_sim_mask( if root_com_state_w: root_com_state_w[i] = set_state_transforms_func(root_com_state_w[i], data[i]) # Get the com pose in the link frame - root_link_pose_w[i] = get_com_pose_in_link_frame_func(root_com_pose_w[i], body_com_pose_b[i, 0]) + root_link_pose_w[i] = get_com_pose_in_link_frame_func(root_com_pose_w[i], body_com_pos_b[i, 0]) if root_link_state_w: root_link_state_w[i] = set_state_transforms_func(root_link_state_w[i], root_link_pose_w[i]) if root_state_w: @@ -707,7 +703,7 @@ def set_root_com_velocity_to_sim_mask( @wp.kernel def set_root_link_velocity_to_sim_index( data: wp.array(dtype=wp.spatial_vectorf), - body_com_pose_b: wp.array2d(dtype=wp.transformf), + body_com_pos_b: wp.array2d(dtype=wp.vec3f), link_pose_w: wp.array(dtype=wp.transformf), env_ids: wp.array(dtype=wp.int32), num_bodies: wp.int32, @@ -726,7 +722,7 @@ def set_root_link_velocity_to_sim_index( Args: data: Input array of root link spatial velocities. Shape is (num_selected_envs,). - body_com_pose_b: Input array of body COM poses in body frame. Shape is + body_com_pos_b: Input array of body COM positions in body frame. Shape is (num_envs, num_bodies). Only the first body (index 0) is used for the root. link_pose_w: Input array of root link poses in world frame. Shape is (num_envs,). env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). @@ -750,7 +746,7 @@ def set_root_link_velocity_to_sim_index( root_link_state_w[env_ids[i]] = set_state_velocities_func(root_link_state_w[env_ids[i]], data[i]) # Get the link velocity in the com frame root_com_velocity_w[env_ids[i]] = get_link_velocity_in_com_frame_func( - root_link_velocity_w[env_ids[i]], link_pose_w[env_ids[i]], body_com_pose_b[env_ids[i], 0] + root_link_velocity_w[env_ids[i]], link_pose_w[env_ids[i]], body_com_pos_b[env_ids[i], 0] ) if root_com_state_w: root_com_state_w[env_ids[i]] = set_state_velocities_func( @@ -766,7 +762,7 @@ def set_root_link_velocity_to_sim_index( @wp.kernel def set_root_link_velocity_to_sim_mask( data: wp.array(dtype=wp.spatial_vectorf), - body_com_pose_b: wp.array2d(dtype=wp.transformf), + body_com_pos_b: wp.array2d(dtype=wp.vec3f), link_pose_w: wp.array(dtype=wp.transformf), env_mask: wp.array(dtype=wp.bool), num_bodies: wp.int32, @@ -785,7 +781,7 @@ def set_root_link_velocity_to_sim_mask( Args: data: Input array of root link spatial velocities. Shape is (num_instances,). - body_com_pose_b: Input array of body COM poses in body frame. Shape is + body_com_pos_b: Input array of body COM positions in body frame. Shape is (num_envs, num_bodies). Only the first body (index 0) is used for the root. link_pose_w: Input array of root link poses in world frame. Shape is (num_envs,). env_mask: Input array of environment mask. Shape is (num_instances,). @@ -810,7 +806,7 @@ def set_root_link_velocity_to_sim_mask( root_link_state_w[i] = set_state_velocities_func(root_link_state_w[i], data[i]) # Get the link velocity in the com frame root_com_velocity_w[i] = get_link_velocity_in_com_frame_func( - root_link_velocity_w[i], link_pose_w[i], body_com_pose_b[i, 0] + root_link_velocity_w[i], link_pose_w[i], body_com_pos_b[i, 0] ) if root_com_state_w: root_com_state_w[i] = set_state_velocities_func(root_com_state_w[i], root_com_velocity_w[i]) @@ -880,7 +876,7 @@ def set_body_link_pose_to_sim( @wp.kernel def set_body_com_pose_to_sim( data: wp.array2d(dtype=wp.transformf), - body_com_pose_b: wp.array2d(dtype=wp.transformf), + body_com_pos_b: wp.array2d(dtype=wp.vec3f), env_ids: wp.array(dtype=wp.int32), body_ids: wp.array(dtype=wp.int32), from_mask: bool, @@ -899,7 +895,7 @@ def set_body_com_pose_to_sim( Args: data: Input array of body COM poses. Shape is (num_envs, num_bodies) or (num_selected_envs, num_selected_bodies) depending on from_mask. - body_com_pose_b: Input array of body COM poses in body frame. Shape is + body_com_pos_b: Input array of body COM positions in body frame. Shape is (num_envs, num_bodies). env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). body_ids: Input array of body indices to write to. Shape is (num_selected_bodies,). @@ -930,7 +926,7 @@ def set_body_com_pose_to_sim( ) # Get the link pose from com pose body_link_pose_w[env_ids[i], body_ids[j]] = get_com_pose_in_link_frame_func( - body_com_pose_w[env_ids[i], body_ids[j]], body_com_pose_b[env_ids[i], body_ids[j]] + body_com_pose_w[env_ids[i], body_ids[j]], body_com_pos_b[env_ids[i], body_ids[j]] ) if body_link_state_w: body_link_state_w[env_ids[i], body_ids[j]] = set_state_transforms_func( @@ -1002,7 +998,7 @@ def set_body_com_velocity_to_sim( @wp.kernel def set_body_link_velocity_to_sim( data: wp.array2d(dtype=wp.spatial_vectorf), - body_com_pose_b: wp.array2d(dtype=wp.transformf), + body_com_pos_b: wp.array2d(dtype=wp.vec3f), body_link_pose_w: wp.array2d(dtype=wp.transformf), env_ids: wp.array(dtype=wp.int32), body_ids: wp.array(dtype=wp.int32), @@ -1023,7 +1019,7 @@ def set_body_link_velocity_to_sim( Args: data: Input array of body link spatial velocities. Shape is (num_envs, num_bodies) or (num_selected_envs, num_selected_bodies) depending on from_mask. - body_com_pose_b: Input array of body COM poses in body frame. Shape is + body_com_pos_b: Input array of body COM positions in body frame. Shape is (num_envs, num_bodies). body_link_pose_w: Input array of body link poses in world frame. Shape is (num_envs, num_bodies). @@ -1060,7 +1056,7 @@ def set_body_link_velocity_to_sim( body_com_velocity_w[env_ids[i], body_ids[j]] = get_link_velocity_in_com_frame_func( body_link_velocity_w[env_ids[i], body_ids[j]], body_link_pose_w[env_ids[i], body_ids[j]], - body_com_pose_b[env_ids[i], body_ids[j]], + body_com_pos_b[env_ids[i], body_ids[j]], ) if body_com_state_w: body_com_state_w[env_ids[i], body_ids[j]] = set_state_velocities_func( @@ -1385,16 +1381,16 @@ def make_dummy_body_com_pose_b( body_com_pos_b: wp.array2d(dtype=wp.vec3f), body_com_pose_b: wp.array2d(dtype=wp.transformf), ): - """Make a dummy body COM pose in body frame. + """Make a body COM pose from position by appending an identity quaternion. - This kernel makes a dummy body COM pose in body frame. + Needed by the ``body_com_pose_b`` property to match the base API that returns + ``wp.transformf`` (pos + quat). Args: body_com_pos_b: Input array of body COM positions in body frame. Shape is (num_envs, num_bodies). body_com_pose_b: Output array where body COM poses are written. Shape is (num_envs, num_bodies). """ i, j = wp.tid() - # Concatenate the position and a unit quaternion body_com_pose_b[i, j] = wp.transformf(body_com_pos_b[i, j], wp.quatf(0.0, 0.0, 0.0, 1.0)) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py index 40f01044b510..80ac693c4842 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py @@ -21,6 +21,7 @@ import isaaclab.sim as sim_utils import isaaclab.utils.string as string_utils from isaaclab.assets.rigid_object.base_rigid_object import BaseRigidObject +from isaaclab.physics import PhysicsEvent from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_newton.assets import kernels as shared_kernels @@ -434,7 +435,7 @@ def write_root_com_pose_to_sim_index( dim=env_ids.shape[0], inputs=[ root_pose, - self.data.body_com_pose_b, + self.data.body_com_pos_b, env_ids, ], outputs=[ @@ -485,7 +486,7 @@ def write_root_com_pose_to_sim_mask( dim=root_pose.shape[0], inputs=[ root_pose, - self.data.body_com_pose_b, + self.data.body_com_pos_b, env_mask, ], outputs=[ @@ -641,7 +642,7 @@ def write_root_link_velocity_to_sim_index( dim=env_ids.shape[0], inputs=[ root_velocity, - self.data.body_com_pose_b, + self.data.body_com_pos_b, self.data.root_link_pose_w, env_ids, 1, @@ -656,9 +657,12 @@ def write_root_link_velocity_to_sim_index( ], device=self.device, ) - self.data._root_link_state_w.timestamp = -1.0 - self.data._root_state_w.timestamp = -1.0 - self.data._root_com_state_w.timestamp = -1.0 + if self.data._root_link_state_w is not None: + self.data._root_link_state_w.timestamp = -1.0 + if self.data._root_state_w is not None: + self.data._root_state_w.timestamp = -1.0 + if self.data._root_com_state_w is not None: + self.data._root_com_state_w.timestamp = -1.0 def write_root_link_velocity_to_sim_mask( self, @@ -693,7 +697,7 @@ def write_root_link_velocity_to_sim_mask( dim=root_velocity.shape[0], inputs=[ root_velocity, - self.data.body_com_pose_b, + self.data.body_com_pos_b, self.data.root_link_pose_w, env_mask, 1, @@ -708,9 +712,12 @@ def write_root_link_velocity_to_sim_mask( ], device=self.device, ) - self.data._root_link_state_w.timestamp = -1.0 - self.data._root_state_w.timestamp = -1.0 - self.data._root_com_state_w.timestamp = -1.0 + if self.data._root_link_state_w is not None: + self.data._root_link_state_w.timestamp = -1.0 + if self.data._root_state_w is not None: + self.data._root_state_w.timestamp = -1.0 + if self.data._root_com_state_w is not None: + self.data._root_com_state_w.timestamp = -1.0 """ Operations - Setters. @@ -986,8 +993,6 @@ def set_inertias_mask( """ def _initialize_impl(self): - # obtain global simulation view - self._physics_sim_view = SimulationManager.get_physics_sim_view() # obtain the first prim in the regex expression (all others are assumed to be a copy of this) template_prim = sim_utils.find_first_matching_prim(self.cfg.prim_path) if template_prim is None: @@ -1030,11 +1035,11 @@ def _initialize_impl(self): root_prim_path = root_prims[0].GetPath().pathString root_prim_path_expr = self.cfg.prim_path + root_prim_path[len(template_prim_path) :] # -- object view - self._root_view = self._physics_sim_view.create_rigid_body_view(root_prim_path_expr.replace(".*", "*")) - - # check if the rigid body was created - if self.root_view._backend is None: - raise RuntimeError(f"Failed to create rigid body at: {self.cfg.prim_path}. Please check PhysX logs.") + self._root_view = ArticulationView( + SimulationManager.get_model(), + root_prim_path_expr.replace(".*", "*"), + verbose=False, + ) # log information about the rigid body logger.info(f"Rigid body initialized at: {self.cfg.prim_path} with root '{root_prim_path_expr}'.") @@ -1045,6 +1050,13 @@ def _initialize_impl(self): # container for data access self._data = RigidObjectData(self.root_view, self.device) + # Register callback to rebind simulation data after a full reset (model/state recreation). + self._physics_ready_handle = SimulationManager.register_callback( + lambda _: self._data._create_simulation_bindings(), + PhysicsEvent.PHYSICS_READY, + name=f"rigid_object_rebind_{self.cfg.prim_path}", + ) + # create buffers self._create_buffers() # process configuration @@ -1054,6 +1066,13 @@ def _initialize_impl(self): # Let the rigid object data know that it is fully instantiated and ready to use. self.data.is_primed = True + def _clear_callbacks(self) -> None: + """Clears all registered callbacks, including the physics-ready rebind handle.""" + super()._clear_callbacks() + if hasattr(self, "_physics_ready_handle") and self._physics_ready_handle is not None: + self._physics_ready_handle.deregister() + self._physics_ready_handle = None + def _create_buffers(self): """Create buffers for storing data.""" # constants diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object_data.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object_data.py index 77ce20385455..7f179d9d5a20 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object_data.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object_data.py @@ -203,7 +203,7 @@ def root_link_vel_w(self) -> wp.array: inputs=[ self.root_com_vel_w, self.root_link_pose_w, - self.body_com_pose_b, + self.body_com_pos_b, ], outputs=[ self._root_link_vel_w.data, @@ -229,7 +229,7 @@ def root_com_pose_w(self) -> wp.array: dim=self._num_instances, inputs=[ self.root_link_pose_w, - self.body_com_pose_b, + self.body_com_pos_b, ], outputs=[ self._root_com_pose_w.data, @@ -261,7 +261,7 @@ def body_mass(self) -> wp.array: Shape is (num_instances, 1, 1), dtype = wp.float32. In torch this resolves to (num_instances, 1, 1). """ - return self._body_mass + return self._sim_bind_body_mass @property def body_inertia(self) -> wp.array: @@ -270,7 +270,7 @@ def body_inertia(self) -> wp.array: Shape is (num_instances, 1, 9), dtype = wp.float32. In torch this resolves to (num_instances, 1, 9). """ - return self._body_inertia + return self._sim_bind_body_inertia @property def body_link_pose_w(self) -> wp.array: @@ -427,6 +427,10 @@ def root_link_lin_vel_b(self) -> wp.array: This quantity is the linear velocity of the actor frame of the root rigid body frame with respect to the rigid body's actor frame. """ + if self._root_link_lin_vel_b is None: + self._root_link_lin_vel_b = TimestampedBuffer( + shape=(self._num_instances,), dtype=wp.vec3f, device=self.device + ) if self._root_link_lin_vel_b.timestamp < self._sim_timestamp: wp.launch( shared_kernels.quat_apply_inverse_1D_kernel, @@ -446,6 +450,10 @@ def root_link_ang_vel_b(self) -> wp.array: This quantity is the angular velocity of the actor frame of the root rigid body frame with respect to the rigid body's actor frame. """ + if self._root_link_ang_vel_b is None: + self._root_link_ang_vel_b = TimestampedBuffer( + shape=(self._num_instances,), dtype=wp.vec3f, device=self.device + ) if self._root_link_ang_vel_b.timestamp < self._sim_timestamp: wp.launch( shared_kernels.quat_apply_inverse_1D_kernel, @@ -465,6 +473,10 @@ def root_com_lin_vel_b(self) -> wp.array: This quantity is the linear velocity of the root rigid body's center of mass frame with respect to the rigid body's actor frame. """ + if self._root_com_lin_vel_b is None: + self._root_com_lin_vel_b = TimestampedBuffer( + shape=(self._num_instances,), dtype=wp.vec3f, device=self.device + ) if self._root_com_lin_vel_b.timestamp < self._sim_timestamp: wp.launch( shared_kernels.quat_apply_inverse_1D_kernel, @@ -484,6 +496,10 @@ def root_com_ang_vel_b(self) -> wp.array: This quantity is the angular velocity of the root rigid body's center of mass frame with respect to the rigid body's actor frame. """ + if self._root_com_ang_vel_b is None: + self._root_com_ang_vel_b = TimestampedBuffer( + shape=(self._num_instances,), dtype=wp.vec3f, device=self.device + ) if self._root_com_ang_vel_b.timestamp < self._sim_timestamp: wp.launch( shared_kernels.quat_apply_inverse_1D_kernel, @@ -681,6 +697,10 @@ def _create_simulation_bindings(self) -> None: .. caution:: This is possible if and only if the properties that we access are strided from newton and not indexed. Newton willing this is the case all the time, but we should pay attention to this if things look off. """ + # Short-hand for the number of instances, number of links, and number of joints. + self._num_instances = self._root_view.count + self._num_bodies = self._root_view.link_count + # -- root properties if self._root_view.is_fixed_base: self._sim_bind_root_link_pose_w = self._root_view.get_root_transforms(SimulationManager.get_state_0())[ @@ -699,7 +719,18 @@ def _create_simulation_bindings(self) -> None: self._sim_bind_body_link_pose_w = self._root_view.get_link_transforms(SimulationManager.get_state_0())[:, 0] self._sim_bind_body_com_vel_w = self._root_view.get_link_velocities(SimulationManager.get_state_0())[:, 0] self._sim_bind_body_mass = self._root_view.get_attribute("body_mass", SimulationManager.get_model())[:, 0] - self._sim_bind_body_inertia = self._root_view.get_attribute("body_inertia", SimulationManager.get_model())[:, 0] + _inertia_mat33 = self._root_view.get_attribute("body_inertia", SimulationManager.get_model())[:, 0] + if _inertia_mat33.ptr is not None and _inertia_mat33.ndim == 4: + self._sim_bind_body_inertia = wp.array( + ptr=_inertia_mat33.ptr, + dtype=wp.float32, + shape=(_inertia_mat33.shape[0], _inertia_mat33.shape[1], 9), + strides=(_inertia_mat33.strides[0], _inertia_mat33.strides[1], _inertia_mat33.strides[3]), + device=_inertia_mat33.device, + copy=False, + ) + else: + self._sim_bind_body_inertia = _inertia_mat33 self._sim_bind_body_external_wrench = self._root_view.get_attribute("body_f", SimulationManager.get_state_0())[ :, 0 ] @@ -725,10 +756,6 @@ def _create_buffers(self) -> None: self._default_root_pose = wp.zeros((self._num_instances,), dtype=wp.transformf, device=self.device) self._default_root_vel = wp.zeros((self._num_instances,), dtype=wp.spatial_vectorf, device=self.device) - # -- Body properties - self._body_mass = wp.clone(self._root_view.get_masses(), device=self.device) - self._body_inertia = wp.clone(self._root_view.get_inertias(), device=self.device) - # Initialize history for finite differencing self._previous_body_com_vel = wp.clone(self._root_view.get_link_velocities(SimulationManager.get_state_0()))[ :, 0 @@ -758,6 +785,9 @@ def _create_buffers(self) -> None: self._body_com_acc_w = TimestampedBuffer( shape=(self._num_instances, 1), dtype=wp.spatial_vectorf, device=self.device ) + self._body_com_pose_b = TimestampedBuffer( + shape=(self._num_instances, 1), dtype=wp.transformf, device=self.device + ) # Empty memory pre-allocations self._root_state_w = None self._root_link_state_w = None @@ -785,7 +815,6 @@ def _create_buffers(self) -> None: self._body_com_ang_vel_w = None self._body_com_lin_acc_w = None self._body_com_ang_acc_w = None - self._body_com_pose_b = None """ Internal helpers. @@ -802,7 +831,7 @@ def _get_pos_from_transform(self, source: wp.array | None, transform: wp.array) """ # Check if we already created the lazy buffer. if source is None: - if transform.is_contiguous(): + if transform.is_contiguous: # Check if the array is contiguous. If so, we can just return a strided array. # Then this update becomes a no-op. return wp.array( @@ -813,34 +842,27 @@ def _get_pos_from_transform(self, source: wp.array | None, transform: wp.array) device=self.device, ) else: - # If the array is no contiguous, we need to create a new array to write to. - source = wp.zeros((transform.shape[0], 3), dtype=wp.vec3f, device=self.device) + # If the array is not contiguous, we need to create a new array to write to. + # Shape matches transform.shape since each element is vec3f (already contains 3 floats) + source = wp.zeros(transform.shape, dtype=wp.vec3f, device=self.device) # If the array is not contiguous, we need to launch the kernel to get the position part of the transform. - if not transform.is_contiguous(): - # Launch the right kernel based on the shape of the source array. - if len(source.shape) > 1: + if not transform.is_contiguous: + # Launch the right kernel based on the shape of the transform array. + if len(transform.shape) > 1: wp.launch( shared_kernels.split_transform_to_pos_2d, - dim=source.shape, - inputs=[ - source, - ], - outputs=[ - source, - ], + dim=transform.shape, + inputs=[transform], + outputs=[source], device=self.device, ) else: wp.launch( shared_kernels.split_transform_to_pos_1d, - dim=source.shape, - inputs=[ - source, - ], - outputs=[ - source, - ], + dim=transform.shape, + inputs=[transform], + outputs=[source], device=self.device, ) return source @@ -856,7 +878,7 @@ def _get_quat_from_transform(self, source: wp.array | None, transform: wp.array) """ # Check if we already created the lazy buffer. if source is None: - if transform.is_contiguous(): + if transform.is_contiguous: # Check if the array is contiguous. If so, we can just return a strided array. # Then this update becomes a no-op. return wp.array( @@ -867,34 +889,27 @@ def _get_quat_from_transform(self, source: wp.array | None, transform: wp.array) device=self.device, ) else: - # If the array is no contiguous, we need to create a new array to write to. - source = wp.zeros((transform.shape[0], 4), dtype=wp.quatf, device=self.device) + # If the array is not contiguous, we need to create a new array to write to. + # Shape matches transform.shape since each element is quatf (already contains 4 floats) + source = wp.zeros(transform.shape, dtype=wp.quatf, device=self.device) # If the array is not contiguous, we need to launch the kernel to get the quaternion part of the transform. - if not transform.is_contiguous(): - # Launch the right kernel based on the shape of the source array. - if len(source.shape) > 1: + if not transform.is_contiguous: + # Launch the right kernel based on the shape of the transform array. + if len(transform.shape) > 1: wp.launch( shared_kernels.split_transform_to_quat_2d, - dim=source.shape, - inputs=[ - source, - ], - outputs=[ - source, - ], + dim=transform.shape, + inputs=[transform], + outputs=[source], device=self.device, ) else: wp.launch( shared_kernels.split_transform_to_quat_1d, - dim=source.shape, - inputs=[ - source, - ], - outputs=[ - source, - ], + dim=transform.shape, + inputs=[transform], + outputs=[source], device=self.device, ) # Return the source array. (no-op if the array is contiguous.) @@ -913,7 +928,7 @@ def _get_top_from_spatial_vector(self, source: wp.array | None, spatial_vector: """ # Check if we already created the lazy buffer. if source is None: - if spatial_vector.is_contiguous(): + if spatial_vector.is_contiguous: # Check if the array is contiguous. If so, we can just return a strided array. # Then this update becomes a no-op. return wp.array( @@ -924,34 +939,27 @@ def _get_top_from_spatial_vector(self, source: wp.array | None, spatial_vector: device=self.device, ) else: - # If the array is no contiguous, we need to create a new array to write to. - source = wp.zeros((spatial_vector.shape[0], 3), dtype=wp.vec3f, device=self.device) + # If the array is not contiguous, we need to create a new array to write to. + # Shape matches spatial_vector.shape since each element is vec3f (already contains 3 floats) + source = wp.zeros(spatial_vector.shape, dtype=wp.vec3f, device=self.device) # If the array is not contiguous, we need to launch the kernel to get the top part of the spatial vector. - if not spatial_vector.is_contiguous(): - # Launch the right kernel based on the shape of the source array. - if len(source.shape) > 1: + if not spatial_vector.is_contiguous: + # Launch the right kernel based on the shape of the spatial_vector array. + if len(spatial_vector.shape) > 1: wp.launch( shared_kernels.split_spatial_vector_to_top_2d, - dim=source.shape, - inputs=[ - source, - ], - outputs=[ - source, - ], + dim=spatial_vector.shape, + inputs=[spatial_vector], + outputs=[source], device=self.device, ) else: wp.launch( shared_kernels.split_spatial_vector_to_top_1d, - dim=source.shape, - inputs=[ - source, - ], - outputs=[ - source, - ], + dim=spatial_vector.shape, + inputs=[spatial_vector], + outputs=[source], device=self.device, ) # Return the source array. (no-op if the array is contiguous.) @@ -970,7 +978,7 @@ def _get_bottom_from_spatial_vector(self, source: wp.array | None, spatial_vecto """ # Check if we already created the lazy buffer. if source is None: - if spatial_vector.is_contiguous(): + if spatial_vector.is_contiguous: # Check if the array is contiguous. If so, we can just return a strided array. # Then this update becomes a no-op. return wp.array( @@ -981,34 +989,27 @@ def _get_bottom_from_spatial_vector(self, source: wp.array | None, spatial_vecto device=self.device, ) else: - # If the array is no contiguous, we need to create a new array to write to. - source = wp.zeros((spatial_vector.shape[0], 3), dtype=wp.vec3f, device=self.device) + # If the array is not contiguous, we need to create a new array to write to. + # Shape matches spatial_vector.shape since each element is vec3f (already contains 3 floats) + source = wp.zeros(spatial_vector.shape, dtype=wp.vec3f, device=self.device) # If the array is not contiguous, we need to launch the kernel to get the bottom part of the spatial vector. - if not spatial_vector.is_contiguous(): - # Launch the right kernel based on the shape of the source array. - if len(source.shape) > 1: + if not spatial_vector.is_contiguous: + # Launch the right kernel based on the shape of the spatial_vector array. + if len(spatial_vector.shape) > 1: wp.launch( shared_kernels.split_spatial_vector_to_bottom_2d, - dim=source.shape, - inputs=[ - source, - ], - outputs=[ - source, - ], + dim=spatial_vector.shape, + inputs=[spatial_vector], + outputs=[source], device=self.device, ) else: wp.launch( shared_kernels.split_spatial_vector_to_bottom_1d, - dim=source.shape, - inputs=[ - source, - ], - outputs=[ - source, - ], + dim=spatial_vector.shape, + inputs=[spatial_vector], + outputs=[source], device=self.device, ) # Return the source array. (no-op if the array is contiguous.) @@ -1027,6 +1028,10 @@ def root_state_w(self) -> wp.array: DeprecationWarning, stacklevel=2, ) + if self._root_state_w is None: + self._root_state_w = TimestampedBuffer( + shape=(self._num_instances,), dtype=shared_kernels.vec13f, device=self.device + ) if self._root_state_w.timestamp < self._sim_timestamp: wp.launch( shared_kernels.concat_root_pose_and_vel_to_state, @@ -1053,6 +1058,10 @@ def root_link_state_w(self) -> wp.array: DeprecationWarning, stacklevel=2, ) + if self._root_link_state_w is None: + self._root_link_state_w = TimestampedBuffer( + shape=(self._num_instances,), dtype=shared_kernels.vec13f, device=self.device + ) if self._root_link_state_w.timestamp < self._sim_timestamp: wp.launch( shared_kernels.concat_root_pose_and_vel_to_state, @@ -1079,6 +1088,10 @@ def root_com_state_w(self) -> wp.array: DeprecationWarning, stacklevel=2, ) + if self._root_com_state_w is None: + self._root_com_state_w = TimestampedBuffer( + shape=(self._num_instances,), dtype=shared_kernels.vec13f, device=self.device + ) if self._root_com_state_w.timestamp < self._sim_timestamp: wp.launch( shared_kernels.concat_root_pose_and_vel_to_state, @@ -1135,6 +1148,10 @@ def body_state_w(self) -> wp.array: stacklevel=2, ) # Access internal buffer directly to avoid cascading deprecation warnings from root_state_w + if self._root_state_w is None: + self._root_state_w = TimestampedBuffer( + shape=(self._num_instances,), dtype=shared_kernels.vec13f, device=self.device + ) if self._root_state_w.timestamp < self._sim_timestamp: wp.launch( shared_kernels.concat_root_pose_and_vel_to_state, @@ -1161,6 +1178,10 @@ def body_link_state_w(self) -> wp.array: stacklevel=2, ) # Access internal buffer directly to avoid cascading deprecation warnings from root_link_state_w + if self._root_link_state_w is None: + self._root_link_state_w = TimestampedBuffer( + shape=(self._num_instances,), dtype=shared_kernels.vec13f, device=self.device + ) if self._root_link_state_w.timestamp < self._sim_timestamp: wp.launch( shared_kernels.concat_root_pose_and_vel_to_state, @@ -1186,6 +1207,10 @@ def body_com_state_w(self) -> wp.array: DeprecationWarning, stacklevel=2, ) + if self._root_com_state_w is None: + self._root_com_state_w = TimestampedBuffer( + shape=(self._num_instances,), dtype=shared_kernels.vec13f, device=self.device + ) if self._root_com_state_w.timestamp < self._sim_timestamp: wp.launch( shared_kernels.concat_root_pose_and_vel_to_state, diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/__init__.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/__init__.py new file mode 100644 index 000000000000..e14e0f6d52c5 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/__init__.pyi new file mode 100644 index 000000000000..8b12ec95e7a2 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/__init__.pyi @@ -0,0 +1,12 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +__all__ = [ + "RigidObjectCollection", + "RigidObjectCollectionData", +] + +from .rigid_object_collection import RigidObjectCollection +from .rigid_object_collection_data import RigidObjectCollectionData diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py new file mode 100644 index 000000000000..b49e3e3289b3 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py @@ -0,0 +1,1366 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import logging +import re +import warnings +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import numpy as np +import torch +import warp as wp +from newton.selection import ArticulationView +from newton.solvers import SolverNotifyFlags + +from pxr import UsdPhysics + +import isaaclab.sim as sim_utils +import isaaclab.utils.string as string_utils +from isaaclab.assets.rigid_object_collection.base_rigid_object_collection import BaseRigidObjectCollection +from isaaclab.utils.wrench_composer import WrenchComposer + +from isaaclab_newton.assets import kernels as shared_kernels +from isaaclab_newton.physics import NewtonManager as SimulationManager + +from .rigid_object_collection_data import RigidObjectCollectionData + +if TYPE_CHECKING: + from isaaclab.assets.rigid_object_collection.rigid_object_collection_cfg import RigidObjectCollectionCfg + +# import logger +logger = logging.getLogger(__name__) + + +class RigidObjectCollection(BaseRigidObjectCollection): + """A rigid object collection class. + + This class represents a collection of rigid objects in the simulation, where the state of the + rigid objects can be accessed and modified using a batched ``(env_ids, object_ids)`` API. + + For each rigid body in the collection, the root prim of the asset must have the `USD RigidBodyAPI`_ + applied to it. This API is used to define the simulation properties of the rigid bodies. On playing the + simulation, the physics engine will automatically register the rigid bodies and create a corresponding + rigid body handle. This handle can be accessed using the :attr:`root_view` attribute. + + Rigid objects in the collection are uniquely identified via the key of the dictionary + :attr:`~isaaclab.assets.RigidObjectCollectionCfg.rigid_objects` in the + :class:`~isaaclab.assets.RigidObjectCollectionCfg` configuration class. + This differs from the :class:`~isaaclab.assets.RigidObject` class, where a rigid object is identified by + the name of the Xform where the `USD RigidBodyAPI`_ is applied. This would not be possible for the rigid + object collection since the :attr:`~isaaclab.assets.RigidObjectCollectionCfg.rigid_objects` dictionary + could contain the same rigid object multiple times, leading to ambiguity. + + .. _`USD RigidBodyAPI`: https://openusd.org/dev/api/class_usd_physics_rigid_body_a_p_i.html + """ + + cfg: RigidObjectCollectionCfg + """Configuration instance for the rigid object.""" + + __backend_name__: str = "newton" + """The name of the backend for the rigid object.""" + + def __init__(self, cfg: RigidObjectCollectionCfg): + """Initialize the rigid object collection. + + Args: + cfg: A configuration instance. + """ + # Note: We never call the parent constructor as it tries to call its own spawning which we don't want. + # check that the config is valid + cfg.validate() + # store inputs + self.cfg = cfg.copy() + # flag for whether the asset is initialized + self._is_initialized = False + # spawn the rigid objects + for rigid_body_cfg in self.cfg.rigid_objects.values(): + # spawn the asset + if rigid_body_cfg.spawn is not None: + rigid_body_cfg.spawn.func( + rigid_body_cfg.prim_path, + rigid_body_cfg.spawn, + translation=rigid_body_cfg.init_state.pos, + orientation=rigid_body_cfg.init_state.rot, + ) + # check that spawn was successful + matching_prims = sim_utils.find_matching_prims(rigid_body_cfg.prim_path) + if len(matching_prims) == 0: + raise RuntimeError(f"Could not find prim with path {rigid_body_cfg.prim_path}.") + # stores object names + self._body_names_list = [] + + # register various callback functions + self._register_callbacks() + self._debug_vis_handle = None + + """ + Properties + """ + + @property + def data(self) -> RigidObjectCollectionData: + return self._data + + @property + def num_instances(self) -> int: + return self._root_view.count // self.num_bodies + + @property + def num_bodies(self) -> int: + """Number of bodies in the rigid object collection.""" + return len(self.body_names) + + @property + def body_names(self) -> list[str]: + """Ordered names of bodies in the rigid object collection.""" + return self._body_names_list + + @property + def root_view(self) -> ArticulationView: + """Root view for the rigid object collection. + + A single :class:`ArticulationView` matching all body types. The 2nd dimension + (matches per world) corresponds to the different body types. + + .. note:: + Use this view with caution. It requires handling of tensors in a specific way. + """ + return self._root_view + + @property + def instantaneous_wrench_composer(self) -> WrenchComposer: + """Instantaneous wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are only valid for the current simulation step. At the end of the simulation step, the wrenches set + to this object are discarded. This is useful to apply forces that change all the time, things like drag forces + for instance. + """ + return self._instantaneous_wrench_composer + + @property + def permanent_wrench_composer(self) -> WrenchComposer: + """Permanent wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are persistent and are applied to the simulation at every step. This is useful to apply forces that + are constant over a period of time, things like the thrust of a motor for instance. + """ + return self._permanent_wrench_composer + + """ + Operations. + """ + + def reset( + self, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + object_ids: slice | torch.Tensor | None = None, + env_mask: wp.array | None = None, + object_mask: wp.array | None = None, + ) -> None: + """Resets all internal buffers of selected environments and objects. + + Args: + env_ids: Environment indices. If None, then all indices are used. + object_ids: Object indices. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + object_mask: Object mask. Not used currently. + """ + # resolve all indices + if env_ids is None: + env_ids = self._ALL_ENV_INDICES + if object_ids is None: + object_ids = self._ALL_BODY_INDICES + # reset external wrench + self._instantaneous_wrench_composer.reset(env_ids) + self._permanent_wrench_composer.reset(env_ids) + + def write_data_to_sim(self) -> None: + """Write external wrench to the simulation. + + .. note:: + We write external wrench to the simulation here since this function is called before the simulation step. + This ensures that the external wrench is applied at every simulation step. + """ + # write external wrench + if self._instantaneous_wrench_composer.active or self._permanent_wrench_composer.active: + if self._instantaneous_wrench_composer.active: + # Compose instantaneous wrench with permanent wrench + self._instantaneous_wrench_composer.add_forces_and_torques_index( + forces=self._permanent_wrench_composer.composed_force, + torques=self._permanent_wrench_composer.composed_torque, + body_ids=self._ALL_BODY_INDICES, + env_ids=self._ALL_ENV_INDICES, + ) + # Apply both instantaneous and permanent wrench to a consolidated 2D buffer + wp.launch( + shared_kernels.update_wrench_array_with_force_and_torque, + dim=(self.num_instances, self.num_bodies), + device=self.device, + inputs=[ + self._instantaneous_wrench_composer.composed_force, + self._instantaneous_wrench_composer.composed_torque, + self._wrench_buffer, + self._ALL_ENV_MASK, + self._ALL_BODY_MASK, + ], + ) + else: + # Apply permanent wrench to a consolidated 2D buffer + wp.launch( + shared_kernels.update_wrench_array_with_force_and_torque, + dim=(self.num_instances, self.num_bodies), + device=self.device, + inputs=[ + self._permanent_wrench_composer.composed_force, + self._permanent_wrench_composer.composed_torque, + self._wrench_buffer, + self._ALL_ENV_MASK, + self._ALL_BODY_MASK, + ], + ) + # Write the wrench buffer directly to the Newton binding (already 2D) + wp.copy(self._data._sim_bind_body_external_wrench, self._wrench_buffer) + self._instantaneous_wrench_composer.reset() + + def update(self, dt: float) -> None: + """Updates the simulation data. + + Args: + dt: The time step size [s]. + """ + self.data.update(dt) + + """ + Operations - Finders. + """ + + def find_bodies( + self, name_keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[torch.Tensor, list[str]]: + """Find bodies in the rigid body collection based on the name keys. + + Please check the :meth:`isaaclab.utils.string_utils.resolve_matching_names` function for more + information on the name matching. + + Args: + name_keys: A regular expression or a list of regular expressions to match the body names. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the body indices and names. + """ + obj_ids, obj_names = string_utils.resolve_matching_names(name_keys, self.body_names, preserve_order) + return torch.tensor(obj_ids, device=self.device, dtype=torch.int32), obj_names + + """ + Operations - Write to simulation. + """ + + def write_body_pose_to_sim_index( + self, + *, + body_poses: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set the body pose over selected environment and body indices into the simulation. + + The body pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + body_poses: Body poses in simulation frame. Shape is (len(env_ids), len(body_ids), 7) + or (len(env_ids), len(body_ids)) with dtype wp.transformf. + body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. + """ + self.write_body_link_pose_to_sim_index(body_poses=body_poses, env_ids=env_ids, body_ids=body_ids) + + def write_body_pose_to_sim_mask( + self, + *, + body_poses: torch.Tensor | wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set the body pose over selected environment mask into the simulation. + + The body pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + body_poses: Body poses in simulation frame. Shape is (num_instances, num_bodies, 7) + or (num_instances, num_bodies) with dtype wp.transformf. + body_mask: Body mask. If None, then all bodies are updated. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + if env_mask is not None: + env_ids = self._resolve_env_mask(env_mask) + else: + env_ids = self._ALL_ENV_INDICES + if body_mask is not None: + body_ids = self._resolve_body_mask(body_mask) + else: + body_ids = self._ALL_BODY_INDICES + self.write_body_link_pose_to_sim_index( + body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True + ) + + def write_body_velocity_to_sim_index( + self, + *, + body_velocities: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set the body velocity over selected environment and body indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the body's center of mass rather than the body's frame. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + body_velocities: Body velocities in simulation frame. + Shape is (len(env_ids), len(body_ids), 6) or (num_instances, num_bodies, 6), + or (len(env_ids), len(body_ids)) / (num_instances, num_bodies) with dtype wp.spatial_vectorf. + body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. + """ + self.write_body_com_velocity_to_sim_index(body_velocities=body_velocities, env_ids=env_ids, body_ids=body_ids) + + def write_body_velocity_to_sim_mask( + self, + *, + body_velocities: torch.Tensor | wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set the body velocity over selected environment mask into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the body's center of mass rather than the body's frame. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + body_velocities: Body velocities in simulation frame. + Shape is (num_instances, num_bodies, 6) + or (num_instances, num_bodies) with dtype wp.spatial_vectorf. + body_mask: Body mask. If None, then all bodies are updated. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + if env_mask is not None: + env_ids = self._resolve_env_mask(env_mask) + else: + env_ids = self._ALL_ENV_INDICES + if body_mask is not None: + body_ids = self._resolve_body_mask(body_mask) + else: + body_ids = self._ALL_BODY_INDICES + self.write_body_com_velocity_to_sim_index( + body_velocities=body_velocities, env_ids=env_ids, body_ids=body_ids, full_data=True + ) + + def write_body_link_pose_to_sim_index( + self, + *, + body_poses: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + full_data: bool = False, + ) -> None: + """Set the body link pose over selected environment and body indices into the simulation. + + The body link pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + body_poses: Body link poses in simulation frame. + Shape is (len(env_ids), len(body_ids), 7) or (num_instances, num_bodies, 7), + or (len(env_ids), len(body_ids)) / (num_instances, num_bodies) with dtype wp.transformf. + body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. + full_data: Whether to expect full data. Defaults to False. + """ + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + if full_data: + self.assert_shape_and_dtype(body_poses, (self.num_instances, self.num_bodies), wp.transformf, "body_poses") + else: + self.assert_shape_and_dtype(body_poses, (env_ids.shape[0], body_ids.shape[0]), wp.transformf, "body_poses") + # Write to consolidated buffer + wp.launch( + shared_kernels.set_body_link_pose_to_sim, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[ + body_poses, + env_ids, + body_ids, + full_data, + ], + outputs=[ + self.data.body_link_pose_w, + None, # body_link_state_w + None, # body_state_w + ], + device=self.device, + ) + # Invalidate dependent timestamps + self.data._body_com_pose_w.timestamp = -1.0 + self.data._body_com_state_w.timestamp = -1.0 + self.data._body_link_state_w.timestamp = -1.0 + self.data._body_state_w.timestamp = -1.0 + + def write_body_link_pose_to_sim_mask( + self, + *, + body_poses: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None = None, + ) -> None: + """Set the body link pose over selected environment mask into the simulation. + + The body link pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + body_poses: Body link poses in simulation frame. Shape is (num_instances, num_bodies, 7) + or (num_instances, num_bodies) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + body_ids: Body indices. If None, then all indices are used. + """ + if env_mask is not None: + env_ids = self._resolve_env_mask(env_mask) + else: + env_ids = self._ALL_ENV_INDICES + self.write_body_link_pose_to_sim_index( + body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True + ) + + def write_body_com_pose_to_sim_index( + self, + *, + body_poses: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + full_data: bool = False, + ) -> None: + """Set the body center of mass pose over selected environment and body indices into the simulation. + + The body center of mass pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + The orientation is the orientation of the principal axes of inertia. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + body_poses: Body center of mass poses in simulation frame. + Shape is (len(env_ids), len(body_ids), 7) or (num_instances, num_bodies, 7), + or (len(env_ids), len(body_ids)) / (num_instances, num_bodies) with dtype wp.transformf. + body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. + full_data: Whether to expect full data. Defaults to False. + """ + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + if full_data: + self.assert_shape_and_dtype(body_poses, (self.num_instances, self.num_bodies), wp.transformf, "body_poses") + else: + self.assert_shape_and_dtype(body_poses, (env_ids.shape[0], body_ids.shape[0]), wp.transformf, "body_poses") + # Write to consolidated buffers (updates both com_pose_w and link_pose_w) + wp.launch( + shared_kernels.set_body_com_pose_to_sim, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[ + body_poses, + self.data.body_com_pose_b, + env_ids, + body_ids, + full_data, + ], + outputs=[ + self.data.body_com_pose_w, + self.data.body_link_pose_w, + None, # body_com_state_w + None, # body_link_state_w + None, # body_state_w + ], + device=self.device, + ) + # Invalidate dependent timestamps + self.data._body_link_state_w.timestamp = -1.0 + self.data._body_state_w.timestamp = -1.0 + self.data._body_com_state_w.timestamp = -1.0 + + def write_body_com_pose_to_sim_mask( + self, + *, + body_poses: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None = None, + ) -> None: + """Set the body center of mass pose over selected environment mask into the simulation. + + The body center of mass pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + The orientation is the orientation of the principal axes of inertia. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + body_poses: Body center of mass poses in simulation frame. Shape is (num_instances, num_bodies, 7) + or (num_instances, num_bodies) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + body_ids: Body indices. If None, then all indices are used. + """ + if env_mask is not None: + env_ids = self._resolve_env_mask(env_mask) + else: + env_ids = self._ALL_ENV_INDICES + self.write_body_com_pose_to_sim_index(body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True) + + def write_body_com_velocity_to_sim_index( + self, + *, + body_velocities: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + full_data: bool = False, + ) -> None: + """Set the body center of mass velocity over selected environment and body indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the body's center of mass rather than the body's frame. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + body_velocities: Body center of mass velocities in simulation frame. + Shape is (len(env_ids), len(body_ids), 6) or (num_instances, num_bodies, 6), + or (len(env_ids), len(body_ids)) / (num_instances, num_bodies) with dtype wp.spatial_vectorf. + body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. + full_data: Whether to expect full data. Defaults to False. + """ + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + if full_data: + self.assert_shape_and_dtype( + body_velocities, (self.num_instances, self.num_bodies), wp.spatial_vectorf, "body_velocities" + ) + else: + self.assert_shape_and_dtype( + body_velocities, (env_ids.shape[0], body_ids.shape[0]), wp.spatial_vectorf, "body_velocities" + ) + # Write to consolidated buffer + wp.launch( + shared_kernels.set_body_com_velocity_to_sim, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[ + body_velocities, + env_ids, + body_ids, + full_data, + ], + outputs=[ + self.data.body_com_vel_w, + self.data.body_com_acc_w, + None, # body_state_w + None, # body_com_state_w + ], + device=self.device, + ) + # Invalidate dependent timestamps + self.data._body_link_vel_w.timestamp = -1.0 + self.data._body_state_w.timestamp = -1.0 + self.data._body_com_state_w.timestamp = -1.0 + self.data._body_link_state_w.timestamp = -1.0 + + def write_body_com_velocity_to_sim_mask( + self, + *, + body_velocities: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None = None, + ) -> None: + """Set the body center of mass velocity over selected environment mask into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the body's center of mass rather than the body's frame. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + body_velocities: Body center of mass velocities in simulation frame. + Shape is (num_instances, num_bodies, 6) + or (num_instances, num_bodies) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + body_ids: Body indices. If None, then all indices are used. + """ + if env_mask is not None: + env_ids = self._resolve_env_mask(env_mask) + else: + env_ids = self._ALL_ENV_INDICES + self.write_body_com_velocity_to_sim_index( + body_velocities=body_velocities, env_ids=env_ids, body_ids=body_ids, full_data=True + ) + + def write_body_link_velocity_to_sim_index( + self, + *, + body_velocities: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + full_data: bool = False, + ) -> None: + """Set the body link velocity over selected environment and body indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the body's frame rather than the body's center of mass. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + body_velocities: Body link velocities in simulation frame. + Shape is (len(env_ids), len(body_ids), 6) or (num_instances, num_bodies, 6), + or (len(env_ids), len(body_ids)) / (num_instances, num_bodies) with dtype wp.spatial_vectorf. + body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. + full_data: Whether to expect full data. Defaults to False. + """ + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + if full_data: + self.assert_shape_and_dtype( + body_velocities, (self.num_instances, self.num_bodies), wp.spatial_vectorf, "body_velocities" + ) + else: + self.assert_shape_and_dtype( + body_velocities, (env_ids.shape[0], body_ids.shape[0]), wp.spatial_vectorf, "body_velocities" + ) + # Access body_com_pose_b and body_link_pose_w to ensure they are current. + wp.launch( + shared_kernels.set_body_link_velocity_to_sim, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[ + body_velocities, + self.data.body_com_pose_b, + self.data.body_link_pose_w, + env_ids, + body_ids, + full_data, + ], + outputs=[ + self.data.body_link_vel_w, + self.data.body_com_vel_w, + self.data.body_com_acc_w, + None, # body_link_state_w + None, # body_state_w + None, # body_com_state_w + ], + device=self.device, + ) + # Invalidate dependent timestamps + self.data._body_link_state_w.timestamp = -1.0 + self.data._body_state_w.timestamp = -1.0 + self.data._body_com_state_w.timestamp = -1.0 + + def write_body_link_velocity_to_sim_mask( + self, + *, + body_velocities: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None = None, + ) -> None: + """Set the body link velocity over selected environment mask into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the body's frame rather than the body's center of mass. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + body_velocities: Body link velocities in simulation frame. Shape is (num_instances, num_bodies, 6) + or (num_instances, num_bodies) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + body_ids: Body indices. If None, then all indices are used. + """ + if env_mask is not None: + env_ids = self._resolve_env_mask(env_mask) + else: + env_ids = self._ALL_ENV_INDICES + self.write_body_link_velocity_to_sim_index( + body_velocities=body_velocities, env_ids=env_ids, body_ids=body_ids, full_data=True + ) + + """ + Operations - Setters. + """ + + def set_masses_index( + self, + *, + masses: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set masses of all bodies using indices. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + masses: Masses of all bodies. Shape is (len(env_ids), len(body_ids)). + body_ids: The body indices to set the masses for. Defaults to None (all bodies). + env_ids: The environment indices to set the masses for. Defaults to None (all environments). + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + self.assert_shape_and_dtype(masses, (env_ids.shape[0], body_ids.shape[0]), wp.float32, "masses") + # Write to consolidated buffer + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[ + masses, + env_ids, + body_ids, + ], + outputs=[ + self.data._sim_bind_body_mass, + ], + device=self.device, + ) + # No copy-back needed — writes go directly to Newton's state via the 2D binding + # Tell the physics engine that some of the body properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + def set_masses_mask( + self, + *, + masses: torch.Tensor | wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set masses of all bodies using masks. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + masses: Masses of all bodies. Shape is (num_instances, num_bodies). + body_mask: Body mask. If None, then all bodies are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + # resolve masks + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if body_mask is None: + body_mask = self._ALL_BODY_MASK + self.assert_shape_and_dtype_mask(masses, (env_mask, body_mask), wp.float32, "masses") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(env_mask.shape[0], body_mask.shape[0]), + inputs=[ + masses, + env_mask, + body_mask, + ], + outputs=[ + self.data._sim_bind_body_mass, + ], + device=self.device, + ) + # No copy-back needed — writes go directly to Newton's state via the 2D binding + # Tell the physics engine that some of the body properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + def set_coms_index( + self, + *, + coms: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set center of mass position of all bodies using indices. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + .. caution:: + Unlike the PhysX version of this method, this method does not set the center of mass orientation. + Only the position is set. This is because Newton considers the center of mass orientation to always be + aligned with the body frame. + + Args: + coms: Center of mass position of all bodies. Shape is (len(env_ids), len(body_ids), 3). + body_ids: The body indices to set the center of mass pose for. Defaults to None (all bodies). + env_ids: The environment indices to set the center of mass pose for. Defaults to None (all environments). + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + self.assert_shape_and_dtype(coms, (env_ids.shape[0], body_ids.shape[0]), wp.vec3f, "coms") + # Write to consolidated buffer + wp.launch( + shared_kernels.write_body_com_position_to_buffer_index, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[ + coms, + env_ids, + body_ids, + ], + outputs=[ + self.data._sim_bind_body_com_pos_b, + ], + device=self.device, + ) + # Invalidate derived buffers that depend on com position + self.data._body_com_pose_b.timestamp = -1.0 + self.data._body_com_pose_w.timestamp = -1.0 + # Tell the physics engine that some of the body properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + def set_coms_mask( + self, + *, + coms: torch.Tensor | wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set center of mass position of all bodies using masks. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + .. caution:: + Unlike the PhysX version of this method, this method does not set the center of mass orientation. + Only the position is set. This is because Newton considers the center of mass orientation to always be + aligned with the body frame. + + Args: + coms: Center of mass position of all bodies. Shape is (num_instances, num_bodies, 3). + body_mask: Body mask. If None, then all bodies are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + # resolve masks + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if body_mask is None: + body_mask = self._ALL_BODY_MASK + self.assert_shape_and_dtype_mask(coms, (env_mask, body_mask), wp.vec3f, "coms") + wp.launch( + shared_kernels.write_body_com_position_to_buffer_mask, + dim=(env_mask.shape[0], body_mask.shape[0]), + inputs=[ + coms, + env_mask, + body_mask, + ], + outputs=[ + self.data._sim_bind_body_com_pos_b, + ], + device=self.device, + ) + # Invalidate derived buffers that depend on com position + self.data._body_com_pose_b.timestamp = -1.0 + self.data._body_com_pose_w.timestamp = -1.0 + # Tell the physics engine that some of the body properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + def set_inertias_index( + self, + *, + inertias: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set inertias of all bodies using indices. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + inertias: Inertias of all bodies. Shape is (len(env_ids), len(body_ids), 9). + body_ids: The body indices to set the inertias for. Defaults to None (all bodies). + env_ids: The environment indices to set the inertias for. Defaults to None (all environments). + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + self.assert_shape_and_dtype(inertias, (env_ids.shape[0], body_ids.shape[0], 9), wp.float32, "inertias") + # Write to consolidated buffer + wp.launch( + shared_kernels.write_body_inertia_to_buffer_index, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[ + inertias, + env_ids, + body_ids, + ], + outputs=[ + self.data._body_inertia, + ], + device=self.device, + ) + # No copy-back needed — writes go directly to Newton's state via the 2D binding + # Tell the physics engine that some of the body properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + def set_inertias_mask( + self, + *, + inertias: torch.Tensor | wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set inertias of all bodies using masks. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + inertias: Inertias of all bodies. Shape is (num_instances, num_bodies, 9). + body_mask: Body mask. If None, then all bodies are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + # resolve masks + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if body_mask is None: + body_mask = self._ALL_BODY_MASK + self.assert_shape_and_dtype_mask(inertias, (env_mask, body_mask), wp.float32, "inertias", trailing_dims=(9,)) + wp.launch( + shared_kernels.write_body_inertia_to_buffer_mask, + dim=(env_mask.shape[0], body_mask.shape[0]), + inputs=[ + inertias, + env_mask, + body_mask, + ], + outputs=[ + self.data._body_inertia, + ], + device=self.device, + ) + # No copy-back needed — writes go directly to Newton's state via the 2D binding + # Tell the physics engine that some of the body properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + """ + Internal helper. + """ + + def _initialize_impl(self): + # clear body names list to prevent double counting on re-initialization + self._body_names_list.clear() + root_prim_path_exprs: list[str] = [] + + for name, rigid_body_cfg in self.cfg.rigid_objects.items(): + # obtain the first prim in the regex expression (all others are assumed to be a copy of this) + template_prim = sim_utils.find_first_matching_prim(rigid_body_cfg.prim_path) + if template_prim is None: + raise RuntimeError(f"Failed to find prim for expression: '{rigid_body_cfg.prim_path}'.") + template_prim_path = template_prim.GetPath().pathString + + # find rigid root prims + root_prims = sim_utils.get_all_matching_child_prims( + template_prim_path, + predicate=lambda prim: prim.HasAPI(UsdPhysics.RigidBodyAPI), + traverse_instance_prims=False, + ) + if len(root_prims) == 0: + raise RuntimeError( + f"Failed to find a rigid body when resolving '{rigid_body_cfg.prim_path}'." + " Please ensure that the prim has 'USD RigidBodyAPI' applied." + ) + if len(root_prims) > 1: + raise RuntimeError( + f"Failed to find a single rigid body when resolving '{rigid_body_cfg.prim_path}'." + f" Found multiple '{root_prims}' under '{template_prim_path}'." + " Please ensure that there is only one rigid body in the prim path tree." + ) + + # check that no rigid object has an articulation root API, which decreases simulation performance + articulation_prims = sim_utils.get_all_matching_child_prims( + template_prim_path, + predicate=lambda prim: prim.HasAPI(UsdPhysics.ArticulationRootAPI), + traverse_instance_prims=False, + ) + if len(articulation_prims) != 0: + if articulation_prims[0].GetAttribute("physxArticulation:articulationEnabled").Get(): + raise RuntimeError( + f"Found an articulation root when resolving '{rigid_body_cfg.prim_path}' in the rigid object" + f" collection. These are located at: '{articulation_prims}' under '{template_prim_path}'." + " Please disable the articulation root in the USD or from code by setting the parameter" + " 'ArticulationRootPropertiesCfg.articulation_enabled' to False in the spawn configuration." + ) + + # resolve root prim back into regex expression + root_prim_path = root_prims[0].GetPath().pathString + root_prim_path_expr = rigid_body_cfg.prim_path + root_prim_path[len(template_prim_path) :] + root_prim_path_exprs.append(root_prim_path_expr.replace(".*", "*")) + self._body_names_list.append(name) + + # Build a single pattern that matches ALL body types by wildcarding the differing path segment + combined_pattern = self._build_combined_pattern(root_prim_path_exprs) + + # Create a single ArticulationView matching all body types. + # The 2nd dimension (matches per world) corresponds to the body types. + self._root_view = ArticulationView( + SimulationManager.get_model(), + combined_pattern, + verbose=False, + ) + + # log information about the rigid body + logger.info(f"Number of instances: {self.num_instances}") + logger.info(f"Number of distinct bodies: {self.num_bodies}") + logger.info(f"Body names: {self.body_names}") + + # container for data access + self._data = RigidObjectCollectionData(self._root_view, self.num_bodies, self.device) + + # create buffers + self._create_buffers() + # process configuration + self._process_cfg() + # update the rigid body data + self.update(0.0) + # Let the rigid object collection data know that it is fully instantiated and ready to use. + self.data.is_primed = True + + def _create_buffers(self): + """Create buffers for storing data.""" + # constants + self._ALL_ENV_INDICES = wp.array( + np.arange(self.num_instances, dtype=np.int32), device=self.device, dtype=wp.int32 + ) + self._ALL_BODY_INDICES = wp.array( + np.arange(self.num_bodies, dtype=np.int32), device=self.device, dtype=wp.int32 + ) + self._ALL_ENV_MASK = wp.ones((self.num_instances,), dtype=wp.bool, device=self.device) + self._ALL_BODY_MASK = wp.ones((self.num_bodies,), dtype=wp.bool, device=self.device) + + # external wrench composer + self._instantaneous_wrench_composer = WrenchComposer(self) + self._permanent_wrench_composer = WrenchComposer(self) + + # Temporary 2D wrench buffer for write_data_to_sim + self._wrench_buffer = wp.zeros( + (self.num_instances, self.num_bodies), dtype=wp.spatial_vectorf, device=self.device + ) + + # set information about rigid body into data + self._data.body_names = self.body_names + + def _process_cfg(self) -> None: + """Post processing of configuration parameters.""" + # default state + # -- body state + default_body_poses = [] + default_body_vels = [] + for rigid_object_cfg in self.cfg.rigid_objects.values(): + default_body_pose = tuple(rigid_object_cfg.init_state.pos) + tuple(rigid_object_cfg.init_state.rot) + default_body_vel = tuple(rigid_object_cfg.init_state.lin_vel) + tuple(rigid_object_cfg.init_state.ang_vel) + default_body_pose = np.tile(np.array(default_body_pose, dtype=np.float32), (self.num_instances, 1)) + default_body_vel = np.tile(np.array(default_body_vel, dtype=np.float32), (self.num_instances, 1)) + default_body_poses.append(default_body_pose) + default_body_vels.append(default_body_vel) + # Stack: each has shape (num_instances, data_size) -> (num_instances, num_bodies, data_size) + default_body_poses = np.stack(default_body_poses, axis=1) + default_body_vels = np.stack(default_body_vels, axis=1) + self.data.default_body_pose = wp.array(default_body_poses, dtype=wp.transformf, device=self.device) + self.data.default_body_vel = wp.array(default_body_vels, dtype=wp.spatial_vectorf, device=self.device) + + def _resolve_env_ids(self, env_ids) -> wp.array: + """Resolve environment indices to a warp array.""" + if isinstance(env_ids, list): + return wp.array(env_ids, dtype=wp.int32, device=self.device) + if (env_ids is None) or (env_ids == slice(None)): + return self._ALL_ENV_INDICES + if isinstance(env_ids, torch.Tensor): + return wp.from_torch(env_ids.to(torch.int32), dtype=wp.int32) + return env_ids + + def _resolve_body_ids(self, body_ids) -> wp.array: + """Resolve body indices to a warp array.""" + if body_ids is None or (body_ids == slice(None)): + return self._ALL_BODY_INDICES + if isinstance(body_ids, slice): + return wp.from_torch( + torch.arange(self.num_bodies, dtype=torch.int32, device=self.device)[body_ids], dtype=wp.int32 + ) + if isinstance(body_ids, list): + return wp.array(body_ids, dtype=wp.int32, device=self.device) + if isinstance(body_ids, torch.Tensor): + return wp.from_torch(body_ids.to(torch.int32), dtype=wp.int32) + return body_ids + + def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array | torch.Tensor: + """Resolve environment mask to indices via torch.nonzero.""" + if env_mask is not None: + if isinstance(env_mask, wp.array): + env_mask = wp.to_torch(env_mask) + env_ids = torch.nonzero(env_mask)[:, 0].to(torch.int32) + else: + env_ids = self._ALL_ENV_INDICES + return env_ids + + def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array | torch.Tensor: + """Resolve body mask to indices via torch.nonzero.""" + if body_mask is not None: + if isinstance(body_mask, wp.array): + body_mask = wp.to_torch(body_mask) + body_ids = torch.nonzero(body_mask)[:, 0].to(torch.int32) + else: + body_ids = self._ALL_BODY_INDICES + return body_ids + + @staticmethod + def _build_combined_pattern(prim_path_exprs: list[str]) -> str: + """Build a single fnmatch pattern that matches all body types. + + Compares path segments across all expressions and wildcards the segments that differ. + For example, given:: + + ["/World/Env_*/DexCube/Cube", "/World/Env_*/DexSphere/Sphere"] + + produces ``"/World/Env_*/*/*"``. + + Args: + prim_path_exprs: List of prim path expressions, one per body type. + + Returns: + A single fnmatch pattern string. + + Raises: + ValueError: If the expressions have different numbers of path segments. + """ + if len(prim_path_exprs) == 1: + return prim_path_exprs[0] + + split_paths = [p.split("/") for p in prim_path_exprs] + lengths = {len(s) for s in split_paths} + if len(lengths) != 1: + raise ValueError( + f"Cannot build combined pattern: path expressions have different segment counts: {prim_path_exprs}" + ) + + combined_segments = [] + for segments in zip(*split_paths): + unique = set(segments) + if len(unique) == 1: + combined_segments.append(segments[0]) + else: + combined_segments.append("*") + return "/".join(combined_segments) + + """ + Internal simulation callbacks. + """ + + def _invalidate_initialize_callback(self, event) -> None: + """Invalidates the scene elements.""" + # call parent + super()._invalidate_initialize_callback(event) + # set all existing views to None to invalidate them + self._root_view = None + + def _on_prim_deletion(self, prim_path: str) -> None: + """Invalidates and deletes the callbacks when the prim is deleted. + + Args: + prim_path: The path to the prim that is being deleted. + + .. note:: + This function is called when the prim is deleted. + """ + if prim_path == "/": + self._clear_callbacks() + return + for prim_path_expr in [obj.prim_path for obj in self.cfg.rigid_objects.values()]: + result = re.match( + pattern="^" + "/".join(prim_path_expr.split("/")[: prim_path.count("/") + 1]) + "$", string=prim_path + ) + if result: + self._clear_callbacks() + return + + """ + Deprecated properties and methods. + """ + + @property + def root_physx_view(self): + """Deprecated property. Please use :attr:`root_view` instead.""" + warnings.warn( + "The `root_physx_view` property will be deprecated in a future release. Please use `root_view` instead.", + DeprecationWarning, + stacklevel=2, + ) + return self.root_view + + def write_body_state_to_sim( + self, + body_states: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + body_ids: slice | torch.Tensor | None = None, + ) -> None: + """Deprecated, same as :meth:`write_body_link_pose_to_sim_index` and + :meth:`write_body_com_velocity_to_sim_index`.""" + warnings.warn( + "The function 'write_body_state_to_sim' will be deprecated in a future release. Please" + " use 'write_body_link_pose_to_sim_index' and 'write_body_com_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + self.write_body_link_pose_to_sim_index(body_poses=body_states[:, :, :7], env_ids=env_ids, body_ids=body_ids) + self.write_body_com_velocity_to_sim_index( + body_velocities=body_states[:, :, 7:], env_ids=env_ids, body_ids=body_ids + ) + + def write_body_com_state_to_sim( + self, + body_states: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + body_ids: slice | torch.Tensor | None = None, + ) -> None: + """Deprecated, same as :meth:`write_body_com_pose_to_sim_index` and + :meth:`write_body_com_velocity_to_sim_index`.""" + warnings.warn( + "The function 'write_body_com_state_to_sim' will be deprecated in a future release. Please" + " use 'write_body_com_pose_to_sim_index' and 'write_body_com_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + self.write_body_com_pose_to_sim_index(body_poses=body_states[:, :, :7], env_ids=env_ids, body_ids=body_ids) + self.write_body_com_velocity_to_sim_index( + body_velocities=body_states[:, :, 7:], env_ids=env_ids, body_ids=body_ids + ) + + def write_body_link_state_to_sim( + self, + body_states: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + body_ids: slice | torch.Tensor | None = None, + ) -> None: + """Deprecated, same as :meth:`write_body_link_pose_to_sim_index` and + :meth:`write_body_link_velocity_to_sim_index`.""" + warnings.warn( + "The function 'write_body_link_state_to_sim' will be deprecated in a future release. Please" + " use 'write_body_link_pose_to_sim_index' and 'write_body_link_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + self.write_body_link_pose_to_sim_index(body_poses=body_states[:, :, :7], env_ids=env_ids, body_ids=body_ids) + self.write_body_link_velocity_to_sim_index( + body_velocities=body_states[:, :, 7:], env_ids=env_ids, body_ids=body_ids + ) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection_data.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection_data.py new file mode 100644 index 000000000000..90d4cac3ecdc --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection_data.py @@ -0,0 +1,803 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations + +import logging +import warnings +import weakref +from typing import TYPE_CHECKING + +import torch +import warp as wp + +from isaaclab.assets.rigid_object_collection.base_rigid_object_collection_data import BaseRigidObjectCollectionData +from isaaclab.utils.buffers import TimestampedBufferWarp as TimestampedBuffer +from isaaclab.utils.math import normalize + +from isaaclab_newton.assets import kernels as shared_kernels +from isaaclab_newton.physics import NewtonManager as SimulationManager + +if TYPE_CHECKING: + from newton.selection import ArticulationView + +# import logger +logger = logging.getLogger(__name__) + + +class RigidObjectCollectionData(BaseRigidObjectCollectionData): + """Data container for a rigid object collection. + + This class contains the data for a rigid object collection in the simulation. The data includes the state of + all the bodies in the collection. The data is stored in the simulation world frame unless otherwise specified. + The data is in the order ``(num_instances, num_objects, data_size)``, where data_size is the size of the data. + + For a rigid body, there are two frames of reference that are used: + + - Actor frame: The frame of reference of the rigid body prim. This typically corresponds to the Xform prim + with the rigid body schema. + - Center of mass frame: The frame of reference of the center of mass of the rigid body. + + Depending on the settings of the simulation, the actor frame and the center of mass frame may be the same. + This needs to be taken into account when interpreting the data. + + The data is lazily updated, meaning that the data is only updated when it is accessed. This is useful + when the data is expensive to compute or retrieve. The data is updated when the timestamp of the buffer + is older than the current simulation timestamp. The timestamp is updated whenever the data is updated. + """ + + __backend_name__: str = "newton" + """The name of the backend for the rigid object collection data.""" + + def __init__(self, root_view: ArticulationView, num_bodies: int, device: str): + """Initializes the rigid object collection data. + + Args: + root_view: A single articulation view matching all body types across all environments. + The view's 2nd dimension (matches per world) corresponds to the body types. + num_bodies: The number of bodies in the collection. + device: The device used for processing. + """ + super().__init__(root_view, num_bodies, device) + self.num_bodies = num_bodies + # Store the view as a weak reference to avoid circular references + self._root_view: ArticulationView = weakref.proxy(root_view) + self.num_instances = self._root_view.count // num_bodies + + # Set initial time stamp + self._sim_timestamp = 0.0 + self._is_primed = False + + # Convert gravity to direction vector + gravity = wp.to_torch(SimulationManager.get_model().gravity)[0] + gravity_dir = torch.tensor((gravity[0], gravity[1], gravity[2]), device=self.device) + gravity_dir = normalize(gravity_dir.unsqueeze(0)).squeeze(0) + + # Initialize constants + self.GRAVITY_VEC_W = wp.from_torch(gravity_dir.repeat(self.num_instances, self.num_bodies, 1), dtype=wp.vec3f) + self.FORWARD_VEC_B = wp.from_torch( + torch.tensor((1.0, 0.0, 0.0), device=self.device).repeat(self.num_instances, self.num_bodies, 1), + dtype=wp.vec3f, + ) + + self._create_simulation_bindings() + self._create_buffers() + + @property + def is_primed(self) -> bool: + """Whether the rigid object collection data is fully instantiated and ready to use.""" + return self._is_primed + + @is_primed.setter + def is_primed(self, value: bool) -> None: + """Set whether the rigid object collection data is fully instantiated and ready to use. + + .. note:: + Once this quantity is set to True, it cannot be changed. + + Args: + value: The primed state. + + Raises: + ValueError: If the rigid object collection data is already primed. + """ + if self._is_primed: + raise ValueError("The rigid object collection data is already primed.") + self._is_primed = value + + def update(self, dt: float) -> None: + """Updates the data for the rigid object collection. + + Args: + dt: The time step for the update. This must be a positive value. + """ + # update the simulation timestamp + self._sim_timestamp += dt + # Trigger an update of the body com acceleration buffer at a higher frequency + # since we do finite differencing. + self.body_com_acc_w + + """ + Names. + """ + + body_names: list[str] = None + """Body names in the order parsed by the simulation view.""" + + """ + Defaults. + """ + + @property + def default_body_pose(self) -> wp.array: + """Default body pose ``[pos, quat]`` in local environment frame. + + The position and quaternion are of the rigid body's actor frame. + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). + """ + return self._default_body_pose + + @default_body_pose.setter + def default_body_pose(self, value: wp.array) -> None: + """Set the default body pose. + + Args: + value: The default body pose. Shape is (num_instances, num_bodies, 7). + + Raises: + ValueError: If the rigid object collection data is already primed. + """ + if self._is_primed: + raise ValueError("The rigid object collection data is already primed.") + self._default_body_pose.assign(value) + + @property + def default_body_vel(self) -> wp.array: + """Default body velocity ``[lin_vel, ang_vel]`` in local environment frame. + + The linear and angular velocities are of the rigid body's center of mass frame. + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). + """ + return self._default_body_vel + + @default_body_vel.setter + def default_body_vel(self, value: wp.array) -> None: + """Set the default body velocity. + + Args: + value: The default body velocity. Shape is (num_instances, num_bodies, 6). + + Raises: + ValueError: If the rigid object collection data is already primed. + """ + if self._is_primed: + raise ValueError("The rigid object collection data is already primed.") + self._default_body_vel.assign(value) + + """ + Body state properties. + """ + + @property + def body_link_pose_w(self) -> wp.array: + """Body link pose ``[pos, quat]`` in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). + This quantity is the pose of the actor frame of the rigid body relative to the world. + The orientation is provided in (x, y, z, w) format. + """ + return self._sim_bind_body_link_pose_w + + @property + def body_link_vel_w(self) -> wp.array: + """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). + This quantity contains the linear and angular velocities of the actor frame of the root + rigid body relative to the world. + """ + if self._body_link_vel_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.get_body_link_vel_from_body_com_vel, + dim=(self.num_instances, self.num_bodies), + inputs=[ + self.body_com_vel_w, + self.body_link_pose_w, + self.body_com_pos_b, + ], + outputs=[ + self._body_link_vel_w.data, + ], + device=self.device, + ) + self._body_link_vel_w.timestamp = self._sim_timestamp + + return self._body_link_vel_w.data + + @property + def body_com_pose_w(self) -> wp.array: + """Body center of mass pose ``[pos, quat]`` in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). + This quantity is the pose of the center of mass frame of the rigid body relative to the world. + The orientation is provided in (x, y, z, w) format. + """ + if self._body_com_pose_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.get_body_com_pose_from_body_link_pose, + dim=(self.num_instances, self.num_bodies), + inputs=[ + self.body_link_pose_w, + self.body_com_pos_b, + ], + outputs=[ + self._body_com_pose_w.data, + ], + device=self.device, + ) + self._body_com_pose_w.timestamp = self._sim_timestamp + + return self._body_com_pose_w.data + + @property + def body_com_vel_w(self) -> wp.array: + """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). + This quantity contains the linear and angular velocities of the root rigid body's center of mass frame + relative to the world. + """ + return self._sim_bind_body_com_vel_w + + @property + def body_com_acc_w(self) -> wp.array: + """Acceleration of all bodies ``[lin_acc, ang_acc]`` in the simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). + This quantity is the acceleration of the rigid bodies' center of mass frame relative to the world. + """ + if self._body_com_acc_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.derive_body_acceleration_from_body_com_velocities, + dim=(self.num_instances, self.num_bodies), + device=self.device, + inputs=[ + self.body_com_vel_w, + SimulationManager.get_dt(), + self._previous_body_com_vel, + ], + outputs=[ + self._body_com_acc_w.data, + ], + ) + self._body_com_acc_w.timestamp = self._sim_timestamp + return self._body_com_acc_w.data + + @property + def body_com_pose_b(self) -> wp.array: + """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. + + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). + This quantity is the pose of the center of mass frame of the rigid body relative to the body's link frame. + The orientation is provided in (x, y, z, w) format. + """ + warnings.warn( + "In Newton, body com pose always has unit quaternion. Consider using body_com_pos_b instead." + "Querying this property requires appending a unit quaternion to the position which is expensive.", + category=UserWarning, + stacklevel=2, + ) + if self._body_com_pose_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.make_dummy_body_com_pose_b, + dim=(self.num_instances, self.num_bodies), + inputs=[ + self.body_com_pos_b, + ], + outputs=[ + self._body_com_pose_b.data, + ], + device=self.device, + ) + self._body_com_pose_b.timestamp = self._sim_timestamp + return self._body_com_pose_b.data + + @property + def body_com_pos_b(self) -> wp.array: + """Center of mass position of all of the bodies in their respective link frames. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + This quantity is the center of mass location relative to its body's link frame. + """ + return self._sim_bind_body_com_pos_b + + @property + def body_mass(self) -> wp.array: + """Mass of all bodies in the simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.float32. + In torch this resolves to (num_instances, num_bodies). + """ + return self._sim_bind_body_mass + + @property + def body_inertia(self) -> wp.array: + """Inertia of all bodies in the simulation world frame. + + Shape is (num_instances, num_bodies, 9), dtype = wp.float32. + In torch this resolves to (num_instances, num_bodies, 9). + """ + return self._body_inertia + + """ + Derived Properties. + """ + + @property + def projected_gravity_b(self) -> wp.array: + """Projection of the gravity direction on base frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + """ + if self._projected_gravity_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_2D_kernel, + dim=(self.num_instances, self.num_bodies), + inputs=[self.GRAVITY_VEC_W, self.body_link_quat_w], + outputs=[self._projected_gravity_b.data], + device=self.device, + ) + self._projected_gravity_b.timestamp = self._sim_timestamp + return self._projected_gravity_b.data + + @property + def heading_w(self) -> wp.array: + """Yaw heading of the base frame (in radians). + + Shape is (num_instances, num_bodies), dtype = wp.float32. In torch this resolves to (num_instances, num_bodies). + + .. note:: + This quantity is computed by assuming that the forward-direction of the base + frame is along x-direction, i.e. :math:`(1, 0, 0)`. + """ + if self._heading_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.body_heading_w, + dim=(self.num_instances, self.num_bodies), + inputs=[self.FORWARD_VEC_B, self.body_link_quat_w], + outputs=[self._heading_w.data], + device=self.device, + ) + self._heading_w.timestamp = self._sim_timestamp + return self._heading_w.data + + @property + def body_link_lin_vel_b(self) -> wp.array: + """Root link linear velocity in base frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + This quantity is the linear velocity of the actor frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + if self._body_link_lin_vel_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_2D_kernel, + dim=(self.num_instances, self.num_bodies), + inputs=[self.body_link_lin_vel_w, self.body_link_quat_w], + outputs=[self._body_link_lin_vel_b.data], + device=self.device, + ) + self._body_link_lin_vel_b.timestamp = self._sim_timestamp + return self._body_link_lin_vel_b.data + + @property + def body_link_ang_vel_b(self) -> wp.array: + """Root link angular velocity in base frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + This quantity is the angular velocity of the actor frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + if self._body_link_ang_vel_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_2D_kernel, + dim=(self.num_instances, self.num_bodies), + inputs=[self.body_link_ang_vel_w, self.body_link_quat_w], + outputs=[self._body_link_ang_vel_b.data], + device=self.device, + ) + self._body_link_ang_vel_b.timestamp = self._sim_timestamp + return self._body_link_ang_vel_b.data + + @property + def body_com_lin_vel_b(self) -> wp.array: + """Root center of mass linear velocity in base frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + This quantity is the linear velocity of the root rigid body's center of mass frame with respect to the + rigid body's actor frame. + """ + if self._body_com_lin_vel_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_2D_kernel, + dim=(self.num_instances, self.num_bodies), + inputs=[self.body_com_lin_vel_w, self.body_link_quat_w], + outputs=[self._body_com_lin_vel_b.data], + device=self.device, + ) + self._body_com_lin_vel_b.timestamp = self._sim_timestamp + return self._body_com_lin_vel_b.data + + @property + def body_com_ang_vel_b(self) -> wp.array: + """Root center of mass angular velocity in base frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + This quantity is the angular velocity of the root rigid body's center of mass frame with respect to the + rigid body's actor frame. + """ + if self._body_com_ang_vel_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_2D_kernel, + dim=(self.num_instances, self.num_bodies), + inputs=[self.body_com_ang_vel_w, self.body_link_quat_w], + outputs=[self._body_com_ang_vel_b.data], + device=self.device, + ) + self._body_com_ang_vel_b.timestamp = self._sim_timestamp + return self._body_com_ang_vel_b.data + + """ + Sliced properties. + """ + + @property + def body_link_pos_w(self) -> wp.array: + """Positions of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + This quantity is the position of the rigid bodies' actor frame relative to the world. + """ + return self._get_pos_from_transform(self.body_link_pose_w) + + @property + def body_link_quat_w(self) -> wp.array: + """Orientation (x, y, z, w) of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to + (num_instances, num_bodies, 4). + This quantity is the orientation of the rigid bodies' actor frame relative to the world. + """ + return self._get_quat_from_transform(self.body_link_pose_w) + + @property + def body_link_lin_vel_w(self) -> wp.array: + """Linear velocity of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + This quantity is the linear velocity of the rigid bodies' actor frame relative to the world. + """ + return self._get_lin_vel_from_spatial_vector(self.body_link_vel_w) + + @property + def body_link_ang_vel_w(self) -> wp.array: + """Angular velocity of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + This quantity is the angular velocity of the rigid bodies' actor frame relative to the world. + """ + return self._get_ang_vel_from_spatial_vector(self.body_link_vel_w) + + @property + def body_com_pos_w(self) -> wp.array: + """Positions of all bodies' center of mass in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + This quantity is the position of the rigid bodies' center of mass frame. + """ + return self._get_pos_from_transform(self.body_com_pose_w) + + @property + def body_com_quat_w(self) -> wp.array: + """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to + (num_instances, num_bodies, 4). + This quantity is the orientation of the principal axes of inertia of the rigid bodies. + """ + return self._get_quat_from_transform(self.body_com_pose_w) + + @property + def body_com_lin_vel_w(self) -> wp.array: + """Linear velocity of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + This quantity is the linear velocity of the rigid bodies' center of mass frame. + """ + return self._get_lin_vel_from_spatial_vector(self.body_com_vel_w) + + @property + def body_com_ang_vel_w(self) -> wp.array: + """Angular velocity of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + This quantity is the angular velocity of the rigid bodies' center of mass frame. + """ + return self._get_ang_vel_from_spatial_vector(self.body_com_vel_w) + + @property + def body_com_lin_acc_w(self) -> wp.array: + """Linear acceleration of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + This quantity is the linear acceleration of the rigid bodies' center of mass frame. + """ + return self._get_lin_vel_from_spatial_vector(self.body_com_acc_w) + + @property + def body_com_ang_acc_w(self) -> wp.array: + """Angular acceleration of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + This quantity is the angular acceleration of the rigid bodies' center of mass frame. + """ + return self._get_ang_vel_from_spatial_vector(self.body_com_acc_w) + + @property + def body_com_quat_b(self) -> wp.array: + """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their + respective link frames. + + Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to + (num_instances, num_bodies, 4). + This quantity is the orientation of the principal axes of inertia relative to its body's link frame. + """ + return self._get_quat_from_transform(self.body_com_pose_b) + + def _create_simulation_bindings(self) -> None: + """Create simulation bindings for the body data. + + For a rigid object collection in Newton, a single :class:`ArticulationView` matches all body + types. Its data is shaped ``(num_envs, num_bodies, ...)``, where the 2nd dimension (matches + per world) corresponds to the different body types. This gives us direct 2D bindings into + Newton's state with no scatter/gather overhead. + """ + state_0 = SimulationManager.get_state_0() + model = SimulationManager.get_model() + + # Root transforms/velocities are (num_envs, num_bodies) — direct 2D bindings + self._sim_bind_body_link_pose_w = self._root_view.get_root_transforms(state_0) + self._sim_bind_body_com_vel_w = self._root_view.get_root_velocities(state_0) + # Attributes have a trailing link dim of 1: (num_envs, num_bodies, 1) -> [:, :, 0] + self._sim_bind_body_com_pos_b = self._root_view.get_attribute("body_com", model)[:, :, 0] + self._sim_bind_body_external_wrench = self._root_view.get_attribute("body_f", state_0)[:, :, 0] + # -- Body mass: (num_envs, num_bodies, 1) float32 → squeeze to (num_envs, num_bodies) + self._sim_bind_body_mass = self._root_view.get_attribute("body_mass", model)[:, :, 0] + # -- Body inertia: (num_envs, num_bodies, 1) mat33f → squeeze, clone to contiguous, convert to float32 + inertia_mat33 = self._root_view.get_attribute("body_inertia", model)[:, :, 0] + self._body_inertia = wp.clone(inertia_mat33).view(wp.float32).reshape((self.num_instances, self.num_bodies, 9)) + + def _create_buffers(self) -> None: + """Create buffers for computing and caching derived quantities.""" + super()._create_buffers() + + # Initialize the lazy buffers. + # -- link frame w.r.t. world frame (computed from com vel) + self._body_link_vel_w = TimestampedBuffer( + (self.num_instances, self.num_bodies), self.device, wp.spatial_vectorf + ) + # -- com frame w.r.t. link frame + self._body_com_pose_b = TimestampedBuffer((self.num_instances, self.num_bodies), self.device, wp.transformf) + # -- com frame w.r.t. world frame + self._body_com_pose_w = TimestampedBuffer((self.num_instances, self.num_bodies), self.device, wp.transformf) + self._body_com_acc_w = TimestampedBuffer((self.num_instances, self.num_bodies), self.device, wp.spatial_vectorf) + # -- combined state (these are cached as they concatenate) + self._body_state_w = TimestampedBuffer( + (self.num_instances, self.num_bodies), self.device, shared_kernels.vec13f + ) + self._body_link_state_w = TimestampedBuffer( + (self.num_instances, self.num_bodies), self.device, shared_kernels.vec13f + ) + self._body_com_state_w = TimestampedBuffer( + (self.num_instances, self.num_bodies), self.device, shared_kernels.vec13f + ) + + # -- Default state + self._default_body_pose = wp.zeros( + (self.num_instances, self.num_bodies), dtype=wp.transformf, device=self.device + ) + self._default_body_vel = wp.zeros( + (self.num_instances, self.num_bodies), dtype=wp.spatial_vectorf, device=self.device + ) + self._default_body_state = None + + # -- Derived properties + self._projected_gravity_b = TimestampedBuffer((self.num_instances, self.num_bodies), self.device, wp.vec3f) + self._heading_w = TimestampedBuffer((self.num_instances, self.num_bodies), self.device, wp.float32) + self._body_link_lin_vel_b = TimestampedBuffer((self.num_instances, self.num_bodies), self.device, wp.vec3f) + self._body_link_ang_vel_b = TimestampedBuffer((self.num_instances, self.num_bodies), self.device, wp.vec3f) + self._body_com_lin_vel_b = TimestampedBuffer((self.num_instances, self.num_bodies), self.device, wp.vec3f) + self._body_com_ang_vel_b = TimestampedBuffer((self.num_instances, self.num_bodies), self.device, wp.vec3f) + + # -- Initialize history for finite differencing + self._previous_body_com_vel = wp.clone(self._sim_bind_body_com_vel_w) + + """ + Helpers. + """ + + def _get_pos_from_transform(self, transform: wp.array) -> wp.array: + """Generates a position array from a transform array.""" + return wp.array( + ptr=transform.ptr, + shape=transform.shape, + dtype=wp.vec3f, + strides=transform.strides, + device=self.device, + ) + + def _get_quat_from_transform(self, transform: wp.array) -> wp.array: + """Generates a quaternion array from a transform array.""" + return wp.array( + ptr=transform.ptr + 3 * 4, + shape=transform.shape, + dtype=wp.quatf, + strides=transform.strides, + device=self.device, + ) + + def _get_lin_vel_from_spatial_vector(self, spatial_vector: wp.array) -> wp.array: + """Generates a linear velocity array from a spatial vector array.""" + return wp.array( + ptr=spatial_vector.ptr, + shape=spatial_vector.shape, + dtype=wp.vec3f, + strides=spatial_vector.strides, + device=self.device, + ) + + def _get_ang_vel_from_spatial_vector(self, spatial_vector: wp.array) -> wp.array: + """Generates an angular velocity array from a spatial vector array.""" + return wp.array( + ptr=spatial_vector.ptr + 3 * 4, + shape=spatial_vector.shape, + dtype=wp.vec3f, + strides=spatial_vector.strides, + device=self.device, + ) + + """ + Deprecated properties. + """ + + @property + def default_body_state(self) -> wp.array: + """Default root state ``[pos, quat, lin_vel, ang_vel]`` in local environment frame. + + The position and quaternion are of the rigid body's actor frame. Meanwhile, the linear and angular velocities + are of the center of mass frame. Shape is (num_instances, num_bodies, 13). + """ + warnings.warn( + "Reading the body state directly is deprecated since IsaacLab 3.0 and will be removed in a future version. " + "Please use the default_body_pose and default_body_vel properties instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._default_body_state is None: + self._default_body_state = wp.zeros( + (self.num_instances, self.num_bodies), dtype=shared_kernels.vec13f, device=self.device + ) + wp.launch( + shared_kernels.concat_body_pose_and_vel_to_state, + dim=(self.num_instances, self.num_bodies), + inputs=[ + self._default_body_pose, + self._default_body_vel, + ], + outputs=[ + self._default_body_state, + ], + device=self.device, + ) + return self._default_body_state + + @property + def body_state_w(self) -> wp.array: + """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_com_vel_w`.""" + warnings.warn( + "The `body_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " + "`body_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_body_pose_and_vel_to_state, + dim=(self.num_instances, self.num_bodies), + inputs=[ + self.body_link_pose_w, + self.body_com_vel_w, + ], + outputs=[ + self._body_state_w.data, + ], + device=self.device, + ) + self._body_state_w.timestamp = self._sim_timestamp + + return self._body_state_w.data + + @property + def body_link_state_w(self) -> wp.array: + """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_link_vel_w`.""" + warnings.warn( + "The `body_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " + "`body_link_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_link_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_body_pose_and_vel_to_state, + dim=(self.num_instances, self.num_bodies), + inputs=[ + self.body_link_pose_w, + self.body_link_vel_w, + ], + outputs=[ + self._body_link_state_w.data, + ], + device=self.device, + ) + self._body_link_state_w.timestamp = self._sim_timestamp + + return self._body_link_state_w.data + + @property + def body_com_state_w(self) -> wp.array: + """Deprecated, same as :attr:`body_com_pose_w` and :attr:`body_com_vel_w`.""" + warnings.warn( + "The `body_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_com_pose_w` and " + "`body_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_com_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_body_pose_and_vel_to_state, + dim=(self.num_instances, self.num_bodies), + inputs=[ + self.body_com_pose_w, + self.body_com_vel_w, + ], + outputs=[ + self._body_com_state_w.data, + ], + device=self.device, + ) + self._body_com_state_w.timestamp = self._sim_timestamp + + return self._body_com_state_w.data diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/cloner/__init__.pyi index 05f017bf6167..f55377295a3e 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/__init__.pyi +++ b/source/isaaclab_newton/isaaclab_newton/cloner/__init__.pyi @@ -4,7 +4,8 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ - "newton_replicate", + "newton_physics_replicate", + "newton_visualizer_prebuild", ] -from .newton_replicate import newton_replicate +from .newton_replicate import newton_physics_replicate, newton_visualizer_prebuild diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py b/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py index 854edc7bb3c9..e3a19a3e6987 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py @@ -5,38 +5,59 @@ from __future__ import annotations +from collections.abc import Callable + import torch import warp as wp from newton import ModelBuilder, solvers +from newton.usd import SchemaResolverNewton, SchemaResolverPhysx + +from pxr import Usd, UsdGeom -from pxr import Usd +from isaaclab.physics.scene_data_requirements import VisualizerPrebuiltArtifacts from isaaclab_newton.physics import NewtonManager -def newton_replicate( +def _build_newton_builder_from_mapping( stage: Usd.Stage, sources: list[str], - destinations: list[str], env_ids: torch.Tensor, mapping: torch.Tensor, positions: torch.Tensor | None = None, quaternions: torch.Tensor | None = None, - device: str = "cpu", up_axis: str = "Z", simplify_meshes: bool = True, -): - """Replicate prims into a Newton ``ModelBuilder`` using a per-source mapping.""" +) -> tuple[ModelBuilder, object]: + """Build a Newton model builder from clone mapping inputs. + + Args: + stage: USD stage containing source assets. + sources: Source prim paths used for cloning. + env_ids: Environment ids for destination worlds. + mapping: Boolean source-to-environment mapping matrix. + positions: Optional per-environment world positions. + quaternions: Optional per-environment orientations in xyzw order. + up_axis: Up axis for the Newton model builder. + simplify_meshes: Whether to run convex-hull mesh approximation. + Returns: + Tuple of the populated Newton model builder and stage metadata returned by ``add_usd``. + """ if positions is None: positions = torch.zeros((mapping.size(1), 3), device=mapping.device, dtype=torch.float32) if quaternions is None: quaternions = torch.zeros((mapping.size(1), 4), device=mapping.device, dtype=torch.float32) quaternions[:, 3] = 1.0 - # load empty stage + schema_resolvers = [SchemaResolverNewton(), SchemaResolverPhysx()] + builder = ModelBuilder(up_axis=up_axis) - stage_info = builder.add_usd(stage, ignore_paths=["/World/envs"] + sources) + stage_info = builder.add_usd( + stage, + ignore_paths=["/World/envs"] + sources, + schema_resolvers=schema_resolvers, + ) # The prototype is built from env_0 in absolute world coordinates. # add_builder xforms are deltas from env_0 so positions don't get double-counted. @@ -50,6 +71,7 @@ def newton_replicate( root_path=src_path, load_visual_shapes=True, skip_mesh_approximation=True, + schema_resolvers=schema_resolvers, ) if simplify_meshes: p.approximate_meshes("convex_hull", keep_visual_shapes=True) @@ -57,12 +79,9 @@ def newton_replicate( # create a separate world for each environment (heterogeneous spawning) # Newton assigns sequential world IDs (0, 1, 2, ...), so we need to track the mapping - newton_world_to_env_id = {} - for col, env_id in enumerate(env_ids.tolist()): + for col, _ in enumerate(env_ids.tolist()): # begin a new world context (Newton assigns world ID = col) builder.begin_world() - newton_world_to_env_id[col] = env_id - # add all active sources for this world delta_pos = (positions[col] - env0_pos).tolist() for row in torch.nonzero(mapping[:, col], as_tuple=True)[0].tolist(): @@ -70,10 +89,24 @@ def newton_replicate( protos[sources[row]], xform=wp.transform(delta_pos, quaternions[col].tolist()), ) - # end the world context builder.end_world() + return builder, stage_info + + +def _rename_builder_labels( + builder: ModelBuilder, sources: list[str], destinations: list[str], env_ids: torch.Tensor, mapping: torch.Tensor +) -> None: + """Rename builder labels/keys from source roots to destination roots. + + Args: + builder: Newton model builder to update in-place. + sources: Source prim root paths. + destinations: Destination prim path templates. + env_ids: Environment ids corresponding to mapping columns. + mapping: Boolean source-to-environment mapping matrix. + """ # per-source, per-world renaming (strict prefix swap), compact style preserved for i, src_path in enumerate(sources): src_prefix_len = len(src_path.rstrip("/")) @@ -92,6 +125,147 @@ def newton_replicate( if world_id in world_roots and labels[k].startswith(src_path): labels[k] = swap(labels[k], world_roots[world_id]) + +def newton_physics_replicate( + stage: Usd.Stage, + sources: list[str], + destinations: list[str], + env_ids: torch.Tensor, + mapping: torch.Tensor, + positions: torch.Tensor | None = None, + quaternions: torch.Tensor | None = None, + device: str = "cpu", + up_axis: str = "Z", + simplify_meshes: bool = True, +): + """Replicate prims into a Newton ``ModelBuilder`` using a per-source mapping. + + Args: + stage: USD stage containing source assets. + sources: Source prim paths used for cloning. + destinations: Destination prim path templates. + env_ids: Environment ids for destination worlds. + mapping: Boolean source-to-environment mapping matrix. + positions: Optional per-environment world positions. + quaternions: Optional per-environment orientations in xyzw order. + device: Device used by the finalized Newton model builder. + up_axis: Up axis for the Newton model builder. + simplify_meshes: Whether to run convex-hull mesh approximation. + + Returns: + Tuple of the populated Newton model builder and stage metadata. + """ + builder, stage_info = _build_newton_builder_from_mapping( + stage=stage, + sources=sources, + env_ids=env_ids, + mapping=mapping, + positions=positions, + quaternions=quaternions, + up_axis=up_axis, + simplify_meshes=simplify_meshes, + ) + _rename_builder_labels(builder, sources, destinations, env_ids, mapping) NewtonManager.set_builder(builder) NewtonManager._num_envs = mapping.size(1) return builder, stage_info + + +def newton_visualizer_prebuild( + stage: Usd.Stage, + sources: list[str], + destinations: list[str], + env_ids: torch.Tensor, + mapping: torch.Tensor, + positions: torch.Tensor | None = None, + quaternions: torch.Tensor | None = None, + device: str = "cpu", + up_axis: str = "Z", + simplify_meshes: bool = True, +): + """Replicate a clone plan into a finalized Newton model/state for visualization. + + Unlike :func:`newton_physics_replicate`, this path does not mutate ``NewtonManager`` and is intended + for prebuilding visualizer-only artifacts that can be consumed by scene data providers. + + Args: + stage: USD stage containing source assets. + sources: Source prim paths used for cloning. + destinations: Destination prim path templates. + env_ids: Environment ids for destination worlds. + mapping: Boolean source-to-environment mapping matrix. + positions: Optional per-environment world positions. + quaternions: Optional per-environment orientations in xyzw order. + device: Device used by the finalized Newton model. + up_axis: Up axis for the Newton model builder. + simplify_meshes: Whether to run convex-hull mesh approximation. + + Returns: + Tuple of finalized Newton model and state. + """ + builder, _ = _build_newton_builder_from_mapping( + stage=stage, + sources=sources, + env_ids=env_ids, + mapping=mapping, + positions=positions, + quaternions=quaternions, + up_axis=up_axis, + simplify_meshes=simplify_meshes, + ) + _rename_builder_labels(builder, sources, destinations, env_ids, mapping) + model = builder.finalize(device=device) + state = model.state() + return model, state + + +def create_newton_visualizer_prebuild_clone_fn( + stage, + set_visualizer_artifact: Callable[[VisualizerPrebuiltArtifacts | None], None], +): + """Create a cloner callback that prebuilds Newton visualizer artifacts. + + Args: + stage: USD stage used by the clone callback. + set_visualizer_artifact: Callback used to store the produced prebuilt artifact. + + Returns: + Clone callback that builds and stores visualizer prebuilt artifacts. + """ + up_axis = UsdGeom.GetStageUpAxis(stage) + + def _visualizer_clone_fn( + stage, + sources, + destinations, + env_ids, + mapping, + positions=None, + quaternions=None, + device="cpu", + ): + """Prebuild Newton model/state and store visualizer artifacts for clone consumers.""" + model, state = newton_visualizer_prebuild( + stage=stage, + sources=sources, + destinations=destinations, + env_ids=env_ids, + mapping=mapping, + positions=positions, + quaternions=quaternions, + device=device, + up_axis=up_axis, + ) + set_visualizer_artifact( + VisualizerPrebuiltArtifacts( + model=model, + state=state, + rigid_body_paths=list(getattr(model, "body_label", None) or getattr(model, "body_key", [])), + articulation_paths=list( + getattr(model, "articulation_label", None) or getattr(model, "articulation_key", []) + ), + num_envs=int(mapping.size(1)), + ) + ) + + return _visualizer_clone_fn diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index 9ed2ca1c792e..6e744ce2048f 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -16,6 +16,7 @@ from newton import Axis, CollisionPipeline, Contacts, Control, Model, ModelBuilder, State, eval_fk from newton.sensors import SensorContact as NewtonContactSensor from newton.solvers import SolverBase, SolverFeatherstone, SolverMuJoCo, SolverNotifyFlags, SolverXPBD +from newton.usd import SchemaResolverNewton, SchemaResolverPhysx from isaaclab.physics import PhysicsEvent, PhysicsManager from isaaclab.sim.utils.stage import get_current_stage @@ -110,8 +111,13 @@ def initialize(cls, sim_context: SimulationContext) -> None: cls._gravity_vector = sim.cfg.gravity # type: ignore[union-attr] # USD/Fabric sync for Omniverse rendering (visualizer) or Newton+RTX (Kit cameras) - viz_str = sim.get_setting("/isaaclab/visualizer") or "" - requested = [v.strip() for v in viz_str.split(",") if v.strip()] + try: + requested = sim.resolve_visualizer_types() + except Exception: + requested = [] + viz_raw = sim.get_setting("/isaaclab/visualizer/types") + if isinstance(viz_raw, str): + requested = [v for part in viz_raw.split(",") for v in part.split() if v] from isaaclab.app.settings_manager import get_settings_manager cameras_enabled = bool(get_settings_manager().get("/isaaclab/cameras_enabled", False)) @@ -186,23 +192,25 @@ def step(cls) -> None: # Notify solver of model changes if cls._model_changes: - for change in cls._model_changes: - cls._solver.notify_model_changed(change) - cls._model_changes = set() + with wp.ScopedDevice(PhysicsManager._device): + for change in cls._model_changes: + cls._solver.notify_model_changed(change) + cls._model_changes = set() - # Step simulation (graphed or not; _graph is None when RTX/Fabric sync is active) + # Step simulation (graphed or not; _graph is None when RTX/Fabric sync is active or on CPU) cfg = PhysicsManager._cfg - if cls._graph is not None: - wp.capture_launch(cls._graph) # type: ignore[arg-type] + if cfg is not None and cfg.use_cuda_graph and cls._graph is not None and "cuda" in PhysicsManager._device: # type: ignore[union-attr] + wp.capture_launch(cls._graph) else: - cls._simulate() + with wp.ScopedDevice(PhysicsManager._device): + cls._simulate() # Debug convergence info if cfg is not None and cfg.debug_mode: # type: ignore[union-attr] convergence_data = cls.get_solver_convergence_steps() + logger.info(f"Solver convergence data: {convergence_data}") if convergence_data["max"] == cls._solver.mjw_model.opt.iterations: logger.warning(f"Solver didn't converge! max_iter={convergence_data['max']}") - PhysicsManager._sim_time += cls._solver_dt * cls._num_substeps @classmethod @@ -281,7 +289,6 @@ def start_simulation(cls) -> None: logger.info(f"Finalizing model on device: {device}") cls._builder.up_axis = Axis.from_string(cls._up_axis) # Set smaller contact margin for manipulation examples (default 10cm is too large) - cls._builder.default_shape_cfg.contact_margin = 0.01 with Timer(name="newton_finalize_builder", msg="Finalize builder took:"): cls._model = cls._builder.finalize(device=device) cls._model.set_gravity(cls._gravity_vector) @@ -315,13 +322,70 @@ def start_simulation(cls) -> None: @classmethod def instantiate_builder_from_stage(cls): - """Create builder from USD stage.""" + """Create builder from USD stage. + + Detects env Xforms (e.g. ``/World/Env_0``, ``/World/Env_1``) and builds + each as a separate Newton world via ``begin_world``/``end_world``. + Falls back to a flat ``add_usd`` when no env Xforms are found. + """ + import re + from pxr import UsdGeom stage = get_current_stage() up_axis = UsdGeom.GetStageUpAxis(stage) + + # Scan /World children for env-like Xforms (Env_0, env_1, ...) + env_pattern = re.compile(r"^[Ee]nv_(\d+)$") + world_prim = stage.GetPrimAtPath("/World") + env_paths: list[tuple[int, str]] = [] + if world_prim and world_prim.IsValid(): + for child in world_prim.GetChildren(): + m = env_pattern.match(child.GetName()) + if m: + env_paths.append((int(m.group(1)), child.GetPath().pathString)) + env_paths.sort(key=lambda x: x[0]) + builder = ModelBuilder(up_axis=up_axis) - builder.add_usd(stage) + + schema_resolvers = [SchemaResolverNewton(), SchemaResolverPhysx()] + + if not env_paths: + # No env Xforms — flat loading + builder.add_usd(stage, schema_resolvers=schema_resolvers) + else: + # Load everything except the env subtrees (ground plane, lights, etc.) + ignore_paths = [path for _, path in env_paths] + builder.add_usd(stage, ignore_paths=ignore_paths, schema_resolvers=schema_resolvers) + + # Build a prototype from the first env (all envs assumed identical) + _, proto_path = env_paths[0] + proto = ModelBuilder(up_axis=up_axis) + proto.add_usd( + stage, + root_path=proto_path, + schema_resolvers=schema_resolvers, + ) + + # Add each env as a separate Newton world + xform_cache = UsdGeom.XformCache() + for _, env_path in env_paths: + builder.begin_world() + world_xform = xform_cache.GetLocalToWorldTransform(stage.GetPrimAtPath(env_path)) + translation = world_xform.ExtractTranslation() + rotation = world_xform.ExtractRotationQuat() + pos = (translation[0], translation[1], translation[2]) + quat = ( + rotation.GetImaginary()[0], + rotation.GetImaginary()[1], + rotation.GetImaginary()[2], + rotation.GetReal(), + ) + builder.add_builder(proto, xform=wp.transform(pos, quat)) + builder.end_world() + + cls._num_envs = len(env_paths) + cls.set_builder(builder) @classmethod @@ -413,16 +477,14 @@ def initialize_solver(cls) -> None: # Initialize contacts and collision pipeline cls._initialize_contacts() - # Ensure we are using a CUDA enabled device device = PhysicsManager._device - assert device.startswith("cuda"), "NewtonManager only supports CUDA enabled devices" # Skip CUDA graph when syncing to USD/Fabric for RTX: capture conflicts with RTX/Replicator # using the legacy stream (cudaErrorStreamCaptureImplicit). use_cuda_graph = cfg.use_cuda_graph and (cls._usdrt_stage is None) # type: ignore[union-attr] with Timer(name="newton_cuda_graph", msg="CUDA graph took:"): - if use_cuda_graph: + if use_cuda_graph and "cuda" in device: with wp.ScopedCapture() as capture: cls._simulate() cls._graph = capture.graph @@ -521,22 +583,20 @@ def add_contact_sensor( shape_names_expr: str | list[str] | None = None, contact_partners_body_expr: str | list[str] | None = None, contact_partners_shape_expr: str | list[str] | None = None, - prune_noncolliding: bool = True, verbose: bool = False, ) -> tuple[str | list[str] | None, str | list[str] | None, str | list[str] | None, str | list[str] | None]: """Add a contact sensor for reporting contacts between bodies/shapes. - Note: Only one contact sensor can be active at a time. + Converts Isaac Lab pattern conventions (``.*`` regex, full USD paths) to + fnmatch globs and delegates to :class:`newton.sensors.SensorContact`. Args: body_names_expr: Expression for body names to sense. shape_names_expr: Expression for shape names to sense. contact_partners_body_expr: Expression for contact partner body names. contact_partners_shape_expr: Expression for contact partner shape names. - prune_noncolliding: Make force matrix sparse using collision pairs. verbose: Print verbose information. """ - # Validate inputs if body_names_expr is None and shape_names_expr is None: raise ValueError("At least one of body_names_expr or shape_names_expr must be provided") if body_names_expr is not None and shape_names_expr is not None: @@ -544,64 +604,63 @@ def add_contact_sensor( if contact_partners_body_expr is not None and contact_partners_shape_expr is not None: raise ValueError("Only one of contact_partners_body_expr or contact_partners_shape_expr must be provided") - # Log sensor configuration sensor_target = body_names_expr or shape_names_expr partner_filter = contact_partners_body_expr or contact_partners_shape_expr or "all bodies/shapes" logger.info(f"Adding contact sensor for {sensor_target} with filter {partner_filter}") - # Create unique key for this sensor - sensor_key = (body_names_expr, shape_names_expr, contact_partners_body_expr, contact_partners_shape_expr) + def _hashable_key(x): + return tuple(x) if isinstance(x, list) else x - def _to_fnmatch_patterns(expr: str | list[str] | None): + def _to_fnmatch(expr: str | list[str] | None) -> str | list[str] | None: + """Convert Isaac Lab regex expressions (``.*``) to fnmatch glob (``*``).""" if expr is None: return None if isinstance(expr, str): return expr.replace(".*", "*") - return [pattern.replace(".*", "*") for pattern in expr] + return [p.replace(".*", "*") for p in expr] + + def _normalize_for_labels(expr: str | list[str] | None, labels: list[str]) -> str | list[str] | None: + """Strip leading path components from *expr* when labels are bare names. - def _normalize_expr_for_model_labels(expr: str | list[str] | None, labels: list[str] | None): + Model labels may be full USD paths (``/World/envs/env_0/Robot/base``) or bare + names (``base``). When the labels are bare names but the user expression + contains slashes, we strip everything up to the last ``/``. + """ if expr is None or not labels: return expr - label_has_paths = any("/" in label for label in labels) - expr_list = [expr] if isinstance(expr, str) else list(expr) - expr_uses_paths = any("/" in pattern for pattern in expr_list) + label_has_paths = any("/" in lbl for lbl in labels) + items = [expr] if isinstance(expr, str) else list(expr) + expr_uses_paths = any("/" in p for p in items) if label_has_paths or not expr_uses_paths: return expr - normalized = [pattern.rsplit("/", 1)[-1] for pattern in expr_list] + normalized = [p.rsplit("/", 1)[-1] for p in items] return normalized[0] if isinstance(expr, str) else normalized - model = cls._model - body_labels = getattr(model, "body_label", None) or getattr(model, "body_key", None) - shape_labels = getattr(model, "shape_label", None) or getattr(model, "shape_key", None) - sensing_obj_bodies = _normalize_expr_for_model_labels(_to_fnmatch_patterns(body_names_expr), body_labels) - sensing_obj_shapes = _normalize_expr_for_model_labels(_to_fnmatch_patterns(shape_names_expr), shape_labels) - counterpart_bodies = _normalize_expr_for_model_labels( - _to_fnmatch_patterns(contact_partners_body_expr), body_labels - ) - counterpart_shapes = _normalize_expr_for_model_labels( - _to_fnmatch_patterns(contact_partners_shape_expr), shape_labels + sensor_key = ( + _hashable_key(body_names_expr), + _hashable_key(shape_names_expr), + _hashable_key(contact_partners_body_expr), + _hashable_key(contact_partners_shape_expr), ) - # Create and store the sensor - # Note: SensorContact constructor requests 'force' attribute from the model - newton_sensor = NewtonContactSensor( - cls._model, - sensing_obj_bodies=sensing_obj_bodies, - sensing_obj_shapes=sensing_obj_shapes, - counterpart_bodies=counterpart_bodies, - counterpart_shapes=counterpart_shapes, - include_total=True, - prune_noncolliding=prune_noncolliding, - verbose=verbose, - ) - cls._newton_contact_sensors[sensor_key] = newton_sensor + body_labels = list(cls._model.body_label) + shape_labels = list(cls._model.shape_label) + + with Timer(name="newton_contact_sensor", msg="Contact sensor construction took:"): + sensor = NewtonContactSensor( + cls._model, + sensing_obj_bodies=_normalize_for_labels(_to_fnmatch(body_names_expr), body_labels), + sensing_obj_shapes=_normalize_for_labels(_to_fnmatch(shape_names_expr), shape_labels), + counterpart_bodies=_normalize_for_labels(_to_fnmatch(contact_partners_body_expr), body_labels), + counterpart_shapes=_normalize_for_labels(_to_fnmatch(contact_partners_shape_expr), shape_labels), + include_total=True, + verbose=verbose, + ) + + cls._newton_contact_sensors[sensor_key] = sensor cls._report_contacts = True - # Regenerate contacts only if they were already created without force attribute - # If solver is not initialized, contacts will be created with force in initialize_solver() - if cls._solver is not None and cls._contacts is not None: - # Only regenerate if contacts don't have force attribute (sensor.update() requires it) - if cls._contacts.force is None: - cls._initialize_contacts() + if cls._solver is not None and cls._contacts is not None and cls._contacts.force is None: + cls._initialize_contacts() return sensor_key diff --git a/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py b/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py index 5dd4e90fd1e8..27342aa68779 100644 --- a/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py +++ b/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py @@ -19,13 +19,12 @@ from isaaclab.renderers import BaseRenderer from isaaclab.sim import SimulationContext from isaaclab.utils.math import convert_camera_frame_orientation_convention -from isaaclab.visualizers import VisualizerCfg from .newton_warp_renderer_cfg import NewtonWarpRendererCfg if TYPE_CHECKING: + from isaaclab.physics import BaseSceneDataProvider from isaaclab.sensors import SensorBase - from isaaclab.sim.scene_data_providers import SceneDataProvider logger = logging.getLogger(__name__) @@ -202,5 +201,5 @@ def cleanup(self, render_data: RenderData | None): if render_data: render_data.sensor = None - def get_scene_data_provider(self) -> SceneDataProvider: - return SimulationContext.instance().initialize_scene_data_provider([VisualizerCfg(visualizer_type="newton")]) + def get_scene_data_provider(self) -> BaseSceneDataProvider: + return SimulationContext.instance().initialize_scene_data_provider() diff --git a/source/isaaclab/isaaclab/sim/scene_data_providers/__init__.py b/source/isaaclab_newton/isaaclab_newton/scene_data_providers/__init__.py similarity index 81% rename from source/isaaclab/isaaclab/sim/scene_data_providers/__init__.py rename to source/isaaclab_newton/isaaclab_newton/scene_data_providers/__init__.py index b19bb3b92008..cf0f30853ead 100644 --- a/source/isaaclab/isaaclab/sim/scene_data_providers/__init__.py +++ b/source/isaaclab_newton/isaaclab_newton/scene_data_providers/__init__.py @@ -3,7 +3,7 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Scene data providers for visualizers and renderers.""" +"""Newton scene data provider backends.""" from isaaclab.utils.module import lazy_export diff --git a/source/isaaclab_newton/isaaclab_newton/scene_data_providers/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/scene_data_providers/__init__.pyi new file mode 100644 index 000000000000..3cb204031738 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/scene_data_providers/__init__.pyi @@ -0,0 +1,10 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +__all__ = [ + "NewtonSceneDataProvider", +] + +from .newton_scene_data_provider import NewtonSceneDataProvider diff --git a/source/isaaclab/isaaclab/sim/scene_data_providers/newton_scene_data_provider.py b/source/isaaclab_newton/isaaclab_newton/scene_data_providers/newton_scene_data_provider.py similarity index 72% rename from source/isaaclab/isaaclab/sim/scene_data_providers/newton_scene_data_provider.py rename to source/isaaclab_newton/isaaclab_newton/scene_data_providers/newton_scene_data_provider.py index ca43170b47a9..f3b10dc40044 100644 --- a/source/isaaclab/isaaclab/sim/scene_data_providers/newton_scene_data_provider.py +++ b/source/isaaclab_newton/isaaclab_newton/scene_data_providers/newton_scene_data_provider.py @@ -14,12 +14,14 @@ from pxr import UsdGeom +from isaaclab.physics.base_scene_data_provider import BaseSceneDataProvider + logger = logging.getLogger(__name__) _ENV_ID_RE = re.compile(r"/World/envs/env_(\d+)") -class NewtonSceneDataProvider: +class NewtonSceneDataProvider(BaseSceneDataProvider): """Scene data provider for Newton physics backend. Provides access to Newton model, state, and USD stage for visualizers and renderers. @@ -28,14 +30,31 @@ class NewtonSceneDataProvider: Newton backend already owns the authoritative model and state. """ - def __init__(self, visualizer_cfgs: list[Any] | None, stage, simulation_context) -> None: + def __init__(self, stage, simulation_context) -> None: + """Initialize the Newton scene data provider. + + Args: + stage: USD stage handle. + simulation_context: Active simulation context. + """ self._simulation_context = simulation_context self._stage = stage self._metadata = {"physics_backend": "newton"} self._num_envs: int | None = None self._warned_once: set[str] = set() + # Determine if usd stage sync is required for selected renderers and visualizers + requirements = self._simulation_context.get_scene_data_requirements() + self._needs_usd_sync = bool(requirements.requires_usd_stage) + def _warn_once(self, key: str, message: str, *args) -> None: + """Emit a warning once per unique key. + + Args: + key: Unique warning key. + message: Warning message format string. + *args: Optional formatting arguments. + """ if key in self._warned_once: return self._warned_once.add(key) @@ -44,6 +63,11 @@ def _warn_once(self, key: str, message: str, *args) -> None: # ---- Environment discovery --------------------------------------------------------------- def get_num_envs(self) -> int: + """Return discovered environment count. + + Returns: + Number of environments discovered from stage prim paths. + """ if self._num_envs is not None and self._num_envs > 0: return self._num_envs discovered = self._determine_num_envs_in_scene() @@ -53,6 +77,11 @@ def get_num_envs(self) -> int: return 0 def _determine_num_envs_in_scene(self) -> int: + """Infer environment count from ``/World/envs/env_`` prim names. + + Returns: + Number of environments inferred from the stage. + """ if self._stage is None: return 0 max_env_id = -1 @@ -68,13 +97,19 @@ def _determine_num_envs_in_scene(self) -> int: # ---- Core provider API ------------------------------------------------------------------- def update(self, env_ids: list[int] | None = None) -> None: - """Sync Newton body transforms to USD Fabric for Kit viewport rendering. + """Sync Newton body transforms to USD Fabric when a Kit viewport is active. Called at render cadence by :meth:`~isaaclab.sim.SimulationContext.update_scene_data_provider`, - after forward kinematics have been evaluated. Delegates to - :meth:`~isaaclab_newton.physics.NewtonManager.sync_transforms_to_usd`, which is a - no-op when no Kit visualizer is active (``_fabric_manager is None``). + after forward kinematics have been evaluated. Only calls + :meth:`~isaaclab_newton.physics.NewtonManager.sync_transforms_to_usd` when a Kit + (or other USD-based) visualizer is in use. When both sim and rendering backend + are Newton (or Rerun), the sync is skipped to avoid unnecessary slowdown. + + Args: + env_ids: Optional environment id selection. Unused in this provider. """ + if not self._needs_usd_sync: + return try: from isaaclab_newton.physics import NewtonManager @@ -83,7 +118,11 @@ def update(self, env_ids: list[int] | None = None) -> None: pass def get_newton_model(self) -> Any | None: - """Return Newton model from NewtonManager.""" + """Return Newton model from ``NewtonManager``. + + Returns: + Newton model object, or ``None`` when unavailable. + """ from isaaclab_newton.physics import NewtonManager return NewtonManager.get_model() @@ -103,22 +142,43 @@ def get_newton_state(self, env_ids: list[int] | None = None) -> Any | None: return NewtonManager.get_state_0() def get_model(self) -> Any | None: - """Alias for get_newton_model (visualizer compatibility).""" + """Alias for :meth:`get_newton_model` for visualizer compatibility. + + Returns: + Newton model object, or ``None`` when unavailable. + """ return self.get_newton_model() def get_state(self, env_ids: list[int] | None = None) -> Any | None: - """Alias for get_newton_state (visualizer compatibility).""" + """Alias for :meth:`get_newton_state` for visualizer compatibility. + + Args: + env_ids: Optional list of environment ids. + + Returns: + Newton state object, or ``None`` when unavailable. + """ return self.get_newton_state(env_ids) def get_usd_stage(self) -> Any | None: - """Return the USD stage handle.""" + """Return the USD stage handle. + + Returns: + USD stage object, or ``None`` when unavailable. + """ if self._stage is not None: return self._stage return getattr(self._simulation_context, "stage", None) def get_metadata(self) -> dict[str, Any]: + """Return provider metadata. + + Returns: + Metadata dictionary with backend and synchronization information. + """ out = dict(self._metadata) out["num_envs"] = self.get_num_envs() + out["needs_usd_sync"] = self._needs_usd_sync return out def get_transforms(self) -> dict[str, Any] | None: @@ -126,9 +186,13 @@ def get_transforms(self) -> dict[str, Any] | None: Reads body_q from the authoritative Newton state and splits it into positions (vec3) and orientations (quaternion xyzw). + + Returns: + Dictionary containing positions and orientations, or ``None`` when unavailable. """ try: import warp as wp + from isaaclab_newton.physics import NewtonManager state = NewtonManager.get_state_0() @@ -148,9 +212,14 @@ def get_transforms(self) -> dict[str, Any] | None: return None def get_velocities(self) -> dict[str, Any] | None: - """Return body velocities from Newton state.""" + """Return body velocities from Newton state. + + Returns: + Dictionary containing linear and angular velocities, or ``None`` when unavailable. + """ try: import warp as wp + from isaaclab_newton.physics import NewtonManager state = NewtonManager.get_state_0() @@ -174,11 +243,20 @@ def get_velocities(self) -> dict[str, Any] | None: return None def get_contacts(self) -> dict[str, Any] | None: - """Contacts not yet supported for Newton provider.""" + """Return contact data for Newton provider. + + Returns: + ``None`` because contacts are not currently implemented in this provider. + """ return None def get_camera_transforms(self) -> dict[str, Any] | None: - """Return per-camera, per-env transforms (positions, orientations).""" + """Return per-camera, per-environment transforms. + + Returns: + Dictionary containing camera order, positions, orientations, and environment count, + or ``None`` when unavailable. + """ if self._stage is None: return None diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor.py b/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor.py index 92edde91c6a3..5517b845438b 100644 --- a/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor.py +++ b/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor.py @@ -309,7 +309,6 @@ def _initialize_impl(self): shape_names_expr=self.cfg.sensor_shape_prim_expr or None, contact_partners_body_expr=self.cfg.filter_prim_paths_expr or None, contact_partners_shape_expr=self.cfg.filter_shape_prim_expr or None, - prune_noncolliding=True, ) self._create_buffers() @@ -355,10 +354,12 @@ def get_name(idx, kind): return shape_labels[idx].split("/")[-1] return "MATCH_ANY" - self._sensor_names = [get_name(idx, kind) for idx, kind in self.contact_view.sensing_objs] + flat_sensing = [obj for world_objs in self.contact_view.sensing_objs for obj in world_objs] + self._sensor_names = [get_name(idx, kind) for idx, kind in flat_sensing] # Assumes the environments are processed in order. self._sensor_names = self._sensor_names[: self._num_sensors] - self._filter_object_names = [get_name(idx, kind) for idx, kind in self.contact_view.counterparts] + flat_counterparts = [obj for world_objs in self.contact_view.counterparts for obj in world_objs] + self._filter_object_names = [get_name(idx, kind) for idx, kind in flat_counterparts] # Number of filter objects (counterparts minus the total column) self._num_filter_objects = max(newton_shape[1] - 1, 0) diff --git a/source/isaaclab_newton/isaaclab_newton/test/mock_interfaces/views/__init__.py b/source/isaaclab_newton/isaaclab_newton/test/mock_interfaces/views/__init__.py index ca4ed19d1514..721f887ba475 100644 --- a/source/isaaclab_newton/isaaclab_newton/test/mock_interfaces/views/__init__.py +++ b/source/isaaclab_newton/isaaclab_newton/test/mock_interfaces/views/__init__.py @@ -5,8 +5,9 @@ """Mock Newton views.""" -from .mock_articulation_view import MockNewtonArticulationView +from .mock_articulation_view import MockNewtonArticulationView, MockNewtonCollectionView __all__ = [ "MockNewtonArticulationView", + "MockNewtonCollectionView", ] diff --git a/source/isaaclab_newton/isaaclab_newton/test/mock_interfaces/views/mock_articulation_view.py b/source/isaaclab_newton/isaaclab_newton/test/mock_interfaces/views/mock_articulation_view.py index 1064eca3e4fa..dc698e97d223 100644 --- a/source/isaaclab_newton/isaaclab_newton/test/mock_interfaces/views/mock_articulation_view.py +++ b/source/isaaclab_newton/isaaclab_newton/test/mock_interfaces/views/mock_articulation_view.py @@ -11,6 +11,151 @@ import warp as wp +class MockNewtonCollectionView: + """Mock Newton ArticulationView for rigid object collection testing. + + This class mimics the interface of a single ``ArticulationView`` whose combined + fnmatch pattern matches **B** body types per world. Its data is shaped + ``(N, B, ...)`` rather than the ``(N, 1, ...)`` convention used for single + articulations or single rigid bodies. + + Data Shapes: + - root_transforms: ``(N, B)`` dtype=wp.transformf + - root_velocities: ``(N, B)`` dtype=wp.spatial_vectorf + - body_com: ``(N, B, 1)`` dtype=wp.vec3f + - body_mass: ``(N, B, 1)`` dtype=wp.float32 + - body_inertia: ``(N, B, 1)`` dtype=wp.mat33f + - body_f: ``(N, B, 1)`` dtype=wp.spatial_vectorf + + Where N = num_envs, B = num_bodies (body types in the collection). + ``count`` returns ``N * B`` so that the data class can compute + ``num_instances = count // num_bodies``. + """ + + def __init__( + self, + num_envs: int = 2, + num_bodies: int = 3, + device: str = "cpu", + body_names: list[str] | None = None, + ): + self._num_envs = num_envs + self._num_bodies = num_bodies + self._device = device + self._noop_setters = False + + self._body_names = body_names if body_names is not None else [f"object_{i}" for i in range(num_bodies)] + + # Internal state (lazily initialised) + self._root_transforms: wp.array | None = None + self._root_velocities: wp.array | None = None + self._attributes: dict[str, wp.array | None] = { + "body_com": None, + "body_mass": None, + "body_inertia": None, + "body_f": None, + } + + # -- Properties -------------------------------------------------------- + + @property + def count(self) -> int: + """Total matched entities (``N * B``).""" + return self._num_envs * self._num_bodies + + @property + def body_names(self) -> list[str]: + return self._body_names + + # -- Lazy init helpers ------------------------------------------------- + + def _ensure_root_transforms(self) -> wp.array: + if self._root_transforms is None: + self._root_transforms = wp.zeros( + (self._num_envs, self._num_bodies), dtype=wp.transformf, device=self._device + ) + return self._root_transforms + + def _ensure_root_velocities(self) -> wp.array: + if self._root_velocities is None: + self._root_velocities = wp.zeros( + (self._num_envs, self._num_bodies), dtype=wp.spatial_vectorf, device=self._device + ) + return self._root_velocities + + def _ensure_attribute(self, name: str) -> wp.array: + if self._attributes[name] is None: + self._attributes[name] = self._create_default_attribute(name) + return self._attributes[name] + + def _create_default_attribute(self, name: str) -> wp.array: + N, B = self._num_envs, self._num_bodies + dev = self._device + if name == "body_com": + return wp.zeros((N, B, 1), dtype=wp.vec3f, device=dev) + elif name == "body_mass": + return wp.zeros((N, B, 1), dtype=wp.float32, device=dev) + elif name == "body_inertia": + return wp.zeros((N, B, 1), dtype=wp.mat33f, device=dev) + elif name == "body_f": + return wp.zeros((N, B, 1), dtype=wp.spatial_vectorf, device=dev) + else: + raise KeyError(f"Unknown attribute: {name}") + + # -- Getters ----------------------------------------------------------- + + def get_root_transforms(self, state) -> wp.array: + return self._ensure_root_transforms() + + def get_root_velocities(self, state) -> wp.array: + return self._ensure_root_velocities() + + def get_attribute(self, name: str, model_or_state) -> wp.array: + return self._ensure_attribute(name) + + # -- Setters ----------------------------------------------------------- + + def set_root_transforms(self, state, transforms: wp.array) -> None: + if self._noop_setters: + return + self._ensure_root_transforms().assign(transforms) + + def set_root_velocities(self, state, velocities: wp.array) -> None: + if self._noop_setters: + return + self._ensure_root_velocities().assign(velocities) + + # -- Mock data injection ----------------------------------------------- + + def set_random_mock_data(self) -> None: + """Set all internal state to random values for testing.""" + N, B = self._num_envs, self._num_bodies + dev = self._device + + # Root transforms + root_tf_np = np.random.randn(N, B, 7).astype(np.float32) + root_tf_np[..., 3:7] /= np.linalg.norm(root_tf_np[..., 3:7], axis=-1, keepdims=True) + self._root_transforms = wp.array(root_tf_np, dtype=wp.transformf, device=dev) + + # Root velocities + root_vel_np = np.random.randn(N, B, 6).astype(np.float32) + self._root_velocities = wp.array(root_vel_np, dtype=wp.spatial_vectorf, device=dev) + + # Attributes (all have trailing link dim of 1) + self._attributes["body_com"] = wp.array( + np.random.randn(N, B, 1, 3).astype(np.float32), dtype=wp.vec3f, device=dev + ) + self._attributes["body_mass"] = wp.array( + (np.random.rand(N, B, 1) * 10 + 0.1).astype(np.float32), dtype=wp.float32, device=dev + ) + self._attributes["body_inertia"] = wp.array( + np.random.randn(N, B, 1, 9).astype(np.float32), dtype=wp.mat33f, device=dev + ) + self._attributes["body_f"] = wp.array( + np.random.randn(N, B, 1, 6).astype(np.float32), dtype=wp.spatial_vectorf, device=dev + ) + + class MockNewtonArticulationView: """Mock Newton ArticulationView for unit testing without simulation. @@ -131,6 +276,11 @@ def body_names(self) -> list[str]: """Names of the bodies.""" return self._body_names + @property + def link_names(self) -> list[str]: + """Alias for body_names (Newton calls bodies 'links').""" + return self._body_names + # -- Lazy Initialization Helpers -- def _ensure_root_transforms(self) -> wp.array: diff --git a/source/isaaclab_newton/setup.py b/source/isaaclab_newton/setup.py index 316b36013609..82e0d1bbe72c 100644 --- a/source/isaaclab_newton/setup.py +++ b/source/isaaclab_newton/setup.py @@ -17,11 +17,12 @@ # Minimum dependencies required prior to installation INSTALL_REQUIRES = [ + # generic + "prettytable==3.3.0", # newton - "warp-lang==1.12.0rc2", # avoids pulling newton dep from newton package that could end up being a dev build "mujoco==3.5.0", "mujoco-warp==3.5.0.2", - "newton==1.0.0rc1", + "newton @ git+https://github.com/newton-physics/newton.git@release-1.0", "PyOpenGL-accelerate==3.1.10", ] diff --git a/source/isaaclab_newton/test/assets/test_articulation.py b/source/isaaclab_newton/test/assets/test_articulation.py new file mode 100644 index 000000000000..b872dd9fb63f --- /dev/null +++ b/source/isaaclab_newton/test/assets/test_articulation.py @@ -0,0 +1,2404 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# ignore private usage of variables warning +# pyright: reportPrivateUsage=none + +"""Launch Isaac Sim Simulator first.""" + +from isaaclab.app import AppLauncher + +HEADLESS = True + +# launch omniverse app +simulation_app = AppLauncher(headless=True).app + +"""Rest everything follows.""" + +import sys + +import pytest +import torch +import warp as wp +from isaaclab_newton.assets import Articulation +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +from isaaclab_newton.physics import NewtonManager as SimulationManager +from newton.solvers import SolverNotifyFlags + +import isaaclab.sim as sim_utils +import isaaclab.utils.math as math_utils +import isaaclab.utils.string as string_utils +from isaaclab.actuators import ActuatorBase, IdealPDActuatorCfg, ImplicitActuatorCfg +from isaaclab.assets import ArticulationCfg +from isaaclab.envs.mdp.terminations import joint_effort_out_of_limit +from isaaclab.managers import SceneEntityCfg +from isaaclab.sim import SimulationCfg, build_simulation_context +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.version import get_isaac_sim_version, has_kit + +## +# Pre-defined configs +## +from isaaclab_assets import ANYMAL_C_CFG, FRANKA_PANDA_CFG # isort:skip +# , SHADOW_HAND_CFG # isort:skip + +SIM_CFGs = { + "humanoid": SimulationCfg( + physics=NewtonCfg( + solver_cfg=MJWarpSolverCfg( + njmax=80, + nconmax=25, + ls_iterations=20, + ls_parallel=False, + cone="pyramidal", + update_data_interval=2, + integrator="implicitfast", + impratio=1, + ), + num_substeps=2, + debug_mode=False, + ), + ), + "anymal": SimulationCfg( + dt=1 / 200, + physics=NewtonCfg( + solver_cfg=MJWarpSolverCfg( + njmax=70, + nconmax=70, + ls_iterations=40, + cone="elliptic", + impratio=100, + ls_parallel=False, + integrator="implicitfast", + ), + num_substeps=2, + debug_mode=True, + ), + ), + "panda": SimulationCfg( + dt=1 / 120, + physics=NewtonCfg( + solver_cfg=MJWarpSolverCfg( + njmax=20, + nconmax=20, + ls_iterations=20, + cone="pyramidal", + impratio=1, + ls_parallel=False, + integrator="implicitfast", + ), + num_substeps=1, + debug_mode=False, + ), + ), + "single_joint_implicit": SimulationCfg( + dt=1 / 120, + physics=NewtonCfg( + solver_cfg=MJWarpSolverCfg( + njmax=20, + nconmax=20, + ls_iterations=20, + cone="pyramidal", + impratio=1, + ls_parallel=False, + integrator="implicitfast", + ), + num_substeps=1, + debug_mode=False, + ), + ), + "single_joint_explicit": SimulationCfg( + dt=1 / 120, + physics=NewtonCfg( + solver_cfg=MJWarpSolverCfg( + njmax=20, + nconmax=20, + ls_iterations=20, + cone="pyramidal", + impratio=1, + ls_parallel=False, + integrator="implicitfast", + ), + num_substeps=1, + debug_mode=False, + ), + ), + "shadow_hand": SimulationCfg( + dt=1 / 120, + physics=NewtonCfg( + solver_cfg=MJWarpSolverCfg( + njmax=70, + nconmax=70, + ls_iterations=40, + cone="elliptic", + impratio=100, + ls_parallel=False, + integrator="implicitfast", + ), + num_substeps=2, + debug_mode=True, + ), + ), +} + + +def generate_articulation_cfg( + articulation_type: str, + stiffness: float | None = 10.0, + damping: float | None = 2.0, + velocity_limit: float | None = None, + effort_limit: float | None = None, + velocity_limit_sim: float | None = None, + effort_limit_sim: float | None = None, +) -> ArticulationCfg: + """Generate an articulation configuration. + + Args: + articulation_type: Type of articulation to generate. + It should be one of: "humanoid", "panda", "anymal", "shadow_hand", "single_joint_implicit", + "single_joint_explicit". + stiffness: Stiffness value for the articulation's actuators. Only currently used for "humanoid". + Defaults to 10.0. + damping: Damping value for the articulation's actuators. Only currently used for "humanoid". + Defaults to 2.0. + velocity_limit: Velocity limit for the actuators. Only currently used for "single_joint_implicit" + and "single_joint_explicit". + effort_limit: Effort limit for the actuators. Only currently used for "single_joint_implicit" + and "single_joint_explicit". + velocity_limit_sim: Velocity limit for the actuators (set into the simulation). + Only currently used for "single_joint_implicit" and "single_joint_explicit". + effort_limit_sim: Effort limit for the actuators (set into the simulation). + Only currently used for "single_joint_implicit" and "single_joint_explicit". + + Returns: + The articulation configuration for the requested articulation type. + + """ + if articulation_type == "humanoid": + articulation_cfg = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/IsaacSim/Humanoid/humanoid_instanceable.usd" + ), + init_state=ArticulationCfg.InitialStateCfg(pos=(0.0, 0.0, 1.34)), + actuators={"body": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=stiffness, damping=damping)}, + ) + elif articulation_type == "panda": + articulation_cfg = FRANKA_PANDA_CFG + elif articulation_type == "anymal": + articulation_cfg = ANYMAL_C_CFG + elif articulation_type == "shadow_hand": + pytest.skip("Shadow hand is not supported in Newton") + # articulation_cfg = SHADOW_HAND_CFG + elif articulation_type == "single_joint_implicit": + articulation_cfg = ArticulationCfg( + # we set 80.0 default for max force because default in USD is 10e10 which makes testing annoying. + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/IsaacSim/SimpleArticulation/revolute_articulation.usd", + joint_drive_props=sim_utils.JointDrivePropertiesCfg(max_effort=80.0, max_velocity=5.0), + ), + actuators={ + "joint": ImplicitActuatorCfg( + joint_names_expr=[".*"], + effort_limit_sim=effort_limit_sim, + velocity_limit_sim=velocity_limit_sim, + effort_limit=effort_limit, + velocity_limit=velocity_limit, + stiffness=2000.0, + damping=100.0, + ), + }, + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.0, 0.0, 0.0), + joint_pos=({"RevoluteJoint": 1.5708}), + rot=(0.7071081, 0, 0, 0.7071055), + ), + ) + elif articulation_type == "single_joint_explicit": + # we set 80.0 default for max force because default in USD is 10e10 which makes testing annoying. + articulation_cfg = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/IsaacSim/SimpleArticulation/revolute_articulation.usd", + joint_drive_props=sim_utils.JointDrivePropertiesCfg(max_effort=80.0, max_velocity=5.0), + ), + actuators={ + "joint": IdealPDActuatorCfg( + joint_names_expr=[".*"], + effort_limit_sim=effort_limit_sim, + velocity_limit_sim=velocity_limit_sim, + effort_limit=effort_limit, + velocity_limit=velocity_limit, + stiffness=0.0, + damping=10.0, + ), + }, + ) + elif articulation_type == "spatial_tendon_test_asset": + # we set 80.0 default for max force because default in USD is 10e10 which makes testing annoying. + articulation_cfg = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/IsaacLab/Tests/spatial_tendons.usd", + ), + actuators={ + "joint": ImplicitActuatorCfg( + joint_names_expr=[".*"], + stiffness=2000.0, + damping=100.0, + ), + }, + ) + else: + raise ValueError( + f"Invalid articulation type: {articulation_type}, valid options are 'humanoid', 'panda', 'anymal'," + " 'shadow_hand', 'single_joint_implicit', 'single_joint_explicit' or 'spatial_tendon_test_asset'." + ) + + return articulation_cfg + + +def fix_reversed_joints(stage): + """Fix reversed joints on the USD stage. + + Some USD assets have joints where physics:body0 is the child and physics:body1 is the parent, + which is the opposite of what Newton expects. This function detects reversed joints by building + a graph of body connections and identifying the root body (the one attached to world via a joint + with a missing body target). Any joint where body0 is closer to the root than body1 is swapped. + """ + from pxr import UsdPhysics + + # First pass: find root bodies (bodies with a joint that has only one target, i.e. attached to world) + root_bodies: set[str] = set() + joints_to_check = [] + for prim in stage.Traverse(): + if not prim.IsA(UsdPhysics.Joint): + continue + body0_targets = prim.GetRelationship("physics:body0").GetTargets() + body1_targets = prim.GetRelationship("physics:body1").GetTargets() + if body0_targets and not body1_targets: + root_bodies.add(str(body0_targets[0])) + elif body1_targets and not body0_targets: + root_bodies.add(str(body1_targets[0])) + elif body0_targets and body1_targets: + joints_to_check.append(prim) + + if not root_bodies: + return + + # Second pass: for each joint with two bodies, ensure body0 is the parent (closer to root) + for prim in joints_to_check: + body0_rel = prim.GetRelationship("physics:body0") + body1_rel = prim.GetRelationship("physics:body1") + body0_path = str(body0_rel.GetTargets()[0]) + body1_path = str(body1_rel.GetTargets()[0]) + + # Determine if we need to swap: body1 is root or ancestor → need to swap + body1_is_parent = body1_path in root_bodies or body0_path.startswith(body1_path + "/") + body0_is_parent = body0_path in root_bodies or body1_path.startswith(body0_path + "/") + + if body0_is_parent or not body1_is_parent: + continue # already correct or ambiguous + + # Swap body0 and body1 + body0_rel.SetTargets(body1_rel.GetTargets()) + body1_rel.SetTargets([body0_path]) + + # Swap local transforms + for attr_suffix in ("localPos", "localRot"): + attr0 = prim.GetAttribute(f"physics:{attr_suffix}0") + attr1 = prim.GetAttribute(f"physics:{attr_suffix}1") + val0, val1 = attr0.Get(), attr1.Get() + if val0 is not None and val1 is not None: + attr0.Set(val1) + attr1.Set(val0) + + +_REVERSED_JOINT_USD_FILES = {"revolute_articulation.usd"} +"""USD filenames with known reversed joint body0/body1 ordering.""" + + +def generate_articulation( + articulation_cfg: ArticulationCfg, num_articulations: int, device: str +) -> tuple[Articulation, torch.tensor]: + """Generate an articulation from a configuration. + + Handles the creation of the articulation, the environment prims and the articulation's environment + translations + + Args: + articulation_cfg: Articulation configuration. + num_articulations: Number of articulations to generate. + device: Device to use for the tensors. + + Returns: + The articulation and environment translations. + + """ + # Generate translations of 2.5 m in x for each articulation + translations = torch.zeros(num_articulations, 3, device=device) + translations[:, 0] = torch.arange(num_articulations) * 2.5 + + # Create Top-level Xforms, one for each articulation + for i in range(num_articulations): + sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=translations[i][:3]) + articulation = Articulation(articulation_cfg.replace(prim_path="/World/Env_.*/Robot")) + + # Fix reversed joints for known-broken USD assets (body0/body1 swapped) + usd_path = getattr(articulation_cfg.spawn, "usd_path", "") + if any(name in usd_path for name in _REVERSED_JOINT_USD_FILES): + import omni.usd + + fix_reversed_joints(omni.usd.get_context().get_stage()) + + return articulation, translations + + +@pytest.fixture +def sim(request): + """Create simulation context with the specified device.""" + device = request.getfixturevalue("device") + if "gravity_enabled" in request.fixturenames: + gravity_enabled = request.getfixturevalue("gravity_enabled") + else: + gravity_enabled = True # default to gravity enabled + if "add_ground_plane" in request.fixturenames: + add_ground_plane = request.getfixturevalue("add_ground_plane") + else: + add_ground_plane = False # default to no ground plane + articulation_type = request.getfixturevalue("articulation_type") + sim_cfg = SIM_CFGs[articulation_type] + sim_cfg.device = device + with build_simulation_context( + device=device, + auto_add_lighting=True, + gravity_enabled=gravity_enabled, + add_ground_plane=add_ground_plane, + sim_cfg=sim_cfg, + ) as sim: + sim._app_control_on_stop_handle = None + yield sim + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.parametrize("articulation_type", ["humanoid"]) +@pytest.mark.isaacsim_ci +def test_initialization_floating_base_non_root(sim, num_articulations, device, add_ground_plane, articulation_type): + """Test initialization for a floating-base with articulation root on a rigid body. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is not fixed base + 3. All buffers have correct shapes + 4. The articulation can be simulated + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type, stiffness=0.0, damping=0.0) + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + + # Check if articulation is initialized + assert articulation.is_initialized + # Check that is fixed base + assert not articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) + assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) + assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 21) + + # -- actuator type + for actuator_name, actuator in articulation.actuators.items(): + is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) + assert actuator.is_implicit_model == is_implicit_model_cfg + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.parametrize("articulation_type", ["anymal"]) +@pytest.mark.isaacsim_ci +def test_initialization_floating_base(sim, num_articulations, device, add_ground_plane, articulation_type): + """Test initialization for a floating-base with articulation root on provided prim path. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is not fixed base + 3. All buffers have correct shapes + 4. The articulation can be simulated + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type, stiffness=0.0, damping=0.0) + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that floating base + assert not articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) + assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) + assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 12) + assert wp.to_torch(articulation.data.body_mass).shape == (num_articulations, articulation.num_bodies) + assert wp.to_torch(articulation.data.body_inertia).shape == (num_articulations, articulation.num_bodies, 9) + + # -- actuator type + for actuator_name, actuator in articulation.actuators.items(): + is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) + assert actuator.is_implicit_model == is_implicit_model_cfg + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("articulation_type", ["panda"]) +@pytest.mark.isaacsim_ci +def test_initialization_fixed_base(sim, num_articulations, device, articulation_type): + """Test initialization for fixed base. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is fixed base + 3. All buffers have correct shapes + 4. The articulation maintains its default state + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, translations = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that fixed base + assert articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) + assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) + assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 9) + assert wp.to_torch(articulation.data.body_mass).shape == (num_articulations, articulation.num_bodies) + assert wp.to_torch(articulation.data.body_inertia).shape == (num_articulations, articulation.num_bodies, 9) + + # -- actuator type + for actuator_name, actuator in articulation.actuators.items(): + is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) + assert actuator.is_implicit_model == is_implicit_model_cfg + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + # check that the root is at the correct state - its default state as it is fixed base + default_root_pose = wp.to_torch(articulation.data.default_root_pose).clone() + default_root_vel = wp.to_torch(articulation.data.default_root_vel).clone() + default_root_pose[:, :3] = default_root_pose[:, :3] + translations + + torch.testing.assert_close(wp.to_torch(articulation.data.root_link_pose_w), default_root_pose) + torch.testing.assert_close(wp.to_torch(articulation.data.root_com_vel_w), default_root_vel) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.parametrize("articulation_type", ["single_joint_implicit"]) +@pytest.mark.isaacsim_ci +def test_initialization_fixed_base_single_joint(sim, num_articulations, device, add_ground_plane, articulation_type): + """Test initialization for fixed base articulation with a single joint. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is fixed base + 3. All buffers have correct shapes + 4. The articulation maintains its default state + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, translations = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that fixed base + assert articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) + assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) + assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 1) + assert wp.to_torch(articulation.data.body_mass).shape == (num_articulations, articulation.num_bodies) + assert wp.to_torch(articulation.data.body_inertia).shape == (num_articulations, articulation.num_bodies, 9) + + # -- actuator type + for actuator_name, actuator in articulation.actuators.items(): + is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) + assert actuator.is_implicit_model == is_implicit_model_cfg + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + # check that the root is at the correct state - its default state as it is fixed base + default_root_pose = wp.to_torch(articulation.data.default_root_pose).clone() + default_root_vel = wp.to_torch(articulation.data.default_root_vel).clone() + default_root_pose[:, :3] = default_root_pose[:, :3] + translations + + torch.testing.assert_close(wp.to_torch(articulation.data.root_link_pose_w), default_root_pose) + torch.testing.assert_close(wp.to_torch(articulation.data.root_com_vel_w), default_root_vel) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("articulation_type", ["shadow_hand"]) +@pytest.mark.isaacsim_ci +def test_initialization_hand_with_tendons(sim, num_articulations, device, articulation_type): + """Test initialization for fixed base articulated hand with tendons. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is fixed base + 3. All buffers have correct shapes + 4. The articulation can be simulated + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that fixed base + assert articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) + assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) + assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 24) + assert wp.to_torch(articulation.data.body_mass).shape == (num_articulations, articulation.num_bodies) + assert wp.to_torch(articulation.data.body_inertia).shape == (num_articulations, articulation.num_bodies, 9) + + # -- actuator type + for actuator_name, actuator in articulation.actuators.items(): + is_implicit_model_cfg = isinstance(articulation_cfg.actuators[actuator_name], ImplicitActuatorCfg) + assert actuator.is_implicit_model == is_implicit_model_cfg + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.parametrize("articulation_type", ["anymal"]) +@pytest.mark.isaacsim_ci +def test_initialization_floating_base_made_fixed_base( + sim, num_articulations, device, add_ground_plane, articulation_type +): + """Test initialization for a floating-base articulation made fixed-base using schema properties. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is fixed base after modification + 3. All buffers have correct shapes + 4. The articulation maintains its default state + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type).copy() + # Fix root link by making it kinematic + articulation_cfg.spawn.articulation_props.fix_root_link = True + articulation, translations = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that is fixed base + assert articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) + assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) + assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 12) + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + # check that the root is at the correct state - its default state as it is fixed base + default_root_pose = wp.to_torch(articulation.data.default_root_pose).clone() + default_root_vel = wp.to_torch(articulation.data.default_root_vel).clone() + default_root_pose[:, :3] = default_root_pose[:, :3] + translations + + torch.testing.assert_close(wp.to_torch(articulation.data.root_link_pose_w), default_root_pose) + torch.testing.assert_close(wp.to_torch(articulation.data.root_com_vel_w), default_root_vel) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.parametrize("articulation_type", ["panda"]) +@pytest.mark.isaacsim_ci +def test_initialization_fixed_base_made_floating_base( + sim, num_articulations, device, add_ground_plane, articulation_type +): + """Test initialization for fixed base made floating-base using schema properties. + + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation is floating base after modification + 3. All buffers have correct shapes + 4. The articulation can be simulated + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + # Unfix root link by making it non-kinematic + articulation_cfg.spawn.articulation_props.fix_root_link = False + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that is floating base + assert not articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) + assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) + assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 9) + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.parametrize("articulation_type", ["panda"]) +@pytest.mark.isaacsim_ci +def test_out_of_range_default_joint_pos(sim, num_articulations, device, add_ground_plane, articulation_type): + """Test that the default joint position from configuration is out of range. + + This test verifies that: + 1. The articulation fails to initialize when joint positions are out of range + 2. The error is properly handled + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + # Create articulation + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type).copy() + articulation_cfg.init_state.joint_pos = { + "panda_joint1": 10.0, + "panda_joint[2, 4]": -20.0, + } + + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + with pytest.raises(ValueError): + sim.reset() + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("articulation_type", ["panda"]) +@pytest.mark.isaacsim_ci +def test_out_of_range_default_joint_vel(sim, device, articulation_type): + """Test that the default joint velocity from configuration is out of range. + + This test verifies that: + 1. The articulation fails to initialize when joint velocities are out of range + 2. The error is properly handled + """ + articulation_cfg = FRANKA_PANDA_CFG.replace(prim_path="/World/Robot") + articulation_cfg.init_state.joint_vel = { + "panda_joint1": 100.0, + "panda_joint[2, 4]": -60.0, + } + articulation = Articulation(articulation_cfg) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + with pytest.raises(ValueError): + sim.reset() + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.parametrize("articulation_type", ["panda"]) +@pytest.mark.isaacsim_ci +def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane, articulation_type): + """Test write_joint_limits_to_sim API and when default pos falls outside of the new limits. + + This test verifies that: + 1. Joint limits can be set correctly + 2. Default positions are preserved when setting new limits + 3. Joint limits can be set with indexing + 4. Invalid joint positions are properly handled + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + # Create articulation + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device) + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + + # Get current default joint pos + default_joint_pos = wp.to_torch(articulation._data.default_joint_pos).clone() + + # Set new joint limits + limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) + limits[..., 0] = (torch.rand(num_articulations, articulation.num_joints, device=device) + 5.0) * -1.0 + limits[..., 1] = torch.rand(num_articulations, articulation.num_joints, device=device) + 5.0 + articulation.write_joint_position_limit_to_sim_index(limits=limits) + + # Check new limits are in place + torch.testing.assert_close(wp.to_torch(articulation._data.joint_pos_limits), limits) + torch.testing.assert_close(wp.to_torch(articulation._data.default_joint_pos), default_joint_pos) + + # Set new joint limits with indexing + env_ids = torch.arange(1, device=device, dtype=torch.int32) + joint_ids = torch.arange(2, device=device, dtype=torch.int32) + limits = torch.zeros(env_ids.shape[0], joint_ids.shape[0], 2, device=device) + limits[..., 0] = (torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) + 5.0) * -1.0 + limits[..., 1] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) + 5.0 + articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) + + # Check new limits are in place + torch.testing.assert_close(wp.to_torch(articulation._data.joint_pos_limits)[env_ids][:, joint_ids], limits) + torch.testing.assert_close(wp.to_torch(articulation._data.default_joint_pos), default_joint_pos) + + # Set new joint limits that invalidate default joint pos + limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) + limits[..., 0] = torch.rand(num_articulations, articulation.num_joints, device=device) * -0.1 + limits[..., 1] = torch.rand(num_articulations, articulation.num_joints, device=device) * 0.1 + articulation.write_joint_position_limit_to_sim_index(limits=limits) + + # Check if all values are within the bounds + default_joint_pos_torch = wp.to_torch(articulation._data.default_joint_pos) + within_bounds = (default_joint_pos_torch >= limits[..., 0]) & (default_joint_pos_torch <= limits[..., 1]) + assert torch.all(within_bounds) + + # Set new joint limits that invalidate default joint pos with indexing + limits = torch.zeros(env_ids.shape[0], joint_ids.shape[0], 2, device=device) + limits[..., 0] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) * -0.1 + limits[..., 1] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) * 0.1 + articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) + + # Check if all values are within the bounds + default_joint_pos_torch = wp.to_torch(articulation._data.default_joint_pos) + within_bounds = (default_joint_pos_torch[env_ids][:, joint_ids] >= limits[..., 0]) & ( + default_joint_pos_torch[env_ids][:, joint_ids] <= limits[..., 1] + ) + assert torch.all(within_bounds) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.parametrize("articulation_type", ["panda"]) +@pytest.mark.isaacsim_ci +def test_joint_effort_limits(sim, num_articulations, device, add_ground_plane, articulation_type): + """Validate joint effort limits via joint_effort_out_of_limit().""" + # Create articulation + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device) + + # Minimal env wrapper exposing scene["robot"] + class _Env: + def __init__(self, art): + self.scene = {"robot": art} + + env = _Env(articulation) + robot_all = SceneEntityCfg(name="robot") + + sim.reset() + assert articulation.is_initialized + + # Case A: no clipping → should NOT terminate + wp.to_torch(articulation._data.computed_torque).zero_() + wp.to_torch(articulation._data.applied_torque).zero_() + out = joint_effort_out_of_limit(env, robot_all) # [N] + assert torch.all(~out) + + # Case B: simulate clipping → should terminate + wp.to_torch(articulation._data.computed_torque).fill_(100.0) # pretend controller commanded 100 + wp.to_torch(articulation._data.applied_torque).fill_(50.0) # pretend actuator clipped to 50 + out = joint_effort_out_of_limit(env, robot_all) # [N] + assert torch.all(out) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("articulation_type", ["anymal"]) +@pytest.mark.isaacsim_ci +def test_external_force_buffer(sim, num_articulations, device, articulation_type): + """Test if external force buffer correctly updates in the force value is zero case. + + This test verifies that: + 1. External forces can be applied correctly + 2. Force buffers are updated properly + 3. Zero forces are handled correctly + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + + # play the simulator + sim.reset() + + # find bodies to apply the force + body_ids, _ = articulation.find_bodies("base") + + # reset root state + articulation.write_root_pose_to_sim_index(root_pose=wp.to_torch(articulation.data.default_root_pose).clone()) + articulation.write_root_velocity_to_sim_index(root_velocity=wp.to_torch(articulation.data.default_root_vel).clone()) + + # reset dof state + joint_pos, joint_vel = ( + wp.to_torch(articulation.data.default_joint_pos), + wp.to_torch(articulation.data.default_joint_vel), + ) + articulation.write_joint_position_to_sim_index(position=joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) + + # reset articulation + articulation.reset() + + # perform simulation + for step in range(5): + # initiate force tensor + external_wrench_b = torch.zeros(articulation.num_instances, len(body_ids), 6, device=sim.device) + + if step == 0 or step == 3: + # set a non-zero force + force = 1 + else: + # set a zero force + force = 0 + + # set force value + external_wrench_b[:, :, 0] = force + external_wrench_b[:, :, 3] = force + + # apply force + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + body_ids=body_ids, + ) + + # check if the articulation's force and torque buffers are correctly updated + for i in range(num_articulations): + assert wp.to_torch(articulation.permanent_wrench_composer.composed_force)[i, 0, 0].item() == force + assert wp.to_torch(articulation.permanent_wrench_composer.composed_torque)[i, 0, 0].item() == force + + # Check if the instantaneous wrench is correctly added to the permanent wrench + articulation.instantaneous_wrench_composer.add_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + body_ids=body_ids, + ) + + # apply action to the articulation + articulation.set_joint_position_target_index(target=wp.to_torch(articulation.data.default_joint_pos).clone()) + articulation.write_data_to_sim() + + # perform step + sim.step() + + # update buffers + articulation.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("articulation_type", ["anymal"]) +@pytest.mark.isaacsim_ci +def test_external_force_on_single_body(sim, num_articulations, device, articulation_type): + """Test application of external force on the base of the articulation. + + This test verifies that: + 1. External forces can be applied to specific bodies + 2. The forces affect the articulation's motion correctly + 3. The articulation responds to the forces as expected + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + # Play the simulator + sim.reset() + + # Find bodies to apply the force + body_ids, _ = articulation.find_bodies("base") + # Sample a large force + external_wrench_b = torch.zeros(articulation.num_instances, len(body_ids), 6, device=sim.device) + external_wrench_b[..., 1] = 100.0 + + # Now we are ready! + for _ in range(5): + # reset root state + articulation.write_root_pose_to_sim_index(root_pose=wp.to_torch(articulation.data.default_root_pose).clone()) + articulation.write_root_velocity_to_sim_index( + root_velocity=wp.to_torch(articulation.data.default_root_vel).clone() + ) + # reset dof state + joint_pos, joint_vel = ( + wp.to_torch(articulation.data.default_joint_pos), + wp.to_torch(articulation.data.default_joint_vel), + ) + articulation.write_joint_position_to_sim_index(position=joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) + # reset articulation + articulation.reset() + # apply force + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], body_ids=body_ids + ) + # perform simulation + for _ in range(100): + # apply action to the articulation + articulation.set_joint_position_target_index( + target=wp.to_torch(articulation.data.default_joint_pos).clone() + ) + articulation.write_data_to_sim() + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + # check condition that the articulations have fallen down + for i in range(num_articulations): + assert wp.to_torch(articulation.data.root_pos_w)[i, 2].item() < 0.2 + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("articulation_type", ["anymal"]) +@pytest.mark.isaacsim_ci +def test_external_force_on_single_body_at_position(sim, num_articulations, device, articulation_type): + """Test application of external force on the base of the articulation at a given position. + + This test verifies that: + 1. External forces can be applied to specific bodies at a given position + 2. External forces can be applied to specific bodies in the global frame + 3. External forces are calculated and composed correctly + 4. The forces affect the articulation's motion correctly + 5. The articulation responds to the forces as expected + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + # Play the simulator + sim.reset() + + # Find bodies to apply the force + body_ids, _ = articulation.find_bodies("base") + # Sample a large force + external_wrench_b = torch.zeros(articulation.num_instances, len(body_ids), 6, device=sim.device) + external_wrench_b[..., 2] = 100.0 + external_wrench_positions_b = torch.zeros(articulation.num_instances, len(body_ids), 3, device=sim.device) + external_wrench_positions_b[..., 1] = 1.0 + + desired_force = torch.zeros(articulation.num_instances, len(body_ids), 3, device=sim.device) + desired_force[..., 2] = 200.0 + desired_torque = torch.zeros(articulation.num_instances, len(body_ids), 3, device=sim.device) + desired_torque[..., 0] = 200.0 + + # Now we are ready! + for i in range(5): + # reset root state + root_pose = wp.to_torch(articulation.data.default_root_pose).clone() + root_pose[0, 0] = 2.5 # space them apart by 2.5m + + articulation.write_root_pose_to_sim_index(root_pose=root_pose) + articulation.write_root_velocity_to_sim_index( + root_velocity=wp.to_torch(articulation.data.default_root_vel).clone() + ) + # reset dof state + joint_pos, joint_vel = ( + wp.to_torch(articulation.data.default_joint_pos), + wp.to_torch(articulation.data.default_joint_vel), + ) + articulation.write_joint_position_to_sim_index(position=joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) + # reset articulation + articulation.reset() + # apply force + is_global = False + + if i % 2 == 0: + body_com_pos_w = wp.to_torch(articulation.data.body_com_pos_w)[:, body_ids, :3] + # is_global = True + external_wrench_positions_b[..., 0] = 0.0 + external_wrench_positions_b[..., 1] = 1.0 + external_wrench_positions_b[..., 2] = 0.0 + external_wrench_positions_b += body_com_pos_w + else: + external_wrench_positions_b[..., 0] = 0.0 + external_wrench_positions_b[..., 1] = 1.0 + external_wrench_positions_b[..., 2] = 0.0 + + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=external_wrench_positions_b, + body_ids=body_ids, + is_global=is_global, + ) + articulation.permanent_wrench_composer.add_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=external_wrench_positions_b, + body_ids=body_ids, + is_global=is_global, + ) + # perform simulation + for _ in range(100): + # apply action to the articulation + articulation.set_joint_position_target_index( + target=wp.to_torch(articulation.data.default_joint_pos).clone() + ) + articulation.write_data_to_sim() + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + # check condition that the articulations have fallen down + for i in range(num_articulations): + assert wp.to_torch(articulation.data.root_pos_w)[i, 2].item() < 0.2 + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("articulation_type", ["anymal"]) +@pytest.mark.isaacsim_ci +def test_external_force_on_multiple_bodies(sim, num_articulations, device, articulation_type): + """Test application of external force on the legs of the articulation. + + This test verifies that: + 1. External forces can be applied to multiple bodies + 2. The forces affect the articulation's motion correctly + 3. The articulation responds to the forces as expected + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + + # Play the simulator + sim.reset() + + # Find bodies to apply the force + body_ids, _ = articulation.find_bodies(".*_SHANK") + # Sample a large force + external_wrench_b = torch.zeros(articulation.num_instances, len(body_ids), 6, device=sim.device) + external_wrench_b[..., 1] = 200.0 + + # Now we are ready! + for _ in range(5): + # reset root state + articulation.write_root_pose_to_sim_index(root_pose=wp.to_torch(articulation.data.default_root_pose).clone()) + articulation.write_root_velocity_to_sim_index( + root_velocity=wp.to_torch(articulation.data.default_root_vel).clone() + ) + # reset dof state + joint_pos, joint_vel = ( + wp.to_torch(articulation.data.default_joint_pos), + wp.to_torch(articulation.data.default_joint_vel), + ) + articulation.write_joint_position_to_sim_index(position=joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) + # reset articulation + articulation.reset() + # apply force + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], body_ids=body_ids + ) + # perform simulation + for _ in range(100): + # apply action to the articulation + articulation.set_joint_position_target_index( + target=wp.to_torch(articulation.data.default_joint_pos).clone() + ) + articulation.write_data_to_sim() + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + # check condition + for i in range(num_articulations): + # since there is a moment applied on the articulation, the articulation should rotate + assert wp.to_torch(articulation.data.root_ang_vel_w)[i, 2].item() > 0.1 + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("articulation_type", ["anymal"]) +@pytest.mark.isaacsim_ci +def test_external_force_on_multiple_bodies_at_position(sim, num_articulations, device, articulation_type): + """Test application of external force on the legs of the articulation at a given position. + + This test verifies that: + 1. External forces can be applied to multiple bodies at a given position + 2. External forces can be applied to multiple bodies in the global frame + 3. External forces are calculated and composed correctly + 4. The forces affect the articulation's motion correctly + 5. The articulation responds to the forces as expected + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + + # Play the simulator + sim.reset() + + # Find bodies to apply the force + body_ids, _ = articulation.find_bodies(".*_SHANK") + # Sample a large force + external_wrench_b = torch.zeros(articulation.num_instances, len(body_ids), 6, device=sim.device) + external_wrench_b[..., 2] = 100.0 + external_wrench_positions_b = torch.zeros(articulation.num_instances, len(body_ids), 3, device=sim.device) + external_wrench_positions_b[..., 1] = 1.0 + + desired_force = torch.zeros(articulation.num_instances, len(body_ids), 3, device=sim.device) + desired_force[..., 2] = 200.0 + desired_torque = torch.zeros(articulation.num_instances, len(body_ids), 3, device=sim.device) + desired_torque[..., 0] = 200.0 + + # Now we are ready! + for i in range(5): + # reset root state + articulation.write_root_pose_to_sim_index(root_pose=wp.to_torch(articulation.data.default_root_pose).clone()) + articulation.write_root_velocity_to_sim_index( + root_velocity=wp.to_torch(articulation.data.default_root_vel).clone() + ) + # reset dof state + joint_pos, joint_vel = ( + wp.to_torch(articulation.data.default_joint_pos), + wp.to_torch(articulation.data.default_joint_vel), + ) + articulation.write_joint_position_to_sim_index(position=joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) + # reset articulation + articulation.reset() + + is_global = False + if i % 2 == 0: + body_com_pos_w = wp.to_torch(articulation.data.body_com_pos_w)[:, body_ids, :3] + is_global = True + external_wrench_positions_b[..., 0] = 0.0 + external_wrench_positions_b[..., 1] = 1.0 + external_wrench_positions_b[..., 2] = 0.0 + external_wrench_positions_b += body_com_pos_w + else: + external_wrench_positions_b[..., 0] = 0.0 + external_wrench_positions_b[..., 1] = 1.0 + external_wrench_positions_b[..., 2] = 0.0 + + # apply force + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=external_wrench_positions_b, + body_ids=body_ids, + is_global=is_global, + ) + articulation.permanent_wrench_composer.add_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=external_wrench_positions_b, + body_ids=body_ids, + is_global=is_global, + ) + # perform simulation + for _ in range(100): + # apply action to the articulation + articulation.set_joint_position_target_index( + target=wp.to_torch(articulation.data.default_joint_pos).clone() + ) + articulation.write_data_to_sim() + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + # check condition + for i in range(num_articulations): + # since there is a moment applied on the articulation, the articulation should rotate + assert torch.abs(wp.to_torch(articulation.data.root_ang_vel_w)[i, 2]).item() > 0.1 + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("articulation_type", ["humanoid"]) +@pytest.mark.isaacsim_ci +def test_loading_gains_from_usd(sim, num_articulations, device, articulation_type): + """Test that gains are loaded from USD file if actuator model has them as None. + + This test verifies that: + 1. Gains are loaded correctly from USD file + 2. Default gains are applied when not specified + 3. The gains match the expected values + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type, stiffness=None, damping=None) + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=sim.device) + + # Play sim + sim.reset() + + # Expected gains + # -- Stiffness values + expected_stiffness = { + ".*_waist.*": 20.0, + ".*_upper_arm.*": 10.0, + "pelvis": 10.0, + ".*_lower_arm": 2.0, + ".*_thigh:0": 10.0, + ".*_thigh:1": 20.0, + ".*_thigh:2": 10.0, + ".*_shin": 5.0, + ".*_foot.*": 2.0, + } + indices_list, _, values_list = string_utils.resolve_matching_names_values( + expected_stiffness, articulation.joint_names + ) + expected_stiffness = torch.zeros(articulation.num_instances, articulation.num_joints, device=articulation.device) + expected_stiffness[:, indices_list] = torch.tensor(values_list, device=articulation.device) + # -- Damping values + expected_damping = { + ".*_waist.*": 5.0, + ".*_upper_arm.*": 5.0, + "pelvis": 5.0, + ".*_lower_arm": 1.0, + ".*_thigh:0": 5.0, + ".*_thigh:1": 5.0, + ".*_thigh:2": 5.0, + ".*_shin": 0.1, + ".*_foot.*": 1.0, + } + indices_list, _, values_list = string_utils.resolve_matching_names_values( + expected_damping, articulation.joint_names + ) + expected_damping = torch.zeros_like(expected_stiffness) + expected_damping[:, indices_list] = torch.tensor(values_list, device=articulation.device) + + # Check that gains are loaded from USD file + torch.testing.assert_close(articulation.actuators["body"].stiffness, expected_stiffness) + torch.testing.assert_close(articulation.actuators["body"].damping, expected_damping) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.parametrize("articulation_type", ["humanoid"]) +@pytest.mark.isaacsim_ci +def test_setting_gains_from_cfg(sim, num_articulations, device, add_ground_plane, articulation_type): + """Test that gains are loaded from the configuration correctly. + + This test verifies that: + 1. Gains are loaded correctly from configuration + 2. The gains match the expected values + 3. The gains are applied correctly to the actuators + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=sim.device + ) + + # Play sim + sim.reset() + + # Expected gains + expected_stiffness = torch.full( + (articulation.num_instances, articulation.num_joints), 10.0, device=articulation.device + ) + expected_damping = torch.full_like(expected_stiffness, 2.0) + + # Check that gains are loaded from USD file + torch.testing.assert_close(articulation.actuators["body"].stiffness, expected_stiffness) + torch.testing.assert_close(articulation.actuators["body"].damping, expected_damping) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("articulation_type", ["humanoid"]) +@pytest.mark.isaacsim_ci +def test_setting_gains_from_cfg_dict(sim, num_articulations, device, articulation_type): + """Test that gains are loaded from the configuration dictionary correctly. + + This test verifies that: + 1. Gains are loaded correctly from configuration dictionary + 2. The gains match the expected values + 3. The gains are applied correctly to the actuators + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + """ + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=sim.device + ) + # Play sim + sim.reset() + + # Expected gains + expected_stiffness = torch.full( + (articulation.num_instances, articulation.num_joints), 10.0, device=articulation.device + ) + expected_damping = torch.full_like(expected_stiffness, 2.0) + + # Check that gains are loaded from USD file + torch.testing.assert_close(articulation.actuators["body"].stiffness, expected_stiffness) + torch.testing.assert_close(articulation.actuators["body"].damping, expected_damping) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("vel_limit_sim", [1e5, None]) +@pytest.mark.parametrize("vel_limit", [1e2, None]) +@pytest.mark.parametrize("add_ground_plane", [False]) +@pytest.mark.parametrize("articulation_type", ["single_joint_implicit"]) +@pytest.mark.isaacsim_ci +def test_setting_velocity_limit_implicit( + sim, num_articulations, device, vel_limit_sim, vel_limit, add_ground_plane, articulation_type +): + """Test setting of velocity limit for implicit actuators. + + This test verifies that: + 1. Velocity limits can be set correctly for implicit actuators + 2. The limits are applied correctly to the simulation + 3. The limits are handled correctly when both sim and non-sim limits are set + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + vel_limit_sim: The velocity limit to set in simulation + vel_limit: The velocity limit to set in actuator + """ + # create simulation + articulation_cfg = generate_articulation_cfg( + articulation_type=articulation_type, + velocity_limit_sim=vel_limit_sim, + velocity_limit=vel_limit, + ) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, + num_articulations=num_articulations, + device=device, + ) + # Play sim + if vel_limit_sim is not None and vel_limit is not None: + with pytest.raises(ValueError): + sim.reset() + return + sim.reset() + + # read the values set into the simulation + newton_vel_limit = wp.to_torch( + articulation.root_view.get_attribute("joint_velocity_limit", SimulationManager.get_model()) + ).to(device)[:, 0, :] + # check data buffer + torch.testing.assert_close(wp.to_torch(articulation.data.joint_velocity_limits), newton_vel_limit) + # check actuator has simulation velocity limit + torch.testing.assert_close(articulation.actuators["joint"].velocity_limit_sim, newton_vel_limit) + # check that both values match for velocity limit + torch.testing.assert_close( + articulation.actuators["joint"].velocity_limit_sim, + articulation.actuators["joint"].velocity_limit, + ) + + if vel_limit_sim is None: + # Case 2: both velocity limit and velocity limit sim are not set + # This is the case where the velocity limit keeps its USD default value + # Case 3: velocity limit sim is not set but velocity limit is set + # For backwards compatibility, we do not set velocity limit to simulation + # Thus, both default to USD default value. + limit = articulation_cfg.spawn.joint_drive_props.max_velocity + else: + # Case 4: only velocity limit sim is set + # In this case, the velocity limit is set to the USD value + limit = vel_limit_sim + + # check max velocity is what we set + expected_velocity_limit = torch.full_like(newton_vel_limit, limit) + torch.testing.assert_close(newton_vel_limit, expected_velocity_limit) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("vel_limit_sim", [1e5, None]) +@pytest.mark.parametrize("vel_limit", [1e2, None]) +@pytest.mark.parametrize("articulation_type", ["single_joint_explicit"]) +@pytest.mark.isaacsim_ci +def test_setting_velocity_limit_explicit(sim, num_articulations, device, vel_limit_sim, vel_limit, articulation_type): + """Test setting of velocity limit for explicit actuators.""" + articulation_cfg = generate_articulation_cfg( + articulation_type=articulation_type, + velocity_limit_sim=vel_limit_sim, + velocity_limit=vel_limit, + ) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, + num_articulations=num_articulations, + device=device, + ) + # Play sim + sim.reset() + + # collect limit init values + newton_vel_limit = wp.to_torch( + articulation.root_view.get_attribute("joint_velocity_limit", SimulationManager.get_model()) + ).to(device)[:, 0, :] + actuator_vel_limit = articulation.actuators["joint"].velocity_limit + actuator_vel_limit_sim = articulation.actuators["joint"].velocity_limit_sim + + # check data buffer for joint_velocity_limits_sim + torch.testing.assert_close(wp.to_torch(articulation.data.joint_velocity_limits), newton_vel_limit) + # check actuator velocity_limit_sim is set to physx + torch.testing.assert_close(actuator_vel_limit_sim, newton_vel_limit) + + if vel_limit is not None: + expected_actuator_vel_limit = torch.full( + (articulation.num_instances, articulation.num_joints), + vel_limit, + device=articulation.device, + ) + # check actuator is set + torch.testing.assert_close(actuator_vel_limit, expected_actuator_vel_limit) + # check physx is not velocity_limit + assert not torch.allclose(actuator_vel_limit, newton_vel_limit) + else: + # check actuator velocity_limit is the same as the PhysX default + torch.testing.assert_close(actuator_vel_limit, newton_vel_limit) + + # simulation velocity limit is set to USD value unless user overrides + if vel_limit_sim is not None: + limit = vel_limit_sim + else: + limit = articulation_cfg.spawn.joint_drive_props.max_velocity + # check physx is set to expected value + expected_vel_limit = torch.full_like(newton_vel_limit, limit) + torch.testing.assert_close(newton_vel_limit, expected_vel_limit) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("effort_limit_sim", [1e5, None]) +@pytest.mark.parametrize("effort_limit", [1e2, 80.0, None]) +@pytest.mark.parametrize("articulation_type", ["single_joint_implicit"]) +@pytest.mark.isaacsim_ci +def test_setting_effort_limit_implicit( + sim, num_articulations, device, effort_limit_sim, effort_limit, articulation_type +): + """Test setting of effort limit for implicit actuators. + + This test verifies the effort limit resolution logic for actuator models implemented in :class:`ActuatorBase`: + - Case 1: If USD value == actuator config value: values match correctly + - Case 2: If USD value != actuator config value: actuator config value is used + - Case 3: If actuator config value is None: USD value is used as default + """ + articulation_cfg = generate_articulation_cfg( + articulation_type=articulation_type, + effort_limit_sim=effort_limit_sim, + effort_limit=effort_limit, + ) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, + num_articulations=num_articulations, + device=device, + ) + # Play sim + if effort_limit_sim is not None and effort_limit is not None: + with pytest.raises(ValueError): + sim.reset() + return + sim.reset() + + # obtain the physx effort limits + newton_effort_limit = wp.to_torch( + articulation.root_view.get_attribute("joint_effort_limit", SimulationManager.get_model()) + ).to(device)[:, 0, :] + + # check that the two are equivalent + torch.testing.assert_close( + articulation.actuators["joint"].effort_limit_sim, + articulation.actuators["joint"].effort_limit, + ) + torch.testing.assert_close(articulation.actuators["joint"].effort_limit_sim, newton_effort_limit) + + # decide the limit based on what is set + if effort_limit_sim is None and effort_limit is None: + limit = articulation_cfg.spawn.joint_drive_props.max_effort + elif effort_limit_sim is not None and effort_limit is None: + limit = effort_limit_sim + elif effort_limit_sim is None and effort_limit is not None: + limit = effort_limit + + # check that the max force is what we set + expected_effort_limit = torch.full_like(newton_effort_limit, limit) + torch.testing.assert_close(newton_effort_limit, expected_effort_limit) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("effort_limit_sim", [1e5, None]) +@pytest.mark.parametrize("effort_limit", [80.0, 1e2, None]) +@pytest.mark.parametrize("articulation_type", ["single_joint_explicit"]) +@pytest.mark.isaacsim_ci +def test_setting_effort_limit_explicit( + sim, num_articulations, device, effort_limit_sim, effort_limit, articulation_type +): + """Test setting of effort limit for explicit actuators. + + This test verifies the effort limit resolution logic for actuator models implemented in :class:`ActuatorBase`: + - Case 1: If USD value == actuator config value: values match correctly + - Case 2: If USD value != actuator config value: actuator config value is used + - Case 3: If actuator config value is None: USD value is used as default + + """ + + articulation_cfg = generate_articulation_cfg( + articulation_type=articulation_type, + effort_limit_sim=effort_limit_sim, + effort_limit=effort_limit, + ) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, + num_articulations=num_articulations, + device=device, + ) + # Play sim + sim.reset() + + # usd default effort limit is set to 80 + usd_default_effort_limit = 80.0 + + # collect limit init values + newton_effort_limit = wp.to_torch( + articulation.root_view.get_attribute("joint_effort_limit", SimulationManager.get_model()) + ).to(device)[:, 0, :] + actuator_effort_limit = articulation.actuators["joint"].effort_limit + actuator_effort_limit_sim = articulation.actuators["joint"].effort_limit_sim + + # check actuator effort_limit_sim is set to physx + torch.testing.assert_close(actuator_effort_limit_sim, newton_effort_limit) + + if effort_limit is not None: + expected_actuator_effort_limit = torch.full_like(actuator_effort_limit, effort_limit) + # check actuator is set + torch.testing.assert_close(actuator_effort_limit, expected_actuator_effort_limit) + + # check physx effort limit does not match the one explicit actuator has + assert not (torch.allclose(actuator_effort_limit, newton_effort_limit)) + else: + # When effort_limit is None, actuator should use USD default values + expected_actuator_effort_limit = torch.full_like(newton_effort_limit, usd_default_effort_limit) + torch.testing.assert_close(actuator_effort_limit, expected_actuator_effort_limit) + + # when using explicit actuators, the limits are set to high unless user overrides + if effort_limit_sim is not None: + limit = effort_limit_sim + else: + limit = ActuatorBase._DEFAULT_MAX_EFFORT_SIM # type: ignore + # check physx internal value matches the expected sim value + expected_effort_limit = torch.full_like(newton_effort_limit, limit) + torch.testing.assert_close(actuator_effort_limit_sim, expected_effort_limit) + torch.testing.assert_close(newton_effort_limit, expected_effort_limit) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("articulation_type", ["humanoid"]) +@pytest.mark.isaacsim_ci +def test_reset(sim, num_articulations, device, articulation_type): + """Test that reset method works properly.""" + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=device + ) + + # Play the simulator + sim.reset() + + # Now we are ready! + # reset articulation + articulation.reset() + + # Reset should zero external forces and torques + assert not articulation._instantaneous_wrench_composer.active + assert not articulation._permanent_wrench_composer.active + assert torch.count_nonzero(wp.to_torch(articulation._instantaneous_wrench_composer.composed_force)) == 0 + assert torch.count_nonzero(wp.to_torch(articulation._instantaneous_wrench_composer.composed_torque)) == 0 + assert torch.count_nonzero(wp.to_torch(articulation._permanent_wrench_composer.composed_force)) == 0 + assert torch.count_nonzero(wp.to_torch(articulation._permanent_wrench_composer.composed_torque)) == 0 + + if num_articulations > 1: + num_bodies = articulation.num_bodies + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=torch.ones((num_articulations, num_bodies, 3), device=device), + torques=torch.ones((num_articulations, num_bodies, 3), device=device), + ) + articulation.instantaneous_wrench_composer.add_forces_and_torques_index( + forces=torch.ones((num_articulations, num_bodies, 3), device=device), + torques=torch.ones((num_articulations, num_bodies, 3), device=device), + ) + articulation.reset(env_ids=torch.tensor([0], device=device)) + assert articulation._instantaneous_wrench_composer.active + assert articulation._permanent_wrench_composer.active + assert ( + torch.count_nonzero(wp.to_torch(articulation._instantaneous_wrench_composer.composed_force)) + == num_bodies * 3 + ) + assert ( + torch.count_nonzero(wp.to_torch(articulation._instantaneous_wrench_composer.composed_torque)) + == num_bodies * 3 + ) + assert ( + torch.count_nonzero(wp.to_torch(articulation._permanent_wrench_composer.composed_force)) == num_bodies * 3 + ) + assert ( + torch.count_nonzero(wp.to_torch(articulation._permanent_wrench_composer.composed_torque)) == num_bodies * 3 + ) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.parametrize("articulation_type", ["panda"]) +@pytest.mark.isaacsim_ci +def test_apply_joint_command(sim, num_articulations, device, add_ground_plane, articulation_type): + """Test applying of joint position target functions correctly for a robotic arm.""" + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=device + ) + + # Play the simulator + sim.reset() + + for _ in range(100): + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + # reset dof state + joint_pos = wp.to_torch(articulation.data.default_joint_pos).clone() + joint_pos[:, 3] = 0.0 + + # apply action to the articulation + articulation.set_joint_position_target_index(target=joint_pos) + articulation.write_data_to_sim() + + for _ in range(100): + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + # Check that current joint position is not the same as default joint position, meaning + # the articulation moved. We can't check that it reached its desired joint position as the gains + # are not properly tuned + assert not torch.allclose(wp.to_torch(articulation.data.joint_pos), joint_pos) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("with_offset", [True, False]) +@pytest.mark.parametrize("articulation_type", ["single_joint_implicit"]) +@pytest.mark.isaacsim_ci +def test_body_root_state(sim, num_articulations, device, with_offset, articulation_type): + """Test for reading the `body_state_w` property. + + This test verifies that: + 1. Body states can be read correctly + 2. States are correct with and without offsets + 3. States are consistent across different devices + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + with_offset: Whether to test with offset + """ + sim._app_control_on_stop_handle = None + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, env_pos = generate_articulation(articulation_cfg, num_articulations, device) + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10, "Possible reference leak for articulation" + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized, "Articulation is not initialized" + # Check that fixed base + assert articulation.is_fixed_base, "Articulation is not a fixed base" + + # Resolve body indices by name (ordering may differ across physics backends) + root_idx = articulation.body_names.index("CenterPivot") + arm_idx = articulation.body_names.index("Arm") + + # change center of mass offset from link frame + if with_offset: + offset = [0.5, 0.0, 0.0] + else: + offset = [0.0, 0.0, 0.0] + + # create com offsets — apply offset to the Arm body + num_bodies = articulation.num_bodies + com = wp.to_torch(articulation.root_view.get_attribute("body_com", SimulationManager.get_model())) + link_offset = [1.0, 0.0, 0.0] # the offset from CenterPivot to Arm frames + new_com = torch.tensor(offset, device=device).repeat(num_articulations, 1, 1) + com[:, 0, arm_idx, :] = new_com.squeeze(-2) + articulation.root_view.set_attribute("body_com", SimulationManager.get_model(), wp.from_torch(com, dtype=wp.vec3f)) + with wp.ScopedDevice(device): + SimulationManager._solver.notify_model_changed(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + # check they are set + torch.testing.assert_close( + wp.to_torch(articulation.root_view.get_attribute("body_com", SimulationManager.get_model())), com + ) + + for i in range(50): + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + # get state properties + root_link_pose_w = wp.to_torch(articulation.data.root_link_pose_w) + root_link_vel_w = wp.to_torch(articulation.data.root_link_vel_w) + root_com_pose_w = wp.to_torch(articulation.data.root_com_pose_w) + root_com_vel_w = wp.to_torch(articulation.data.root_com_vel_w) + body_link_pose_w = wp.to_torch(articulation.data.body_link_pose_w) + body_link_vel_w = wp.to_torch(articulation.data.body_link_vel_w) + body_com_pose_w = wp.to_torch(articulation.data.body_com_pose_w) + body_com_vel_w = wp.to_torch(articulation.data.body_com_vel_w) + + if with_offset: + # get joint state + joint_pos = wp.to_torch(articulation.data.joint_pos).unsqueeze(-1) + joint_vel = wp.to_torch(articulation.data.joint_vel).unsqueeze(-1) + + # LINK state + # angular velocity should be the same for both COM and link frames + torch.testing.assert_close(root_com_vel_w[..., 3:], root_link_vel_w[..., 3:]) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) + + # lin_vel arm + lin_vel_gt = torch.zeros(num_articulations, num_bodies, 3, device=device) + vx = -(link_offset[0]) * joint_vel * torch.sin(joint_pos) + vy = torch.zeros(num_articulations, 1, 1, device=device) + vz = (link_offset[0]) * joint_vel * torch.cos(joint_pos) + lin_vel_gt[:, arm_idx, :] = torch.cat([vx, vy, vz], dim=-1).squeeze(-2) + + # linear velocity of root link should be zero + torch.testing.assert_close(lin_vel_gt[:, root_idx, :], root_link_vel_w[..., :3], atol=1e-3, rtol=1e-1) + # linear velocity of pendulum link should be + torch.testing.assert_close(lin_vel_gt, body_link_vel_w[..., :3], atol=1e-3, rtol=1e-1) + + # ang_vel + torch.testing.assert_close(root_com_vel_w[..., 3:], root_link_vel_w[..., 3:]) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) + + # COM state + # position and orientation shouldn't match for the _state_com_w but everything else will + pos_gt = torch.zeros(num_articulations, num_bodies, 3, device=device) + px = (link_offset[0] + offset[0]) * torch.cos(joint_pos) + py = torch.zeros(num_articulations, 1, 1, device=device) + pz = (link_offset[0] + offset[0]) * torch.sin(joint_pos) + pos_gt[:, arm_idx, :] = torch.cat([px, py, pz], dim=-1).squeeze(-2) + pos_gt += env_pos.unsqueeze(-2).repeat(1, num_bodies, 1) + torch.testing.assert_close(pos_gt[:, root_idx, :], root_com_pose_w[..., :3], atol=1e-3, rtol=1e-1) + torch.testing.assert_close(pos_gt, body_com_pose_w[..., :3], atol=1e-3, rtol=1e-1) + + # orientation + com_quat_b = wp.to_torch(articulation.data.body_com_quat_b) + com_quat_w = math_utils.quat_mul(body_link_pose_w[..., 3:], com_quat_b) + torch.testing.assert_close(com_quat_w, body_com_pose_w[..., 3:]) + torch.testing.assert_close(com_quat_w[:, root_idx, :], root_com_pose_w[..., 3:]) + + # angular velocity should be the same for both COM and link frames + torch.testing.assert_close(root_com_vel_w[..., 3:], root_link_vel_w[..., 3:]) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) + else: + # single joint center of masses are at link frames so they will be the same + torch.testing.assert_close(root_link_pose_w, root_com_pose_w) + torch.testing.assert_close(root_com_vel_w, root_link_vel_w) + torch.testing.assert_close(body_link_pose_w, body_com_pose_w) + torch.testing.assert_close(body_com_vel_w, body_link_vel_w) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("with_offset", [True, False]) +@pytest.mark.parametrize("state_location", ["com", "link"]) +@pytest.mark.parametrize("gravity_enabled", [False]) +@pytest.mark.parametrize("articulation_type", ["anymal"]) +@pytest.mark.isaacsim_ci +def test_write_root_state( + sim, num_articulations, device, with_offset, state_location, gravity_enabled, articulation_type +): + """Test the setters for root_state using both the link frame and center of mass as reference frame. + + This test verifies that: + 1. Root states can be written correctly + 2. States are correct with and without offsets + 3. States can be written for both COM and link frames + 4. States are consistent across different devices + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + with_offset: Whether to test with offset + state_location: Whether to test COM or link frame + """ + sim._app_control_on_stop_handle = None + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, env_pos = generate_articulation(articulation_cfg, num_articulations, device) + env_idx = torch.tensor([x for x in range(num_articulations)], device=device, dtype=torch.int32) + + # Play sim + sim.reset() + + # Resolve root body index by name (ordering may differ across physics backends) + root_idx = articulation.find_bodies("base")[0][0] + + # change center of mass offset from link frame + if with_offset: + offset = torch.tensor([1.0, 0.0, 0.0]).repeat(num_articulations, 1, 1) + else: + offset = torch.tensor([0.0, 0.0, 0.0]).repeat(num_articulations, 1, 1) + + # create com offsets + com = wp.to_torch(articulation.root_view.get_attribute("body_com", SimulationManager.get_model())) + new_com = offset + com[:, 0, root_idx, :] = new_com.squeeze(-2) + articulation.root_view.set_attribute("body_com", SimulationManager.get_model(), wp.from_torch(com, dtype=wp.vec3f)) + with wp.ScopedDevice(device): + SimulationManager._solver.notify_model_changed(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + # check they are set + torch.testing.assert_close( + wp.to_torch(articulation.root_view.get_attribute("body_com", SimulationManager.get_model())), com + ) + + rand_state = torch.zeros(num_articulations, 13, device=device) + rand_state[..., :7] = wp.to_torch(articulation.data.default_root_pose) + rand_state[..., :3] += env_pos + # make quaternion a unit vector + rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) + + env_idx = env_idx.to(device) + for i in range(10): + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + if state_location == "com": + if i % 2 == 0: + articulation.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7]) + articulation.write_root_com_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) + else: + articulation.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_idx) + articulation.write_root_com_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_idx) + elif state_location == "link": + if i % 2 == 0: + articulation.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7]) + articulation.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) + else: + articulation.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_idx) + articulation.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_idx) + + if state_location == "com": + torch.testing.assert_close(rand_state[..., :7], wp.to_torch(articulation.data.root_com_pose_w)) + torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(articulation.data.root_com_vel_w)) + elif state_location == "link": + torch.testing.assert_close(rand_state[..., :7], wp.to_torch(articulation.data.root_link_pose_w)) + torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(articulation.data.root_link_vel_w)) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("articulation_type", ["single_joint_implicit"]) +@pytest.mark.xfail(reason="Newton body_parent_f uses different convention than PhysX get_link_incoming_joint_force") +@pytest.mark.isaacsim_ci +def test_body_incoming_joint_wrench_b_single_joint(sim, num_articulations, device, articulation_type): + """Test the data.body_incoming_joint_wrench_b buffer is populated correctly and statically correct for single joint. + + This test verifies that: + 1. The body incoming joint wrench buffer has correct shape + 2. The wrench values are statically correct for a single joint + 3. The wrench values match expected values from gravity and external forces + + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=device + ) + + # Play the simulator + sim.reset() + # apply external force + external_force_vector_b = torch.zeros((num_articulations, articulation.num_bodies, 3), device=device) + external_force_vector_b[:, 1, 1] = 10.0 # 10 N in Y direction + external_torque_vector_b = torch.zeros((num_articulations, articulation.num_bodies, 3), device=device) + external_torque_vector_b[:, 1, 2] = 10.0 # 10 Nm in z direction + + # apply action to the articulation + joint_pos = torch.ones_like(wp.to_torch(articulation.data.joint_pos)) * 1.5708 / 2.0 + articulation.write_joint_position_to_sim_index( + position=torch.ones_like(wp.to_torch(articulation.data.joint_pos)), + ) + articulation.write_joint_velocity_to_sim_index( + velocity=torch.zeros_like(wp.to_torch(articulation.data.joint_vel)), + ) + articulation.set_joint_position_target_index(target=joint_pos) + articulation.write_data_to_sim() + for _ in range(50): + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_force_vector_b, torques=external_torque_vector_b + ) + articulation.write_data_to_sim() + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + # check shape + assert wp.to_torch(articulation.data.body_incoming_joint_wrench_b).shape == ( + num_articulations, + articulation.num_bodies, + 6, + ) + + # calculate expected static + mass = wp.to_torch(articulation.data.body_mass).to("cpu") + pos_w = wp.to_torch(articulation.data.body_pos_w) + quat_w = wp.to_torch(articulation.data.body_quat_w) + + mass_link2 = mass[:, 1].view(num_articulations, -1) + gravity = torch.tensor(sim.cfg.gravity, device="cpu").repeat(num_articulations, 1).view((num_articulations, 3)) + + # NOTE: the com and link pose for single joint are colocated + weight_vector_w = mass_link2 * gravity + # expected wrench from link mass and external wrench + arm_idx = articulation.body_names.index("Arm") + root_idx = articulation.body_names.index("CenterPivot") + expected_wrench = torch.zeros((num_articulations, 6), device=device) + expected_wrench[:, :3] = math_utils.quat_apply( + math_utils.quat_conjugate(quat_w[:, root_idx, :]), + weight_vector_w.to(device) + math_utils.quat_apply(quat_w[:, arm_idx, :], external_force_vector_b[:, 1, :]), + ) + expected_wrench[:, 3:] = math_utils.quat_apply( + math_utils.quat_conjugate(quat_w[:, root_idx, :]), + torch.cross( + pos_w[:, arm_idx, :].to(device) - pos_w[:, root_idx, :].to(device), + weight_vector_w.to(device) + math_utils.quat_apply(quat_w[:, arm_idx, :], external_force_vector_b[:, 1, :]), + dim=-1, + ) + + math_utils.quat_apply(quat_w[:, arm_idx, :], external_torque_vector_b[:, 1, :]), + ) + + # check value of last joint wrench + actual_wrench = wp.to_torch(articulation.data.body_incoming_joint_wrench_b)[:, arm_idx, :].squeeze(1) + torch.testing.assert_close( + expected_wrench, + actual_wrench, + atol=1e-2, + rtol=1e-3, + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("articulation_type", ["humanoid"]) +@pytest.mark.isaacsim_ci +def test_setting_articulation_root_prim_path(sim, device, articulation_type): + """Test that the articulation root prim path can be set explicitly.""" + sim._app_control_on_stop_handle = None + # Create articulation + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation_cfg.articulation_root_prim_path = "/torso" + articulation, _ = generate_articulation(articulation_cfg, 1, device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation._is_initialized + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("articulation_type", ["humanoid"]) +@pytest.mark.isaacsim_ci +def test_setting_invalid_articulation_root_prim_path(sim, device, articulation_type): + """Test that the articulation root prim path can be set explicitly.""" + sim._app_control_on_stop_handle = None + # Create articulation + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation_cfg.articulation_root_prim_path = "/non_existing_prim_path" + articulation, _ = generate_articulation(articulation_cfg, 1, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + with pytest.raises((RuntimeError, KeyError)): + sim.reset() + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("gravity_enabled", [False]) +@pytest.mark.parametrize("articulation_type", ["anymal"]) +@pytest.mark.isaacsim_ci +def test_write_joint_state_data_consistency(sim, num_articulations, device, gravity_enabled, articulation_type): + """Test the setters for root_state using both the link frame and center of mass as reference frame. + + This test verifies that after write_joint_state_to_sim operations: + 1. state, com_state, link_state value consistency + 2. body_pose, link + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + sim._app_control_on_stop_handle = None + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, env_pos = generate_articulation(articulation_cfg, num_articulations, device) + env_idx = torch.tensor([x for x in range(num_articulations)]) + + # Play sim + sim.reset() + + limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) + limits[..., 0] = (torch.rand(num_articulations, articulation.num_joints, device=device) + 5.0) * -1.0 + limits[..., 1] = torch.rand(num_articulations, articulation.num_joints, device=device) + 5.0 + articulation.write_joint_position_limit_to_sim_index(limits=limits) + + from torch.distributions import Uniform + + joint_pos_limits = wp.to_torch(articulation.data.joint_pos_limits) + joint_vel_limits = wp.to_torch(articulation.data.joint_vel_limits) + pos_dist = Uniform(joint_pos_limits[..., 0], joint_pos_limits[..., 1]) + vel_dist = Uniform(-joint_vel_limits, joint_vel_limits) + + original_body_link_pose_w = wp.to_torch(articulation.data.body_link_pose_w).clone() + original_body_com_vel_w = wp.to_torch(articulation.data.body_com_vel_w).clone() + + rand_joint_pos = pos_dist.sample() + rand_joint_vel = vel_dist.sample() + + articulation.write_joint_position_to_sim_index(position=rand_joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=rand_joint_vel) + # make sure valued updated + body_link_pose_w = wp.to_torch(articulation.data.body_link_pose_w) + body_com_vel_w = wp.to_torch(articulation.data.body_com_vel_w) + original_body_states = torch.cat([original_body_link_pose_w, original_body_com_vel_w], dim=-1) + body_state_w = torch.cat([body_link_pose_w, body_com_vel_w], dim=-1) + assert torch.count_nonzero(original_body_states[:, 1:] != body_state_w[:, 1:]) > ( + len(original_body_states[:, 1:]) / 2 + ) + # validate body - link consistency + body_link_vel_w = wp.to_torch(articulation.data.body_link_vel_w) + torch.testing.assert_close(body_link_pose_w, wp.to_torch(articulation.data.body_link_pose_w)) + # skip lin_vel because it differs from link frame, this should be fine because we are only checking + # if velocity update is triggered, which can be determined by comparing angular velocity + torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) + + # validate link - com conistency + body_com_pos_b = wp.to_torch(articulation.data.body_com_pos_b) + body_com_quat_b = wp.to_torch(articulation.data.body_com_quat_b) + expected_com_pos, expected_com_quat = math_utils.combine_frame_transforms( + body_link_pose_w[..., :3].view(-1, 3), + body_link_pose_w[..., 3:].view(-1, 4), + body_com_pos_b.view(-1, 3), + body_com_quat_b.view(-1, 4), + ) + body_com_pos_w = wp.to_torch(articulation.data.body_com_pos_w) + body_com_quat_w = wp.to_torch(articulation.data.body_com_quat_w) + torch.testing.assert_close(expected_com_pos.view(len(env_idx), -1, 3), body_com_pos_w) + torch.testing.assert_close(expected_com_quat.view(len(env_idx), -1, 4), body_com_quat_w) + + # validate body - com consistency + body_com_lin_vel_w = wp.to_torch(articulation.data.body_com_lin_vel_w) + body_com_ang_vel_w = wp.to_torch(articulation.data.body_com_ang_vel_w) + torch.testing.assert_close(body_com_vel_w[..., :3], body_com_lin_vel_w) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_com_ang_vel_w) + + # validate pos_w, quat_w, pos_b, quat_b is consistent with pose_w and pose_b + expected_com_pose_w = torch.cat((body_com_pos_w, body_com_quat_w), dim=2) + expected_com_pose_b = torch.cat((body_com_pos_b, body_com_quat_b), dim=2) + body_pos_w = wp.to_torch(articulation.data.body_pos_w) + body_quat_w = wp.to_torch(articulation.data.body_quat_w) + expected_body_pose_w = torch.cat((body_pos_w, body_quat_w), dim=2) + body_link_pos_w = wp.to_torch(articulation.data.body_link_pos_w) + body_link_quat_w = wp.to_torch(articulation.data.body_link_quat_w) + expected_body_link_pose_w = torch.cat((body_link_pos_w, body_link_quat_w), dim=2) + body_com_pose_w = wp.to_torch(articulation.data.body_com_pose_w) + body_com_pose_b = wp.to_torch(articulation.data.body_com_pose_b) + body_pose_w = wp.to_torch(articulation.data.body_pose_w) + body_link_pose_w_fresh = wp.to_torch(articulation.data.body_link_pose_w) + torch.testing.assert_close(body_com_pose_w, expected_com_pose_w) + torch.testing.assert_close(body_com_pose_b, expected_com_pose_b) + torch.testing.assert_close(body_pose_w, expected_body_pose_w) + torch.testing.assert_close(body_link_pose_w_fresh, expected_body_link_pose_w) + + # validate pose_w is consistent with individual properties + body_vel_w = wp.to_torch(articulation.data.body_vel_w) + body_com_vel_w_fresh = wp.to_torch(articulation.data.body_com_vel_w) + torch.testing.assert_close(body_pose_w, body_link_pose_w) + torch.testing.assert_close(body_vel_w, body_com_vel_w) + torch.testing.assert_close(body_link_pose_w_fresh, body_link_pose_w) + torch.testing.assert_close(body_com_pose_w, wp.to_torch(articulation.data.body_com_pose_w)) + torch.testing.assert_close(body_vel_w, body_com_vel_w_fresh) + + +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("articulation_type", ["shadow_hand"]) +@pytest.mark.skip(reason="Spatial tendons are not supported in Newton yet.") +def test_spatial_tendons(sim, num_articulations, device, articulation_type): + """Test spatial tendons apis. + This test verifies that: + 1. The articulation is properly initialized + 2. The articulation has spatial tendons + 3. All buffers have correct shapes + 4. The articulation can be simulated + Args: + sim: The simulation fixture + num_articulations: Number of articulations to test + device: The device to run the simulation on + """ + # skip test if Isaac Sim version is less than 5.0 + if has_kit() and get_isaac_sim_version().major < 5: + pytest.skip("Spatial tendons are not supported in Isaac Sim < 5.0. Please update to Isaac Sim 5.0 or later.") + return + articulation_cfg = generate_articulation_cfg(articulation_type="spatial_tendon_test_asset") + articulation, _ = generate_articulation(articulation_cfg, num_articulations, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(articulation) < 10 + + # Play sim + sim.reset() + # Check if articulation is initialized + assert articulation.is_initialized + # Check that fixed base + assert articulation.is_fixed_base + # Check buffers that exists and have correct shapes + assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) + assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) + assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 3) + assert wp.to_torch(articulation.data.body_mass).shape == (num_articulations, articulation.num_bodies) + assert wp.to_torch(articulation.data.body_inertia).shape == (num_articulations, articulation.num_bodies, 9) + assert articulation.num_spatial_tendons == 1 + + articulation.set_spatial_tendon_stiffness_index(stiffness=10.0) + articulation.set_spatial_tendon_limit_stiffness_index(limit_stiffness=10.0) + articulation.set_spatial_tendon_damping_index(damping=10.0) + articulation.set_spatial_tendon_offset_index(offset=10.0) + + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update articulation + articulation.update(sim.cfg.dt) + + +@pytest.mark.parametrize("add_ground_plane", [True]) +@pytest.mark.parametrize("num_articulations", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("articulation_type", ["panda"]) +@pytest.mark.isaacsim_ci +def test_write_joint_frictions_to_sim(sim, num_articulations, device, add_ground_plane, articulation_type): + """Test applying of joint position target functions correctly for a robotic arm.""" + articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type) + articulation, _ = generate_articulation( + articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=device + ) + + # Play the simulator + sim.reset() + + for _ in range(100): + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + # apply action to the articulation + friction = torch.rand(num_articulations, articulation.num_joints, device=device) + + # The static friction must be set first to be sure the dynamic friction is not greater than static + # when both are set. + articulation.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff=friction, + ) + articulation.write_data_to_sim() + + for _ in range(100): + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + joint_friction_coeff_sim = wp.to_torch( + articulation.root_view.get_attribute("joint_friction", SimulationManager.get_model()) + )[:, 0, :] + assert torch.allclose(joint_friction_coeff_sim, friction) + + # Reset simulator to ensure a clean state for the alternative API path + sim.reset() + + # Warm up a few steps to populate buffers + for _ in range(100): + sim.step() + articulation.update(sim.cfg.dt) + + # New random coefficients + friction_2 = torch.rand(num_articulations, articulation.num_joints, device=device) + + # Use the combined setter to write all three at once + articulation.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff=friction_2, + ) + articulation.write_data_to_sim() + + # Step to let sim ingest new params and refresh data buffers + for _ in range(100): + sim.step() + articulation.update(sim.cfg.dt) + + joint_friction_coeff_sim_2 = wp.to_torch( + articulation.root_view.get_attribute("joint_friction", SimulationManager.get_model()) + )[:, 0, :] + + # Validate values propagated + assert torch.allclose(joint_friction_coeff_sim_2, friction_2) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "--maxfail=1"]) diff --git a/source/isaaclab_newton/test/assets/test_rigid_object.py b/source/isaaclab_newton/test/assets/test_rigid_object.py new file mode 100644 index 000000000000..9cc1717d8c5f --- /dev/null +++ b/source/isaaclab_newton/test/assets/test_rigid_object.py @@ -0,0 +1,1285 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# ignore private usage of variables warning +# pyright: reportPrivateUsage=none + + +"""Launch Isaac Sim Simulator first.""" + +from isaaclab.app import AppLauncher + +# launch omniverse app +simulation_app = AppLauncher(headless=True).app + +"""Rest everything follows.""" + +import sys +from typing import Literal + +import pytest +import torch +import warp as wp +from flaky import flaky +from isaaclab_newton.assets import RigidObject +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +from isaaclab_newton.physics import NewtonManager as SimulationManager +from newton.solvers import SolverNotifyFlags + +import isaaclab.sim as sim_utils +from isaaclab.assets import RigidObjectCfg +from isaaclab.sim import SimulationCfg, build_simulation_context +from isaaclab.sim.spawners import materials +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR +from isaaclab.utils.math import ( + combine_frame_transforms, + default_orientation, + quat_apply_inverse, + quat_inv, + quat_mul, + quat_rotate, + random_orientation, +) + +NEWTON_SIM_CFG = SimulationCfg( + physics=NewtonCfg( + solver_cfg=MJWarpSolverCfg(), + ), +) + + +def _newton_sim_context(device, gravity_enabled=True, dt=None, **kwargs): + """Helper to create a Newton simulation context with the correct device. + + When sim_cfg is provided to build_simulation_context, the device, gravity_enabled, and dt + kwargs are ignored. This helper applies them to the shared NEWTON_SIM_CFG before calling. + """ + NEWTON_SIM_CFG.device = device + NEWTON_SIM_CFG.gravity = (0.0, 0.0, -9.81) if gravity_enabled else (0.0, 0.0, 0.0) + if dt is not None: + NEWTON_SIM_CFG.dt = dt + return build_simulation_context(device=device, sim_cfg=NEWTON_SIM_CFG, **kwargs) + + +def generate_cubes_scene( + num_cubes: int = 1, + height=1.0, + api: Literal["none", "rigid_body", "articulation_root"] = "rigid_body", + kinematic_enabled: bool = False, + device: str = "cuda:0", +) -> tuple[RigidObject, torch.Tensor]: + """Generate a scene with the provided number of cubes. + + Args: + num_cubes: Number of cubes to generate. + height: Height of the cubes. + api: The type of API that the cubes should have. + kinematic_enabled: Whether the cubes are kinematic. + device: Device to use for the simulation. + + Returns: + A tuple containing the rigid object representing the cubes and the origins of the cubes. + + """ + origins = torch.tensor([(i * 1.0, 0, height) for i in range(num_cubes)]).to(device) + # Create Top-level Xforms, one for each cube + for i, origin in enumerate(origins): + sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=origin) + + # Resolve spawn configuration + if api == "none": + # since no rigid body properties defined, this is just a static collider + spawn_cfg = sim_utils.CuboidCfg( + size=(0.1, 0.1, 0.1), + collision_props=sim_utils.CollisionPropertiesCfg(), + ) + elif api == "rigid_body": + spawn_cfg = sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", + rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=kinematic_enabled), + ) + elif api == "articulation_root": + spawn_cfg = sim_utils.UsdFileCfg( + usd_path=f"{ISAACLAB_NUCLEUS_DIR}/Tests/RigidObject/Cube/dex_cube_instanceable_with_articulation_root.usd", + rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=kinematic_enabled), + ) + else: + raise ValueError(f"Unknown api: {api}") + + # Create rigid object + cube_object_cfg = RigidObjectCfg( + prim_path="/World/Env_.*/Object", + spawn=spawn_cfg, + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, height)), + ) + cube_object = RigidObject(cfg=cube_object_cfg) + + return cube_object, origins + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_initialization(num_cubes, device): + """Test initialization for prim with rigid body API at the provided prim path.""" + with _newton_sim_context(device, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + # Generate cubes scene + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(cube_object) < 10 + + # Play sim + sim.reset() + + # Check if object is initialized + assert cube_object.is_initialized + assert len(cube_object.body_names) == 1 + + # Check buffers that exists and have correct shapes + assert wp.to_torch(cube_object.data.root_pos_w).shape == (num_cubes, 3) + assert wp.to_torch(cube_object.data.root_quat_w).shape == (num_cubes, 4) + assert wp.to_torch(cube_object.data.body_mass).shape == (num_cubes, 1) + assert wp.to_torch(cube_object.data.body_inertia).shape == (num_cubes, 1, 9) + + # Simulate physics + for _ in range(2): + # perform rendering + sim.step() + # update object + cube_object.update(sim.cfg.dt) + + +@pytest.mark.skip(reason="Newton does not support kinematic rigid bodies") +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_initialization_with_kinematic_enabled(num_cubes, device): + """Test that initialization for prim with kinematic flag enabled.""" + with _newton_sim_context(device, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + # Generate cubes scene + cube_object, origins = generate_cubes_scene(num_cubes=num_cubes, kinematic_enabled=True, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(cube_object) < 10 + + # Play sim + sim.reset() + + # Check if object is initialized + assert cube_object.is_initialized + assert len(cube_object.body_names) == 1 + + # Check buffers that exists and have correct shapes + assert wp.to_torch(cube_object.data.root_pos_w).shape == (num_cubes, 3) + assert wp.to_torch(cube_object.data.root_quat_w).shape == (num_cubes, 4) + + # Simulate physics + for _ in range(2): + # perform rendering + sim.step() + # update object + cube_object.update(sim.cfg.dt) + # check that the object is kinematic + default_root_pose = wp.to_torch(cube_object.data.default_root_pose).clone() + default_root_vel = wp.to_torch(cube_object.data.default_root_vel).clone() + default_root_pose[:, :3] += origins + torch.testing.assert_close(wp.to_torch(cube_object.data.root_link_pose_w), default_root_pose) + torch.testing.assert_close(wp.to_torch(cube_object.data.root_com_vel_w), default_root_vel) + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_initialization_with_no_rigid_body(num_cubes, device): + """Test that initialization fails when no rigid body is found at the provided prim path.""" + with _newton_sim_context(device, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + # Generate cubes scene + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, api="none", device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(cube_object) < 10 + + # Play sim + with pytest.raises(RuntimeError): + sim.reset() + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_initialization_with_articulation_root(num_cubes, device): + """Test that initialization fails when an articulation root is found at the provided prim path.""" + with _newton_sim_context(device, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + # Generate cubes scene + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, api="articulation_root", device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(cube_object) < 10 + + # Play sim + with pytest.raises(RuntimeError): + sim.reset() + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_external_force_buffer(device): + """Test if external force buffer correctly updates in the force value is zero case. + + In this test, we apply a non-zero force, then a zero force, then finally a non-zero force + to an object. We check if the force buffer is properly updated at each step. + """ + + # Generate cubes scene + with _newton_sim_context(device, add_ground_plane=True, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, origins = generate_cubes_scene(num_cubes=1, device=device) + + # play the simulator + sim.reset() + + # find bodies to apply the force + body_ids, body_names = cube_object.find_bodies(".*") + + # reset object + cube_object.reset() + + # perform simulation + for step in range(5): + # initiate force tensor + external_wrench_b = torch.zeros(cube_object.num_instances, len(body_ids), 6, device=sim.device) + + if step == 0 or step == 3: + # set a non-zero force + force = 1 + else: + # set a zero force + force = 0 + + # set force value + external_wrench_b[:, :, 0] = force + external_wrench_b[:, :, 3] = force + + # apply force + cube_object.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + body_ids=body_ids, + ) + + # check if the cube's force and torque buffers are correctly updated + for i in range(cube_object.num_instances): + assert wp.to_torch(cube_object._permanent_wrench_composer.composed_force)[i, 0, 0].item() == force + assert wp.to_torch(cube_object._permanent_wrench_composer.composed_torque)[i, 0, 0].item() == force + + # Check if the instantaneous wrench is correctly added to the permanent wrench + cube_object.permanent_wrench_composer.add_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + body_ids=body_ids, + ) + + # apply action to the object + cube_object.write_data_to_sim() + + # perform step + sim.step() + + # update buffers + cube_object.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_cubes", [2, 4]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_external_force_on_single_body(num_cubes, device): + """Test application of external force on the base of the object. + + In this test, we apply a force equal to the weight of an object on the base of + one of the objects. We check that the object does not move. For the other object, + we do not apply any force and check that it falls down. + + We validate that this works when we apply the force in the global frame and in the local frame. + """ + # Generate cubes scene + with _newton_sim_context(device, add_ground_plane=True, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, origins = generate_cubes_scene(num_cubes=num_cubes, device=device) + + # Play the simulator + sim.reset() + + # Find bodies to apply the force + body_ids, body_names = cube_object.find_bodies(".*") + + # Sample a force equal to the weight of the object + external_wrench_b = torch.zeros(cube_object.num_instances, len(body_ids), 6, device=sim.device) + # Every 2nd cube should have a force applied to it + external_wrench_b[0::2, :, 2] = 9.81 * wp.to_torch(cube_object.data.body_mass)[0] + + # Now we are ready! + for i in range(5): + # reset root state + root_pose = wp.to_torch(cube_object.data.default_root_pose).clone() + root_vel = wp.to_torch(cube_object.data.default_root_vel).clone() + + # need to shift the position of the cubes otherwise they will be on top of each other + root_pose[:, :3] = origins + cube_object.write_root_pose_to_sim_index(root_pose=root_pose) + cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) + + # reset object + cube_object.reset() + + is_global = False + if i % 2 == 0: + is_global = True + positions = wp.to_torch(cube_object.data.body_com_pos_w)[:, body_ids, :3] + else: + positions = None + + # apply force + cube_object.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=positions, + body_ids=body_ids, + is_global=is_global, + ) + # perform simulation + for _ in range(5): + # apply action to the object + cube_object.write_data_to_sim() + + # perform step + sim.step() + + # update buffers + cube_object.update(sim.cfg.dt) + + # First object should still be at the same Z position (1.0) + torch.testing.assert_close( + wp.to_torch(cube_object.data.root_pos_w)[0::2, 2], torch.ones(num_cubes // 2, device=sim.device) + ) + # Second object should have fallen, so it's Z height should be less than initial height of 1.0 + assert torch.all(wp.to_torch(cube_object.data.root_pos_w)[1::2, 2] < 1.0) + + +@pytest.mark.skip(reason="Newton wrench composer at-position force composition differs from PhysX") +@pytest.mark.parametrize("num_cubes", [2, 4]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_external_force_on_single_body_at_position(num_cubes, device): + """Test application of external force on the base of the object at a specific position. + + In this test, we apply a force equal to the weight of an object on the base of + one of the objects at 1m in the Y direction, we check that the object rotates around it's X axis. + For the other object, we do not apply any force and check that it falls down. + + We validate that this works when we apply the force in the global frame and in the local frame. + """ + # Generate cubes scene + with _newton_sim_context(device, add_ground_plane=True, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, origins = generate_cubes_scene(num_cubes=num_cubes, device=device) + + # Play the simulator + sim.reset() + + # Find bodies to apply the force + body_ids, body_names = cube_object.find_bodies(".*") + + # Sample a force equal to the weight of the object + external_wrench_b = torch.zeros(cube_object.num_instances, len(body_ids), 6, device=sim.device) + external_wrench_positions_b = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=sim.device) + # Every 2nd cube should have a force applied to it + external_wrench_b[0::2, :, 2] = 500.0 + external_wrench_positions_b[0::2, :, 1] = 1.0 + + # Desired force and torque + desired_force = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=sim.device) + desired_force[0::2, :, 2] = 1000.0 + desired_torque = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=sim.device) + desired_torque[0::2, :, 0] = 1000.0 + # Now we are ready! + for i in range(5): + # reset root state + root_pose = wp.to_torch(cube_object.data.default_root_pose).clone() + root_vel = wp.to_torch(cube_object.data.default_root_vel).clone() + + # need to shift the position of the cubes otherwise they will be on top of each other + root_pose[:, :3] = origins + cube_object.write_root_pose_to_sim_index(root_pose=root_pose) + cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) + + # reset object + cube_object.reset() + + is_global = False + if i % 2 == 0: + is_global = True + body_com_pos_w = wp.to_torch(cube_object.data.body_com_pos_w)[:, body_ids, :3] + external_wrench_positions_b[..., 0] = 0.0 + external_wrench_positions_b[..., 1] = 1.0 + external_wrench_positions_b[..., 2] = 0.0 + external_wrench_positions_b += body_com_pos_w + else: + external_wrench_positions_b[..., 0] = 0.0 + external_wrench_positions_b[..., 1] = 1.0 + external_wrench_positions_b[..., 2] = 0.0 + + # apply force + cube_object.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=external_wrench_positions_b, + body_ids=body_ids, + is_global=is_global, + ) + cube_object.permanent_wrench_composer.add_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=external_wrench_positions_b, + body_ids=body_ids, + is_global=is_global, + ) + torch.testing.assert_close( + wp.to_torch(cube_object._permanent_wrench_composer.composed_force)[:, 0, :], + desired_force[:, 0, :], + rtol=1e-6, + atol=1e-7, + ) + torch.testing.assert_close( + wp.to_torch(cube_object._permanent_wrench_composer.composed_torque)[:, 0, :], + desired_torque[:, 0, :], + rtol=1e-6, + atol=1e-7, + ) + # perform simulation + for _ in range(5): + # apply action to the object + cube_object.write_data_to_sim() + + # perform step + sim.step() + + # update buffers + cube_object.update(sim.cfg.dt) + + # The first object should be rotating around it's X axis + assert torch.all(torch.abs(wp.to_torch(cube_object.data.root_ang_vel_b)[0::2, 0]) > 0.1) + # Second object should have fallen, so it's Z height should be less than initial height of 1.0 + assert torch.all(wp.to_torch(cube_object.data.root_pos_w)[1::2, 2] < 1.0) + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_set_rigid_object_state(num_cubes, device): + """Test setting the state of the rigid object. + + In this test, we set the state of the rigid object to a random state and check + that the object is in that state after simulation. We set gravity to zero as + we don't want any external forces acting on the object to ensure state remains static. + """ + # Turn off gravity for this test as we don't want any external forces acting on the object + # to ensure state remains static + with _newton_sim_context(device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + # Generate cubes scene + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, device=device) + + # Play the simulator + sim.reset() + + state_types = ["root_pos_w", "root_quat_w", "root_lin_vel_w", "root_ang_vel_w"] + + # Set each state type individually as they are dependent on each other + for state_type_to_randomize in state_types: + state_dict = { + "root_pos_w": torch.zeros_like(wp.to_torch(cube_object.data.root_pos_w), device=sim.device), + "root_quat_w": default_orientation(num=num_cubes, device=sim.device), + "root_lin_vel_w": torch.zeros_like(wp.to_torch(cube_object.data.root_lin_vel_w), device=sim.device), + "root_ang_vel_w": torch.zeros_like(wp.to_torch(cube_object.data.root_ang_vel_w), device=sim.device), + } + + # Now we are ready! + for _ in range(5): + # reset object + cube_object.reset() + + # Set random state + if state_type_to_randomize == "root_quat_w": + state_dict[state_type_to_randomize] = random_orientation(num=num_cubes, device=sim.device) + else: + state_dict[state_type_to_randomize] = torch.randn(num_cubes, 3, device=sim.device) + + # perform simulation + for _ in range(5): + root_pose = torch.cat( + [state_dict["root_pos_w"], state_dict["root_quat_w"]], + dim=-1, + ) + root_vel = torch.cat( + [state_dict["root_lin_vel_w"], state_dict["root_ang_vel_w"]], + dim=-1, + ) + # reset root state + cube_object.write_root_pose_to_sim_index(root_pose=root_pose) + cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) + + sim.step() + + # assert that set root quantities are equal to the ones set in the state_dict + for key, expected_value in state_dict.items(): + value = wp.to_torch(getattr(cube_object.data, key)) + # Newton reads state directly from sim (not cached), so post-step drift + # from velocity integration causes larger differences than PhysX + torch.testing.assert_close(value, expected_value, rtol=1e-1, atol=1e-1) + + cube_object.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_reset_rigid_object(num_cubes, device): + """Test resetting the state of the rigid object.""" + with _newton_sim_context(device, gravity_enabled=True, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + # Generate cubes scene + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, device=device) + + # Play the simulator + sim.reset() + + for i in range(5): + # perform rendering + sim.step() + + # update object + cube_object.update(sim.cfg.dt) + + # Move the object to a random position + root_pose = wp.to_torch(cube_object.data.default_root_pose).clone() + root_pose[:, :3] = torch.randn(num_cubes, 3, device=sim.device) + + # Random orientation + root_pose[:, 3:7] = random_orientation(num=num_cubes, device=sim.device) + cube_object.write_root_pose_to_sim_index(root_pose=root_pose) + root_vel = wp.to_torch(cube_object.data.default_root_vel).clone() + cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) + + if i % 2 == 0: + # reset object + cube_object.reset() + + # Reset should zero external forces and torques + assert not cube_object._instantaneous_wrench_composer.active + assert not cube_object._permanent_wrench_composer.active + assert torch.count_nonzero(wp.to_torch(cube_object._instantaneous_wrench_composer.composed_force)) == 0 + assert torch.count_nonzero(wp.to_torch(cube_object._instantaneous_wrench_composer.composed_torque)) == 0 + assert torch.count_nonzero(wp.to_torch(cube_object._permanent_wrench_composer.composed_force)) == 0 + assert torch.count_nonzero(wp.to_torch(cube_object._permanent_wrench_composer.composed_torque)) == 0 + + +@pytest.mark.skip(reason="Newton friction/restitution/material not yet supported") +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_rigid_body_set_material_properties(num_cubes, device): + """Test getting and setting material properties of rigid object.""" + with _newton_sim_context(device, gravity_enabled=True, add_ground_plane=True, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + # Generate cubes scene + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, device=device) + + # Play sim + sim.reset() + + # Set material properties + static_friction = torch.FloatTensor(num_cubes, 1).uniform_(0.4, 0.8) + dynamic_friction = torch.FloatTensor(num_cubes, 1).uniform_(0.4, 0.8) + restitution = torch.FloatTensor(num_cubes, 1).uniform_(0.0, 0.2) + + materials = torch.cat([static_friction, dynamic_friction, restitution], dim=-1) + + indices = torch.tensor(range(num_cubes), dtype=torch.int32) + # Add friction to cube + cube_object.root_view.set_material_properties( + wp.from_torch(materials, dtype=wp.float32), wp.from_torch(indices, dtype=wp.int32) + ) + + # Simulate physics + # perform rendering + sim.step() + # update object + cube_object.update(sim.cfg.dt) + + # Get material properties + materials_to_check = wp.to_torch(cube_object.root_view.get_material_properties()) + + # Check if material properties are set correctly + torch.testing.assert_close(materials_to_check.reshape(num_cubes, 3), materials) + + +@pytest.mark.skip(reason="Newton friction/restitution/material not yet supported") +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_rigid_body_no_friction(num_cubes, device): + """Test that a rigid object with no friction will maintain it's velocity when sliding across a plane.""" + with _newton_sim_context(device, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + # Generate cubes scene + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=0.0, device=device) + + # Create ground plane with no friction + cfg = sim_utils.GroundPlaneCfg( + physics_material=materials.RigidBodyMaterialCfg( + static_friction=0.0, + dynamic_friction=0.0, + restitution=0.0, + ) + ) + cfg.func("/World/GroundPlane", cfg) + + # Play sim + sim.reset() + + # Set material friction properties to be all zero + static_friction = torch.zeros(num_cubes, 1) + dynamic_friction = torch.zeros(num_cubes, 1) + restitution = torch.FloatTensor(num_cubes, 1).uniform_(0.0, 0.2) + + cube_object_materials = torch.cat([static_friction, dynamic_friction, restitution], dim=-1) + indices = torch.tensor(range(num_cubes), dtype=torch.int32) + + cube_object.root_view.set_material_properties( + wp.from_torch(cube_object_materials, dtype=wp.float32), wp.from_torch(indices, dtype=wp.int32) + ) + + # Set initial velocity + # Initial velocity in X to get the block moving + initial_velocity = torch.zeros((num_cubes, 6), device=sim.cfg.device) + initial_velocity[:, 0] = 0.1 + + cube_object.write_root_velocity_to_sim_index(root_velocity=initial_velocity) + + # Simulate physics + for _ in range(5): + # perform rendering + sim.step() + # update object + cube_object.update(sim.cfg.dt) + + # Non-deterministic when on GPU, so we use different tolerances + if device == "cuda:0": + tolerance = 1e-2 + else: + tolerance = 1e-5 + + torch.testing.assert_close( + wp.to_torch(cube_object.data.root_lin_vel_w), initial_velocity[:, :3], rtol=1e-5, atol=tolerance + ) + + +@pytest.mark.skip(reason="Newton friction/restitution/material not yet supported") +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_rigid_body_with_static_friction(num_cubes, device): + """Test that static friction applied to rigid object works as expected. + + This test works by applying a force to the object and checking if the object moves or not based on the + mu (coefficient of static friction) value set for the object. We set the static friction to be non-zero and + apply a force to the object. When the force applied is below mu, the object should not move. When the force + applied is above mu, the object should move. + """ + with _newton_sim_context(device, dt=0.01, add_ground_plane=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=0.03125, device=device) + + # Create ground plane + static_friction_coefficient = 0.5 + cfg = sim_utils.GroundPlaneCfg( + physics_material=materials.RigidBodyMaterialCfg( + static_friction=static_friction_coefficient, + dynamic_friction=static_friction_coefficient, # This shouldn't be required but is due to a bug in PhysX + ) + ) + cfg.func("/World/GroundPlane", cfg) + + # Play sim + sim.reset() + + # Set static friction to be non-zero + # Dynamic friction also needs to be zero due to a bug in PhysX + static_friction = torch.Tensor([[static_friction_coefficient]] * num_cubes) + dynamic_friction = torch.Tensor([[static_friction_coefficient]] * num_cubes) + restitution = torch.zeros(num_cubes, 1) + + cube_object_materials = torch.cat([static_friction, dynamic_friction, restitution], dim=-1) + + indices = torch.tensor(range(num_cubes), dtype=torch.int32) + + # Add friction to cube + cube_object.root_view.set_material_properties( + wp.from_torch(cube_object_materials, dtype=wp.float32), wp.from_torch(indices, dtype=wp.int32) + ) + + # let everything settle + for _ in range(100): + sim.step() + cube_object.update(sim.cfg.dt) + cube_object.write_root_velocity_to_sim_index(root_velocity=torch.zeros((num_cubes, 6), device=sim.device)) + cube_mass = wp.to_torch(cube_object.data.body_mass) + gravity_magnitude = abs(sim.cfg.gravity[2]) + # 2 cases: force applied is below and above mu + # below mu: block should not move as the force applied is <= mu + # above mu: block should move as the force applied is > mu + for force in "below_mu", "above_mu": + # set initial velocity to zero + cube_object.write_root_velocity_to_sim_index(root_velocity=torch.zeros((num_cubes, 6), device=sim.device)) + + external_wrench_b = torch.zeros((num_cubes, 1, 6), device=sim.device) + if force == "below_mu": + external_wrench_b[..., 0] = static_friction_coefficient * cube_mass * gravity_magnitude * 0.99 + else: + external_wrench_b[..., 0] = static_friction_coefficient * cube_mass * gravity_magnitude * 1.01 + + cube_object.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + ) + + # Get root state + initial_root_pos = wp.to_torch(cube_object.data.root_pos_w).clone() + # Simulate physics + for _ in range(200): + # apply the wrench + cube_object.write_data_to_sim() + sim.step() + # update object + cube_object.update(sim.cfg.dt) + if force == "below_mu": + # Assert that the block has not moved + torch.testing.assert_close( + wp.to_torch(cube_object.data.root_pos_w), initial_root_pos, rtol=2e-3, atol=2e-3 + ) + if force == "above_mu": + assert (wp.to_torch(cube_object.data.root_pos_w)[..., 0] - initial_root_pos[..., 0] > 0.02).all() + + +@pytest.mark.skip(reason="Newton friction/restitution/material not yet supported") +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_rigid_body_with_restitution(num_cubes, device): + """Test that restitution when applied to rigid object works as expected. + + This test works by dropping a block from a height and checking if the block bounces or not based on the + restitution value set for the object. We set the restitution to be non-zero and drop the block from a height. + When the restitution is 0, the block should not bounce. When the restitution is between 0 and 1, the block + should bounce with less energy. + """ + for expected_collision_type in "partially_elastic", "inelastic": + with _newton_sim_context(device, add_ground_plane=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=1.0, device=device) + + # Set static friction to be non-zero + if expected_collision_type == "inelastic": + restitution_coefficient = 0.0 + elif expected_collision_type == "partially_elastic": + restitution_coefficient = 0.5 + + # Create ground plane such that has a restitution of 1.0 (perfectly elastic collision) + cfg = sim_utils.GroundPlaneCfg( + physics_material=materials.RigidBodyMaterialCfg( + restitution=restitution_coefficient, + ) + ) + cfg.func("/World/GroundPlane", cfg) + + indices = torch.tensor(range(num_cubes), dtype=torch.int32) + + # Play sim + sim.reset() + + root_pose = torch.zeros(num_cubes, 7, device=sim.device) + root_pose[:, 3] = 1.0 # To make orientation a quaternion + for i in range(num_cubes): + root_pose[i, 1] = 1.0 * i + root_pose[:, 2] = 1.0 # Set an initial drop height + root_vel = torch.zeros(num_cubes, 6, device=sim.device) + root_vel[:, 2] = -1.0 # Set an initial downward velocity + + cube_object.write_root_pose_to_sim_index(root_pose=root_pose) + cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) + + static_friction = torch.zeros(num_cubes, 1) + dynamic_friction = torch.zeros(num_cubes, 1) + restitution = torch.Tensor([[restitution_coefficient]] * num_cubes) + + cube_object_materials = torch.cat([static_friction, dynamic_friction, restitution], dim=-1) + + # Add restitution to cube + cube_object.root_view.set_material_properties( + wp.from_torch(cube_object_materials, dtype=wp.float32), wp.from_torch(indices, dtype=wp.int32) + ) + + curr_z_velocity = wp.to_torch(cube_object.data.root_lin_vel_w)[:, 2].clone() + + for _ in range(100): + sim.step() + + # update object + cube_object.update(sim.cfg.dt) + curr_z_velocity = wp.to_torch(cube_object.data.root_lin_vel_w)[:, 2].clone() + + if expected_collision_type == "inelastic": + # assert that the block has not bounced by checking that the z velocity is less than or equal to 0 + assert (curr_z_velocity <= 0.0).all() + + if torch.all(curr_z_velocity <= 0.0): + # Still in the air + prev_z_velocity = curr_z_velocity + else: + # collision has happened, exit the for loop + break + + if expected_collision_type == "partially_elastic": + # Assert that the block has lost some energy by checking that the z velocity is less + assert torch.all(torch.le(abs(curr_z_velocity), abs(prev_z_velocity))) + assert (curr_z_velocity > 0.0).all() + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.isaacsim_ci +def test_rigid_body_set_mass(num_cubes, device): + """Test getting and setting mass of rigid object.""" + with _newton_sim_context(device, gravity_enabled=False, add_ground_plane=True, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + # Create a scene with random cubes + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, height=1.0, device=device) + + # Play sim + sim.reset() + + # Get masses before increasing + original_masses = wp.to_torch(cube_object.data.body_mass) + + assert original_masses.shape == (num_cubes, 1) + + # Randomize mass of the object + masses = original_masses + torch.zeros(num_cubes, 1, device=device).uniform_(4, 8) + + # Set masses using Newton API + cube_object.set_masses_index(masses=wp.from_torch(masses, dtype=wp.float32)) + + torch.testing.assert_close(wp.to_torch(cube_object.data.body_mass), masses) + + # Simulate physics + # perform rendering + sim.step() + # update object + cube_object.update(sim.cfg.dt) + + masses_to_check = wp.to_torch(cube_object.data.body_mass) + + # Check if mass is set correctly + torch.testing.assert_close(masses, masses_to_check) + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("gravity_enabled", [True, False]) +@pytest.mark.isaacsim_ci +def test_gravity_vec_w(num_cubes, device, gravity_enabled): + """Test that gravity vector direction is set correctly for the rigid object.""" + with _newton_sim_context(device, gravity_enabled=gravity_enabled) as sim: + sim._app_control_on_stop_handle = None + # Create a scene with random cubes + cube_object, _ = generate_cubes_scene(num_cubes=num_cubes, device=device) + + # Obtain gravity direction + if gravity_enabled: + gravity_dir = (0.0, 0.0, -1.0) + else: + gravity_dir = (0.0, 0.0, 0.0) + + # Play sim + sim.reset() + + # Check that gravity is set correctly + assert wp.to_torch(cube_object.data.GRAVITY_VEC_W)[0, 0] == gravity_dir[0] + assert wp.to_torch(cube_object.data.GRAVITY_VEC_W)[0, 1] == gravity_dir[1] + assert wp.to_torch(cube_object.data.GRAVITY_VEC_W)[0, 2] == gravity_dir[2] + + # Simulate physics + for _ in range(2): + # perform rendering + sim.step() + # update object + cube_object.update(sim.cfg.dt) + + # Expected gravity value is the acceleration of the body + gravity = torch.zeros(num_cubes, 1, 6, device=device) + if gravity_enabled: + gravity[:, :, 2] = -9.81 + # Check the body accelerations are correct + torch.testing.assert_close(wp.to_torch(cube_object.data.body_acc_w), gravity) + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("with_offset", [True, False]) +@pytest.mark.isaacsim_ci +@flaky(max_runs=3, min_passes=1) +def test_body_root_state_properties(num_cubes, device, with_offset): + """Test the root_com_state_w, root_link_state_w, body_com_state_w, and body_link_state_w properties.""" + with _newton_sim_context(device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + # Create a scene with random cubes + cube_object, env_pos = generate_cubes_scene(num_cubes=num_cubes, height=0.0, device=device) + + # Play sim + sim.reset() + + # Check if cube_object is initialized + assert cube_object.is_initialized + + # change center of mass offset from link frame + if with_offset: + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_cubes, 1) + else: + offset = torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_cubes, 1) + + # Set center of mass offset via Newton API (position only, no quaternion) + com_pos = offset.unsqueeze(1) # (N, 1, 3) + cube_object.set_coms_index(coms=wp.from_torch(com_pos, dtype=wp.vec3f)) + with wp.ScopedDevice(device): + SimulationManager._solver.notify_model_changed(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + # check center of mass has been set + torch.testing.assert_close(wp.to_torch(cube_object.data.body_com_pos_b).squeeze(1), offset) + + # random z spin velocity + spin_twist = torch.zeros(6, device=device) + spin_twist[5] = torch.randn(1, device=device) + + # Simulate physics + for _ in range(100): + # spin the object around Z axis (com) + cube_object.write_root_velocity_to_sim_index(root_velocity=spin_twist.repeat(num_cubes, 1)) + # perform rendering + sim.step() + # update object + cube_object.update(sim.cfg.dt) + + # get state properties + root_link_pose_w = wp.to_torch(cube_object.data.root_link_pose_w) + root_link_vel_w = wp.to_torch(cube_object.data.root_link_vel_w) + root_com_pose_w = wp.to_torch(cube_object.data.root_com_pose_w) + root_com_vel_w = wp.to_torch(cube_object.data.root_com_vel_w) + body_link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) + body_link_vel_w = wp.to_torch(cube_object.data.body_link_vel_w) + body_com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) + body_com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) + + # if offset is [0,0,0] all root_state_%_w will match and all body_%_w will match + if not with_offset: + torch.testing.assert_close(root_link_pose_w, root_com_pose_w) + torch.testing.assert_close(root_com_vel_w, root_link_vel_w) + torch.testing.assert_close(root_link_pose_w, root_link_pose_w) + torch.testing.assert_close(root_com_vel_w, root_link_vel_w) + torch.testing.assert_close(body_link_pose_w, body_com_pose_w) + torch.testing.assert_close(body_com_vel_w, body_link_vel_w) + torch.testing.assert_close(body_link_pose_w, body_link_pose_w) + torch.testing.assert_close(body_com_vel_w, body_link_vel_w) + else: + # cubes are spinning around center of mass + # position will not match + # center of mass position will be constant (i.e. spinning around com) + _tol = dict(atol=2e-3, rtol=2e-3) + torch.testing.assert_close(env_pos + offset, root_com_pose_w[..., :3], **_tol) + torch.testing.assert_close(env_pos + offset, body_com_pose_w[..., :3].squeeze(-2), **_tol) + # link position will be moving but should stay constant away from center of mass + root_link_state_pos_rel_com = quat_apply_inverse( + root_link_pose_w[..., 3:], + root_link_pose_w[..., :3] - root_com_pose_w[..., :3], + ) + torch.testing.assert_close(-offset, root_link_state_pos_rel_com, **_tol) + body_link_state_pos_rel_com = quat_apply_inverse( + body_link_pose_w[..., 3:], + body_link_pose_w[..., :3] - body_com_pose_w[..., :3], + ) + torch.testing.assert_close(-offset, body_link_state_pos_rel_com.squeeze(-2), **_tol) + + # orientation of com will be a constant rotation from link orientation + com_quat_b = wp.to_torch(cube_object.data.body_com_quat_b) + com_quat_w = quat_mul(body_link_pose_w[..., 3:], com_quat_b) + torch.testing.assert_close(com_quat_w, body_com_pose_w[..., 3:], **_tol) + torch.testing.assert_close(com_quat_w.squeeze(-2), root_com_pose_w[..., 3:], **_tol) + + # orientation of link will match root state will always match + torch.testing.assert_close(root_link_pose_w[..., 3:], root_link_pose_w[..., 3:], **_tol) + torch.testing.assert_close(body_link_pose_w[..., 3:], body_link_pose_w[..., 3:], **_tol) + + # lin_vel will not match + # center of mass vel will be constant (i.e. spinning around com) + torch.testing.assert_close(torch.zeros_like(root_com_vel_w[..., :3]), root_com_vel_w[..., :3], **_tol) + torch.testing.assert_close(torch.zeros_like(body_com_vel_w[..., :3]), body_com_vel_w[..., :3], **_tol) + # link frame will be moving, and should be equal to input angular velocity cross offset + lin_vel_rel_root_gt = quat_apply_inverse(root_link_pose_w[..., 3:], root_link_vel_w[..., :3]) + lin_vel_rel_body_gt = quat_apply_inverse(body_link_pose_w[..., 3:], body_link_vel_w[..., :3]) + lin_vel_rel_gt = torch.linalg.cross(spin_twist.repeat(num_cubes, 1)[..., 3:], -offset) + torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_root_gt, **_tol) + torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_body_gt.squeeze(-2), **_tol) + + # ang_vel will always match + torch.testing.assert_close(root_com_vel_w[..., 3:], root_com_vel_w[..., 3:]) + torch.testing.assert_close(root_com_vel_w[..., 3:], root_link_vel_w[..., 3:]) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_com_vel_w[..., 3:]) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("with_offset", [True, False]) +@pytest.mark.parametrize("state_location", ["com", "link"]) +@pytest.mark.isaacsim_ci +def test_write_root_state(num_cubes, device, with_offset, state_location): + """Test the setters for root_state using both the link frame and center of mass as reference frame.""" + with _newton_sim_context(device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + # Create a scene with random cubes + cube_object, env_pos = generate_cubes_scene(num_cubes=num_cubes, height=0.0, device=device) + env_idx = torch.tensor([x for x in range(num_cubes)], dtype=torch.int32, device=device) + + # Play sim + sim.reset() + + # Check if cube_object is initialized + assert cube_object.is_initialized + + # change center of mass offset from link frame + if with_offset: + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_cubes, 1) + else: + offset = torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_cubes, 1) + + # Set center of mass offset via Newton API (position only) + com_pos = offset.unsqueeze(1) # (N, 1, 3) + cube_object.set_coms_index(coms=wp.from_torch(com_pos, dtype=wp.vec3f)) + with wp.ScopedDevice(device): + SimulationManager._solver.notify_model_changed(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + # check center of mass has been set + torch.testing.assert_close(wp.to_torch(cube_object.data.body_com_pos_b).squeeze(1), offset) + + rand_state = torch.zeros(num_cubes, 13, device=device) + rand_state[..., :7] = wp.to_torch(cube_object.data.default_root_pose) + rand_state[..., :3] += env_pos + # make quaternion a unit vector + rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) + + for i in range(10): + # perform step + sim.step() + # update buffers + cube_object.update(sim.cfg.dt) + + if state_location == "com": + if i % 2 == 0: + cube_object.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7]) + cube_object.write_root_com_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) + else: + cube_object.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_idx) + cube_object.write_root_com_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_idx) + elif state_location == "link": + if i % 2 == 0: + cube_object.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7]) + cube_object.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) + else: + cube_object.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_idx) + cube_object.write_root_link_velocity_to_sim_index( + root_velocity=rand_state[..., 7:], env_ids=env_idx + ) + + if state_location == "com": + torch.testing.assert_close(rand_state[..., :7], wp.to_torch(cube_object.data.root_com_pose_w)) + torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(cube_object.data.root_com_vel_w)) + elif state_location == "link": + torch.testing.assert_close(rand_state[..., :7], wp.to_torch(cube_object.data.root_link_pose_w)) + torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(cube_object.data.root_link_vel_w)) + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("with_offset", [True]) +@pytest.mark.parametrize("state_location", ["com", "link", "root"]) +@pytest.mark.isaacsim_ci +def test_write_state_functions_data_consistency(num_cubes, device, with_offset, state_location): + """Test the setters for root_state using both the link frame and center of mass as reference frame.""" + with _newton_sim_context(device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + # Create a scene with random cubes + cube_object, env_pos = generate_cubes_scene(num_cubes=num_cubes, height=0.0, device=device) + + # Play sim + sim.reset() + + # Check if cube_object is initialized + assert cube_object.is_initialized + + # change center of mass offset from link frame + if with_offset: + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_cubes, 1) + else: + offset = torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_cubes, 1) + + # Set center of mass offset via Newton API (position only) + com_pos = offset.unsqueeze(1) # (N, 1, 3) + cube_object.set_coms_index(coms=wp.from_torch(com_pos, dtype=wp.vec3f)) + with wp.ScopedDevice(device): + SimulationManager._solver.notify_model_changed(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + # check center of mass has been set + torch.testing.assert_close(wp.to_torch(cube_object.data.body_com_pos_b).squeeze(1), offset) + + rand_state = torch.rand(num_cubes, 13, device=device) + rand_state[..., :3] += env_pos + # make quaternion a unit vector + rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) + + # perform step + sim.step() + # update buffers + cube_object.update(sim.cfg.dt) + + if state_location == "com": + cube_object.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7]) + cube_object.write_root_com_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) + elif state_location == "link": + cube_object.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7]) + cube_object.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) + elif state_location == "root": + cube_object.write_root_pose_to_sim_index(root_pose=rand_state[..., :7]) + cube_object.write_root_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) + + if state_location == "com": + root_com_pose_w = wp.to_torch(cube_object.data.root_com_pose_w) + root_com_vel_w = wp.to_torch(cube_object.data.root_com_vel_w) + body_com_pose_b = wp.to_torch(cube_object.data.body_com_pose_b) + expected_root_link_pos, expected_root_link_quat = combine_frame_transforms( + root_com_pose_w[:, :3], + root_com_pose_w[:, 3:], + quat_rotate(quat_inv(body_com_pose_b[:, 0, 3:7]), -body_com_pose_b[:, 0, :3]), + quat_inv(body_com_pose_b[:, 0, 3:7]), + ) + expected_root_link_pose = torch.cat((expected_root_link_pos, expected_root_link_quat), dim=1) + root_link_pose_w = wp.to_torch(cube_object.data.root_link_pose_w) + root_link_vel_w = wp.to_torch(cube_object.data.root_link_vel_w) + # test both root_pose and root_link successfully updated when root_com updates + torch.testing.assert_close(expected_root_link_pose, root_link_pose_w) + # skip lin_vel because it differs from link frame, this should be fine because we are only checking + # if velocity update is triggered, which can be determined by comparing angular velocity + torch.testing.assert_close(root_com_vel_w[:, 3:], root_link_vel_w[:, 3:]) + torch.testing.assert_close(expected_root_link_pose, root_link_pose_w) + torch.testing.assert_close(root_com_vel_w[:, 3:], wp.to_torch(cube_object.data.root_com_vel_w)[:, 3:]) + elif state_location == "link": + root_link_pose_w = wp.to_torch(cube_object.data.root_link_pose_w) + root_link_vel_w = wp.to_torch(cube_object.data.root_link_vel_w) + body_com_pose_b = wp.to_torch(cube_object.data.body_com_pose_b) + expected_com_pos, expected_com_quat = combine_frame_transforms( + root_link_pose_w[:, :3], + root_link_pose_w[:, 3:], + body_com_pose_b[:, 0, :3], + body_com_pose_b[:, 0, 3:7], + ) + expected_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1) + root_com_pose_w = wp.to_torch(cube_object.data.root_com_pose_w) + root_com_vel_w = wp.to_torch(cube_object.data.root_com_vel_w) + # test both root_pose and root_com successfully updated when root_link updates + torch.testing.assert_close(expected_com_pose, root_com_pose_w) + # skip lin_vel because it differs from link frame, this should be fine because we are only checking + # if velocity update is triggered, which can be determined by comparing angular velocity + torch.testing.assert_close(root_link_vel_w[:, 3:], root_com_vel_w[:, 3:]) + torch.testing.assert_close(root_link_pose_w, wp.to_torch(cube_object.data.root_link_pose_w)) + torch.testing.assert_close(root_link_vel_w[:, 3:], wp.to_torch(cube_object.data.root_com_vel_w)[:, 3:]) + elif state_location == "root": + root_link_pose_w = wp.to_torch(cube_object.data.root_link_pose_w) + root_com_vel_w = wp.to_torch(cube_object.data.root_com_vel_w) + body_com_pose_b = wp.to_torch(cube_object.data.body_com_pose_b) + expected_com_pos, expected_com_quat = combine_frame_transforms( + root_link_pose_w[:, :3], + root_link_pose_w[:, 3:], + body_com_pose_b[:, 0, :3], + body_com_pose_b[:, 0, 3:7], + ) + expected_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1) + root_com_pose_w = wp.to_torch(cube_object.data.root_com_pose_w) + root_link_vel_w = wp.to_torch(cube_object.data.root_link_vel_w) + # test both root_com and root_link successfully updated when root_pose updates + torch.testing.assert_close(expected_com_pose, root_com_pose_w) + torch.testing.assert_close(root_com_vel_w, wp.to_torch(cube_object.data.root_com_vel_w)) + torch.testing.assert_close(root_link_pose_w, wp.to_torch(cube_object.data.root_link_pose_w)) + torch.testing.assert_close(root_com_vel_w[:, 3:], root_link_vel_w[:, 3:]) + + +@pytest.mark.skip(reason="PhysX-specific warmup test") +@pytest.mark.isaacsim_ci +def test_warmup_attach_stage_not_called_for_cpu(): + """Regression test: attach_stage() must not be called for CPU in _warmup_and_create_views(). + + Bug (commit 0ba9c5cb3b): ``PhysxManager._warmup_and_create_views()`` called + ``_physx_sim.attach_stage()`` unconditionally before ``force_load_physics_from_usd()``. + These are two alternative initialization patterns; combining them causes + double-initialization that corrupts the CPU MBP broadphase, producing + non-deterministic collision failures (objects passing through surfaces). + + Fix: guard ``attach_stage()`` with ``if is_gpu:`` — it is only required by the + GPU pipeline, which needs explicit stage attachment before the physics load step. + The CPU pipeline attaches implicitly via ``force_load_physics_from_usd()``. + + This test verifies the guard is in place by monkeypatching ``attach_stage`` on + the PhysX simulation interface and asserting it is *not* called during CPU warmup. + The simulation test itself (1 cube falling onto a ground plane) is intentionally + omitted here because the MBP corruption is non-deterministic and depends on scene + complexity (multiple dynamic actors on a mesh collider), making it unreliable as a + unit test assertion. + """ + from unittest.mock import MagicMock + + from isaaclab_physx.physics import PhysxManager + + with _newton_sim_context("cpu", add_ground_plane=True, dt=0.01, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + generate_cubes_scene(num_cubes=1, height=1.0, device="cpu") + + # IPhysxSimulation is a C++ binding whose attributes are read-only, so we cannot + # assign to _physx_sim.attach_stage directly. Instead, replace the class-level + # reference with a MagicMock that wraps the real object so all other calls still + # work, then restore it in the finally block. + original_physx_sim = PhysxManager._physx_sim + spy = MagicMock(wraps=original_physx_sim) + PhysxManager._physx_sim = spy + try: + sim.reset() + finally: + PhysxManager._physx_sim = original_physx_sim + + assert spy.attach_stage.call_count == 0, ( + f"attach_stage() was called {spy.attach_stage.call_count} time(s) during CPU warmup. " + f"This indicates the CPU MBP broadphase double-initialization regression is present: " + f"attach_stage() + force_load_physics_from_usd() must not be combined for CPU." + ) diff --git a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py new file mode 100644 index 000000000000..f5b23db07178 --- /dev/null +++ b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py @@ -0,0 +1,907 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# ignore private usage of variables warning +# pyright: reportPrivateUsage=none + + +"""Launch Isaac Sim Simulator first.""" + +from isaaclab.app import AppLauncher + +# launch omniverse app +simulation_app = AppLauncher(headless=True).app + +"""Rest everything follows.""" + +import sys + +import pytest +import torch +import warp as wp +from isaaclab_newton.assets import RigidObjectCollection +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +from isaaclab_newton.physics import NewtonManager as SimulationManager +from newton.solvers import SolverNotifyFlags + +import isaaclab.sim as sim_utils +from isaaclab.assets import RigidObjectCfg, RigidObjectCollectionCfg +from isaaclab.sim import SimulationCfg, build_simulation_context +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.math import ( + combine_frame_transforms, + default_orientation, + quat_apply_inverse, + quat_inv, + quat_mul, + quat_rotate, + random_orientation, + subtract_frame_transforms, +) + +NEWTON_SIM_CFG = SimulationCfg( + physics=NewtonCfg( + solver_cfg=MJWarpSolverCfg(), + ), +) + + +def _newton_sim_context(device, gravity_enabled=True, dt=None, **kwargs): + """Helper to create a Newton simulation context with the correct device. + + When sim_cfg is provided to build_simulation_context, the device, gravity_enabled, and dt + kwargs are ignored. This helper applies them to the shared NEWTON_SIM_CFG before calling. + """ + NEWTON_SIM_CFG.device = device + NEWTON_SIM_CFG.gravity = (0.0, 0.0, -9.81) if gravity_enabled else (0.0, 0.0, 0.0) + if dt is not None: + NEWTON_SIM_CFG.dt = dt + return build_simulation_context(device=device, sim_cfg=NEWTON_SIM_CFG, **kwargs) + + +def generate_cubes_scene( + num_envs: int = 1, + num_cubes: int = 1, + height=1.0, + has_api: bool = True, + kinematic_enabled: bool = False, + device: str = "cuda:0", +) -> tuple[RigidObjectCollection, torch.Tensor]: + """Generate a scene with the provided number of cubes. + + Args: + num_envs: Number of envs to generate. + num_cubes: Number of cubes to generate. + height: Height of the cubes. + has_api: Whether the cubes have a rigid body API on them. + kinematic_enabled: Whether the cubes are kinematic. + device: Device to use for the simulation. + + Returns: + A tuple containing the rigid object collection representing the cubes and the origins of the cubes. + + """ + origins = torch.tensor([(i * 3.0, 0, height) for i in range(num_envs)]).to(device) + # Create Top-level Xforms, one for each cube + for i, origin in enumerate(origins): + sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=origin) + + # Resolve spawn configuration + if has_api: + spawn_cfg = sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", + rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=kinematic_enabled), + ) + else: + # since no rigid body properties defined, this is just a static collider + spawn_cfg = sim_utils.CuboidCfg( + size=(0.1, 0.1, 0.1), + collision_props=sim_utils.CollisionPropertiesCfg(), + ) + + # create the rigid object configs + cube_config_dict = {} + for i in range(num_cubes): + cube_object_cfg = RigidObjectCfg( + prim_path=f"/World/Env_.*/Object_{i}", + spawn=spawn_cfg, + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 3 * i, height)), + ) + cube_config_dict[f"cube_{i}"] = cube_object_cfg + # create the rigid object collection + cube_object_collection_cfg = RigidObjectCollectionCfg(rigid_objects=cube_config_dict) + cube_object_collection = RigidObjectCollection(cfg=cube_object_collection_cfg) + + return cube_object_collection, origins + + +@pytest.mark.parametrize("num_envs", [1, 2]) +@pytest.mark.parametrize("num_cubes", [1, 3]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_initialization(num_envs, num_cubes, device): + """Test initialization for prim with rigid body API at the provided prim path.""" + with _newton_sim_context(device, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(object_collection) < 10 + + # Play sim + sim.reset() + + # Check if object is initialized + assert object_collection.is_initialized + assert len(object_collection.body_names) == num_cubes + + # Check buffers that exist and have correct shapes + assert wp.to_torch(object_collection.data.body_link_pos_w).shape == (num_envs, num_cubes, 3) + assert wp.to_torch(object_collection.data.body_link_quat_w).shape == (num_envs, num_cubes, 4) + assert wp.to_torch(object_collection.data.body_mass).shape == (num_envs, num_cubes) + assert wp.to_torch(object_collection.data.body_inertia).shape == (num_envs, num_cubes, 9) + + # Simulate physics + for _ in range(2): + sim.step() + object_collection.update(sim.cfg.dt) + + +@pytest.mark.skip(reason="Newton doesn't support kinematic rigid bodies yet") +@pytest.mark.parametrize("num_envs", [1, 2]) +@pytest.mark.parametrize("num_cubes", [1, 3]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_initialization_with_kinematic_enabled(num_envs, num_cubes, device): + """Test that initialization for prim with kinematic flag enabled.""" + with _newton_sim_context(device, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + object_collection, origins = generate_cubes_scene( + num_envs=num_envs, num_cubes=num_cubes, kinematic_enabled=True, device=device + ) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(object_collection) < 10 + + # Play sim + sim.reset() + + # Check if object is initialized + assert object_collection.is_initialized + assert len(object_collection.body_names) == num_cubes + + # Check buffers that exist and have correct shapes + assert wp.to_torch(object_collection.data.body_link_pos_w).shape == (num_envs, num_cubes, 3) + assert wp.to_torch(object_collection.data.body_link_quat_w).shape == (num_envs, num_cubes, 4) + + # Simulate physics + for _ in range(2): + sim.step() + object_collection.update(sim.cfg.dt) + # check that the object is kinematic + default_body_pose = wp.to_torch(object_collection.data.default_body_pose).clone() + default_body_vel = wp.to_torch(object_collection.data.default_body_vel).clone() + default_body_pose[..., :3] += origins.unsqueeze(1) + torch.testing.assert_close(wp.to_torch(object_collection.data.body_link_pose_w), default_body_pose) + torch.testing.assert_close(wp.to_torch(object_collection.data.body_link_vel_w), default_body_vel) + + +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_initialization_with_no_rigid_body(num_cubes, device): + """Test that initialization fails when no rigid body is found at the provided prim path.""" + with _newton_sim_context(device, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + object_collection, _ = generate_cubes_scene(num_cubes=num_cubes, has_api=False, device=device) + + # Check that the framework doesn't hold excessive strong references. + assert sys.getrefcount(object_collection) < 10 + + # Play sim + with pytest.raises(RuntimeError): + sim.reset() + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_external_force_buffer(device): + """Test if external force buffer correctly updates in the force value is zero case.""" + with _newton_sim_context(device, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + num_envs = 2 + num_cubes = 1 + object_collection, origins = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) + sim.reset() + + # find objects to apply the force + object_ids, object_names = object_collection.find_bodies(".*") + # reset object + object_collection.reset() + + # perform simulation + for step in range(5): + # initiate force tensor + external_wrench_b = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=sim.device) + + # decide if zero or non-zero force + if step == 0 or step == 3: + force = 1.0 + else: + force = 0.0 + + # apply force to the object + external_wrench_b[:, :, 0] = force + external_wrench_b[:, :, 3] = force + + object_collection.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + body_ids=object_ids, + env_ids=None, + ) + + # check if the object collection's force and torque buffers are correctly updated + for i in range(num_envs): + assert wp.to_torch(object_collection._permanent_wrench_composer.composed_force)[i, 0, 0].item() == force + assert ( + wp.to_torch(object_collection._permanent_wrench_composer.composed_torque)[i, 0, 0].item() == force + ) + + object_collection.instantaneous_wrench_composer.add_forces_and_torques_index( + body_ids=object_ids, + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + ) + + # apply action to the object collection + object_collection.write_data_to_sim() + sim.step() + object_collection.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_envs", [1, 2]) +@pytest.mark.parametrize("num_cubes", [1, 4]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_external_force_on_single_body(num_envs, num_cubes, device): + """Test application of external force on the base of the object.""" + with _newton_sim_context(device, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + object_collection, origins = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) + sim.reset() + + # find objects to apply the force + object_ids, object_names = object_collection.find_bodies(".*") + + # Sample a force equal to the weight of the object + external_wrench_b = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=sim.device) + # Every 2nd cube should have a force applied to it + external_wrench_b[:, 0::2, 2] = 9.81 * wp.to_torch(object_collection.data.body_mass)[:, 0::2] + + for i in range(5): + # reset object state + body_pose = wp.to_torch(object_collection.data.default_body_pose).clone() + body_vel = wp.to_torch(object_collection.data.default_body_vel).clone() + # need to shift the position of the cubes otherwise they will be on top of each other + body_pose[..., :2] += origins.unsqueeze(1)[..., :2] + object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) + object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) + # reset object + object_collection.reset() + + is_global = False + if i % 2 == 0: + positions = wp.to_torch(object_collection.data.body_link_pos_w)[:, object_ids, :3] + is_global = True + else: + positions = None + + # apply force + object_collection.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=positions, + body_ids=object_ids, + env_ids=None, + is_global=is_global, + ) + for _ in range(10): + # write data to sim + object_collection.write_data_to_sim() + # step sim + sim.step() + # update object collection + object_collection.update(sim.cfg.dt) + + # First object should still be at the same Z position (1.0) + torch.testing.assert_close( + wp.to_torch(object_collection.data.body_link_pos_w)[:, 0::2, 2], + torch.ones_like(wp.to_torch(object_collection.data.body_link_pos_w)[:, 0::2, 2]), + ) + # Second object should have fallen, so it's Z height should be less than initial height of 1.0 + assert torch.all(wp.to_torch(object_collection.data.body_link_pos_w)[:, 1::2, 2] < 1.0) + + +@pytest.mark.parametrize("num_envs", [1, 2]) +@pytest.mark.parametrize("num_cubes", [1, 4]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_external_force_on_single_body_at_position(num_envs, num_cubes, device): + """Test application of external force on the base of the object at a specific position. + + In this test, we apply a force equal to the weight of an object on the base of + one of the objects at 1m in the Y direction, we check that the object rotates around it's X axis. + For the other object, we do not apply any force and check that it falls down. + """ + with _newton_sim_context(device, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + object_collection, origins = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) + sim.reset() + + # find objects to apply the force + object_ids, object_names = object_collection.find_bodies(".*") + + # Sample a force equal to the weight of the object + external_wrench_b = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=sim.device) + external_wrench_positions_b = torch.zeros( + object_collection.num_instances, len(object_ids), 3, device=sim.device + ) + # Every 2nd cube should have a force applied to it + external_wrench_b[:, 0::2, 2] = 500.0 + external_wrench_positions_b[:, 0::2, 1] = 1.0 + + # Desired force and torque + for i in range(5): + # reset object state + body_pose = wp.to_torch(object_collection.data.default_body_pose).clone() + body_vel = wp.to_torch(object_collection.data.default_body_vel).clone() + # need to shift the position of the cubes otherwise they will be on top of each other + body_pose[..., :2] += origins.unsqueeze(1)[..., :2] + object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) + object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) + # reset object + object_collection.reset() + + is_global = False + if i % 2 == 0: + body_com_pos_w = wp.to_torch(object_collection.data.body_link_pos_w)[:, object_ids, :3] + external_wrench_positions_b[..., 0] = 0.0 + external_wrench_positions_b[..., 1] = 1.0 + external_wrench_positions_b[..., 2] = 0.0 + external_wrench_positions_b += body_com_pos_w + is_global = True + else: + external_wrench_positions_b[..., 0] = 0.0 + external_wrench_positions_b[..., 1] = 1.0 + external_wrench_positions_b[..., 2] = 0.0 + + # apply force + object_collection.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=external_wrench_positions_b, + body_ids=object_ids, + env_ids=None, + is_global=is_global, + ) + object_collection.permanent_wrench_composer.add_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], + positions=external_wrench_positions_b, + body_ids=object_ids, + is_global=is_global, + ) + + for _ in range(10): + # write data to sim + object_collection.write_data_to_sim() + # step sim + sim.step() + # update object collection + object_collection.update(sim.cfg.dt) + + # First object should be rotating around it's X axis + assert torch.all(wp.to_torch(object_collection.data.body_com_ang_vel_b)[:, 0::2, 0] > 0.1) + # Second object should have fallen, so it's Z height should be less than initial height of 1.0 + assert torch.all(wp.to_torch(object_collection.data.body_link_pos_w)[:, 1::2, 2] < 1.0) + + +@pytest.mark.parametrize("num_envs", [1, 3]) +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_set_object_state(num_envs, num_cubes, device): + """Test setting the state of the object. + + .. note:: + Turn off gravity for this test as we don't want any external forces acting on the object + to ensure state remains static + """ + with _newton_sim_context(device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + object_collection, origins = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) + sim.reset() + + state_types = ["body_link_pos_w", "body_link_quat_w", "body_com_lin_vel_w", "body_com_ang_vel_w"] + + # Set each state type individually as they are dependent on each other + for state_type_to_randomize in state_types: + state_dict = { + "body_link_pos_w": torch.zeros_like( + wp.to_torch(object_collection.data.body_link_pos_w), device=sim.device + ), + "body_link_quat_w": default_orientation(num=num_cubes * num_envs, device=sim.device).view( + num_envs, num_cubes, 4 + ), + "body_com_lin_vel_w": torch.zeros_like( + wp.to_torch(object_collection.data.body_com_lin_vel_w), device=sim.device + ), + "body_com_ang_vel_w": torch.zeros_like( + wp.to_torch(object_collection.data.body_com_ang_vel_w), device=sim.device + ), + } + + for _ in range(5): + # reset object + object_collection.reset() + + # Set random state + if state_type_to_randomize == "body_link_quat_w": + state_dict[state_type_to_randomize] = random_orientation( + num=num_cubes * num_envs, device=sim.device + ).view(num_envs, num_cubes, 4) + else: + state_dict[state_type_to_randomize] = torch.randn(num_envs, num_cubes, 3, device=sim.device) + # make sure objects do not overlap + if state_type_to_randomize == "body_link_pos_w": + state_dict[state_type_to_randomize][..., :2] += origins.unsqueeze(1)[..., :2] + + # perform simulation + for _ in range(5): + body_pose = torch.cat( + [state_dict["body_link_pos_w"], state_dict["body_link_quat_w"]], + dim=-1, + ) + body_vel = torch.cat( + [state_dict["body_com_lin_vel_w"], state_dict["body_com_ang_vel_w"]], + dim=-1, + ) + # reset object state + object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) + object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) + sim.step() + + # assert that set object quantities are equal to the ones set in the state_dict + for key, expected_value in state_dict.items(): + value = wp.to_torch(getattr(object_collection.data, key)) + # Newton reads state directly from sim (not cached), so post-step drift + # from velocity integration causes larger differences than PhysX + torch.testing.assert_close(value, expected_value, rtol=1e-1, atol=1e-1) + + object_collection.update(sim.cfg.dt) + + +@pytest.mark.parametrize("num_envs", [1, 3]) +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_reset_object_collection(num_envs, num_cubes, device): + """Test resetting the state of the rigid object.""" + with _newton_sim_context(device, gravity_enabled=True, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) + sim.reset() + + for i in range(5): + sim.step() + object_collection.update(sim.cfg.dt) + + # Move the object to a random position + body_pose = wp.to_torch(object_collection.data.default_body_pose).clone() + body_pose[..., :3] = torch.randn(num_envs, num_cubes, 3, device=sim.device) + # Random orientation + body_pose[..., 3:7] = random_orientation(num=num_cubes, device=sim.device) + object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) + body_vel = wp.to_torch(object_collection.data.default_body_vel).clone() + object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) + + if i % 2 == 0: + object_collection.reset() + + # Reset should zero external forces and torques + assert not object_collection._instantaneous_wrench_composer.active + assert not object_collection._permanent_wrench_composer.active + assert ( + torch.count_nonzero(wp.to_torch(object_collection._instantaneous_wrench_composer.composed_force)) + == 0 + ) + assert ( + torch.count_nonzero(wp.to_torch(object_collection._instantaneous_wrench_composer.composed_torque)) + == 0 + ) + assert ( + torch.count_nonzero(wp.to_torch(object_collection._permanent_wrench_composer.composed_force)) == 0 + ) + assert ( + torch.count_nonzero(wp.to_torch(object_collection._permanent_wrench_composer.composed_torque)) == 0 + ) + + +@pytest.mark.skip(reason="Newton doesn't support friction/restitution/material properties yet") +@pytest.mark.parametrize("num_envs", [1, 3]) +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_set_material_properties(num_envs, num_cubes, device): + """Test getting and setting material properties of rigid object.""" + with _newton_sim_context(device, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) + sim.reset() + + # Set material properties + static_friction = torch.FloatTensor(num_envs, num_cubes, 1).uniform_(0.4, 0.8) + dynamic_friction = torch.FloatTensor(num_envs, num_cubes, 1).uniform_(0.4, 0.8) + restitution = torch.FloatTensor(num_envs, num_cubes, 1).uniform_(0.0, 0.2) + + materials = torch.cat([static_friction, dynamic_friction, restitution], dim=-1) + + # Add friction to cube + indices = torch.tensor(range(num_cubes * num_envs), dtype=torch.int) + object_collection.root_view.set_material_properties( + object_collection.reshape_data_to_view_3d(wp.from_torch(materials, dtype=wp.float32), 3, device="cpu"), + wp.from_torch(indices, dtype=wp.int32), + ) + + # Perform simulation + sim.step() + object_collection.update(sim.cfg.dt) + + # Get material properties + materials_to_check = object_collection.root_view.get_material_properties() + + # Check if material properties are set correctly + torch.testing.assert_close( + wp.to_torch(object_collection.reshape_view_to_data_3d(materials_to_check, 3, device="cpu")), materials + ) + + +@pytest.mark.parametrize("num_envs", [1, 3]) +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("gravity_enabled", [True, False]) +def test_gravity_vec_w(num_envs, num_cubes, device, gravity_enabled): + """Test that gravity vector direction is set correctly for the rigid object.""" + with _newton_sim_context(device, gravity_enabled=gravity_enabled, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + object_collection, _ = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) + + # Obtain gravity direction + gravity_dir = (0.0, 0.0, -1.0) if gravity_enabled else (0.0, 0.0, 0.0) + + sim.reset() + + # Check if gravity vector is set correctly + gravity_vec = wp.to_torch(object_collection.data.GRAVITY_VEC_W) + assert gravity_vec[0, 0, 0] == gravity_dir[0] + assert gravity_vec[0, 0, 1] == gravity_dir[1] + assert gravity_vec[0, 0, 2] == gravity_dir[2] + + # Perform simulation + for _ in range(2): + sim.step() + object_collection.update(sim.cfg.dt) + + # Expected gravity value is the acceleration of the body + gravity = torch.zeros(num_envs, num_cubes, 6, device=device) + if gravity_enabled: + gravity[..., 2] = -9.81 + + # Check the body accelerations are correct + torch.testing.assert_close(wp.to_torch(object_collection.data.body_com_acc_w), gravity) + + +@pytest.mark.parametrize("num_envs", [1, 4]) +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("with_offset", [True, False]) +def test_object_state_properties(num_envs, num_cubes, device, with_offset): + """Test the object_com_state_w and object_link_state_w properties.""" + with _newton_sim_context(device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, env_pos = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, height=0.0, device=device) + + sim.reset() + + # check if cube_object is initialized + assert cube_object.is_initialized + + # change center of mass offset from link frame + offset = ( + torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) + if with_offset + else torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) + ) + + # Set center of mass offset via Newton API (position only, shape (E, B, 3)) + cube_object.set_coms_index(coms=wp.from_torch(offset, dtype=wp.vec3f)) + # Flush the model change immediately so it takes effect before the next step + with wp.ScopedDevice(device): + SimulationManager._solver.notify_model_changed(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + # check center of mass has been set + torch.testing.assert_close(wp.to_torch(cube_object.data.body_com_pos_b), offset) + + # random z spin velocity + spin_twist = torch.zeros(6, device=device) + spin_twist[5] = torch.randn(1, device=device) + + # initial spawn point + init_com = wp.to_torch(cube_object.data.body_com_pose_w)[..., :3] + + for i in range(10): + # spin the object around Z axis (com) + cube_object.write_body_com_velocity_to_sim_index(body_velocities=spin_twist.repeat(num_envs, num_cubes, 1)) + sim.step() + cube_object.update(sim.cfg.dt) + + # get state properties + object_link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) + object_link_vel_w = wp.to_torch(cube_object.data.body_link_vel_w) + object_com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) + object_com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) + + # if offset is [0,0,0] all object_state_%_w will match and all body_%_w will match + if not with_offset: + torch.testing.assert_close(object_link_pose_w, object_com_pose_w) + torch.testing.assert_close(object_com_vel_w, object_link_vel_w) + else: + # cubes are spinning around center of mass + # position will not match + # center of mass position will be constant (i.e. spinning around com) + torch.testing.assert_close(init_com, object_com_pose_w[..., :3]) + + # link position will be moving but should stay constant away from center of mass + object_link_state_pos_rel_com = quat_apply_inverse( + object_link_pose_w[..., 3:], + object_link_pose_w[..., :3] - object_com_pose_w[..., :3], + ) + + torch.testing.assert_close(-offset, object_link_state_pos_rel_com) + + # orientation of com will be a constant rotation from link orientation + com_quat_b = wp.to_torch(cube_object.data.body_com_quat_b) + com_quat_w = quat_mul(object_link_pose_w[..., 3:], com_quat_b) + torch.testing.assert_close(com_quat_w, object_com_pose_w[..., 3:]) + + # orientation of link will match object state will always match + torch.testing.assert_close(object_link_pose_w[..., 3:], object_link_pose_w[..., 3:]) + + # lin_vel will not match + # center of mass vel will be constant (i.e. spinning around com) + torch.testing.assert_close( + torch.zeros_like(object_com_vel_w[..., :3]), + object_com_vel_w[..., :3], + atol=1e-4, + rtol=1e-3, + ) + + # link frame will be moving, and should be equal to input angular velocity cross offset + lin_vel_rel_object_gt = quat_apply_inverse(object_link_pose_w[..., 3:], object_link_vel_w[..., :3]) + lin_vel_rel_gt = torch.linalg.cross(spin_twist.repeat(num_envs, num_cubes, 1)[..., 3:], -offset) + torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_object_gt, atol=1e-4, rtol=1e-3) + + # ang_vel will always match + torch.testing.assert_close(object_com_vel_w[..., 3:], object_com_vel_w[..., 3:]) + torch.testing.assert_close(object_com_vel_w[..., 3:], object_link_vel_w[..., 3:]) + + +@pytest.mark.parametrize("num_envs", [1, 3]) +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("with_offset", [True, False]) +@pytest.mark.parametrize("state_location", ["com", "link"]) +def test_write_object_state(num_envs, num_cubes, device, with_offset, state_location): + """Test the setters for object_state using both the link frame and center of mass as reference frame.""" + with _newton_sim_context(device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + # Create a scene with random cubes + cube_object, env_pos = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, height=0.0, device=device) + env_ids = torch.tensor([x for x in range(num_envs)], dtype=torch.int32) + object_ids = torch.tensor([x for x in range(num_cubes)], dtype=torch.int32) + + sim.reset() + + # Check if cube_object is initialized + assert cube_object.is_initialized + + # change center of mass offset from link frame + offset = ( + torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) + if with_offset + else torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) + ) + + # Set center of mass offset via Newton API (position only, shape (E, B, 3)) + cube_object.set_coms_index(coms=wp.from_torch(offset, dtype=wp.vec3f)) + # Flush the model change immediately so it takes effect before the next step + with wp.ScopedDevice(device): + SimulationManager._solver.notify_model_changed(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + # check center of mass has been set + torch.testing.assert_close(wp.to_torch(cube_object.data.body_com_pos_b), offset) + + rand_state = torch.zeros(num_envs, num_cubes, 13, device=device) + rand_state[..., :7] = wp.to_torch(cube_object.data.default_body_pose) + rand_state[..., :3] += wp.to_torch(cube_object.data.body_link_pos_w) + # make quaternion a unit vector + rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) + + env_ids = env_ids.to(device) + object_ids = object_ids.to(device) + for i in range(10): + sim.step() + cube_object.update(sim.cfg.dt) + + if state_location == "com": + if i % 2 == 0: + cube_object.write_body_com_pose_to_sim_index(body_poses=rand_state[..., :7]) + cube_object.write_body_com_velocity_to_sim_index(body_velocities=rand_state[..., 7:]) + else: + cube_object.write_body_com_pose_to_sim_index( + body_poses=rand_state[..., :7], env_ids=env_ids, body_ids=object_ids + ) + cube_object.write_body_com_velocity_to_sim_index( + body_velocities=rand_state[..., 7:], env_ids=env_ids, body_ids=object_ids + ) + elif state_location == "link": + if i % 2 == 0: + cube_object.write_body_link_pose_to_sim_index(body_poses=rand_state[..., :7]) + cube_object.write_body_link_velocity_to_sim_index(body_velocities=rand_state[..., 7:]) + else: + cube_object.write_body_link_pose_to_sim_index( + body_poses=rand_state[..., :7], env_ids=env_ids, body_ids=object_ids + ) + cube_object.write_body_link_velocity_to_sim_index( + body_velocities=rand_state[..., 7:], env_ids=env_ids, body_ids=object_ids + ) + + if state_location == "com": + torch.testing.assert_close(rand_state[..., :7], wp.to_torch(cube_object.data.body_com_pose_w)) + torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(cube_object.data.body_com_vel_w)) + elif state_location == "link": + torch.testing.assert_close(rand_state[..., :7], wp.to_torch(cube_object.data.body_link_pose_w)) + torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(cube_object.data.body_link_vel_w)) + + +@pytest.mark.parametrize("num_envs", [1, 3]) +@pytest.mark.parametrize("num_cubes", [1, 2]) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("with_offset", [True]) +@pytest.mark.parametrize("state_location", ["com", "link", "root"]) +def test_write_object_state_functions_data_consistency(num_envs, num_cubes, device, with_offset, state_location): + """Test the setters for object_state using both the link frame and center of mass as reference frame.""" + with _newton_sim_context(device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + # Create a scene with random cubes + cube_object, env_pos = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, height=0.0, device=device) + env_ids = torch.tensor([x for x in range(num_envs)], dtype=torch.int32) + object_ids = torch.tensor([x for x in range(num_cubes)], dtype=torch.int32) + + sim.reset() + + # Check if cube_object is initialized + assert cube_object.is_initialized + + # change center of mass offset from link frame + offset = ( + torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) + if with_offset + else torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) + ) + + # Set center of mass offset via Newton API (position only, shape (E, B, 3)) + cube_object.set_coms_index(coms=wp.from_torch(offset, dtype=wp.vec3f)) + # Flush the model change immediately so it takes effect before the next step + with wp.ScopedDevice(device): + SimulationManager._solver.notify_model_changed(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + # check center of mass has been set + torch.testing.assert_close(wp.to_torch(cube_object.data.body_com_pos_b), offset) + + rand_state = torch.rand(num_envs, num_cubes, 13, device=device) + rand_state[..., :3] += wp.to_torch(cube_object.data.body_link_pos_w) + # make quaternion a unit vector + rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) + + env_ids = env_ids.to(device) + object_ids = object_ids.to(device) + sim.step() + cube_object.update(sim.cfg.dt) + + body_link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) + body_com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) + object_link_to_com_pos, object_link_to_com_quat = subtract_frame_transforms( + body_link_pose_w[..., :3].view(-1, 3), + body_link_pose_w[..., 3:7].view(-1, 4), + body_com_pose_w[..., :3].view(-1, 3), + body_com_pose_w[..., 3:7].view(-1, 4), + ) + + if state_location == "com": + cube_object.write_body_com_pose_to_sim_index( + body_poses=rand_state[..., :7], env_ids=env_ids, body_ids=object_ids + ) + cube_object.write_body_com_velocity_to_sim_index( + body_velocities=rand_state[..., 7:], env_ids=env_ids, body_ids=object_ids + ) + elif state_location == "link": + cube_object.write_body_link_pose_to_sim_index( + body_poses=rand_state[..., :7], env_ids=env_ids, body_ids=object_ids + ) + cube_object.write_body_link_velocity_to_sim_index( + body_velocities=rand_state[..., 7:], env_ids=env_ids, body_ids=object_ids + ) + elif state_location == "root": + cube_object.write_body_link_pose_to_sim_index( + body_poses=rand_state[..., :7], env_ids=env_ids, body_ids=object_ids + ) + cube_object.write_body_com_velocity_to_sim_index( + body_velocities=rand_state[..., 7:], env_ids=env_ids, body_ids=object_ids + ) + + if state_location == "com": + com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) + com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) + expected_root_link_pos, expected_root_link_quat = combine_frame_transforms( + com_pose_w[..., :3].view(-1, 3), + com_pose_w[..., 3:].view(-1, 4), + quat_rotate(quat_inv(object_link_to_com_quat), -object_link_to_com_pos), + quat_inv(object_link_to_com_quat), + ) + expected_object_link_pose = torch.cat((expected_root_link_pos, expected_root_link_quat), dim=1).view( + num_envs, -1, 7 + ) + link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) + link_vel_w = wp.to_torch(cube_object.data.body_link_vel_w) + # test both root_pose and root_link successfully updated when root_com updates + torch.testing.assert_close(expected_object_link_pose, link_pose_w) + # skip lin_vel because it differs from link frame, this should be fine because we are only checking + # if velocity update is triggered, which can be determined by comparing angular velocity + torch.testing.assert_close(com_vel_w[..., 3:], link_vel_w[..., 3:]) + torch.testing.assert_close(expected_object_link_pose, link_pose_w) + torch.testing.assert_close(com_vel_w[..., 3:], wp.to_torch(cube_object.data.body_com_vel_w)[..., 3:]) + elif state_location == "link": + link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) + link_vel_w = wp.to_torch(cube_object.data.body_link_vel_w) + expected_com_pos, expected_com_quat = combine_frame_transforms( + link_pose_w[..., :3].view(-1, 3), + link_pose_w[..., 3:].view(-1, 4), + object_link_to_com_pos, + object_link_to_com_quat, + ) + expected_object_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1).view(num_envs, -1, 7) + com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) + com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) + # test both root_pose and root_com successfully updated when root_link updates + torch.testing.assert_close(expected_object_com_pose, com_pose_w) + # skip lin_vel because it differs from link frame, this should be fine because we are only checking + # if velocity update is triggered, which can be determined by comparing angular velocity + torch.testing.assert_close(link_vel_w[..., 3:], com_vel_w[..., 3:]) + torch.testing.assert_close(link_pose_w, wp.to_torch(cube_object.data.body_link_pose_w)) + torch.testing.assert_close(link_vel_w[..., 3:], wp.to_torch(cube_object.data.body_com_vel_w)[..., 3:]) + elif state_location == "root": + body_link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) + body_com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) + expected_object_com_pos, expected_object_com_quat = combine_frame_transforms( + body_link_pose_w[..., :3].view(-1, 3), + body_link_pose_w[..., 3:].view(-1, 4), + object_link_to_com_pos, + object_link_to_com_quat, + ) + expected_object_com_pose = torch.cat((expected_object_com_pos, expected_object_com_quat), dim=1).view( + num_envs, -1, 7 + ) + com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) + com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) + link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) + link_vel_w = wp.to_torch(cube_object.data.body_link_vel_w) + # test both root_com and root_link successfully updated when root_pose updates + torch.testing.assert_close(expected_object_com_pose, com_pose_w) + torch.testing.assert_close(body_com_vel_w, com_vel_w) + torch.testing.assert_close(body_link_pose_w, link_pose_w) + torch.testing.assert_close(body_com_vel_w[..., 3:], link_vel_w[..., 3:]) diff --git a/source/isaaclab_newton/test/cloner/__init__.py b/source/isaaclab_newton/test/cloner/__init__.py new file mode 100644 index 000000000000..460a30569089 --- /dev/null +++ b/source/isaaclab_newton/test/cloner/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause diff --git a/source/isaaclab_newton/test/physics/__init__.py b/source/isaaclab_newton/test/physics/__init__.py new file mode 100644 index 000000000000..ad6daf8a5502 --- /dev/null +++ b/source/isaaclab_newton/test/physics/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Shared utilities for physics and sensor tests.""" diff --git a/source/isaaclab_newton/test/physics/physics_test_utils.py b/source/isaaclab_newton/test/physics/physics_test_utils.py new file mode 100644 index 000000000000..65177a9c8a93 --- /dev/null +++ b/source/isaaclab_newton/test/physics/physics_test_utils.py @@ -0,0 +1,295 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Shared utilities for physics and sensor testing. + +This module provides common functionality used across multiple test files: +- Shape type definitions and helpers +- Simulation configuration +- Shape spawning utilities +""" + +import time +from contextlib import contextmanager +from enum import Enum, auto + +import pytest +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg + +import isaaclab.sim as sim_utils +from isaaclab.assets import RigidObjectCfg +from isaaclab.sim import SimulationCfg + +## +# Newton Configuration +## + + +def make_sim_cfg( + use_mujoco_contacts: bool = False, device: str = "cuda:0", gravity: tuple = (0.0, 0.0, -9.81) +) -> SimulationCfg: + """Create simulation configuration with specified collision pipeline. + + Args: + use_mujoco_contacts: If True, use MuJoCo contact pipeline. If False, use Newton contact pipeline. + device: Device to run simulation on ("cuda:0" or "cpu"). + gravity: Gravity vector (x, y, z). + + Returns: + SimulationCfg configured for the specified collision pipeline. + """ + solver_cfg = MJWarpSolverCfg( + njmax=200, + ls_iterations=20, + cone="elliptic", + ls_parallel=False, + integrator="implicitfast", + use_mujoco_contacts=use_mujoco_contacts, + ) + + newton_cfg = NewtonCfg( + solver_cfg=solver_cfg, + num_substeps=1, + debug_mode=False, + use_cuda_graph=False, + ) + + return SimulationCfg( + dt=1.0 / 120.0, + device=device, + gravity=gravity, + create_stage_in_memory=False, + physics=newton_cfg, + ) + + +COLLISION_PIPELINES = [ + pytest.param(False, id="newton_contacts"), + pytest.param(True, id="mujoco_contacts"), +] + + +## +# Shape Types +## + + +class ShapeType(Enum): + SPHERE = auto() + BOX = auto() + CAPSULE = auto() + CYLINDER = auto() + CONE = auto() + MESH_SPHERE = auto() + MESH_BOX = auto() + MESH_CAPSULE = auto() + MESH_CYLINDER = auto() + MESH_CONE = auto() + + +PRIMITIVE_SHAPES = [ShapeType.SPHERE, ShapeType.BOX, ShapeType.CAPSULE, ShapeType.CYLINDER, ShapeType.CONE] +MESH_SHAPES = [ + ShapeType.MESH_SPHERE, + ShapeType.MESH_BOX, + ShapeType.MESH_CAPSULE, + ShapeType.MESH_CYLINDER, + ShapeType.MESH_CONE, +] +ALL_SHAPES = PRIMITIVE_SHAPES + MESH_SHAPES + +STABLE_SHAPES = [ + ShapeType.SPHERE, + ShapeType.BOX, + ShapeType.CAPSULE, + ShapeType.CYLINDER, + ShapeType.MESH_SPHERE, + ShapeType.MESH_BOX, + ShapeType.MESH_CAPSULE, + ShapeType.MESH_CYLINDER, +] + +BOX_SHAPES = [ShapeType.BOX, ShapeType.MESH_BOX] +SPHERE_SHAPES = [ShapeType.SPHERE, ShapeType.MESH_SPHERE] + + +def shape_type_to_str(shape_type: ShapeType) -> str: + return shape_type.name.lower() + + +def is_mesh_shape(shape_type: ShapeType) -> bool: + return shape_type in MESH_SHAPES + + +## +# Shape Configuration Factory +## + + +def create_shape_cfg( + shape_type: ShapeType, + prim_path: str, + pos: tuple = (0, 0, 1.0), + disable_gravity: bool = True, + activate_contact_sensors: bool = True, +) -> RigidObjectCfg: + """Create RigidObjectCfg for a shape type.""" + rigid_props = sim_utils.RigidBodyPropertiesCfg( + disable_gravity=disable_gravity, + linear_damping=0.0, + angular_damping=0.0, + ) + collision_props = sim_utils.CollisionPropertiesCfg(collision_enabled=True) + mass_props = sim_utils.MassPropertiesCfg(mass=1.0) + + spawner_map = { + ShapeType.SPHERE: lambda: sim_utils.SphereCfg( + radius=0.25, + rigid_props=rigid_props, + collision_props=collision_props, + mass_props=mass_props, + activate_contact_sensors=activate_contact_sensors, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.4, 0.4, 0.8)), + ), + ShapeType.BOX: lambda: sim_utils.CuboidCfg( + size=(0.5, 0.5, 0.5), + rigid_props=rigid_props, + collision_props=collision_props, + mass_props=mass_props, + activate_contact_sensors=activate_contact_sensors, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.4, 0.8, 0.4)), + ), + ShapeType.CAPSULE: lambda: sim_utils.CapsuleCfg( + radius=0.15, + height=0.3, + axis="Z", + rigid_props=rigid_props, + collision_props=collision_props, + mass_props=mass_props, + activate_contact_sensors=activate_contact_sensors, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.8, 0.4, 0.4)), + ), + ShapeType.CYLINDER: lambda: sim_utils.CylinderCfg( + radius=0.2, + height=0.4, + axis="Z", + rigid_props=rigid_props, + collision_props=collision_props, + mass_props=mass_props, + activate_contact_sensors=activate_contact_sensors, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.8, 0.8, 0.4)), + ), + ShapeType.CONE: lambda: sim_utils.ConeCfg( + radius=0.25, + height=0.5, + axis="Z", + rigid_props=rigid_props, + collision_props=collision_props, + mass_props=mass_props, + activate_contact_sensors=activate_contact_sensors, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.8, 0.4, 0.8)), + ), + ShapeType.MESH_SPHERE: lambda: sim_utils.MeshSphereCfg( + radius=0.25, + rigid_props=rigid_props, + collision_props=collision_props, + mass_props=mass_props, + activate_contact_sensors=activate_contact_sensors, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.2, 0.2, 0.6)), + ), + ShapeType.MESH_BOX: lambda: sim_utils.MeshCuboidCfg( + size=(0.5, 0.5, 0.5), + rigid_props=rigid_props, + collision_props=collision_props, + mass_props=mass_props, + activate_contact_sensors=activate_contact_sensors, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.2, 0.6, 0.2)), + ), + ShapeType.MESH_CAPSULE: lambda: sim_utils.MeshCapsuleCfg( + radius=0.15, + height=0.3, + axis="Z", + rigid_props=rigid_props, + collision_props=collision_props, + mass_props=mass_props, + activate_contact_sensors=activate_contact_sensors, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.6, 0.2, 0.2)), + ), + ShapeType.MESH_CYLINDER: lambda: sim_utils.MeshCylinderCfg( + radius=0.2, + height=0.4, + axis="Z", + rigid_props=rigid_props, + collision_props=collision_props, + mass_props=mass_props, + activate_contact_sensors=activate_contact_sensors, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.6, 0.6, 0.2)), + ), + ShapeType.MESH_CONE: lambda: sim_utils.MeshConeCfg( + radius=0.25, + height=0.5, + axis="Z", + rigid_props=rigid_props, + collision_props=collision_props, + mass_props=mass_props, + activate_contact_sensors=activate_contact_sensors, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.6, 0.2, 0.6)), + ), + } + + return RigidObjectCfg( + prim_path=prim_path, + spawn=spawner_map[shape_type](), + init_state=RigidObjectCfg.InitialStateCfg(pos=pos), + ) + + +def get_shape_extent(shape_type: ShapeType) -> float: + """Get the XY extent (radius/half-size) of a shape for positioning.""" + extents = { + ShapeType.SPHERE: 0.25, + ShapeType.BOX: 0.25, + ShapeType.CAPSULE: 0.15, + ShapeType.CYLINDER: 0.20, + ShapeType.CONE: 0.25, + ShapeType.MESH_SPHERE: 0.25, + ShapeType.MESH_BOX: 0.25, + ShapeType.MESH_CAPSULE: 0.15, + ShapeType.MESH_CYLINDER: 0.20, + ShapeType.MESH_CONE: 0.25, + } + return extents[shape_type] + + +def get_shape_height(shape_type: ShapeType) -> float: + """Get the height of a shape for stacking calculations.""" + heights = { + ShapeType.SPHERE: 0.5, + ShapeType.BOX: 0.5, + ShapeType.CAPSULE: 0.6, + ShapeType.CYLINDER: 0.4, + ShapeType.CONE: 0.5, + ShapeType.MESH_SPHERE: 0.5, + ShapeType.MESH_BOX: 0.5, + ShapeType.MESH_CAPSULE: 0.6, + ShapeType.MESH_CYLINDER: 0.4, + ShapeType.MESH_CONE: 0.5, + } + return heights[shape_type] + + +@contextmanager +def phase_timer(label: str): + """Context manager that prints elapsed wall-clock time for a test phase.""" + t0 = time.perf_counter() + yield + elapsed = time.perf_counter() - t0 + print(f" [{label}] {elapsed:.2f}s") + + +def perform_sim_step(sim, scene, dt: float): + """Perform a single simulation step and update scene.""" + scene.write_data_to_sim() + sim.step(render=True) + scene.update(dt=dt) diff --git a/source/isaaclab_newton/test/sensors/__init__.py b/source/isaaclab_newton/test/sensors/__init__.py new file mode 100644 index 000000000000..460a30569089 --- /dev/null +++ b/source/isaaclab_newton/test/sensors/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause diff --git a/source/isaaclab_newton/test/sensors/test_contact_sensor.py b/source/isaaclab_newton/test/sensors/test_contact_sensor.py new file mode 100644 index 000000000000..78fde123c12e --- /dev/null +++ b/source/isaaclab_newton/test/sensors/test_contact_sensor.py @@ -0,0 +1,783 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Tests to verify contact sensor functionality using Newton physics. + +This test suite verifies that: +1. Contact detection is accurate (no false positives or true negatives) +2. Contact forces are reported correctly +3. Contact filtering works properly +4. Contact time tracking is accurate + +Uses proper collision scenarios (falling, stacking, horizontal collision) instead of +teleporting objects into interpenetrating states. +""" + +# pyright: reportPrivateUsage=none + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import math + +import pytest +import torch +import warp as wp +from physics.physics_test_utils import ( + COLLISION_PIPELINES, + STABLE_SHAPES, + ShapeType, + create_shape_cfg, + get_shape_extent, + get_shape_height, + make_sim_cfg, + perform_sim_step, + shape_type_to_str, +) + +import isaaclab.sim as sim_utils +from isaaclab.assets import Articulation, RigidObject, RigidObjectCfg +from isaaclab.scene import InteractiveScene, InteractiveSceneCfg +from isaaclab.sensors import ContactSensor, ContactSensorCfg +from isaaclab.sim import build_simulation_context +from isaaclab.terrains import TerrainImporterCfg +from isaaclab.utils import configclass + +from isaaclab_assets.robots.allegro import ALLEGRO_HAND_CFG + +## +# Scene Configuration +## + + +@configclass +class ContactSensorTestSceneCfg(InteractiveSceneCfg): + """Configuration for contact sensor test scenes.""" + + terrain: TerrainImporterCfg | None = None + object_a: RigidObjectCfg | None = None + object_b: RigidObjectCfg | None = None + object_c: RigidObjectCfg | None = None + contact_sensor_a: ContactSensorCfg | None = None + contact_sensor_b: ContactSensorCfg | None = None + + +SIM_DT = 1.0 / 120.0 + + +# =================================================================== +# Priority 1: Contact Detection Accuracy +# =================================================================== + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("use_mujoco_contacts", COLLISION_PIPELINES) +@pytest.mark.parametrize("shape_type", STABLE_SHAPES, ids=[shape_type_to_str(s) for s in STABLE_SHAPES]) +def test_contact_lifecycle(device: str, use_mujoco_contacts: bool, shape_type: ShapeType): + """Test full contact detection lifecycle with varied heights across environments. + + 8 environments (2 groups x 4 envs) with objects at different heights. + + Verifies: + - No contact initially while objects are falling + - Contact detected after landing (timing validated against physics) + - Lower drops land before higher drops + - Contact stops when objects are lifted + """ + num_envs = 8 + num_groups = 2 + envs_per_group = num_envs // num_groups + + base_heights = [0.5, 1.5] + object_offset = get_shape_height(shape_type) / 2 + + gravity_mag = 9.81 + total_fall_steps = 100 + lift_steps = 30 + + sim_cfg = make_sim_cfg(use_mujoco_contacts=use_mujoco_contacts, device=device, gravity=(0.0, 0.0, -gravity_mag)) + + with build_simulation_context(sim_cfg=sim_cfg, auto_add_lighting=True, add_ground_plane=True) as sim: + scene_cfg = ContactSensorTestSceneCfg(num_envs=num_envs, env_spacing=5.0, lazy_sensor_update=False) + scene_cfg.object_a = create_shape_cfg( + shape_type, + "{ENV_REGEX_NS}/Object", + pos=(0.0, 0.0, 3.0), + disable_gravity=False, + activate_contact_sensors=True, + ) + scene_cfg.contact_sensor_a = ContactSensorCfg( + prim_path="{ENV_REGEX_NS}/Object", + update_period=0.0, + history_length=1, + track_air_time=True, + ) + + scene = InteractiveScene(scene_cfg) + sim.reset() + scene.reset() + + contact_sensor: ContactSensor = scene["contact_sensor_a"] + obj: RigidObject = scene["object_a"] + + root_pose = wp.to_torch(obj.data.root_link_pose_w).clone() + for group_idx, base_height in enumerate(base_heights): + for i in range(envs_per_group): + env_idx = group_idx * envs_per_group + i + root_pose[env_idx, 2] = base_height + object_offset + obj.write_root_pose_to_sim_index(root_pose=root_pose) + + expected_land_ticks = [int(math.sqrt(2 * h / gravity_mag) / SIM_DT) for h in base_heights] + + contact_detected = [False] * num_envs + contact_tick = [-1] * num_envs + + for _ in range(5): + perform_sim_step(sim, scene, SIM_DT) + + forces = torch.norm(wp.to_torch(contact_sensor.data.net_forces_w), dim=-1) + for env_idx in range(num_envs): + assert forces[env_idx].max().item() < 0.01, f"Env {env_idx}: No contact should be detected while in air." + + for tick in range(5, total_fall_steps): + perform_sim_step(sim, scene, SIM_DT) + forces = torch.norm(wp.to_torch(contact_sensor.data.net_forces_w), dim=-1) + for env_idx in range(num_envs): + if forces[env_idx].max().item() > 0.1 and not contact_detected[env_idx]: + contact_detected[env_idx] = True + contact_tick[env_idx] = tick + + for env_idx in range(num_envs): + group_idx = env_idx // envs_per_group + assert contact_detected[env_idx], ( + f"Env {env_idx} (group {group_idx}, h={base_heights[group_idx]}m): Contact should be detected" + ) + + for env_idx in range(num_envs): + group_idx = env_idx // envs_per_group + expected_tick = expected_land_ticks[group_idx] + tolerance_ticks = int(0.3 * expected_tick) + 10 + assert abs(contact_tick[env_idx] - expected_tick) < tolerance_ticks, ( + f"Env {env_idx}: Contact at tick {contact_tick[env_idx]}, expected ~{expected_tick} ± {tolerance_ticks}" + ) + + group_land_times = [] + for group_idx in range(num_groups): + group_ticks = [contact_tick[group_idx * envs_per_group + i] for i in range(envs_per_group)] + group_land_times.append(sum(group_ticks) / len(group_ticks)) + + for i in range(num_groups - 1): + assert group_land_times[i] < group_land_times[i + 1], ( + f"Group {i} should land before Group {i + 1}. " + f"Avg ticks: {group_land_times[i]:.1f} vs {group_land_times[i + 1]:.1f}" + ) + + velocity = torch.zeros(num_envs, 6, device=device) + velocity[:, 2] = 5.0 + obj.write_root_velocity_to_sim_index(root_velocity=velocity) + + no_contact_detected = [False] * num_envs + for step in range(lift_steps): + perform_sim_step(sim, scene, SIM_DT) + if step > 10: + forces = torch.norm(wp.to_torch(contact_sensor.data.net_forces_w), dim=-1) + for env_idx in range(num_envs): + if forces[env_idx].max().item() < 0.01: + no_contact_detected[env_idx] = True + + for env_idx in range(num_envs): + assert no_contact_detected[env_idx], f"Env {env_idx}: Contact should stop after lift." + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("use_mujoco_contacts", COLLISION_PIPELINES) +@pytest.mark.parametrize("shape_type", STABLE_SHAPES, ids=[shape_type_to_str(s) for s in STABLE_SHAPES]) +def test_horizontal_collision_detects_contact(device: str, use_mujoco_contacts: bool, shape_type: ShapeType): + """Test horizontal collision detection with varied velocities and separations. + + 8 environments (2 groups x 4 envs) with different collision speeds. + + Verifies: + - Contact detected for all collision configurations + - Both objects in each pair detect contact + """ + collision_steps = 90 + num_envs = 8 + num_groups = 2 + envs_per_group = num_envs // num_groups + extent = get_shape_extent(shape_type) + + group_configs = [ + (2.0, 0.6 + 2 * extent), + (4.0, 0.8 + 2 * extent), + ] + + sim_cfg = make_sim_cfg(use_mujoco_contacts=use_mujoco_contacts, device=device, gravity=(0.0, 0.0, 0.0)) + + with build_simulation_context(sim_cfg=sim_cfg, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + + max_separation = max(cfg[1] for cfg in group_configs) + scene_cfg = ContactSensorTestSceneCfg(num_envs=num_envs, env_spacing=5.0, lazy_sensor_update=False) + scene_cfg.object_a = create_shape_cfg( + shape_type, + "{ENV_REGEX_NS}/ObjectA", + pos=(-max_separation / 2, 0.0, 0.5), + disable_gravity=True, + activate_contact_sensors=True, + ) + scene_cfg.object_b = create_shape_cfg( + shape_type, + "{ENV_REGEX_NS}/ObjectB", + pos=(max_separation / 2, 0.0, 0.5), + disable_gravity=True, + activate_contact_sensors=True, + ) + scene_cfg.contact_sensor_a = ContactSensorCfg( + prim_path="{ENV_REGEX_NS}/ObjectA", + update_period=0.0, + history_length=3, + ) + scene_cfg.contact_sensor_b = ContactSensorCfg( + prim_path="{ENV_REGEX_NS}/ObjectB", + update_period=0.0, + history_length=3, + ) + + scene = InteractiveScene(scene_cfg) + sim.reset() + scene.reset() + + object_a: RigidObject = scene["object_a"] + object_b: RigidObject = scene["object_b"] + sensor_a: ContactSensor = scene["contact_sensor_a"] + sensor_b: ContactSensor = scene["contact_sensor_b"] + + pose_a = wp.to_torch(object_a.data.root_link_pose_w).clone() + pose_b = wp.to_torch(object_b.data.root_link_pose_w).clone() + for group_idx, (_, separation) in enumerate(group_configs): + for i in range(envs_per_group): + env_idx = group_idx * envs_per_group + i + pose_a[env_idx, 0] = -separation / 2 + pose_b[env_idx, 0] = separation / 2 + object_a.write_root_pose_to_sim_index(root_pose=pose_a) + object_b.write_root_pose_to_sim_index(root_pose=pose_b) + + velocity = torch.zeros(num_envs, 6, device=device) + for group_idx, (vel, _) in enumerate(group_configs): + for i in range(envs_per_group): + velocity[group_idx * envs_per_group + i, 0] = vel + object_a.write_root_velocity_to_sim_index(root_velocity=velocity) + + contact_detected_a = [False] * num_envs + contact_detected_b = [False] * num_envs + + for tick in range(collision_steps): + perform_sim_step(sim, scene, SIM_DT) + forces_a = torch.norm(wp.to_torch(sensor_a.data.net_forces_w), dim=-1) + forces_b = torch.norm(wp.to_torch(sensor_b.data.net_forces_w), dim=-1) + for env_idx in range(num_envs): + if forces_a[env_idx].max().item() > 0.1: + contact_detected_a[env_idx] = True + if forces_b[env_idx].max().item() > 0.1: + contact_detected_b[env_idx] = True + + for env_idx in range(num_envs): + group_idx = env_idx // envs_per_group + vel, sep = group_configs[group_idx] + assert contact_detected_a[env_idx], f"Env {env_idx} (v={vel}m/s): Object A should detect contact" + assert contact_detected_b[env_idx], f"Env {env_idx} (v={vel}m/s): Object B should detect contact" + + +# =================================================================== +# Priority 2: Net Forces +# =================================================================== + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("use_mujoco_contacts", COLLISION_PIPELINES) +def test_resting_object_contact_force(device: str, use_mujoco_contacts: bool): + """Test that resting object contact force equals weight and points upward. + + Two objects (light=2kg and heavy=4kg) rest on ground. + + Verifies: + - Force magnitude ~ mass x gravity + - Force direction is upward (positive Z) + - Heavier object has proportionally larger force + """ + settle_steps = 120 + num_envs = 4 + mass_a, mass_b = 2.0, 4.0 + gravity_magnitude = 9.81 + expected_force_a = mass_a * gravity_magnitude + expected_force_b = mass_b * gravity_magnitude + + sim_cfg = make_sim_cfg( + use_mujoco_contacts=use_mujoco_contacts, device=device, gravity=(0.0, 0.0, -gravity_magnitude) + ) + + with build_simulation_context(sim_cfg=sim_cfg, auto_add_lighting=True, add_ground_plane=True) as sim: + sim._app_control_on_stop_handle = None + + scene_cfg = ContactSensorTestSceneCfg(num_envs=num_envs, env_spacing=5.0, lazy_sensor_update=False) + rigid_props = sim_utils.RigidBodyPropertiesCfg(disable_gravity=False, linear_damping=0.5, angular_damping=0.5) + + scene_cfg.object_a = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/BoxA", + spawn=sim_utils.CuboidCfg( + size=(0.3, 0.3, 0.3), + rigid_props=rigid_props, + collision_props=sim_utils.CollisionPropertiesCfg(collision_enabled=True), + mass_props=sim_utils.MassPropertiesCfg(mass=mass_a), + activate_contact_sensors=True, + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(-0.5, 0.0, 0.5)), + ) + scene_cfg.object_b = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/BoxB", + spawn=sim_utils.CuboidCfg( + size=(0.3, 0.3, 0.3), + rigid_props=rigid_props, + collision_props=sim_utils.CollisionPropertiesCfg(collision_enabled=True), + mass_props=sim_utils.MassPropertiesCfg(mass=mass_b), + activate_contact_sensors=True, + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.5, 0.0, 0.5)), + ) + scene_cfg.contact_sensor_a = ContactSensorCfg( + prim_path="{ENV_REGEX_NS}/BoxA", update_period=0.0, history_length=1 + ) + scene_cfg.contact_sensor_b = ContactSensorCfg( + prim_path="{ENV_REGEX_NS}/BoxB", update_period=0.0, history_length=1 + ) + + scene = InteractiveScene(scene_cfg) + sim.reset() + scene.reset() + + sensor_a: ContactSensor = scene["contact_sensor_a"] + sensor_b: ContactSensor = scene["contact_sensor_b"] + + for _ in range(settle_steps): + perform_sim_step(sim, scene, SIM_DT) + + forces_a = wp.to_torch(sensor_a.data.net_forces_w) + forces_b = wp.to_torch(sensor_b.data.net_forces_w) + force_mags_a = torch.norm(forces_a, dim=-1) + force_mags_b = torch.norm(forces_b, dim=-1) + + for env_idx in range(num_envs): + fa = force_mags_a[env_idx].max().item() + fb = force_mags_b[env_idx].max().item() + + assert abs(fa - expected_force_a) < 0.2 * expected_force_a, ( + f"Env {env_idx}: BoxA ({mass_a}kg) force should be ~{expected_force_a:.2f} N. Got {fa:.2f} N" + ) + assert abs(fb - expected_force_b) < 0.2 * expected_force_b, ( + f"Env {env_idx}: BoxB ({mass_b}kg) force should be ~{expected_force_b:.2f} N. Got {fb:.2f} N" + ) + assert fb > fa, f"Env {env_idx}: Heavier BoxB should have larger force. A: {fa:.2f}, B: {fb:.2f}" + + assert forces_a[env_idx, 0, 2].item() > 0.1, f"Env {env_idx}: BoxA Z force should be positive" + assert forces_b[env_idx, 0, 2].item() > 0.1, f"Env {env_idx}: BoxB Z force should be positive" + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("use_mujoco_contacts", COLLISION_PIPELINES) +def test_higher_drop_produces_larger_impact_force(device: str, use_mujoco_contacts: bool): + """Test that dropping from higher produces larger peak impact force. + + 8 environments with heights from 0.3m to 3.0m. + + Verifies: + - Peak impact force generally increases with height + - Overall trend: highest/lowest force ratio > 1.5 + """ + num_envs = 8 + min_height, max_height = 0.3, 3.0 + drop_heights = [min_height + (max_height - min_height) * i / (num_envs - 1) for i in range(num_envs)] + gravity_mag = 9.81 + object_radius = 0.25 + + sim_cfg = make_sim_cfg(use_mujoco_contacts=use_mujoco_contacts, device=device, gravity=(0.0, 0.0, -gravity_mag)) + + with build_simulation_context(sim_cfg=sim_cfg, auto_add_lighting=True, add_ground_plane=True) as sim: + sim._app_control_on_stop_handle = None + + scene_cfg = ContactSensorTestSceneCfg(num_envs=num_envs, env_spacing=5.0, lazy_sensor_update=False) + scene_cfg.object_a = create_shape_cfg( + ShapeType.SPHERE, + "{ENV_REGEX_NS}/Sphere", + pos=(0.0, 0.0, max_height + object_radius), + disable_gravity=False, + activate_contact_sensors=True, + ) + scene_cfg.contact_sensor_a = ContactSensorCfg( + prim_path="{ENV_REGEX_NS}/Sphere", + update_period=0.0, + history_length=1, + ) + + scene = InteractiveScene(scene_cfg) + sim.reset() + scene.reset() + + obj: RigidObject = scene["object_a"] + contact_sensor: ContactSensor = scene["contact_sensor_a"] + + root_pose = wp.to_torch(obj.data.root_link_pose_w).clone() + for env_idx in range(num_envs): + root_pose[env_idx, 2] = drop_heights[env_idx] + object_radius + obj.write_root_pose_to_sim_index(root_pose=root_pose) + + total_steps = int(((2 * max_height / gravity_mag) ** 0.5 + 0.5) / SIM_DT) + peak_forces = [0.0] * num_envs + contact_detected = [False] * num_envs + + for _ in range(total_steps): + perform_sim_step(sim, scene, SIM_DT) + force_magnitudes = torch.norm(wp.to_torch(contact_sensor.data.net_forces_w), dim=-1) + for env_idx in range(num_envs): + f = force_magnitudes[env_idx].max().item() + if f > 0.1: + contact_detected[env_idx] = True + peak_forces[env_idx] = max(peak_forces[env_idx], f) + + for env_idx in range(num_envs): + assert contact_detected[env_idx], f"Env {env_idx} (h={drop_heights[env_idx]:.2f}m): No contact" + + violations = [] + for i in range(num_envs - 1): + if peak_forces[i + 1] < peak_forces[i] * 0.95: + violations.append( + f"Env {i} (h={drop_heights[i]:.2f}m, F={peak_forces[i]:.2f}N) -> " + f"Env {i + 1} (h={drop_heights[i + 1]:.2f}m, F={peak_forces[i + 1]:.2f}N)" + ) + assert len(violations) <= 2, "Peak force should increase with height. Violations:\n" + "\n".join(violations) + + force_ratio = peak_forces[-1] / peak_forces[0] if peak_forces[0] > 0 else 0 + assert force_ratio > 1.5, f"Force ratio (highest/lowest) should be > 1.5. Got {force_ratio:.2f}" + + +# =================================================================== +# Priority 3: Filtering +# =================================================================== + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize( + "use_mujoco_contacts", + [ + pytest.param(False, id="newton_contacts", marks=pytest.mark.flaky(reruns=3)), + pytest.param(True, id="mujoco_contacts"), + ], +) +def test_filter_enables_force_matrix(device: str, use_mujoco_contacts: bool): + """Test that filter_prim_paths_expr filters contacts and enables force_matrix_w. + + Object A rests on ground, Object B stacked on A. + Sensor on A is filtered for B only (not ground). + + Verifies: + - force_matrix_w reports only filtered contact (A-B) + - net_forces_w reports total contact (ground + B) + - force_matrix < net_forces (ground contact excluded from matrix) + """ + settle_steps = 180 + num_envs = 4 + mass_b = 2.0 + gravity = 9.81 + expected_force_from_b = mass_b * gravity + + sim_cfg = make_sim_cfg(use_mujoco_contacts=use_mujoco_contacts, device=device, gravity=(0.0, 0.0, -gravity)) + + with build_simulation_context(sim_cfg=sim_cfg, auto_add_lighting=True, add_ground_plane=True) as sim: + sim._app_control_on_stop_handle = None + + scene_cfg = ContactSensorTestSceneCfg(num_envs=num_envs, env_spacing=5.0, lazy_sensor_update=False) + + rigid_props_a = sim_utils.RigidBodyPropertiesCfg(disable_gravity=False, linear_damping=0.5, angular_damping=0.5) + scene_cfg.object_a = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/ObjectA", + spawn=sim_utils.CuboidCfg( + size=(0.5, 0.5, 0.3), + rigid_props=rigid_props_a, + collision_props=sim_utils.CollisionPropertiesCfg(collision_enabled=True), + mass_props=sim_utils.MassPropertiesCfg(mass=5.0), + activate_contact_sensors=True, + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)), + ) + + rigid_props_b = sim_utils.RigidBodyPropertiesCfg(disable_gravity=False, linear_damping=2.0, angular_damping=2.0) + scene_cfg.object_b = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/ObjectB", + spawn=sim_utils.CuboidCfg( + size=(0.3, 0.3, 0.3), + rigid_props=rigid_props_b, + collision_props=sim_utils.CollisionPropertiesCfg(collision_enabled=True), + mass_props=sim_utils.MassPropertiesCfg(mass=mass_b), + activate_contact_sensors=True, + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 0.8)), + ) + + scene_cfg.contact_sensor_a = ContactSensorCfg( + prim_path="{ENV_REGEX_NS}/ObjectA", + update_period=0.0, + history_length=1, + filter_prim_paths_expr=["{ENV_REGEX_NS}/ObjectB"], + ) + + scene = InteractiveScene(scene_cfg) + sim.reset() + scene.reset() + + contact_sensor: ContactSensor = scene["contact_sensor_a"] + + for _ in range(settle_steps): + perform_sim_step(sim, scene, SIM_DT) + + force_matrix_raw = contact_sensor.data.force_matrix_w + net_forces_raw = contact_sensor.data.net_forces_w + + assert force_matrix_raw is not None, "force_matrix_w should not be None when filter is set" + + force_matrix = wp.to_torch(force_matrix_raw) + net_forces = wp.to_torch(net_forces_raw) + + for env_idx in range(num_envs): + matrix_force = torch.norm(force_matrix[env_idx]).item() + net_force = torch.norm(net_forces[env_idx]).item() + + tolerance = 0.3 * expected_force_from_b + assert abs(matrix_force - expected_force_from_b) < tolerance, ( + f"Env {env_idx}: force_matrix should be ~{expected_force_from_b:.2f} N. Got: {matrix_force:.2f} N" + ) + assert matrix_force < net_force, ( + f"Env {env_idx}: force_matrix (B only) should be < net_forces (all). " + f"Matrix: {matrix_force:.2f} N, Net: {net_force:.2f} N" + ) + + +# =================================================================== +# Priority 4: Articulated System +# =================================================================== + +ALLEGRO_FINGERTIP_OFFSETS = { + "index": (-0.052, -0.252, 0.052), + "middle": (-0.001, -0.252, 0.052), + "ring": (0.054, -0.252, 0.052), + "thumb": (-0.168, -0.039, 0.080), +} + +ALLEGRO_FINGER_LINKS = { + "index": "index_link_3", + "middle": "middle_link_3", + "ring": "ring_link_3", + "thumb": "thumb_link_3", +} + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize( + "use_mujoco_contacts", + [ + pytest.param( + False, + id="newton_contacts", + marks=pytest.mark.xfail( + reason="Newton contact pipeline reports inaccurate per-finger forces in articulated systems" + ), + ), + pytest.param(True, id="mujoco_contacts"), + ], +) +@pytest.mark.parametrize( + "drop_shape", + [ + pytest.param(ShapeType.SPHERE, id="sphere"), + pytest.param(ShapeType.MESH_BOX, id="mesh_box"), + ], +) +def test_finger_contact_sensor_isolation(device: str, use_mujoco_contacts: bool, drop_shape: ShapeType): + """Test contact sensor on Allegro hand fingers detects localized contacts. + + Uses 4 environments, each dropping a ball on a different finger + (env 0 -> index, env 1 -> middle, env 2 -> ring, env 3 -> thumb). + Verifies that in each env the target finger has the highest peak force. + """ + drop_steps = 120 + num_envs = 4 + hand_pos = (0.0, 0.0, 0.5) + finger_names = ["index", "middle", "ring", "thumb"] + + sim_cfg = make_sim_cfg(use_mujoco_contacts=use_mujoco_contacts, device=device, gravity=(0.0, 0.0, 0.0)) + + with build_simulation_context(sim_cfg=sim_cfg, add_ground_plane=True, add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + + scene_cfg = ContactSensorTestSceneCfg(num_envs=num_envs, env_spacing=1.0, lazy_sensor_update=False) + + scene_cfg.hand = ALLEGRO_HAND_CFG.copy() + scene_cfg.hand.prim_path = "{ENV_REGEX_NS}/Hand" + scene_cfg.hand.init_state.pos = hand_pos + + for finger in finger_names: + setattr( + scene_cfg, + f"contact_sensor_{finger}", + ContactSensorCfg( + prim_path=f"{{ENV_REGEX_NS}}/Hand/{ALLEGRO_FINGER_LINKS[finger]}", + update_period=0.0, + history_length=1, + ), + ) + + drop_rigid_props = sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, linear_damping=0.0, angular_damping=0.0 + ) + drop_collision_props = sim_utils.CollisionPropertiesCfg(collision_enabled=True) + drop_mass_props = sim_utils.MassPropertiesCfg(mass=1.0) + drop_visual = sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0)) + + spawn_map = { + ShapeType.SPHERE: lambda: sim_utils.SphereCfg( + radius=0.035, + rigid_props=drop_rigid_props, + collision_props=drop_collision_props, + mass_props=drop_mass_props, + visual_material=drop_visual, + activate_contact_sensors=True, + ), + ShapeType.MESH_BOX: lambda: sim_utils.MeshCuboidCfg( + size=(0.05, 0.05, 0.05), + rigid_props=drop_rigid_props, + collision_props=drop_collision_props, + mass_props=drop_mass_props, + visual_material=drop_visual, + activate_contact_sensors=True, + ), + } + + default_offset = ALLEGRO_FINGERTIP_OFFSETS["index"] + default_drop_pos = ( + hand_pos[0] + default_offset[0], + hand_pos[1] + default_offset[1], + hand_pos[2] + default_offset[2] + 0.10, + ) + + scene_cfg.drop_object = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/DropObject", + spawn=spawn_map[drop_shape](), + init_state=RigidObjectCfg.InitialStateCfg(pos=default_drop_pos), + ) + + scene = InteractiveScene(scene_cfg) + sim.reset() + scene.reset() + + hand: Articulation = scene["hand"] + drop_object: RigidObject = scene["drop_object"] + finger_sensors = {f: scene[f"contact_sensor_{f}"] for f in finger_names} + + # Newton's articulation.reset() doesn't write default_joint_pos to sim (unlike + # ManagerBasedEnv's reset_scene_to_default event). Without this, joints start at 0.0 + # which is below thumb_joint_0's lower limit (0.279 rad), causing violent oscillation. + default_jpos = wp.to_torch(hand.data.default_joint_pos).clone() + default_jvel = wp.to_torch(hand.data.default_joint_vel).clone() + hand.write_joint_position_to_sim_index(position=default_jpos) + hand.write_joint_velocity_to_sim_index(velocity=default_jvel) + hand.set_joint_position_target_index(target=default_jpos) + + hand_world_pos = wp.to_torch(hand.data.root_link_pose_w)[:, :3] + drop_pose = wp.to_torch(drop_object.data.root_link_pose_w).clone() + for env_idx, finger in enumerate(finger_names): + offset = ALLEGRO_FINGERTIP_OFFSETS[finger] + drop_pose[env_idx, 0] = hand_world_pos[env_idx, 0] + offset[0] + drop_pose[env_idx, 1] = hand_world_pos[env_idx, 1] + offset[1] + drop_pose[env_idx, 2] = hand_world_pos[env_idx, 2] + offset[2] + 0.10 + drop_object.write_root_pose_to_sim_index(root_pose=drop_pose) + + for _ in range(30): + perform_sim_step(sim, scene, SIM_DT) + + hand_world_pos = wp.to_torch(hand.data.root_link_pose_w)[:, :3] + drop_object.reset() + drop_pose = wp.to_torch(drop_object.data.root_link_pose_w).clone() + for env_idx, finger in enumerate(finger_names): + offset = ALLEGRO_FINGERTIP_OFFSETS[finger] + drop_pose[env_idx, 0] = hand_world_pos[env_idx, 0] + offset[0] + drop_pose[env_idx, 1] = hand_world_pos[env_idx, 1] + offset[1] + drop_pose[env_idx, 2] = hand_world_pos[env_idx, 2] + offset[2] + 0.10 + drop_object.write_root_pose_to_sim_index(root_pose=drop_pose) + + initial_velocity = torch.zeros(num_envs, 6, device=device) + initial_velocity[:, 2] = -1.5 + drop_object.write_root_velocity_to_sim_index(root_velocity=initial_velocity) + + peak_forces = {f: [0.0] * num_envs for f in finger_names} + + for _ in range(drop_steps): + perform_sim_step(sim, scene, SIM_DT) + for finger_name, sensor in finger_sensors.items(): + if sensor.data.net_forces_w is not None: + forces = wp.to_torch(sensor.data.net_forces_w) + for env_idx in range(num_envs): + f = torch.norm(forces[env_idx]).item() + peak_forces[finger_name][env_idx] = max(peak_forces[finger_name][env_idx], f) + + for env_idx, target_finger in enumerate(finger_names): + target_peak = peak_forces[target_finger][env_idx] + assert target_peak > 0.5, ( + f"Env {env_idx}: Target finger '{target_finger}' peak force: {target_peak:.4f} N (expected > 0.5)" + ) + + for other_finger in finger_names: + if other_finger != target_finger: + assert target_peak >= peak_forces[other_finger][env_idx], ( + f"Env {env_idx}: '{target_finger}' (peak={target_peak:.4f}N) should have " + f"highest force, but '{other_finger}' had " + f"{peak_forces[other_finger][env_idx]:.4f}N" + ) + + +# =================================================================== +# Utility +# =================================================================== + + +def test_sensor_print(): + """Test that contact sensor print/repr works correctly.""" + sim_cfg = make_sim_cfg(use_mujoco_contacts=False, device="cuda:0", gravity=(0.0, 0.0, -9.81)) + + with build_simulation_context(sim_cfg=sim_cfg, auto_add_lighting=True, add_ground_plane=True) as sim: + sim._app_control_on_stop_handle = None + + scene_cfg = ContactSensorTestSceneCfg(num_envs=4, env_spacing=5.0, lazy_sensor_update=False) + scene_cfg.object_a = create_shape_cfg( + ShapeType.BOX, + "{ENV_REGEX_NS}/Object", + pos=(0.0, 0.0, 2.0), + disable_gravity=False, + activate_contact_sensors=True, + ) + scene_cfg.contact_sensor_a = ContactSensorCfg( + prim_path="{ENV_REGEX_NS}/Object", + update_period=0.0, + history_length=3, + track_air_time=True, + ) + + scene = InteractiveScene(scene_cfg) + sim.reset() + scene.reset() + + sensor_str = str(scene["contact_sensor_a"]) + assert len(sensor_str) > 0 + print(sensor_str) diff --git a/source/isaaclab_ov/config/extension.toml b/source/isaaclab_ov/config/extension.toml index f6a3ce55a2f4..f32bb96fc7b9 100644 --- a/source/isaaclab_ov/config/extension.toml +++ b/source/isaaclab_ov/config/extension.toml @@ -1,5 +1,5 @@ [package] -version = "0.1.0" +version = "0.1.1" title = "Omniverse renderers for IsaacLab" description = "Extension providing Omniverse renderers (OVRTX, ovphysx, etc.) for tiled camera rendering." readme = "docs/README.md" diff --git a/source/isaaclab_ov/docs/CHANGELOG.rst b/source/isaaclab_ov/docs/CHANGELOG.rst index 6013c0b9d13f..ed93de19af99 100644 --- a/source/isaaclab_ov/docs/CHANGELOG.rst +++ b/source/isaaclab_ov/docs/CHANGELOG.rst @@ -1,6 +1,17 @@ Changelog --------- +0.1.1 (2026-03-07) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added ``ovrtx>=0.2.0,<0.3.0`` as a declared dependency, installable from the + public NVIDIA package index (``pypi.nvidia.com``). +* Added ``ov`` to the list of valid sub-packages for selective installation via + ``./isaaclab.sh -i ov``. + 0.1.0 (2026-03-04) ~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py index 4b4ab6bf8775..8b9d8239cd90 100644 --- a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py +++ b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py @@ -120,6 +120,7 @@ def __init__(self, cfg: OVRTXRendererCfg): self._initialized_scene = False self._sensor_ref: weakref.ref[object] | None = None self._exported_usd_path: str | None = None + self._camera_rel_path: str | None = None def prepare_stage(self, stage: Any, num_envs: int) -> None: """Export the USD stage for OVRTX before create_render_data. @@ -133,7 +134,7 @@ def prepare_stage(self, stage: Any, num_envs: int) -> None: use_cloning = self.cfg.use_cloning logger.info("Preparing stage for export (%d envs, cloning=%s)...", num_envs, use_cloning) - create_cloning_attributes(stage, "Camera", num_envs, use_cloning) + create_cloning_attributes(stage, num_envs, use_cloning) export_path = "/tmp/stage_before_ovrtx.usda" export_stage_for_ovrtx(stage, export_path, num_envs, use_cloning) @@ -154,6 +155,12 @@ def initialize(self, sensor: SensorBase): num_envs = sensor.num_instances data_types = sensor.cfg.data_types if sensor.cfg.data_types else ["rgb"] + env_0_prefix = "/World/envs/env_0/" + first_cam_path = sensor._view.prims[0].GetPath().pathString + if not first_cam_path.startswith(env_0_prefix): + raise RuntimeError(f"Expected camera prim under '{env_0_prefix}', got '{first_cam_path}'") + self._camera_rel_path = first_cam_path.removeprefix(env_0_prefix) + usd_scene_path = self._exported_usd_path use_cloning = self.cfg.use_cloning @@ -176,6 +183,7 @@ def initialize(self, sensor: SensorBase): height=height, num_envs=num_envs, data_types=data_types, + camera_rel_path=self._camera_rel_path, ) self._render_product_paths.append(render_product_path) @@ -195,7 +203,7 @@ def initialize(self, sensor: SensorBase): self._initialized_scene = True - camera_paths = [f"/World/envs/env_{i}/Camera" for i in range(num_envs)] + camera_paths = [f"/World/envs/env_{i}/{self._camera_rel_path}" for i in range(num_envs)] self._camera_binding = self._renderer.bind_attribute( prim_paths=camera_paths, attribute_name="omni:xform", @@ -237,7 +245,7 @@ def _update_scene_partitions_after_clone(self, usd_file_path: str, num_envs: int logger.info("Writing scene partitions for %d environments...", num_envs) partition_tokens = [f"env_{i}" for i in range(num_envs)] env_prim_paths = [f"/World/envs/env_{i}" for i in range(num_envs)] - camera_prim_paths = [f"/World/envs/env_{i}/Camera" for i in range(num_envs)] + camera_prim_paths = [f"/World/envs/env_{i}/{self._camera_rel_path}" for i in range(num_envs)] try: self._renderer.write_attribute( @@ -262,11 +270,8 @@ def _setup_object_bindings(self): """Setup OVRTX bindings for scene objects to sync with Newton physics.""" try: from isaaclab.sim import SimulationContext - from isaaclab.visualizers import VisualizerCfg - provider = SimulationContext.instance().initialize_scene_data_provider( - [VisualizerCfg(visualizer_type="newton")] - ) + provider = SimulationContext.instance().initialize_scene_data_provider() newton_model = provider.get_newton_model() if newton_model is None: logger.info("Newton model not available, skipping object bindings") @@ -280,7 +285,7 @@ def _setup_object_bindings(self): object_paths = [] newton_indices = [] for idx, path in enumerate(all_body_paths): - if "/World/envs/" in path and "Camera" not in path and "GroundPlane" not in path: + if "/World/envs/" in path and self._camera_rel_path not in path and "GroundPlane" not in path: object_paths.append(path) newton_indices.append(idx) @@ -336,11 +341,8 @@ def update_transforms(self) -> None: try: from isaaclab.sim import SimulationContext - from isaaclab.visualizers import VisualizerCfg - provider = SimulationContext.instance().initialize_scene_data_provider( - [VisualizerCfg(visualizer_type="newton")] - ) + provider = SimulationContext.instance().initialize_scene_data_provider() newton_state = provider.get_newton_state() if newton_state is None: return diff --git a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_usd.py b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_usd.py index 212e86ff6c30..28cdfbc6646c 100644 --- a/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_usd.py +++ b/source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_usd.py @@ -11,7 +11,7 @@ from pathlib import Path from typing import TYPE_CHECKING -from pxr import Sdf, Usd +from pxr import Sdf, Usd, UsdGeom if TYPE_CHECKING: from .ovrtx_renderer_cfg import OVRTXRendererCfg @@ -94,6 +94,7 @@ def inject_cameras_into_usd( height: int, num_envs: int, data_types: list[str], + camera_rel_path: str = "Camera", ) -> tuple[str, str]: """Inject camera and render product definitions into an existing USD file. @@ -107,6 +108,8 @@ def inject_cameras_into_usd( height: Tile height from sensor config. num_envs: Number of environments from scene. data_types: Data types from sensor config. + camera_rel_path: Camera prim path relative to the env root (e.g. ``"Camera"`` + or ``"Robot/head_cam"``). """ with open(usd_scene_path) as f: original_usd = f.read() @@ -114,7 +117,7 @@ def inject_cameras_into_usd( data_types = data_types if data_types else ["rgb"] tiled_width, tiled_height = _tiled_resolution(num_envs, width, height) - camera_paths = [f"/World/envs/env_{i}/Camera" for i in range(num_envs)] + camera_paths = [f"/World/envs/env_{i}/{camera_rel_path}" for i in range(num_envs)] render_product_name = "RenderProduct" render_product_path = f"/Render/{render_product_name}" @@ -140,18 +143,20 @@ def inject_cameras_into_usd( return temp_path, render_product_path -def create_cloning_attributes( - stage, camera_prim_name: str = "Camera", num_envs: int = 1, use_cloning: bool = True -) -> int: +def create_cloning_attributes(stage, num_envs: int = 1, use_cloning: bool = True) -> int: """Create OVRTX cloning attributes (scene partition, xform) on env_0 only. Only env_0 is exported for OVRTX; env_1..env_{n-1} are deactivated before export. OVRTX clones env_0 internally and _update_scene_partitions_after_clone sets partition attributes on the clones. So we only need to set attributes on env_0 here. + Camera prims are discovered by USD type (``UsdGeom.Camera``) rather than by + name, so this works regardless of where the camera is placed in the hierarchy. + Args: stage: USD stage to modify. - camera_prim_name: Name of the camera prim under each env (e.g. "Camera"). + num_envs: Number of environments. + use_cloning: Whether OVRTX cloning is enabled. Returns: Total number of objects (non-camera prims) that received partition attributes. @@ -167,15 +172,13 @@ def create_cloning_attributes( attr = env_prim.CreateAttribute("primvars:omni:scenePartition", Sdf.ValueTypeNames.Token) attr.Set(partition_name) for prim in Usd.PrimRange(env_prim): - if prim.GetPath() == env_prim.GetPath() or "Camera" in prim.GetPath().pathString: + if prim.GetPath() == env_prim.GetPath(): continue - obj_attr = prim.CreateAttribute("primvars:omni:scenePartition", Sdf.ValueTypeNames.Token) - obj_attr.Set(partition_name) - total_objects += 1 - camera_path = f"{env_path}/{camera_prim_name}" - camera_prim = stage.GetPrimAtPath(camera_path) - if camera_prim.IsValid(): - camera_prim.CreateAttribute("omni:scenePartition", Sdf.ValueTypeNames.Token).Set(partition_name) + if prim.IsA(UsdGeom.Camera): + prim.CreateAttribute("omni:scenePartition", Sdf.ValueTypeNames.Token).Set(partition_name) + else: + prim.CreateAttribute("primvars:omni:scenePartition", Sdf.ValueTypeNames.Token).Set(partition_name) + total_objects += 1 return total_objects diff --git a/source/isaaclab_ov/setup.py b/source/isaaclab_ov/setup.py index 967b25f4da1c..854afb6526b1 100644 --- a/source/isaaclab_ov/setup.py +++ b/source/isaaclab_ov/setup.py @@ -7,11 +7,13 @@ from setuptools import setup -INSTALL_REQUIRES = [] +INSTALL_REQUIRES = [ + "ovrtx>=0.2.0,<0.3.0", +] setup( name="isaaclab_ov", - version="0.1.0", + version="0.1.1", author="Isaac Lab Project Developers", maintainer="Isaac Lab Project Developers", url="https://github.com/isaac-sim/IsaacLab", diff --git a/source/isaaclab_physx/config/extension.toml b/source/isaaclab_physx/config/extension.toml index efab73db623a..f75689975e98 100644 --- a/source/isaaclab_physx/config/extension.toml +++ b/source/isaaclab_physx/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "0.5.5" +version = "0.5.7" # Description title = "PhysX simulation interfaces for IsaacLab core package" diff --git a/source/isaaclab_physx/docs/CHANGELOG.rst b/source/isaaclab_physx/docs/CHANGELOG.rst index b6807e9c9c40..a3747a469d72 100644 --- a/source/isaaclab_physx/docs/CHANGELOG.rst +++ b/source/isaaclab_physx/docs/CHANGELOG.rst @@ -1,6 +1,16 @@ Changelog --------- +0.5.7 (2026-03-06) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Made several PhysX articulation tests more backend-agnostic by relaxing + PhysX-specific assumptions in ``test_articulation.py``. + + 0.5.6 (2026-03-03) ~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 16df16b19807..89337502f04f 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -3596,7 +3596,8 @@ def _initialize_impl(self): # Find all articulation root prims in the first environment. first_env_root_prims = get_all_matching_child_prims( first_env_matching_prim_path, - predicate=lambda prim: prim.HasAPI(UsdPhysics.ArticulationRootAPI), + predicate=lambda prim: prim.HasAPI(UsdPhysics.ArticulationRootAPI) + and prim.GetAttribute("physxArticulation:articulationEnabled").Get(), traverse_instance_prims=False, ) if len(first_env_root_prims) == 0: diff --git a/source/isaaclab_physx/isaaclab_physx/scene_data_providers/__init__.py b/source/isaaclab_physx/isaaclab_physx/scene_data_providers/__init__.py new file mode 100644 index 000000000000..0a3fe4d79fb4 --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/scene_data_providers/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""PhysX scene data provider backends.""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab/isaaclab/sim/scene_data_providers/__init__.pyi b/source/isaaclab_physx/isaaclab_physx/scene_data_providers/__init__.pyi similarity index 83% rename from source/isaaclab/isaaclab/sim/scene_data_providers/__init__.pyi rename to source/isaaclab_physx/isaaclab_physx/scene_data_providers/__init__.pyi index 7ad60162991c..d1612d1f3bbd 100644 --- a/source/isaaclab/isaaclab/sim/scene_data_providers/__init__.pyi +++ b/source/isaaclab_physx/isaaclab_physx/scene_data_providers/__init__.pyi @@ -6,9 +6,7 @@ __all__ = [ "NewtonSceneDataProvider", "PhysxSceneDataProvider", - "SceneDataProvider", ] from .newton_scene_data_provider import NewtonSceneDataProvider from .physx_scene_data_provider import PhysxSceneDataProvider -from .scene_data_provider import SceneDataProvider diff --git a/source/isaaclab/isaaclab/sim/scene_data_providers/physx_scene_data_provider.py b/source/isaaclab_physx/isaaclab_physx/scene_data_providers/physx_scene_data_provider.py similarity index 75% rename from source/isaaclab/isaaclab/sim/scene_data_providers/physx_scene_data_provider.py rename to source/isaaclab_physx/isaaclab_physx/scene_data_providers/physx_scene_data_provider.py index 5c71f55a3531..5560575d2cf9 100644 --- a/source/isaaclab/isaaclab/sim/scene_data_providers/physx_scene_data_provider.py +++ b/source/isaaclab_physx/isaaclab_physx/scene_data_providers/physx_scene_data_provider.py @@ -9,18 +9,47 @@ import logging import re +import time from collections import deque from typing import Any +import warp as wp + from pxr import UsdGeom +from isaaclab.physics.base_scene_data_provider import BaseSceneDataProvider + logger = logging.getLogger(__name__) # Path pattern for env prims: /World/envs/env_/... _ENV_ID_RE = re.compile(r"/World/envs/env_(\d+)") -class PhysxSceneDataProvider: +@wp.kernel(enable_backward=False) +def _set_body_q_kernel( + positions: wp.array(dtype=wp.vec3), + orientations: wp.array(dtype=wp.quatf), + body_q: wp.array(dtype=wp.transformf), +): + """Write pose arrays into Newton ``body_q`` in one-to-one index order.""" + i = wp.tid() + body_q[i] = wp.transformf(positions[i], orientations[i]) + + +@wp.kernel(enable_backward=False) +def _set_body_q_subset_kernel( + positions: wp.array(dtype=wp.vec3), + orientations: wp.array(dtype=wp.quatf), + body_indices: wp.array(dtype=wp.int32), + body_q: wp.array(dtype=wp.transformf), +): + """Write pose arrays into selected Newton ``body_q`` indices.""" + i = wp.tid() + bi = body_indices[i] + body_q[bi] = wp.transformf(positions[i], orientations[i]) + + +class PhysxSceneDataProvider(BaseSceneDataProvider): """Scene data provider for Omni PhysX backend. Supports: @@ -28,7 +57,6 @@ class PhysxSceneDataProvider: - camera poses & intrinsics - USD stage handles - Newton model/state handles - - TODO: mesh data access """ # ---- Environment discovery / metadata ------------------------------------------------- @@ -64,7 +92,13 @@ def _determine_num_envs_in_scene(self) -> int: max_env_id = max(max_env_id, int(match.group(1))) return max_env_id + 1 if max_env_id >= 0 else 0 - def __init__(self, visualizer_cfgs: list[Any] | None, stage, simulation_context) -> None: + def __init__(self, stage, simulation_context) -> None: + """Initialize the PhysX scene data provider. + + Args: + stage: USD stage handle. + simulation_context: Active simulation context. + """ from isaacsim.core.simulation_manager import SimulationManager self._simulation_context = simulation_context @@ -80,8 +114,13 @@ def __init__(self, visualizer_cfgs: list[Any] | None, stage, simulation_context) # Single source of truth: discovered from stage and cached once available. self._num_envs: int | None = None - viz_types = {getattr(cfg, "visualizer_type", None) for cfg in (visualizer_cfgs or [])} - self._needs_newton_sync = bool({"newton", "rerun", "viser"} & viz_types) + # Determine if newton model sync is required for selected renderers and visualizers + requirements = self._simulation_context.get_scene_data_requirements() + self._needs_newton_sync = bool(requirements.requires_newton_model) + + # Benchmark/debug override: force USD traversal fallback even when prebuilt + # visualizer artifacts are available from the cloner path. + self._force_usd_fallback_for_newton_model_build = False # Fixed metadata for visualizers. get_metadata() returns this plus num_envs so visualizers # can .get("num_envs", 0), .get("physics_backend", ...) etc. without the provider exposing many methods. @@ -103,10 +142,21 @@ def __init__(self, visualizer_cfgs: list[Any] | None, stage, simulation_context) self._filtered_body_indices: list[int] = [] self._rigid_body_paths: list[str] = [] self._articulation_paths: list[str] = [] - self._set_body_q_kernel = None # env_id -> list of body indices (in Newton body_key order) self._env_id_to_body_indices: dict[int, list[int]] = {} + # Reused pose buffers (MR perf): avoid per-call allocations in _read_poses_from_best_source. + self._pose_buf_num_bodies = 0 + self._positions_buf = None + self._orientations_buf = None + self._covered_buf = None + self._xform_mask_buf = None + # View index order as device tensors for vectorized scatter in _apply_view_poses. + self._view_order_tensors: dict[str, Any] = {} + # Last full-model build source for tests/debugging ("prebuilt", "usd_fallback", "error"). + self._last_newton_model_build_source: str | None = None + self._last_newton_model_build_elapsed_ms: float | None = None + # Initialize Newton pipeline only if needed for visualization if self._needs_newton_sync: self._build_newton_model_from_usd() @@ -137,9 +187,82 @@ def _refresh_newton_model_if_needed(self) -> None: self._setup_rigid_body_view() self._setup_articulation_view() + def _model_body_paths(self, model) -> list[str]: + """Return body paths/keys from a Newton model. + + Args: + model: Newton model object. + + Returns: + Body paths/keys from the model, or an empty list when unavailable. + """ + if model is None: + return [] + return list(getattr(model, "body_label", None) or getattr(model, "body_key", [])) + + def _model_articulation_paths(self, model) -> list[str]: + """Return articulation paths/keys from a Newton model. + + Args: + model: Newton model object. + + Returns: + Articulation paths/keys from the model, or an empty list when unavailable. + """ + if model is None: + return [] + return list(getattr(model, "articulation_label", None) or getattr(model, "articulation_key", [])) + + def _try_use_prebuilt_newton_artifact(self) -> bool: + """Use scene-time prebuilt Newton visualizer artifact when available. + + Returns: + ``True`` when a valid prebuilt artifact was consumed, otherwise ``False``. + """ + if self._force_usd_fallback_for_newton_model_build: + return False + artifact = self._simulation_context.get_scene_data_visualizer_prebuilt_artifact() + if not artifact: + return False + + model = artifact.model + state = artifact.state + if model is None or state is None: + return False + + self._newton_model = model + self._newton_state = state + self._rigid_body_paths = list(artifact.rigid_body_paths) or self._model_body_paths(model) + self._articulation_paths = list(artifact.articulation_paths) or self._model_articulation_paths(model) + + self._xform_views.clear() + self._view_body_index_map = {} + self._view_order_tensors.clear() + self._pose_buf_num_bodies = 0 + self._positions_buf = None + self._orientations_buf = None + self._covered_buf = None + self._xform_mask_buf = None + self._env_id_to_body_indices = {} + self._num_envs_at_last_newton_build = int(artifact.num_envs) + self._filtered_newton_model = None + self._filtered_newton_state = None + self._filtered_env_ids_key = None + self._filtered_body_indices = [] + return True + def _build_newton_model_from_usd(self) -> None: """Build Newton model from USD and cache body/articulation paths.""" + # TODO: Deprecate this USD-traversal fallback once cloner/prebuilt coverage + # is complete for full and partial visualization model-build paths. + start_t = time.perf_counter() try: + if self._try_use_prebuilt_newton_artifact(): + self._last_newton_model_build_source = "prebuilt" + return + self._last_newton_model_build_source = ( + "usd_fallback_forced" if self._force_usd_fallback_for_newton_model_build else "usd_fallback" + ) from newton import ModelBuilder builder = ModelBuilder(up_axis=self._up_axis) @@ -152,11 +275,17 @@ def _build_newton_model_from_usd(self) -> None: self._newton_state = self._newton_model.state() # Extract scene structure from Newton model (single source of truth) - self._rigid_body_paths = list(self._newton_model.body_label) - self._articulation_paths = list(self._newton_model.articulation_label) + self._rigid_body_paths = self._model_body_paths(self._newton_model) + self._articulation_paths = self._model_articulation_paths(self._newton_model) self._xform_views.clear() self._view_body_index_map = {} + self._view_order_tensors.clear() + self._pose_buf_num_bodies = 0 + self._positions_buf = None + self._orientations_buf = None + self._covered_buf = None + self._xform_mask_buf = None self._env_id_to_body_indices = {} self._num_envs_at_last_newton_build = self.get_num_envs() # Invalidate any filtered model when full model changes. @@ -165,24 +294,47 @@ def _build_newton_model_from_usd(self) -> None: self._filtered_env_ids_key = None self._filtered_body_indices = [] except ModuleNotFoundError as exc: + self._last_newton_model_build_source = "error" logger.error( "[PhysxSceneDataProvider] Newton module not available. " "Install the Newton backend to use newton/rerun/viser visualizers." ) logger.debug(f"[PhysxSceneDataProvider] Newton import error: {exc}") except Exception as exc: + self._last_newton_model_build_source = "error" logger.error(f"[PhysxSceneDataProvider] Failed to build Newton model from USD: {exc}") self._newton_model = None self._newton_state = None self._rigid_body_paths = [] self._articulation_paths = [] self._num_envs_at_last_newton_build = None + finally: + elapsed_ms = (time.perf_counter() - start_t) * 1000.0 + self._last_newton_model_build_elapsed_ms = elapsed_ms + try: + num_envs = self.get_num_envs() + except Exception: + num_envs = -1 + logger.debug( + "[PhysxSceneDataProvider] Newton model build source=%s num_envs=%d elapsed_ms=%.2f", + self._last_newton_model_build_source, + num_envs, + elapsed_ms, + ) def _build_filtered_newton_model(self, env_ids: list[int]) -> None: - """Build Newton model/state for a subset of envs.""" + """Build Newton model/state for a subset of environments. + + Args: + env_ids: Environment ids to include in the subset model. + """ + # TODO: Deprecate this USD-traversal fallback once cloner/prebuilt coverage + # is complete for full and partial visualization model-build paths. try: from newton import ModelBuilder + # Newton model building from USD with partial visualization does not currently use cloner, + # and falls back to slower USD-stage traversal. builder = ModelBuilder(up_axis=self._up_axis) builder.add_usd(self._stage, ignore_paths=[r"/World/envs/.*"]) for env_id in env_ids: @@ -193,7 +345,7 @@ def _build_filtered_newton_model(self, env_ids: list[int]) -> None: self._filtered_newton_state = self._filtered_newton_model.state() full_index_by_path = {path: i for i, path in enumerate(self._rigid_body_paths)} - filtered_paths = list(self._filtered_newton_model.body_key) + filtered_paths = self._model_body_paths(self._filtered_newton_model) self._filtered_body_indices = [] missing = [] for path in filtered_paths: @@ -267,12 +419,12 @@ def _setup_articulation_view(self) -> None: # ---- Pose/velocity read pipeline ------------------------------------------------------ - def _warn_once(self, key: str, message: str, *args) -> None: + def _warn_once(self, key: str, message: str, *args, level=logging.WARNING) -> None: """Log a warning only once for a given key.""" if key in self._warned_once: return self._warned_once.add(key) - logger.warning(message, *args) + logger.log(level, message, *args) def _get_view_world_poses(self, view: Any): """Read world poses from a PhysX view. @@ -318,6 +470,10 @@ def _cache_view_index_map(self, view, key: str) -> None: if all(idx is not None for idx in order): self._view_body_index_map[key] = order # type: ignore[arg-type] + # Cache as device tensor for vectorized scatter in _apply_view_poses. + import torch + + self._view_order_tensors[key] = torch.tensor(order, dtype=torch.long, device=self._device) def _split_env_relative_path(self, path: str) -> tuple[int | None, str]: """Extract (env_id, relative_path) from a prim path.""" @@ -354,11 +510,6 @@ def _apply_view_poses(self, view: Any, view_key: str, positions: Any, orientatio order = self._view_body_index_map.get(view_key) if not order: - self._warn_once( - f"missing-index-map-{view_key}", - "[PhysxSceneDataProvider] Missing index map for %s; cannot scatter transforms.", - view_key, - ) return 0 # Normalize returned arrays to torch tensors across backends (torch/warp/other). @@ -376,7 +527,20 @@ def _apply_view_poses(self, view: Any, view_key: str, positions: Any, orientatio pos = pos.to(device=self._device, dtype=torch.float32) quat = quat.to(device=self._device, dtype=torch.float32) - # Scatter view outputs into the canonical Newton body order. + # Vectorized scatter when we have a cached order tensor (view fully covers bodies). + order_t = self._view_order_tensors.get(view_key) + if order_t is not None: + uncovered_mask = ~covered + if uncovered_mask.any(): + newton_indices = uncovered_mask.nonzero(as_tuple=True)[0] + view_indices = order_t[newton_indices] + positions[newton_indices] = pos[view_indices] + orientations[newton_indices] = quat[view_indices] + covered[newton_indices] = True + return newton_indices.numel() + return 0 + + # Fallback: Python loop when view does not fully cover or cache missing. count = 0 for newton_idx, view_idx in enumerate(order): if view_idx is not None and not covered[newton_idx]: @@ -426,6 +590,7 @@ def _apply_xform_poses(self, positions: Any, orientations: Any, covered: Any, xf "xform-fallback-failures", "[PhysxSceneDataProvider] Xform fallback failed for %d body paths.", len(self._xform_view_failures), + level=logging.DEBUG, ) return count @@ -450,11 +615,21 @@ def _read_poses_from_best_source(self) -> tuple[Any, Any, str, Any] | None: logger.warning(f"Body count mismatch: body_key={num_bodies}, state={self._newton_state.body_q.shape[0]}") return None - # Allocate outputs in Newton body order. - positions = torch.zeros((num_bodies, 3), dtype=torch.float32, device=self._device) - orientations = torch.zeros((num_bodies, 4), dtype=torch.float32, device=self._device) - covered = torch.zeros(num_bodies, dtype=torch.bool, device=self._device) - xform_mask = torch.zeros(num_bodies, dtype=torch.bool, device=self._device) + # Reuse buffers when size unchanged to avoid per-call allocations (MR perf). + if num_bodies != self._pose_buf_num_bodies or self._positions_buf is None: + self._pose_buf_num_bodies = num_bodies + self._positions_buf = torch.zeros((num_bodies, 3), dtype=torch.float32, device=self._device) + self._orientations_buf = torch.zeros((num_bodies, 4), dtype=torch.float32, device=self._device) + self._covered_buf = torch.zeros(num_bodies, dtype=torch.bool, device=self._device) + self._xform_mask_buf = torch.zeros(num_bodies, dtype=torch.bool, device=self._device) + else: + self._covered_buf.zero_() + self._xform_mask_buf.zero_() + + positions = self._positions_buf + orientations = self._orientations_buf + covered = self._covered_buf + xform_mask = self._xform_mask_buf # Apply sources in preferred order: articulation, rigid bodies, then USD fallback. articulation_count = self._apply_view_poses( @@ -466,6 +641,7 @@ def _read_poses_from_best_source(self) -> tuple[Any, Any, str, Any] | None: self._warn_once( "rigid-source-unused", "[PhysxSceneDataProvider] RigidBodyView did not provide any body transforms; using fallback sources.", + level=logging.DEBUG, ) if not covered.all(): @@ -490,51 +666,12 @@ def _read_poses_from_best_source(self) -> tuple[Any, Any, str, Any] | None: return positions, orientations, source, xform_mask def _get_set_body_q_kernel(self): - """Get or create the Warp kernel for writing transforms to Newton state.""" - if self._set_body_q_kernel is not None: - return self._set_body_q_kernel - try: - import warp as wp - - @wp.kernel(enable_backward=False) - def _set_body_q( - positions: wp.array(dtype=wp.vec3), - orientations: wp.array(dtype=wp.quatf), - body_q: wp.array(dtype=wp.transformf), - ): - i = wp.tid() - body_q[i] = wp.transformf(positions[i], orientations[i]) - - self._set_body_q_kernel = _set_body_q - return self._set_body_q_kernel - except Exception as exc: - logger.warning(f"[PhysxSceneDataProvider] Warp unavailable for Newton state sync: {exc}") - return None + """Return module-level Warp kernel for writing transforms to Newton state.""" + return _set_body_q_kernel def _get_set_body_q_subset_kernel(self): - """Kernel that writes only body_q at given indices.""" - kernel = getattr(self, "_set_body_q_subset_kernel", None) - if kernel is not None: - return kernel - try: - import warp as wp - - @wp.kernel(enable_backward=False) - def _set_body_q_subset( - positions: wp.array(dtype=wp.vec3), - orientations: wp.array(dtype=wp.quatf), - body_indices: wp.array(dtype=wp.int32), - body_q: wp.array(dtype=wp.transformf), - ): - i = wp.tid() - bi = body_indices[i] - body_q[bi] = wp.transformf(positions[i], orientations[i]) - - self._set_body_q_subset_kernel = _set_body_q_subset - return self._set_body_q_subset_kernel - except Exception as exc: - logger.debug(f"Warp subset kernel: {exc}") - return None + """Return module-level Warp kernel for subset writes.""" + return _set_body_q_subset_kernel # ---- Newton state sync ---------------------------------------------------------------- @@ -548,8 +685,6 @@ def update(self, env_ids: list[int] | None = None) -> None: return try: - import warp as wp - # Re-check env count in case stage population completed after provider construction. self._refresh_newton_model_if_needed() @@ -606,11 +741,22 @@ def update(self, env_ids: list[int] | None = None) -> None: ) def get_newton_model(self) -> Any | None: - """Return Newton model when sync is enabled.""" + """Return Newton model when sync is enabled. + + Returns: + Newton model object, or ``None`` when unavailable. + """ return self._newton_model if self._needs_newton_sync else None def get_newton_model_for_env_ids(self, env_ids: list[int] | None) -> Any | None: - """Return Newton model for a subset of envs if requested.""" + """Return Newton model for selected environments. + + Args: + env_ids: Optional environment ids. ``None`` returns full model. + + Returns: + Full or filtered Newton model, or ``None`` when unavailable. + """ if not self._needs_newton_sync: return None if env_ids is None: @@ -686,6 +832,8 @@ def _create_subset_state(self, body_q_subset): body_q_subset = wp.from_torch(body_q_subset, dtype=wp.transformf) class _SubsetState: + """Minimal state carrier with ``body_q`` field for subset rendering.""" + pass s = _SubsetState() @@ -695,13 +843,22 @@ class _SubsetState: # ---- Public provider API --------------------------------------------------------------- def get_usd_stage(self) -> Any: - """Return the USD stage handle.""" + """Return USD stage handle. + + Returns: + USD stage object. + """ if self._stage is not None: return self._stage return getattr(self._simulation_context, "stage", None) def get_camera_transforms(self) -> dict[str, Any] | None: - """Return per-camera, per-env transforms (positions, orientations).""" + """Return per-camera, per-environment transforms. + + Returns: + Dictionary containing camera order, positions, orientations, and environment count, + or ``None`` when unavailable. + """ if self._stage is None: return None @@ -768,13 +925,21 @@ def get_camera_transforms(self) -> dict[str, Any] | None: return {"order": shared_paths, "positions": positions, "orientations": orientations, "num_envs": num_envs} def get_metadata(self) -> dict[str, Any]: - """Return metadata for visualizers (num_envs, physics_backend, etc.).""" + """Return provider metadata for visualizers and renderers. + + Returns: + Metadata dictionary with backend and environment count. + """ out = dict(self._metadata) out["num_envs"] = self.get_num_envs() return out def get_transforms(self) -> dict[str, Any] | None: - """Return merged body transforms from available PhysX views.""" + """Return merged body transforms from available PhysX views. + + Returns: + Dictionary with positions/orientations, or ``None`` when unavailable. + """ try: result = self._read_poses_from_best_source() if result is None: @@ -792,7 +957,11 @@ def get_transforms(self) -> dict[str, Any] | None: return None def get_velocities(self) -> dict[str, Any] | None: - """Return linear/angular velocities from available PhysX views.""" + """Return linear/angular velocities from available PhysX views. + + Returns: + Dictionary with linear/angular velocities, or ``None`` when unavailable. + """ for source, view in ( ("articulation_view", self._articulation_view), ("rigid_body_view", self._rigid_body_view), @@ -803,5 +972,9 @@ def get_velocities(self) -> dict[str, Any] | None: return None def get_contacts(self) -> dict[str, Any] | None: - """Contacts not yet supported for PhysX provider.""" + """Return contact data for PhysX provider. + + Returns: + ``None`` because contacts are not currently implemented in this provider. + """ return None diff --git a/source/isaaclab_physx/test/assets/test_articulation.py b/source/isaaclab_physx/test/assets/test_articulation.py index 15a16671514e..083c61092c09 100644 --- a/source/isaaclab_physx/test/assets/test_articulation.py +++ b/source/isaaclab_physx/test/assets/test_articulation.py @@ -1680,18 +1680,22 @@ def test_body_root_state(sim, num_articulations, device, with_offset): # Check that fixed base assert articulation.is_fixed_base, "Articulation is not a fixed base" + # Resolve body indices by name (ordering may differ across physics backends) + root_idx = articulation.body_names.index("CenterPivot") + arm_idx = articulation.body_names.index("Arm") + # change center of mass offset from link frame if with_offset: offset = [0.5, 0.0, 0.0] else: offset = [0.0, 0.0, 0.0] - # create com offsets + # create com offsets — apply offset to the Arm body num_bodies = articulation.num_bodies com = wp.to_torch(articulation.root_view.get_coms()) link_offset = [1.0, 0.0, 0.0] # the offset from CenterPivot to Arm frames new_com = torch.tensor(offset, device=device).repeat(num_articulations, 1, 1) - com[:, 1, :3] = new_com.squeeze(-2) + com[:, arm_idx, :3] = new_com.squeeze(-2) articulation.root_view.set_coms( wp.from_torch(com.cpu(), dtype=wp.float32), wp.from_torch(env_idx.cpu(), dtype=wp.int32) ) @@ -1730,10 +1734,10 @@ def test_body_root_state(sim, num_articulations, device, with_offset): vx = -(link_offset[0]) * joint_vel * torch.sin(joint_pos) vy = torch.zeros(num_articulations, 1, 1, device=device) vz = (link_offset[0]) * joint_vel * torch.cos(joint_pos) - lin_vel_gt[:, 1, :] = torch.cat([vx, vy, vz], dim=-1).squeeze(-2) + lin_vel_gt[:, arm_idx, :] = torch.cat([vx, vy, vz], dim=-1).squeeze(-2) # linear velocity of root link should be zero - torch.testing.assert_close(lin_vel_gt[:, 0, :], root_link_vel_w[..., :3], atol=1e-3, rtol=1e-1) + torch.testing.assert_close(lin_vel_gt[:, root_idx, :], root_link_vel_w[..., :3], atol=1e-3, rtol=1e-1) # linear velocity of pendulum link should be torch.testing.assert_close(lin_vel_gt, body_link_vel_w[..., :3], atol=1e-3, rtol=1e-1) @@ -1747,16 +1751,16 @@ def test_body_root_state(sim, num_articulations, device, with_offset): px = (link_offset[0] + offset[0]) * torch.cos(joint_pos) py = torch.zeros(num_articulations, 1, 1, device=device) pz = (link_offset[0] + offset[0]) * torch.sin(joint_pos) - pos_gt[:, 1, :] = torch.cat([px, py, pz], dim=-1).squeeze(-2) + pos_gt[:, arm_idx, :] = torch.cat([px, py, pz], dim=-1).squeeze(-2) pos_gt += env_pos.unsqueeze(-2).repeat(1, num_bodies, 1) - torch.testing.assert_close(pos_gt[:, 0, :], root_com_pose_w[..., :3], atol=1e-3, rtol=1e-1) + torch.testing.assert_close(pos_gt[:, root_idx, :], root_com_pose_w[..., :3], atol=1e-3, rtol=1e-1) torch.testing.assert_close(pos_gt, body_com_pose_w[..., :3], atol=1e-3, rtol=1e-1) # orientation com_quat_b = wp.to_torch(articulation.data.body_com_quat_b) com_quat_w = math_utils.quat_mul(body_link_pose_w[..., 3:], com_quat_b) torch.testing.assert_close(com_quat_w, body_com_pose_w[..., 3:]) - torch.testing.assert_close(com_quat_w[:, 0, :], root_com_pose_w[..., 3:]) + torch.testing.assert_close(com_quat_w[:, root_idx, :], root_com_pose_w[..., 3:]) # angular velocity should be the same for both COM and link frames torch.testing.assert_close(root_com_vel_w[..., 3:], root_link_vel_w[..., 3:]) diff --git a/source/isaaclab_rl/config/extension.toml b/source/isaaclab_rl/config/extension.toml index 79f171d3398c..35ce26490606 100644 --- a/source/isaaclab_rl/config/extension.toml +++ b/source/isaaclab_rl/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "0.4.7" +version = "0.5.0" # Description title = "Isaac Lab RL" diff --git a/source/isaaclab_rl/docs/CHANGELOG.rst b/source/isaaclab_rl/docs/CHANGELOG.rst index 6fe0be78d0a5..9666e7214e81 100644 --- a/source/isaaclab_rl/docs/CHANGELOG.rst +++ b/source/isaaclab_rl/docs/CHANGELOG.rst @@ -1,6 +1,31 @@ Changelog --------- +0.5.0 (2026-3-04) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ +* Added function to handle deprecated RSL-RL configurations and automatically convert them to the new format compatible + with RSL-RL 4.0 and 5.0. +* Added new configuration classes "MLPModelCfg", "RNNModelCfg", and "CNNModelCfg", and "DistributionCfg" for the new + versions of RSL-RL. +* Added "check_for_nan" and "share_cnn_encoders" parameters to the configuration classes for RSL-RL 5.0. +* Added recurrent configurations for the "Isaac-Velocity-Flat-Anymal-D-v0" task for RSL-RL. to run RSL-RL CI. + +Changed +^^^^^^^ +* Adapted RSL-RL's train.py and play.py scripts to work with both old and the new versions of RSL-RL. + +Deprecated +^^^^^^^^^^ +* Deprecated old configuration classes "RslRlDistillationStudentTeacherCfg", + "RslRlDistillationStudentTeacherRecurrentCfg", "RslRlPpoActorCriticCfg", and "RslRlPpoActorCriticRecurrentCfg" in + favor of the new "MLPModelCfg", "RNNModelCfg", and "CNNModelCfg" configuration classes for RSL-RL 4.0. +* Deprecated old parameters "stochastic", "init_noise_std", "noise_std_type", amd "state_dependent_std" in favor of the + new "DistributionCfg" configuration class for RSL-RL 5.0. + + 0.4.7 (2025-12-29) ~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab_rl/isaaclab_rl/rsl_rl/__init__.pyi b/source/isaaclab_rl/isaaclab_rl/rsl_rl/__init__.pyi index 16010e41754f..3eae1d119b5b 100644 --- a/source/isaaclab_rl/isaaclab_rl/rsl_rl/__init__.pyi +++ b/source/isaaclab_rl/isaaclab_rl/rsl_rl/__init__.pyi @@ -10,11 +10,15 @@ __all__ = [ "RslRlDistillationStudentTeacherRecurrentCfg", "export_policy_as_jit", "export_policy_as_onnx", + "handle_deprecated_rsl_rl_cfg", "RslRlBaseRunnerCfg", + "RslRlCNNModelCfg", + "RslRlMLPModelCfg", "RslRlOnPolicyRunnerCfg", "RslRlPpoActorCriticCfg", "RslRlPpoActorCriticRecurrentCfg", "RslRlPpoAlgorithmCfg", + "RslRlRNNModelCfg", "RslRlRndCfg", "RslRlSymmetryCfg", "RslRlVecEnvWrapper", @@ -29,11 +33,15 @@ from .distillation_cfg import ( from .exporter import export_policy_as_jit, export_policy_as_onnx from .rl_cfg import ( RslRlBaseRunnerCfg, + RslRlCNNModelCfg, + RslRlMLPModelCfg, RslRlOnPolicyRunnerCfg, RslRlPpoActorCriticCfg, RslRlPpoActorCriticRecurrentCfg, RslRlPpoAlgorithmCfg, + RslRlRNNModelCfg, ) from .rnd_cfg import RslRlRndCfg from .symmetry_cfg import RslRlSymmetryCfg +from .utils import handle_deprecated_rsl_rl_cfg from .vecenv_wrapper import RslRlVecEnvWrapper diff --git a/source/isaaclab_rl/isaaclab_rl/rsl_rl/distillation_cfg.py b/source/isaaclab_rl/isaaclab_rl/rsl_rl/distillation_cfg.py index 1a631eeffdaa..8e81aae4d1f2 100644 --- a/source/isaaclab_rl/isaaclab_rl/rsl_rl/distillation_cfg.py +++ b/source/isaaclab_rl/isaaclab_rl/rsl_rl/distillation_cfg.py @@ -10,16 +10,79 @@ from isaaclab.utils import configclass -from .rl_cfg import RslRlBaseRunnerCfg +from .rl_cfg import RslRlBaseRunnerCfg, RslRlMLPModelCfg + +############################ +# Algorithm configurations # +############################ + + +@configclass +class RslRlDistillationAlgorithmCfg: + """Configuration for the distillation algorithm.""" + + class_name: str = "Distillation" + """The algorithm class name. Default is Distillation.""" + + num_learning_epochs: int = MISSING + """The number of updates performed with each sample.""" + + learning_rate: float = MISSING + """The learning rate for the student policy.""" + + gradient_length: int = MISSING + """The number of environment steps the gradient flows back.""" + + max_grad_norm: None | float = None + """The maximum norm the gradient is clipped to.""" + + optimizer: Literal["adam", "adamw", "sgd", "rmsprop"] = "adam" + """The optimizer to use for the student policy.""" + + loss_type: Literal["mse", "huber"] = "mse" + """The loss type to use for the student policy.""" + ######################### -# Policy configurations # +# Runner configurations # ######################### +@configclass +class RslRlDistillationRunnerCfg(RslRlBaseRunnerCfg): + """Configuration of the runner for distillation algorithms.""" + + class_name: str = "DistillationRunner" + """The runner class name. Default is DistillationRunner.""" + + student: RslRlMLPModelCfg = MISSING + """The student configuration.""" + + teacher: RslRlMLPModelCfg = MISSING + """The teacher configuration.""" + + algorithm: RslRlDistillationAlgorithmCfg = MISSING + """The algorithm configuration.""" + + policy: RslRlDistillationStudentTeacherCfg = MISSING + """The policy configuration. + + For rsl-rl >= 4.0.0, this configuration is deprecated. Please use `student` and `teacher` model configurations + instead. + """ + + +############################# +# Deprecated configurations # +############################# + + @configclass class RslRlDistillationStudentTeacherCfg: - """Configuration for the distillation student-teacher networks.""" + """Configuration for the distillation student-teacher networks. + + For rsl-rl >= 4.0.0, this configuration is deprecated. Please use `RslRlMLPModelCfg` instead. + """ class_name: str = "StudentTeacher" """The policy class name. Default is StudentTeacher.""" @@ -48,7 +111,10 @@ class RslRlDistillationStudentTeacherCfg: @configclass class RslRlDistillationStudentTeacherRecurrentCfg(RslRlDistillationStudentTeacherCfg): - """Configuration for the distillation student-teacher recurrent networks.""" + """Configuration for the distillation student-teacher recurrent networks. + + For rsl-rl >= 4.0.0, this configuration is deprecated. Please use `RslRlRNNModelCfg` instead. + """ class_name: str = "StudentTeacherRecurrent" """The policy class name. Default is StudentTeacherRecurrent.""" @@ -64,53 +130,3 @@ class RslRlDistillationStudentTeacherRecurrentCfg(RslRlDistillationStudentTeache teacher_recurrent: bool = MISSING """Whether the teacher network is recurrent too.""" - - -############################ -# Algorithm configurations # -############################ - - -@configclass -class RslRlDistillationAlgorithmCfg: - """Configuration for the distillation algorithm.""" - - class_name: str = "Distillation" - """The algorithm class name. Default is Distillation.""" - - num_learning_epochs: int = MISSING - """The number of updates performed with each sample.""" - - learning_rate: float = MISSING - """The learning rate for the student policy.""" - - gradient_length: int = MISSING - """The number of environment steps the gradient flows back.""" - - max_grad_norm: None | float = None - """The maximum norm the gradient is clipped to.""" - - optimizer: Literal["adam", "adamw", "sgd", "rmsprop"] = "adam" - """The optimizer to use for the student policy.""" - - loss_type: Literal["mse", "huber"] = "mse" - """The loss type to use for the student policy.""" - - -######################### -# Runner configurations # -######################### - - -@configclass -class RslRlDistillationRunnerCfg(RslRlBaseRunnerCfg): - """Configuration of the runner for distillation algorithms.""" - - class_name: str = "DistillationRunner" - """The runner class name. Default is DistillationRunner.""" - - policy: RslRlDistillationStudentTeacherCfg = MISSING - """The policy configuration.""" - - algorithm: RslRlDistillationAlgorithmCfg = MISSING - """The algorithm configuration.""" diff --git a/source/isaaclab_rl/isaaclab_rl/rsl_rl/rl_cfg.py b/source/isaaclab_rl/isaaclab_rl/rsl_rl/rl_cfg.py index 7be991174dec..0b9f14bf6b2e 100644 --- a/source/isaaclab_rl/isaaclab_rl/rsl_rl/rl_cfg.py +++ b/source/isaaclab_rl/isaaclab_rl/rsl_rl/rl_cfg.py @@ -14,48 +14,91 @@ from .symmetry_cfg import RslRlSymmetryCfg ######################### -# Policy configurations # +# Model configurations # ######################### @configclass -class RslRlPpoActorCriticCfg: - """Configuration for the PPO actor-critic networks.""" +class RslRlMLPModelCfg: + """Configuration for the MLP model.""" - class_name: str = "ActorCritic" - """The policy class name. Default is ActorCritic.""" + class_name: str = "MLPModel" + """The model class name. Default is MLPModel.""" - init_noise_std: float = MISSING - """The initial noise standard deviation for the policy.""" + hidden_dims: list[int] = MISSING + """The hidden dimensions of the MLP network.""" - noise_std_type: Literal["scalar", "log"] = "scalar" - """The type of noise standard deviation for the policy. Default is scalar.""" + activation: str = MISSING + """The activation function for the MLP network.""" - state_dependent_std: bool = False - """Whether to use state-dependent standard deviation for the policy. Default is False.""" + obs_normalization: bool = False + """Whether to normalize the observation for the model. Default is False.""" - actor_obs_normalization: bool = MISSING - """Whether to normalize the observation for the actor network.""" + distribution_cfg: DistributionCfg | None = None + """The configuration for the output distribution. Default is None, in which case no distribution is used.""" - critic_obs_normalization: bool = MISSING - """Whether to normalize the observation for the critic network.""" + @configclass + class DistributionCfg: + """Configuration for the output distribution.""" - actor_hidden_dims: list[int] = MISSING - """The hidden dimensions of the actor network.""" + class_name: str = MISSING + """The distribution class name.""" - critic_hidden_dims: list[int] = MISSING - """The hidden dimensions of the critic network.""" + @configclass + class GaussianDistributionCfg(DistributionCfg): + """Configuration for the Gaussian output distribution.""" - activation: str = MISSING - """The activation function for the actor and critic networks.""" + class_name: str = "GaussianDistribution" + """The distribution class name. Default is GaussianDistribution.""" + + init_std: float = MISSING + """The initial standard deviation of the output distribution.""" + + std_type: Literal["scalar", "log"] = "scalar" + """The parameterization type of the output distribution's standard deviation. Default is scalar.""" + + @configclass + class HeteroscedasticGaussianDistributionCfg(GaussianDistributionCfg): + """Configuration for the heteroscedastic Gaussian output distribution.""" + + class_name: str = "HeteroscedasticGaussianDistribution" + """The distribution class name. Default is HeteroscedasticGaussianDistribution.""" + + stochastic: bool = MISSING + """Whether the model output is stochastic. + + For rsl-rl >= 5.0.0, this configuration is is deprecated. Please use `distribution_cfg` instead and set it to None + for deterministic output or to a valid configuration class, e.g., `GaussianDistributionCfg` for stochastic output. + """ + + init_noise_std: float = MISSING + """The initial noise standard deviation for the model. + + For rsl-rl >= 5.0.0, this configuration is is deprecated. Please use `distribution_cfg` instead and use the + `init_std` field of the distribution configuration to specify the initial noise standard deviation. + """ + + noise_std_type: Literal["scalar", "log"] = "scalar" + """The type of noise standard deviation for the model. Default is scalar. + + For rsl-rl >= 5.0.0, this configuration is is deprecated. Please use `distribution_cfg` instead and use the + `std_type` field of the distribution configuration to specify the type of noise standard deviation. + """ + + state_dependent_std: bool = False + """Whether to use state-dependent standard deviation for the policy. Default is False. + + For rsl-rl >= 5.0.0, this configuration is is deprecated. Please use `distribution_cfg` instead and use + the `HeteroscedasticGaussianDistributionCfg` if state-dependent standard deviation is desired. + """ @configclass -class RslRlPpoActorCriticRecurrentCfg(RslRlPpoActorCriticCfg): - """Configuration for the PPO actor-critic networks with recurrent layers.""" +class RslRlRNNModelCfg(RslRlMLPModelCfg): + """Configuration for RNN model.""" - class_name: str = "ActorCriticRecurrent" - """The policy class name. Default is ActorCriticRecurrent.""" + class_name: str = "RNNModel" + """The model class name. Default is RNNModel.""" rnn_type: str = MISSING """The type of RNN to use. Either "lstm" or "gru".""" @@ -67,6 +110,49 @@ class RslRlPpoActorCriticRecurrentCfg(RslRlPpoActorCriticCfg): """The number of RNN layers.""" +@configclass +class RslRlCNNModelCfg(RslRlMLPModelCfg): + """Configuration for CNN model.""" + + class_name: str = "CNNModel" + """The model class name. Default is CNNModel.""" + + @configclass + class CNNCfg: + output_channels: tuple[int] | list[int] = MISSING + """The number of output channels for each convolutional layer for the CNN.""" + + kernel_size: int | tuple[int] | list[int] = MISSING + """The kernel size for the CNN.""" + + stride: int | tuple[int] | list[int] = 1 + """The stride for the CNN.""" + + dilation: int | tuple[int] | list[int] = 1 + """The dilation for the CNN.""" + + padding: Literal["none", "zeros", "reflect", "replicate", "circular"] = "none" + """The padding for the CNN.""" + + norm: Literal["none", "batch", "layer"] | tuple[str] | list[str] = "none" + """The normalization for the CNN.""" + + activation: str = MISSING + """The activation function for the CNN.""" + + max_pool: bool | tuple[bool] | list[bool] = False + """Whether to use max pooling for the CNN.""" + + global_pool: Literal["none", "max", "avg"] = "none" + """The global pooling for the CNN.""" + + flatten: bool = True + """Whether to flatten the output of the CNN.""" + + cnn_cfg: CNNCfg = MISSING + """The configuration for the CNN(s).""" + + ############################ # Algorithm configurations # ############################ @@ -106,6 +192,9 @@ class RslRlPpoAlgorithmCfg: max_grad_norm: float = MISSING """The maximum gradient norm.""" + optimizer: Literal["adam", "adamw", "sgd", "rmsprop"] = "adam" + """The optimizer to use.""" + value_loss_coef: float = MISSING """The coefficient for the value loss.""" @@ -122,6 +211,9 @@ class RslRlPpoAlgorithmCfg: Otherwise, the advantage is normalized over the entire collected trajectories. """ + share_cnn_encoders: bool = False + """Whether to share the CNN networks between actor and critic, in case CNNModels are used. Defaults to False.""" + rnd_cfg: RslRlRndCfg | None = None """The RND configuration. Default is None, in which case RND is not used.""" @@ -150,10 +242,11 @@ class RslRlBaseRunnerCfg: max_iterations: int = MISSING """The maximum number of iterations.""" - empirical_normalization: bool | None = None + empirical_normalization: bool = MISSING """This parameter is deprecated and will be removed in the future. - Use `actor_obs_normalization` and `critic_obs_normalization` instead. + For rsl-rl < 4.0.0, use `actor_obs_normalization` and `critic_obs_normalization` of the policy instead. + For rsl-rl >= 4.0.0, use `obs_normalization` of the model instead. """ obs_groups: dict[str, list[str]] = MISSING @@ -168,11 +261,11 @@ class RslRlBaseRunnerCfg: .. code-block:: python obs_groups = { - "policy": ["policy", "images"], + "actor": ["policy", "images"], "critic": ["policy", "privileged"], } - This way, the policy will receive the "policy" and "images" observations, and the critic will + This way, the actor will receive the "policy" and "images" observations, and the critic will receive the "policy" and "privileged" observations. For more details, please check ``vec_env.py`` in the rsl_rl library. @@ -185,6 +278,9 @@ class RslRlBaseRunnerCfg: This clipping is performed inside the :class:`RslRlVecEnvWrapper` wrapper. """ + check_for_nan: bool = True + """Whether to check for NaN values coming from the environment.""" + save_interval: int = MISSING """The number of iterations between saves.""" @@ -234,8 +330,78 @@ class RslRlOnPolicyRunnerCfg(RslRlBaseRunnerCfg): class_name: str = "OnPolicyRunner" """The runner class name. Default is OnPolicyRunner.""" - policy: RslRlPpoActorCriticCfg = MISSING - """The policy configuration.""" + actor: RslRlMLPModelCfg = MISSING + """The actor configuration.""" + + critic: RslRlMLPModelCfg = MISSING + """The critic configuration.""" algorithm: RslRlPpoAlgorithmCfg = MISSING """The algorithm configuration.""" + + policy: RslRlPpoActorCriticCfg = MISSING + """The policy configuration. + + For rsl-rl >= 4.0.0, this configuration is is deprecated. Please use `actor` and `critic` model configurations + instead. + """ + + +############################# +# Deprecated configurations # +############################# + + +@configclass +class RslRlPpoActorCriticCfg: + """Configuration for the PPO actor-critic networks. + + For rsl-rl >= 4.0.0, this configuration is deprecated. Please use `RslRlMLPModelCfg` instead. + """ + + class_name: str = "ActorCritic" + """The policy class name. Default is ActorCritic.""" + + init_noise_std: float = MISSING + """The initial noise standard deviation for the policy.""" + + noise_std_type: Literal["scalar", "log"] = "scalar" + """The type of noise standard deviation for the policy. Default is scalar.""" + + state_dependent_std: bool = False + """Whether to use state-dependent standard deviation for the policy. Default is False.""" + + actor_obs_normalization: bool = MISSING + """Whether to normalize the observation for the actor network.""" + + critic_obs_normalization: bool = MISSING + """Whether to normalize the observation for the critic network.""" + + actor_hidden_dims: list[int] = MISSING + """The hidden dimensions of the actor network.""" + + critic_hidden_dims: list[int] = MISSING + """The hidden dimensions of the critic network.""" + + activation: str = MISSING + """The activation function for the actor and critic networks.""" + + +@configclass +class RslRlPpoActorCriticRecurrentCfg(RslRlPpoActorCriticCfg): + """Configuration for the PPO actor-critic networks with recurrent layers. + + For rsl-rl >= 4.0.0, this configuration is deprecated. Please use `RslRlRNNModelCfg` instead. + """ + + class_name: str = "ActorCriticRecurrent" + """The policy class name. Default is ActorCriticRecurrent.""" + + rnn_type: str = MISSING + """The type of RNN to use. Either "lstm" or "gru".""" + + rnn_hidden_dim: int = MISSING + """The dimension of the RNN layers.""" + + rnn_num_layers: int = MISSING + """The number of RNN layers.""" diff --git a/source/isaaclab_rl/isaaclab_rl/rsl_rl/utils.py b/source/isaaclab_rl/isaaclab_rl/rsl_rl/utils.py new file mode 100644 index 000000000000..3d97765a667a --- /dev/null +++ b/source/isaaclab_rl/isaaclab_rl/rsl_rl/utils.py @@ -0,0 +1,267 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from dataclasses import MISSING +from typing import TYPE_CHECKING + +from packaging import version + +if TYPE_CHECKING: + from isaaclab_rl.rsl_rl import RslRlBaseRunnerCfg + + +_V4_0_0 = version.parse("4.0.0") +_V5_0_0 = version.parse("5.0.0") +_MODEL_CFG_NAMES = ("actor", "critic", "student", "teacher") + + +def handle_deprecated_rsl_rl_cfg(agent_cfg: RslRlBaseRunnerCfg, installed_version) -> RslRlBaseRunnerCfg: + """Handle deprecated RSL-RL configurations across version boundaries. + + This function mutates ``agent_cfg`` to keep configurations compatible with the installed ``rsl-rl`` version: + + - For ``rsl-rl < 4.0.0``, ``policy`` is required; new model configs (``actor``, ``critic``, ``student``, + ``teacher``) are ignored and cleared. + - For ``rsl-rl >= 4.0.0``, deprecated ``policy`` can be used to infer missing model configs, then ``policy`` is + cleared. + - For ``rsl-rl >= 5.0.0``, legacy stochastic parameters are migrated to ``distribution_cfg`` when needed; for + ``4.0.0 <= rsl-rl < 5.0.0``, those legacy parameters are validated instead. + + Raises: + ValueError: If required legacy parameters are missing for the selected ``rsl-rl`` version. + """ + installed_version = version.parse(installed_version) + + # Handle configurations for rsl-rl < 4.0.0 + if installed_version < _V4_0_0: + # exit if no policy configuration is present + if not hasattr(agent_cfg, "policy") or _is_missing(agent_cfg.policy): + raise ValueError( + "The `policy` configuration is required for rsl-rl < 4.0.0. Please specify the `policy` configuration" + " or update rsl-rl." + ) + + # handle deprecated obs_normalization argument + if _has_non_missing_attr(agent_cfg, "empirical_normalization"): + _handle_empirical_normalization(agent_cfg.policy, agent_cfg) + + # remove optimizer argument for PPO only available in rsl-rl >= 4.0.0 + from isaaclab_rl.rsl_rl import RslRlPpoAlgorithmCfg + + if hasattr(agent_cfg.algorithm, "optimizer") and isinstance(agent_cfg.algorithm, RslRlPpoAlgorithmCfg): + if agent_cfg.algorithm.optimizer != "adam": + print( + "[WARNING]: The `optimizer` parameter for PPO is only available for rsl-rl >= 4.0.0. Consider" + " updating rsl-rl to use this feature. Defaulting to `adam` optimizer." + ) + del agent_cfg.algorithm.optimizer + + # warn about model configurations only used in rsl-rl >= 4.0.0 + for model_name in _MODEL_CFG_NAMES: + if _has_non_missing_attr(agent_cfg, model_name): + _clear_new_model_cfg(agent_cfg, model_name) + + # Handle configurations for rsl-rl >= 4.0.0 + else: + # Handle deprecated policy configuration + if _has_non_missing_attr(agent_cfg, "policy"): + print( + "[WARNING]: The `policy` configuration is deprecated for rsl-rl >= 4.0.0. Please use, e.g., `actor` and" + " `critic` model configurations instead." + ) + + # handle deprecated obs_normalization argument + if _has_non_missing_attr(agent_cfg, "empirical_normalization"): + _handle_empirical_normalization(agent_cfg.policy, agent_cfg) + + # import old and new config classes + from isaaclab_rl.rsl_rl import ( + RslRlDistillationStudentTeacherCfg, + RslRlDistillationStudentTeacherRecurrentCfg, + RslRlMLPModelCfg, + RslRlPpoActorCriticCfg, + RslRlPpoActorCriticRecurrentCfg, + RslRlRNNModelCfg, + ) + + # set actor model configuration if missing + if hasattr(agent_cfg, "actor") and _is_missing(agent_cfg.actor): + print("[WARNING]: The `policy` configuration is used to infer the `actor` model configuration.") + if type(agent_cfg.policy) is RslRlPpoActorCriticCfg: + agent_cfg.actor = RslRlMLPModelCfg( + hidden_dims=agent_cfg.policy.actor_hidden_dims, + activation=agent_cfg.policy.activation, + obs_normalization=agent_cfg.policy.actor_obs_normalization, + stochastic=True, + init_noise_std=agent_cfg.policy.init_noise_std, + noise_std_type=agent_cfg.policy.noise_std_type, + state_dependent_std=agent_cfg.policy.state_dependent_std, + ) + elif type(agent_cfg.policy) is RslRlPpoActorCriticRecurrentCfg: + agent_cfg.actor = RslRlRNNModelCfg( + hidden_dims=agent_cfg.policy.actor_hidden_dims, + activation=agent_cfg.policy.activation, + obs_normalization=agent_cfg.policy.actor_obs_normalization, + stochastic=True, + init_noise_std=agent_cfg.policy.init_noise_std, + noise_std_type=agent_cfg.policy.noise_std_type, + state_dependent_std=agent_cfg.policy.state_dependent_std, + rnn_type=agent_cfg.policy.rnn_type, + rnn_hidden_dim=agent_cfg.policy.rnn_hidden_dim, + rnn_num_layers=agent_cfg.policy.rnn_num_layers, + ) + # set critic model configuration if missing + if hasattr(agent_cfg, "critic") and _is_missing(agent_cfg.critic): + print("[WARNING]: The `policy` configuration is used to infer the `critic` model configuration.") + if type(agent_cfg.policy) is RslRlPpoActorCriticCfg: + agent_cfg.critic = RslRlMLPModelCfg( + hidden_dims=agent_cfg.policy.critic_hidden_dims, + activation=agent_cfg.policy.activation, + obs_normalization=agent_cfg.policy.critic_obs_normalization, + stochastic=False, + ) + elif type(agent_cfg.policy) is RslRlPpoActorCriticRecurrentCfg: + agent_cfg.critic = RslRlRNNModelCfg( + hidden_dims=agent_cfg.policy.critic_hidden_dims, + activation=agent_cfg.policy.activation, + obs_normalization=agent_cfg.policy.critic_obs_normalization, + stochastic=False, + rnn_type=agent_cfg.policy.rnn_type, + rnn_hidden_dim=agent_cfg.policy.rnn_hidden_dim, + rnn_num_layers=agent_cfg.policy.rnn_num_layers, + ) + # set student model configuration if missing + if hasattr(agent_cfg, "student") and _is_missing(agent_cfg.student): + print("[WARNING]: The `policy` configuration is used to infer the `student` model configuration.") + if type(agent_cfg.policy) is RslRlDistillationStudentTeacherCfg: + agent_cfg.student = RslRlMLPModelCfg( + hidden_dims=agent_cfg.policy.student_hidden_dims, + activation=agent_cfg.policy.activation, + obs_normalization=agent_cfg.policy.student_obs_normalization, + stochastic=True, + init_noise_std=agent_cfg.policy.init_noise_std, + noise_std_type=agent_cfg.policy.noise_std_type, + ) + elif type(agent_cfg.policy) is RslRlDistillationStudentTeacherRecurrentCfg: + agent_cfg.student = RslRlRNNModelCfg( + hidden_dims=agent_cfg.policy.student_hidden_dims, + activation=agent_cfg.policy.activation, + obs_normalization=agent_cfg.policy.student_obs_normalization, + stochastic=True, + init_noise_std=agent_cfg.policy.init_noise_std, + noise_std_type=agent_cfg.policy.noise_std_type, + rnn_type=agent_cfg.policy.rnn_type, + rnn_hidden_dim=agent_cfg.policy.rnn_hidden_dim, + rnn_num_layers=agent_cfg.policy.rnn_num_layers, + ) + # set teacher model configuration if missing + if hasattr(agent_cfg, "teacher") and _is_missing(agent_cfg.teacher): + print("[WARNING]: The `policy` configuration is used to infer the `teacher` model configuration.") + if type(agent_cfg.policy) is RslRlDistillationStudentTeacherCfg: + agent_cfg.teacher = RslRlMLPModelCfg( + hidden_dims=agent_cfg.policy.teacher_hidden_dims, + activation=agent_cfg.policy.activation, + obs_normalization=agent_cfg.policy.teacher_obs_normalization, + stochastic=True, + init_noise_std=0.0, + ) + elif type(agent_cfg.policy) is RslRlDistillationStudentTeacherRecurrentCfg: + agent_cfg.teacher = RslRlRNNModelCfg( + hidden_dims=agent_cfg.policy.teacher_hidden_dims, + activation=agent_cfg.policy.activation, + obs_normalization=agent_cfg.policy.teacher_obs_normalization, + stochastic=True, + init_noise_std=0.0, + rnn_type=agent_cfg.policy.rnn_type, + rnn_hidden_dim=agent_cfg.policy.rnn_hidden_dim, + rnn_num_layers=agent_cfg.policy.rnn_num_layers, + ) + + # remove deprecated policy configuration + agent_cfg.policy = MISSING + + # Handle new distribution configuration + if installed_version < _V5_0_0: + for model_name in _MODEL_CFG_NAMES: + if _has_non_missing_attr(agent_cfg, model_name): + _validate_old_stochastic_cfg(getattr(agent_cfg, model_name)) + else: # rsl-rl >= 5.0.0 + # import new distribution config classes + from isaaclab_rl.rsl_rl import RslRlMLPModelCfg + + for model_name in _MODEL_CFG_NAMES: + if _has_non_missing_attr(agent_cfg, model_name): + _update_distribution_cfg(getattr(agent_cfg, model_name), RslRlMLPModelCfg) + + return agent_cfg + + +def _is_missing(value) -> bool: + return isinstance(value, type(MISSING)) + + +def _has_non_missing_attr(obj, attr_name: str) -> bool: + return hasattr(obj, attr_name) and not _is_missing(getattr(obj, attr_name)) + + +def _handle_empirical_normalization(policy_cfg, agent_cfg): + print( + "[WARNING]: The `empirical_normalization` parameter is deprecated. Please set `actor_obs_normalization` and" + " `critic_obs_normalization` as part of the `policy` configuration instead." + ) + if _is_missing(policy_cfg.actor_obs_normalization): + policy_cfg.actor_obs_normalization = agent_cfg.empirical_normalization + if _is_missing(policy_cfg.critic_obs_normalization): + policy_cfg.critic_obs_normalization = agent_cfg.empirical_normalization + agent_cfg.empirical_normalization = MISSING + + +def _clear_new_model_cfg(agent_cfg, model_name: str): + print( + f"[WARNING]: The `{model_name}` model configuration is only used for rsl-rl >= 4.0.0. Consider updating rsl-rl" + " or use the `policy` configuration for rsl-rl < 4.0.0." + ) + setattr(agent_cfg, model_name, MISSING) + + +def _validate_old_stochastic_cfg(model_cfg): + if not hasattr(model_cfg, "stochastic") or _is_missing(model_cfg.stochastic): + raise ValueError( + "Please parameterize the output distribution using the old parameters `stochastic`, `init_noise_std`," + " `noise_std_type`, and `state_dependent_std` or update rsl-rl." + ) + # remove new distribution configuration + if hasattr(model_cfg, "distribution_cfg"): + del model_cfg.distribution_cfg + + +def _update_distribution_cfg(model_cfg, rsl_rl_mlp_model_cfg_cls): + if model_cfg.distribution_cfg is not None: + pass # new distribution configuration is used, no need to handle deprecated configurations + elif model_cfg.stochastic is True: # distribution config is None but stochastic output is requested + print( + "[WARNING]: The `distribution_cfg` configuration is now used to specify the output distribution for" + " stochastic policies. Consider updating the configuration to use `distribution_cfg` instead of" + " `stochastic`, `init_noise_std`, `noise_std_type`, and `state_dependent_std` parameters." + ) + if model_cfg.state_dependent_std is False: # gaussian distribution + model_cfg.distribution_cfg = rsl_rl_mlp_model_cfg_cls.GaussianDistributionCfg( + init_std=model_cfg.init_noise_std, std_type=model_cfg.noise_std_type + ) + elif model_cfg.state_dependent_std is True: # heteroscedastic gaussian distribution + model_cfg.distribution_cfg = rsl_rl_mlp_model_cfg_cls.HeteroscedasticGaussianDistributionCfg( + init_std=model_cfg.init_noise_std, std_type=model_cfg.noise_std_type + ) + # remove deprecated stochastic parameters + if hasattr(model_cfg, "stochastic"): + del model_cfg.stochastic + if hasattr(model_cfg, "init_noise_std"): + del model_cfg.init_noise_std + if hasattr(model_cfg, "noise_std_type"): + del model_cfg.noise_std_type + if hasattr(model_cfg, "state_dependent_std"): + del model_cfg.state_dependent_std diff --git a/source/isaaclab_rl/setup.py b/source/isaaclab_rl/setup.py index 5476836c6c23..df42b2728154 100644 --- a/source/isaaclab_rl/setup.py +++ b/source/isaaclab_rl/setup.py @@ -46,7 +46,7 @@ "gym", "standard-distutils", ], # rl-games still needs gym :( - "rsl-rl": ["rsl-rl-lib==3.1.2", "onnxscript>=0.5"], # linux aarch 64 requires manual onnxscript installation + "rsl-rl": ["rsl-rl-lib==5.0.1", "onnxscript>=0.5"], # linux aarch 64 requires manual onnxscript installation } # Add the names with hyphens as aliases for convenience EXTRAS_REQUIRE["rl_games"] = EXTRAS_REQUIRE["rl-games"] diff --git a/source/isaaclab_rl/test/test_rsl_rl_cfg_deprecation.py b/source/isaaclab_rl/test/test_rsl_rl_cfg_deprecation.py new file mode 100644 index 000000000000..ad8e79190ce4 --- /dev/null +++ b/source/isaaclab_rl/test/test_rsl_rl_cfg_deprecation.py @@ -0,0 +1,533 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Tests for handle_deprecated_rsl_rl_cfg across rsl-rl version boundaries.""" + +"""Launch Isaac Sim Simulator first.""" + +from isaaclab.app import AppLauncher + +app_launcher = AppLauncher(headless=True) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +from dataclasses import MISSING + +import pytest + +from isaaclab_rl.rsl_rl import ( + RslRlDistillationAlgorithmCfg, + RslRlDistillationRunnerCfg, + RslRlDistillationStudentTeacherCfg, + RslRlDistillationStudentTeacherRecurrentCfg, + RslRlMLPModelCfg, + RslRlOnPolicyRunnerCfg, + RslRlPpoActorCriticCfg, + RslRlPpoActorCriticRecurrentCfg, + RslRlPpoAlgorithmCfg, + RslRlRNNModelCfg, +) +from isaaclab_rl.rsl_rl.utils import _is_missing, handle_deprecated_rsl_rl_cfg + + +def _ppo_algo(): + return RslRlPpoAlgorithmCfg( + num_learning_epochs=5, + num_mini_batches=4, + learning_rate=1e-3, + schedule="adaptive", + gamma=0.99, + lam=0.95, + entropy_coef=0.01, + desired_kl=0.01, + max_grad_norm=1.0, + value_loss_coef=1.0, + use_clipped_value_loss=True, + clip_param=0.2, + ) + + +def _distillation_algo(): + return RslRlDistillationAlgorithmCfg(num_learning_epochs=5, learning_rate=1e-3, gradient_length=1) + + +def _ppo_mlp_policy(): + return RslRlPpoActorCriticCfg( + init_noise_std=1.0, + actor_obs_normalization=False, + critic_obs_normalization=False, + actor_hidden_dims=[256, 256], + critic_hidden_dims=[128, 128], + activation="elu", + ) + + +def _ppo_rnn_policy(): + return RslRlPpoActorCriticRecurrentCfg( + init_noise_std=1.0, + actor_obs_normalization=False, + critic_obs_normalization=False, + actor_hidden_dims=[256, 256], + critic_hidden_dims=[128, 128], + activation="elu", + rnn_type="lstm", + rnn_hidden_dim=256, + rnn_num_layers=1, + ) + + +def _distillation_mlp_policy(): + return RslRlDistillationStudentTeacherCfg( + init_noise_std=1.0, + noise_std_type="scalar", + student_obs_normalization=False, + teacher_obs_normalization=True, + student_hidden_dims=[256, 256], + teacher_hidden_dims=[128, 128], + activation="elu", + ) + + +def _distillation_rnn_policy(): + return RslRlDistillationStudentTeacherRecurrentCfg( + init_noise_std=1.0, + noise_std_type="scalar", + student_obs_normalization=False, + teacher_obs_normalization=True, + student_hidden_dims=[256, 256], + teacher_hidden_dims=[128, 128], + activation="elu", + rnn_type="gru", + rnn_hidden_dim=128, + rnn_num_layers=2, + teacher_recurrent=True, + ) + + +def _mlp_model(): + return RslRlMLPModelCfg( + hidden_dims=[256, 256], + activation="elu", + stochastic=True, + init_noise_std=1.0, + noise_std_type="scalar", + state_dependent_std=False, + ) + + +def _rnn_model(): + return RslRlRNNModelCfg( + hidden_dims=[256, 256], + activation="elu", + stochastic=True, + init_noise_std=1.0, + noise_std_type="scalar", + state_dependent_std=False, + rnn_type="lstm", + rnn_hidden_dim=256, + rnn_num_layers=1, + ) + + +def _on_policy_runner(**kw): + return RslRlOnPolicyRunnerCfg( + num_steps_per_env=24, + max_iterations=100, + save_interval=50, + experiment_name="test", + **kw, + ) + + +def _distillation_runner(**kw): + return RslRlDistillationRunnerCfg( + num_steps_per_env=24, + max_iterations=100, + save_interval=50, + experiment_name="test", + **kw, + ) + + +# =================================================================== +# rsl-rl < 4.0.0 +# =================================================================== +class TestBelow4: + def test_no_policy_raises(self): + with pytest.raises(ValueError, match="policy"): + handle_deprecated_rsl_rl_cfg(_on_policy_runner(algorithm=_ppo_algo()), "3.0.0") + + def test_distillation_no_policy_raises(self): + with pytest.raises(ValueError, match="policy"): + handle_deprecated_rsl_rl_cfg(_distillation_runner(algorithm=_distillation_algo()), "3.0.0") + + def test_preserves_policy(self): + cfg = _on_policy_runner(policy=_ppo_mlp_policy(), algorithm=_ppo_algo()) + handle_deprecated_rsl_rl_cfg(cfg, "3.0.0") + assert not _is_missing(cfg.policy) + assert cfg.policy.actor_hidden_dims == [256, 256] + + def test_removes_optimizer_silently_for_adam(self, capsys): + cfg = _on_policy_runner(policy=_ppo_mlp_policy(), algorithm=_ppo_algo()) + handle_deprecated_rsl_rl_cfg(cfg, "3.0.0") + assert "optimizer" not in capsys.readouterr().out + + def test_removes_optimizer_with_warning_for_non_adam(self, capsys): + algo = _ppo_algo() + algo.optimizer = "adamw" + cfg = _on_policy_runner(policy=_ppo_mlp_policy(), algorithm=algo) + handle_deprecated_rsl_rl_cfg(cfg, "3.0.0") + assert "optimizer" in capsys.readouterr().out.lower() + + def test_distillation_optimizer_untouched(self): + cfg = _distillation_runner(policy=_distillation_mlp_policy(), algorithm=_distillation_algo()) + handle_deprecated_rsl_rl_cfg(cfg, "3.0.0") + assert cfg.algorithm.optimizer == "adam" + + def test_empirical_normalization_migrates(self): + p = _ppo_mlp_policy() + p.actor_obs_normalization = MISSING + p.critic_obs_normalization = MISSING + cfg = _on_policy_runner(policy=p, algorithm=_ppo_algo(), empirical_normalization=True) + handle_deprecated_rsl_rl_cfg(cfg, "3.0.0") + assert cfg.policy.actor_obs_normalization is True + assert cfg.policy.critic_obs_normalization is True + assert _is_missing(cfg.empirical_normalization) + + def test_empirical_normalization_does_not_overwrite(self): + cfg = _on_policy_runner(policy=_ppo_mlp_policy(), algorithm=_ppo_algo(), empirical_normalization=True) + handle_deprecated_rsl_rl_cfg(cfg, "3.0.0") + assert cfg.policy.actor_obs_normalization is False + assert cfg.policy.critic_obs_normalization is False + + def test_clears_model_cfgs(self): + critic = _mlp_model() + critic.stochastic = False + cfg = _on_policy_runner(policy=_ppo_mlp_policy(), algorithm=_ppo_algo(), actor=_mlp_model(), critic=critic) + handle_deprecated_rsl_rl_cfg(cfg, "3.0.0") + assert _is_missing(cfg.actor) + assert _is_missing(cfg.critic) + + def test_distillation_clears_model_cfgs(self): + cfg = _distillation_runner( + policy=_distillation_mlp_policy(), + algorithm=_distillation_algo(), + student=_mlp_model(), + teacher=_mlp_model(), + ) + handle_deprecated_rsl_rl_cfg(cfg, "3.0.0") + assert _is_missing(cfg.student) + assert _is_missing(cfg.teacher) + + +# =================================================================== +# 4.0.0 <= rsl-rl < 5.0.0 +# =================================================================== +class TestV4: + # PPO tests + def test_infers_mlp_actor_critic(self): + p = RslRlPpoActorCriticCfg( + actor_hidden_dims=[512], + critic_hidden_dims=[64], + activation="elu", + actor_obs_normalization=False, + critic_obs_normalization=False, + init_noise_std=0.5, + noise_std_type="log", + state_dependent_std=True, + ) + cfg = _on_policy_runner(policy=p, algorithm=_ppo_algo()) + handle_deprecated_rsl_rl_cfg(cfg, "4.0.0") + + assert isinstance(cfg.actor, RslRlMLPModelCfg) and not isinstance(cfg.actor, RslRlRNNModelCfg) + assert cfg.actor.hidden_dims == [512] + assert cfg.actor.stochastic is True + assert cfg.actor.init_noise_std == 0.5 + assert cfg.actor.noise_std_type == "log" + assert cfg.actor.state_dependent_std is True + + assert isinstance(cfg.critic, RslRlMLPModelCfg) and not isinstance(cfg.critic, RslRlRNNModelCfg) + assert cfg.critic.hidden_dims == [64] + assert cfg.critic.stochastic is False + + def test_infers_rnn_actor_critic(self): + p = RslRlPpoActorCriticRecurrentCfg( + init_noise_std=1.0, + actor_obs_normalization=False, + critic_obs_normalization=False, + actor_hidden_dims=[256, 256], + critic_hidden_dims=[128, 128], + activation="elu", + rnn_type="gru", + rnn_hidden_dim=128, + rnn_num_layers=3, + ) + cfg = _on_policy_runner(policy=p, algorithm=_ppo_algo()) + handle_deprecated_rsl_rl_cfg(cfg, "4.0.0") + + assert isinstance(cfg.actor, RslRlRNNModelCfg) + assert cfg.actor.rnn_type == "gru" + assert cfg.actor.rnn_hidden_dim == 128 + assert cfg.actor.rnn_num_layers == 3 + assert isinstance(cfg.critic, RslRlRNNModelCfg) + + def test_clears_policy_after_inference(self): + cfg = _on_policy_runner(policy=_ppo_mlp_policy(), algorithm=_ppo_algo()) + handle_deprecated_rsl_rl_cfg(cfg, "4.0.0") + assert _is_missing(cfg.policy) + + def test_skips_existing_actor(self): + actor = _mlp_model() + actor.hidden_dims = [999] + cfg = _on_policy_runner(policy=_ppo_mlp_policy(), algorithm=_ppo_algo(), actor=actor) + handle_deprecated_rsl_rl_cfg(cfg, "4.0.0") + assert cfg.actor.hidden_dims == [999] + + def test_skips_existing_critic(self): + critic = _mlp_model() + critic.stochastic = False + critic.hidden_dims = [777] + cfg = _on_policy_runner(policy=_ppo_mlp_policy(), algorithm=_ppo_algo(), critic=critic) + handle_deprecated_rsl_rl_cfg(cfg, "4.0.0") + assert cfg.critic.hidden_dims == [777] + + def test_empirical_norm_then_inference(self): + p = _ppo_mlp_policy() + p.actor_obs_normalization = MISSING + p.critic_obs_normalization = MISSING + cfg = _on_policy_runner(policy=p, algorithm=_ppo_algo(), empirical_normalization=True) + handle_deprecated_rsl_rl_cfg(cfg, "4.0.0") + assert cfg.actor.obs_normalization is True + assert cfg.critic.obs_normalization is True + assert _is_missing(cfg.empirical_normalization) + + # Distillation tests + def test_distillation_infers_mlp_student_teacher(self): + p = RslRlDistillationStudentTeacherCfg( + student_hidden_dims=[512], + teacher_hidden_dims=[64], + activation="elu", + student_obs_normalization=False, + teacher_obs_normalization=True, + init_noise_std=0.8, + noise_std_type="log", + ) + cfg = _distillation_runner(policy=p, algorithm=_distillation_algo()) + handle_deprecated_rsl_rl_cfg(cfg, "4.0.0") + + assert isinstance(cfg.student, RslRlMLPModelCfg) + assert cfg.student.hidden_dims == [512] + assert cfg.student.init_noise_std == 0.8 + assert cfg.student.noise_std_type == "log" + assert isinstance(cfg.teacher, RslRlMLPModelCfg) + assert cfg.teacher.hidden_dims == [64] + assert cfg.teacher.init_noise_std == 0.0 # hardcoded + assert _is_missing(cfg.policy) + + def test_distillation_infers_rnn_student_teacher(self): + cfg = _distillation_runner(policy=_distillation_rnn_policy(), algorithm=_distillation_algo()) + handle_deprecated_rsl_rl_cfg(cfg, "4.0.0") + + assert isinstance(cfg.student, RslRlRNNModelCfg) + assert cfg.student.rnn_type == "gru" + assert isinstance(cfg.teacher, RslRlRNNModelCfg) + assert cfg.teacher.init_noise_std == 0.0 + + def test_distillation_skips_existing_student_teacher(self): + student = _mlp_model() + student.hidden_dims = [999] + teacher = _mlp_model() + teacher.hidden_dims = [777] + cfg = _distillation_runner( + policy=_distillation_mlp_policy(), + algorithm=_distillation_algo(), + student=student, + teacher=teacher, + ) + handle_deprecated_rsl_rl_cfg(cfg, "4.0.0") + assert cfg.student.hidden_dims == [999] + assert cfg.teacher.hidden_dims == [777] + + # Stochastic tests + def test_validates_stochastic_on_existing_models(self): + critic = _mlp_model() + critic.stochastic = False + cfg = _on_policy_runner(algorithm=_ppo_algo(), actor=_mlp_model(), critic=critic) + handle_deprecated_rsl_rl_cfg(cfg, "4.0.0") # should not raise + + def test_missing_stochastic_raises(self): + a = _mlp_model() + a.stochastic = MISSING + cfg = _on_policy_runner(algorithm=_ppo_algo(), actor=a) + with pytest.raises(ValueError, match="stochastic"): + handle_deprecated_rsl_rl_cfg(cfg, "4.0.0") + + def test_removes_distribution_cfg(self): + a = _mlp_model() + a.distribution_cfg = RslRlMLPModelCfg.GaussianDistributionCfg(init_std=1.0) + cfg = _on_policy_runner(algorithm=_ppo_algo(), actor=a) + handle_deprecated_rsl_rl_cfg(cfg, "4.0.0") + assert not hasattr(cfg.actor, "distribution_cfg") + + def test_no_policy_no_models_is_noop(self): + cfg = _on_policy_runner(algorithm=_ppo_algo()) + handle_deprecated_rsl_rl_cfg(cfg, "4.0.0") + assert _is_missing(cfg.actor) and _is_missing(cfg.critic) + + +# =================================================================== +# rsl-rl >= 5.0.0 +# =================================================================== +class TestV5: + # Distribution tests + def test_gaussian_from_stochastic(self): + a = _mlp_model() + a.init_noise_std = 0.5 + a.noise_std_type = "log" + cfg = _on_policy_runner(algorithm=_ppo_algo(), actor=a) + handle_deprecated_rsl_rl_cfg(cfg, "5.0.0") + + d = cfg.actor.distribution_cfg + assert isinstance(d, RslRlMLPModelCfg.GaussianDistributionCfg) + assert not isinstance(d, RslRlMLPModelCfg.HeteroscedasticGaussianDistributionCfg) + assert d.init_std == 0.5 + assert d.std_type == "log" + + def test_heteroscedastic_from_stochastic(self): + a = _mlp_model() + a.state_dependent_std = True + a.init_noise_std = 0.3 + cfg = _on_policy_runner(algorithm=_ppo_algo(), actor=a) + handle_deprecated_rsl_rl_cfg(cfg, "5.0.0") + + assert isinstance(cfg.actor.distribution_cfg, RslRlMLPModelCfg.HeteroscedasticGaussianDistributionCfg) + assert cfg.actor.distribution_cfg.init_std == 0.3 + + def test_keeps_existing_distribution_cfg(self): + a = _mlp_model() + a.state_dependent_std = True + a.distribution_cfg = RslRlMLPModelCfg.HeteroscedasticGaussianDistributionCfg(init_std=2.0, std_type="log") + cfg = _on_policy_runner(algorithm=_ppo_algo(), actor=a) + handle_deprecated_rsl_rl_cfg(cfg, "5.0.0") + + assert isinstance(cfg.actor.distribution_cfg, RslRlMLPModelCfg.HeteroscedasticGaussianDistributionCfg) + assert cfg.actor.distribution_cfg.init_std == 2.0 + assert cfg.actor.distribution_cfg.std_type == "log" + + def test_non_stochastic_no_distribution(self): + a = _mlp_model() + a.stochastic = False + cfg = _on_policy_runner(algorithm=_ppo_algo(), actor=a) + handle_deprecated_rsl_rl_cfg(cfg, "5.0.0") + assert cfg.actor.distribution_cfg is None + + def test_removes_deprecated_params(self): + cfg = _on_policy_runner(algorithm=_ppo_algo(), actor=_mlp_model()) + handle_deprecated_rsl_rl_cfg(cfg, "5.0.0") + assert not hasattr(cfg.actor, "stochastic") + assert not hasattr(cfg.actor, "init_noise_std") + assert not hasattr(cfg.actor, "noise_std_type") + assert not hasattr(cfg.actor, "state_dependent_std") + + def test_migrates_rnn_models(self): + a = _rnn_model() + a.state_dependent_std = True + c = _rnn_model() + c.stochastic = False + cfg = _on_policy_runner(algorithm=_ppo_algo(), actor=a, critic=c) + handle_deprecated_rsl_rl_cfg(cfg, "5.0.0") + + assert isinstance(cfg.actor.distribution_cfg, RslRlMLPModelCfg.HeteroscedasticGaussianDistributionCfg) + assert cfg.critic.distribution_cfg is None + + # Full conversion pipeline tests + def test_mlp_policy_full_pipeline(self): + p = _ppo_mlp_policy() + p.init_noise_std = 0.7 + p.noise_std_type = "log" + cfg = _on_policy_runner(policy=p, algorithm=_ppo_algo()) + handle_deprecated_rsl_rl_cfg(cfg, "5.0.0") + + assert _is_missing(cfg.policy) + assert isinstance(cfg.actor.distribution_cfg, RslRlMLPModelCfg.GaussianDistributionCfg) + assert cfg.actor.distribution_cfg.init_std == 0.7 + assert cfg.actor.distribution_cfg.std_type == "log" + assert cfg.critic.distribution_cfg is None + + def test_mlp_policy_state_dependent_pipeline(self): + p = _ppo_mlp_policy() + p.init_noise_std = 0.5 + p.state_dependent_std = True + cfg = _on_policy_runner(policy=p, algorithm=_ppo_algo()) + handle_deprecated_rsl_rl_cfg(cfg, "5.0.0") + + assert isinstance(cfg.actor.distribution_cfg, RslRlMLPModelCfg.HeteroscedasticGaussianDistributionCfg) + + def test_rnn_policy_full_pipeline(self): + p = _ppo_rnn_policy() + p.init_noise_std = 1.5 + p.rnn_type = "gru" + p.rnn_hidden_dim = 64 + p.rnn_num_layers = 3 + cfg = _on_policy_runner(policy=p, algorithm=_ppo_algo()) + handle_deprecated_rsl_rl_cfg(cfg, "5.0.0") + + assert _is_missing(cfg.policy) + assert isinstance(cfg.actor, RslRlRNNModelCfg) + assert cfg.actor.rnn_type == "gru" + assert isinstance(cfg.actor.distribution_cfg, RslRlMLPModelCfg.GaussianDistributionCfg) + assert cfg.actor.distribution_cfg.init_std == 1.5 + assert isinstance(cfg.critic, RslRlRNNModelCfg) + assert cfg.critic.distribution_cfg is None + + def test_distillation_mlp_policy_full_pipeline(self): + p = _distillation_mlp_policy() + p.init_noise_std = 0.9 + p.noise_std_type = "log" + cfg = _distillation_runner(policy=p, algorithm=_distillation_algo()) + handle_deprecated_rsl_rl_cfg(cfg, "5.0.0") + + assert _is_missing(cfg.policy) + assert isinstance(cfg.student.distribution_cfg, RslRlMLPModelCfg.GaussianDistributionCfg) + assert cfg.student.distribution_cfg.init_std == 0.9 + assert cfg.student.distribution_cfg.std_type == "log" + assert isinstance(cfg.teacher.distribution_cfg, RslRlMLPModelCfg.GaussianDistributionCfg) + assert cfg.teacher.distribution_cfg.init_std == 0.0 + + def test_distillation_rnn_policy_full_pipeline(self): + p = _distillation_rnn_policy() + p.init_noise_std = 1.2 + cfg = _distillation_runner(policy=p, algorithm=_distillation_algo()) + handle_deprecated_rsl_rl_cfg(cfg, "5.0.0") + + assert _is_missing(cfg.policy) + assert isinstance(cfg.student, RslRlRNNModelCfg) + assert isinstance(cfg.student.distribution_cfg, RslRlMLPModelCfg.GaussianDistributionCfg) + assert isinstance(cfg.teacher, RslRlRNNModelCfg) + assert cfg.teacher.distribution_cfg.init_std == 0.0 + + def test_no_policy_no_models_is_noop(self): + cfg = _on_policy_runner(algorithm=_ppo_algo()) + handle_deprecated_rsl_rl_cfg(cfg, "5.0.0") + assert _is_missing(cfg.actor) and _is_missing(cfg.critic) + + def test_empirical_norm_rnn_policy_full_pipeline(self): + p = _ppo_rnn_policy() + p.init_noise_std = 2.0 + p.state_dependent_std = True + p.actor_obs_normalization = MISSING + p.critic_obs_normalization = MISSING + cfg = _on_policy_runner(policy=p, algorithm=_ppo_algo(), empirical_normalization=True) + handle_deprecated_rsl_rl_cfg(cfg, "5.0.0") + + assert _is_missing(cfg.policy) and _is_missing(cfg.empirical_normalization) + assert cfg.actor.obs_normalization is True + assert isinstance(cfg.actor.distribution_cfg, RslRlMLPModelCfg.HeteroscedasticGaussianDistributionCfg) + + # Edge tests + def test_returns_same_object(self): + cfg = _on_policy_runner(policy=_ppo_mlp_policy(), algorithm=_ppo_algo()) + assert handle_deprecated_rsl_rl_cfg(cfg, "5.0.0") is cfg diff --git a/source/isaaclab_tasks/config/extension.toml b/source/isaaclab_tasks/config/extension.toml index 63aad3cb76fa..b13dfdc73d68 100644 --- a/source/isaaclab_tasks/config/extension.toml +++ b/source/isaaclab_tasks/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "1.5.1" +version = "1.5.5" # Description title = "Isaac Lab Environments" diff --git a/source/isaaclab_tasks/docs/CHANGELOG.rst b/source/isaaclab_tasks/docs/CHANGELOG.rst index 642ba17b7e2b..1c64ea3cc530 100644 --- a/source/isaaclab_tasks/docs/CHANGELOG.rst +++ b/source/isaaclab_tasks/docs/CHANGELOG.rst @@ -1,6 +1,51 @@ Changelog --------- +1.5.5 (2026-03-10) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Simplified Dexsuite gravity randomization to use the unified + :class:`~isaaclab.envs.mdp.randomize_physics_scene_gravity` term, removing the + backend-specific ``GravityRandomizationCfg`` preset. + +Added +^^^^^ + +* Added Dexsuite multi-hand dexterous manipulation environments with Kuka Allegro configurations for lift and + reorientation tasks. + + +1.5.4 (2026-03-07) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Simplified all task MDP ``__init__.py`` files to call ``lazy_export()`` without + arguments. Fallback packages are now inferred from ``__init__.pyi`` stubs. + +Added +^^^^^ + +* Added ``from isaaclab.envs.mdp import *`` wildcard re-exports to all task MDP + ``__init__.pyi`` stubs, fixing broken type hints for base MDP symbols. + +* Added ``test_lazy_export_stubs.py`` to enforce that ``lazy_export()`` is called + without arguments across the codebase. + +1.5.3 (2026-03-08) +~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Fixed ``TypeError: 'NoneType' object is not iterable`` in + :func:`~isaaclab_tasks.utils.hydra.apply_overrides` when a preset value is + ``None`` (e.g. ``default = None`` in a :class:`~isaaclab_tasks.utils.PresetCfg`). + 1.5.2 (2026-03-05) ~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_env_cfg.py index 2a3b7045186e..dfbea0e203a7 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_env_cfg.py @@ -13,6 +13,8 @@ from isaaclab.sim import SimulationCfg from isaaclab.utils import configclass +from isaaclab_tasks.utils.presets import MultiBackendRendererCfg + from isaaclab_assets.robots.cartpole import CARTPOLE_CFG @@ -31,7 +33,7 @@ class CartpoleRGBCameraEnvCfg(DirectRLEnvCfg): cart_dof_name = "slider_to_cart" pole_dof_name = "cart_to_pole" - # camera + # camera: default=RTX, newton_renderer=Warp. Override: env.tiled_camera.renderer_cfg=newton_renderer tiled_camera: TiledCameraCfg = TiledCameraCfg( prim_path="/World/envs/env_.*/Camera", offset=TiledCameraCfg.OffsetCfg(pos=(-5.0, 0.0, 2.0), rot=(0.0, 0.0, 0.0, 1.0), convention="world"), @@ -41,6 +43,7 @@ class CartpoleRGBCameraEnvCfg(DirectRLEnvCfg): ), width=100, height=100, + renderer_cfg=MultiBackendRendererCfg(), ) write_image_to_file = False diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py index e097a98836d1..987f82a510e9 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py @@ -135,9 +135,15 @@ def _apply_action(self) -> None: def _get_observations(self) -> dict: if self.cfg.asymmetric_obs: - self.fingertip_force_sensors = wp.to_torch(self.hand.data.body_incoming_joint_wrench_b)[ - :, self.finger_bodies - ] + # Newton does not implement body_incoming_joint_wrench_b; fall back to zeros. + try: + self.fingertip_force_sensors = wp.to_torch(self.hand.data.body_incoming_joint_wrench_b)[ + :, self.finger_bodies + ] + except NotImplementedError: + self.fingertip_force_sensors = torch.zeros( + self.num_envs, len(self.finger_bodies), 6, dtype=torch.float32, device=self.device + ) if self.cfg.obs_type == "openai": obs = self.compute_reduced_observations() diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/__init__.py index c0e2adb96556..0041b165cb0a 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/__init__.py @@ -76,3 +76,14 @@ "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_vision_cfg.yaml", }, ) + +gym.register( + id="Isaac-Repose-Cube-Shadow-Vision-Benchmark-Direct-v0", + entry_point=f"{__name__}.shadow_hand_vision_env:ShadowHandVisionEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.shadow_hand_vision_env_cfg:ShadowHandVisionBenchmarkEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:ShadowHandVisionFFPPORunnerCfg", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_vision_cfg.yaml", + }, +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py index 75a7b6d04a20..84a2e8df2a54 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py @@ -13,27 +13,81 @@ from isaaclab.sensors import save_images_to_file from isaaclab.utils import configclass +# Number of output channels for each supported camera data type. +_DATA_TYPE_CHANNELS: dict[str, int] = { + "rgb": 3, + "depth": 1, + "semantic_segmentation": 3, + "albedo": 3, + "simple_shading_constant_diffuse": 3, + "simple_shading_diffuse_mdl": 3, + "simple_shading_full_mdl": 3, +} + +# Data types whose channels should receive ImageNet normalization in the CNN forward pass. +_IMAGENET_NORM_TYPES: frozenset[str] = frozenset( + { + "rgb", + "semantic_segmentation", + "albedo", + "simple_shading_constant_diffuse", + "simple_shading_diffuse_mdl", + "simple_shading_full_mdl", + } +) + + +def _conv_out(size: int, kernel: int, stride: int, padding: int = 0) -> int: + """Compute the spatial output size of a single convolutional layer.""" + return (size + 2 * padding - kernel) // stride + 1 + class FeatureExtractorNetwork(nn.Module): """CNN architecture used to regress keypoint positions of the in-hand cube from image data.""" - def __init__(self): + def __init__( + self, + num_channel: int = 7, + data_types: list[str] | None = None, + height: int = 120, + width: int = 120, + ): + """Initialize the CNN. + + Args: + num_channel: Total number of input channels across all data types. + data_types: Ordered list of camera data types that form the channel stack. + Used to determine which channel ranges receive ImageNet normalization. + Defaults to ``["rgb", "depth", "semantic_segmentation"]``. + height: Input image height [px]. Used to compute :class:`~torch.nn.LayerNorm` + spatial dimensions. Default is ``120``. + width: Input image width [px]. Used to compute :class:`~torch.nn.LayerNorm` + spatial dimensions. Default is ``120``. + """ super().__init__() - num_channel = 7 + if data_types is None: + data_types = ["rgb", "depth", "semantic_segmentation"] + + # Compute spatial sizes after each conv to build resolution-adaptive LayerNorms. + h1, w1 = _conv_out(height, 6, 2), _conv_out(width, 6, 2) + h2, w2 = _conv_out(h1, 4, 2), _conv_out(w1, 4, 2) + h3, w3 = _conv_out(h2, 4, 2), _conv_out(w2, 4, 2) + h4, w4 = _conv_out(h3, 3, 2), _conv_out(w3, 3, 2) + self.cnn = nn.Sequential( nn.Conv2d(num_channel, 16, kernel_size=6, stride=2, padding=0), nn.ReLU(), - nn.LayerNorm([16, 58, 58]), + nn.LayerNorm([16, h1, w1]), nn.Conv2d(16, 32, kernel_size=4, stride=2, padding=0), nn.ReLU(), - nn.LayerNorm([32, 28, 28]), + nn.LayerNorm([32, h2, w2]), nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=0), nn.ReLU(), - nn.LayerNorm([64, 13, 13]), + nn.LayerNorm([64, h3, w3]), nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=0), nn.ReLU(), - nn.LayerNorm([128, 6, 6]), - nn.AvgPool2d(6), + nn.LayerNorm([128, h4, w4]), + nn.AdaptiveAvgPool2d(1), ) self.linear = nn.Sequential( @@ -46,10 +100,19 @@ def __init__(self): ] ) + # Pre-compute channel ranges that require ImageNet normalization. + self._imagenet_norm_ranges: list[tuple[int, int]] = [] + channel_idx = 0 + for dt in data_types: + n_ch = _DATA_TYPE_CHANNELS.get(dt, 3) + if dt in _IMAGENET_NORM_TYPES: + self._imagenet_norm_ranges.append((channel_idx, channel_idx + n_ch)) + channel_idx += n_ch + def forward(self, x): - x = x.permute(0, 3, 1, 2) - x[:, 0:3, :, :] = self.data_transforms(x[:, 0:3, :, :]) - x[:, 4:7, :, :] = self.data_transforms(x[:, 4:7, :, :]) + x = x.permute(0, 3, 1, 2).clone() + for start, end in self._imagenet_norm_ranges: + x[:, start:end, :, :] = self.data_transforms(x[:, start:end, :, :]) cnn_x = self.cnn(x) out = self.linear(cnn_x.view(-1, 128)) return out @@ -60,7 +123,7 @@ class FeatureExtractorCfg: """Configuration for the feature extractor model.""" train: bool = True - """If True, the feature extractor model is trained during the rollout process. Default is False.""" + """If True, the feature extractor model is trained during the rollout process. Default is True.""" load_checkpoint: bool = False """If True, the feature extractor model is loaded from a checkpoint. Default is False.""" @@ -68,29 +131,64 @@ class FeatureExtractorCfg: write_image_to_file: bool = False """If True, the images from the camera sensor are written to file. Default is False.""" + enabled: bool = True + """If True, the CNN forward pass is executed each step. + + Set to False to bypass the network entirely and return zero embeddings. This is useful + for benchmarking rendering throughput without CNN inference overhead. Default is True. + """ + class FeatureExtractor: """Class for extracting features from image data. - It uses a CNN to regress keypoint positions from normalized RGB, depth, and segmentation images. - If the train flag is set to True, the CNN is trained during the rollout process. + It uses a CNN to regress keypoint positions from normalized images. + If :attr:`FeatureExtractorCfg.train` is ``True``, the CNN is trained during rollouts. + If :attr:`FeatureExtractorCfg.enabled` is ``False``, the network is bypassed and zero + embeddings are returned (useful for benchmarking rendering throughput). + + The input data types (and therefore the CNN's input channel count) are determined by + the camera's ``data_types`` at construction time, passed via the ``data_types`` argument. + This means changing the camera preset (e.g. ``presets=rgb``) automatically reconfigures + the CNN without requiring a separate environment config class. """ - def __init__(self, cfg: FeatureExtractorCfg, device: str, log_dir: str | None = None): + def __init__( + self, + cfg: FeatureExtractorCfg, + device: str, + data_types: list[str], + log_dir: str | None = None, + height: int = 120, + width: int = 120, + ): """Initialize the feature extractor model. Args: cfg: Configuration for the feature extractor model. device: Device to run the model on. + data_types: Ordered list of camera data types that form the CNN input channel + stack. Should match the resolved :attr:`~isaaclab.sensors.TiledCameraCfg.data_types` + of the tiled camera. Total input channels are derived from + :data:`_DATA_TYPE_CHANNELS`. log_dir: Directory to save checkpoints. Default is None, which uses the local "logs" folder resolved relative to this file. + height: Camera image height [px]. Must match the tiled camera + :attr:`~isaaclab.sensors.TiledCameraCfg.height`. Default is ``120``. + width: Camera image width [px]. Must match the tiled camera + :attr:`~isaaclab.sensors.TiledCameraCfg.width`. Default is ``120``. """ - self.cfg = cfg self.device = device + self.data_types = data_types - # Feature extractor model - self.feature_extractor = FeatureExtractorNetwork() + # Compute total input channels from the camera data types. + num_channel = sum(_DATA_TYPE_CHANNELS.get(dt, 3) for dt in data_types) + + # Feature extractor model. + self.feature_extractor = FeatureExtractorNetwork( + num_channel=num_channel, data_types=data_types, height=height, width=width + ) self.feature_extractor.to(self.device) self.step_count = 0 @@ -115,66 +213,96 @@ def __init__(self, cfg: FeatureExtractorCfg, device: str, log_dir: str | None = else: self.feature_extractor.eval() - def _preprocess_images( - self, rgb_img: torch.Tensor, depth_img: torch.Tensor, segmentation_img: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Preprocesses the input images. + def _preprocess_images(self, camera_output: dict[str, torch.Tensor]) -> torch.Tensor: + """Preprocesses and concatenates camera images into a single tensor. + + Each data type in :attr:`FeatureExtractorCfg.data_types` is extracted from + ``camera_output``, normalized, and concatenated along the channel dimension. Args: - rgb_img (torch.Tensor): RGB image tensor. Shape: (N, H, W, 3). - depth_img (torch.Tensor): Depth image tensor. Shape: (N, H, W, 1). - segmentation_img (torch.Tensor): Segmentation image tensor. Shape: (N, H, W, 3) + camera_output: Dictionary mapping data type names to image tensors. Returns: - tuple[torch.Tensor, torch.Tensor, torch.Tensor]: Preprocessed RGB, depth, and segmentation + Concatenated preprocessed image tensor of shape (N, H, W, C). """ - rgb_img = rgb_img / 255.0 - # process depth image - depth_img[depth_img == float("inf")] = 0 - depth_img /= 5.0 - depth_img /= torch.max(depth_img) - # process segmentation image - segmentation_img = segmentation_img / 255.0 - mean_tensor = torch.mean(segmentation_img, dim=(1, 2), keepdim=True) - segmentation_img -= mean_tensor - return rgb_img, depth_img, segmentation_img - - def _save_images(self, rgb_img: torch.Tensor, depth_img: torch.Tensor, segmentation_img: torch.Tensor): - """Writes image buffers to file. + tensors = [] + for dt in self.data_types: + img = camera_output[dt].float() + if dt == "rgb": + img = img / 255.0 + elif dt == "depth": + img[img == float("inf")] = 0 + img /= 5.0 + max_val = img.max() + if max_val > 0: + img /= max_val + elif dt == "semantic_segmentation": + img = img[..., :3] / 255.0 + mean_tensor = torch.mean(img, dim=(1, 2), keepdim=True) + img = img - mean_tensor + else: + # albedo and simple_shading_* are RGB-like 3-channel outputs. + img = img[..., :3] / 255.0 + tensors.append(img) + return torch.cat(tensors, dim=-1) + + def _save_images(self, camera_output: dict[str, torch.Tensor]): + """Writes configured camera data buffers to file as normalized float images. + + Raw camera tensors are converted to float ``[0, 1]`` before saving so that + :func:`~isaaclab.sensors.save_images_to_file` (which delegates to + ``torchvision.utils.save_image``) receives the expected float input. Args: - rgb_img (torch.Tensor): RGB image tensor. Shape: (N, H, W, 3). - depth_img (torch.Tensor): Depth image tensor. Shape: (N, H, W, 1). - segmentation_img (torch.Tensor): Segmentation image tensor. Shape: (N, H, W, 3). + camera_output: Dictionary mapping data type names to image tensors. """ - save_images_to_file(rgb_img, "shadow_hand_rgb.png") - save_images_to_file(depth_img, "shadow_hand_depth.png") - save_images_to_file(segmentation_img, "shadow_hand_segmentation.png") + for dt in self.data_types: + if dt not in camera_output: + continue + img = camera_output[dt].float() + if dt == "depth": + img = img.clone() + img[img == float("inf")] = 0 + max_val = img.max() + if max_val > 0: + img = img / max_val + else: + # rgb, semantic_segmentation, albedo, and simple_shading_* are uint8 [0, 255] + img = img[..., :3] / 255.0 + save_images_to_file(img, f"shadow_hand_{dt}.png") def step( - self, rgb_img: torch.Tensor, depth_img: torch.Tensor, segmentation_img: torch.Tensor, gt_pose: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: - """Extracts the features using the images and trains the model if the train flag is set to True. + self, camera_output: dict[str, torch.Tensor], gt_pose: torch.Tensor + ) -> tuple[torch.Tensor | None, torch.Tensor]: + """Extracts features and optionally trains the CNN. + + Image saving (when :attr:`FeatureExtractorCfg.write_image_to_file` is ``True``) always + runs first, regardless of whether the network is enabled. When + :attr:`FeatureExtractorCfg.enabled` is ``False``, the network is then bypassed and + zero embeddings are returned without any further image preprocessing. Args: - rgb_img (torch.Tensor): RGB image tensor. Shape: (N, H, W, 3). - depth_img (torch.Tensor): Depth image tensor. Shape: (N, H, W, 1). - segmentation_img (torch.Tensor): Segmentation image tensor. Shape: (N, H, W, 3). - gt_pose (torch.Tensor): Ground truth pose tensor (position and corners). Shape: (N, 27). + camera_output: Dictionary mapping data type names to image tensors from the + tiled camera sensor. + gt_pose: Ground truth pose tensor (position and keypoint corners). Shape: (N, 27). Returns: - tuple[torch.Tensor, torch.Tensor]: Pose loss and predicted pose. + tuple[torch.Tensor | None, torch.Tensor]: Pose loss (``None`` when not training + or when the network is disabled) and the predicted pose embedding of shape + (N, 27). """ + if self.cfg.write_image_to_file: + self._save_images(camera_output) - rgb_img, depth_img, segmentation_img = self._preprocess_images(rgb_img, depth_img, segmentation_img) + if not self.cfg.enabled: + batch_size = next(iter(camera_output.values())).shape[0] + return None, torch.zeros(batch_size, 27, dtype=torch.float32, device=self.device) - if self.cfg.write_image_to_file: - self._save_images(rgb_img, depth_img, segmentation_img) + img_input = self._preprocess_images(camera_output) if self.cfg.train: with torch.enable_grad(): with torch.inference_mode(False): - img_input = torch.cat((rgb_img, depth_img, segmentation_img), dim=-1) self.optimizer.zero_grad() predicted_pose = self.feature_extractor(img_input) @@ -193,6 +321,5 @@ def step( return pose_loss, predicted_pose else: - img_input = torch.cat((rgb_img, depth_img, segmentation_img), dim=-1) predicted_pose = self.feature_extractor(img_input) return None, predicted_pose diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_env_cfg.py index b2fd847bc765..7faf28a78c12 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_env_cfg.py @@ -162,7 +162,7 @@ class ShadowHandRobotCfg(PresetCfg): # discards the root body's native USD orientation, so we must re-apply it here as a # spawn rotation. PhysX or USD does not have this issue. Remove once Newton fixes root joint # transform handling in import_usd.py. - rot=(-0.7071, 0.0, 0.0, 0.7071), + rot=(0.0, 0.0, 0.0, 1.0), joint_pos={".*": 0.0}, ), actuators={ @@ -232,7 +232,7 @@ class ObjectCfg(PresetCfg): scale=(0.9, 0.9, 0.9), ), init_state=ArticulationCfg.InitialStateCfg( - pos=(0.0, 0.36, 0.535), rot=(0.0, 0.0, 0.0, 1.0), joint_pos={}, joint_vel={} + pos=(0.0, -0.36, 0.535), rot=(0.0, 0.0, 0.0, 1.0), joint_pos={}, joint_vel={} ), actuators={}, articulation_root_prim_path="", @@ -240,6 +240,23 @@ class ObjectCfg(PresetCfg): default = physx +@configclass +class ShadowHandSceneCfg(PresetCfg): + """Scene configuration presets for the shadow hand environment. + + PhysX supports ``clone_in_fabric=True`` for faster scene cloning via the Fabric layer. + Newton does not support Fabric cloning, so ``clone_in_fabric`` must be ``False``. + """ + + physx: InteractiveSceneCfg = InteractiveSceneCfg( + num_envs=8192, env_spacing=0.75, replicate_physics=True, clone_in_fabric=True + ) + newton: InteractiveSceneCfg = InteractiveSceneCfg( + num_envs=8192, env_spacing=0.75, replicate_physics=True, clone_in_fabric=False + ) + default: InteractiveSceneCfg = physx + + @configclass class PhysicsCfg(PresetCfg): physx = PhysxCfg( @@ -328,10 +345,8 @@ class ShadowHandEnvCfg(DirectRLEnvCfg): ) }, ) - # scene - scene: InteractiveSceneCfg = InteractiveSceneCfg( - num_envs=8192, env_spacing=0.75, replicate_physics=True, clone_in_fabric=True - ) + # scene — use ShadowHandSceneCfg so that presets=newton disables clone_in_fabric automatically + scene: ShadowHandSceneCfg = ShadowHandSceneCfg() # reset reset_position_noise = 0.01 # range of position at reset diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env.py index 74c3ce4747a7..a3bc9fa09435 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env.py @@ -23,14 +23,57 @@ if TYPE_CHECKING: from .shadow_hand_vision_env_cfg import ShadowHandVisionEnvCfg +# Data types supported by the Newton Warp renderer. +_WARP_SUPPORTED_DATA_TYPES: frozenset[str] = frozenset({"rgb", "depth"}) + + +def _validate_cfg(cfg: ShadowHandVisionEnvCfg) -> None: + """Validate preset combination compatibility before scene setup. + + Raises: + ValueError: If the Warp renderer is paired with an unsupported data type, or if a + depth-only camera is used without disabling the feature extractor. + """ + renderer_type = getattr(cfg.tiled_camera.renderer_cfg, "renderer_type", None) + if renderer_type == "newton_warp": + unsupported = set(cfg.tiled_camera.data_types) - _WARP_SUPPORTED_DATA_TYPES + if unsupported: + raise ValueError( + f"Warp renderer only supports data types {sorted(_WARP_SUPPORTED_DATA_TYPES)}, " + f"but the camera is configured with unsupported types: {sorted(unsupported)}. " + f"Choose a compatible preset, e.g. presets=warp,rgb or presets=warp,depth." + ) + + if set(cfg.tiled_camera.data_types) == {"depth"} and cfg.feature_extractor.enabled: + raise ValueError( + "Depth-only camera data type is intended for benchmarking only. " + "The keypoint-regression CNN cannot be meaningfully trained from depth alone. " + "Disable the feature extractor with 'feature_extractor.enabled=False' " + "(e.g. use Isaac-Repose-Cube-Shadow-Vision-Benchmark-Direct-v0), " + "or choose a data type that includes colour, e.g. presets=rgb." + ) + class ShadowHandVisionEnv(InHandManipulationEnv): cfg: ShadowHandVisionEnvCfg def __init__(self, cfg: ShadowHandVisionEnvCfg, render_mode: str | None = None, **kwargs): + # Validate renderer/data-type compatibility before any scene setup so the error + # surfaces immediately rather than crashing deep inside the renderer. + _validate_cfg(cfg) + super().__init__(cfg, render_mode, **kwargs) - # Use the log directory from the configuration - self.feature_extractor = FeatureExtractor(self.cfg.feature_extractor, self.device, self.cfg.log_dir) + # Derive CNN input data types from the resolved camera config so that any camera + # preset (e.g. presets=rgb, presets=albedo) automatically configures the right + # network input channels without requiring a separate env config class. + self.feature_extractor = FeatureExtractor( + self.cfg.feature_extractor, + self.device, + self.cfg.tiled_camera.data_types, + self.cfg.log_dir, + height=self.cfg.tiled_camera.height, + width=self.cfg.tiled_camera.width, + ) # hide goal cubes self.goal_pos[:, :] = torch.tensor([-0.2, 0.1, 0.6], device=self.device) # keypoints buffer @@ -40,7 +83,7 @@ def __init__(self, cfg: ShadowHandVisionEnvCfg, render_mode: str | None = None, def _setup_scene(self): # add hand, in-hand object, and goal object self.hand = Articulation(self.cfg.robot_cfg) - self.object = RigidObject(self.cfg.object_cfg) + self.object: Articulation | RigidObject = self.cfg.object_cfg.class_type(self.cfg.object_cfg) self._tiled_camera = TiledCamera(self.cfg.tiled_camera) # clone and replicate (no need to filter for this environment) self.scene.clone_environments(copy_from_source=False) @@ -60,9 +103,7 @@ def _compute_image_observations(self): # train CNN to regress on keypoint positions pose_loss, embeddings = self.feature_extractor.step( - self._tiled_camera.data.output["rgb"], - self._tiled_camera.data.output["depth"], - self._tiled_camera.data.output["semantic_segmentation"][..., :3], + self._tiled_camera.data.output, object_pose, ) @@ -80,10 +121,11 @@ def _compute_image_observations(self): dim=-1, ) - # log pose loss from CNN training - if "log" not in self.extras: - self.extras["log"] = dict() - self.extras["log"]["pose_loss"] = pose_loss + # log pose loss from CNN training (None when disabled or in inference mode) + if pose_loss is not None: + if "log" not in self.extras: + self.extras["log"] = dict() + self.extras["log"]["pose_loss"] = pose_loss return obs @@ -120,10 +162,15 @@ def _get_observations(self) -> dict: # vision observations from CMM image_obs = self._compute_image_observations() obs = torch.cat((state_obs, image_obs), dim=-1) - # asymmetric critic states - self.fingertip_force_sensors = wp.to_torch(self.hand.root_view.get_link_incoming_joint_force())[ - :, self.finger_bodies - ] + # asymmetric critic states — Newton does not implement body_incoming_joint_wrench_b + try: + self.fingertip_force_sensors = wp.to_torch(self.hand.data.body_incoming_joint_wrench_b)[ + :, self.finger_bodies + ] + except NotImplementedError: + self.fingertip_force_sensors = torch.zeros( + self.num_envs, len(self.finger_bodies), 6, dtype=torch.float32, device=self.device + ) state = self._compute_states() observations = {"policy": obs, "critic": state} diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env_cfg.py index b5536577285d..cee05599cc45 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env_cfg.py @@ -10,31 +10,112 @@ from isaaclab.sensors import TiledCameraCfg from isaaclab.utils import configclass +from isaaclab_tasks.utils import PresetCfg +from isaaclab_tasks.utils.presets import MultiBackendRendererCfg + from .feature_extractor import FeatureExtractorCfg from .shadow_hand_env_cfg import ShadowHandEnvCfg +@configclass +class _ShadowHandBaseTiledCameraCfg(TiledCameraCfg): + """Base tiled camera configuration for the shadow hand vision environment. + + This is an internal config used by :class:`ShadowHandVisionTiledCameraCfg` presets and + by derived env configs that hard-code a specific data type. It embeds + :class:`~isaaclab_tasks.utils.MultiBackendRendererCfg` so the renderer backend can + still be selected via the ``presets`` CLI argument. + """ + + prim_path: str = "/World/envs/env_.*/Camera" + offset: TiledCameraCfg.OffsetCfg = TiledCameraCfg.OffsetCfg( + pos=(0, -0.35, 1.0), rot=(0.0, 0.7071, 0.0, 0.7071), convention="world" + ) + data_types: list[str] = [] + spawn: sim_utils.PinholeCameraCfg = sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 20.0) + ) + width: int = 120 + height: int = 120 + renderer_cfg: MultiBackendRendererCfg = MultiBackendRendererCfg() + + +@configclass +class ShadowHandVisionTiledCameraCfg(PresetCfg): + """Camera data-type presets for the shadow hand vision environment. + + Each preset selects which image modalities are captured. The selected data types must + match :attr:`FeatureExtractorCfg.data_types` so the CNN receives the expected channels. + + Select a data-type preset via the ``presets`` CLI argument, e.g.:: + + presets = rgb # RGB only (3 channels) + presets = albedo # albedo (3 channels) + presets = simple_shading_constant_diffuse # simple shading, constant diffuse (3 channels) + + Renderer and data-type presets can be combined:: + + presets = newton_renderer, rgb + """ + + default: _ShadowHandBaseTiledCameraCfg = _ShadowHandBaseTiledCameraCfg( + data_types=["rgb", "depth", "semantic_segmentation"] + ) + """Default: RGB + depth + semantic segmentation (7 CNN input channels).""" + + full: _ShadowHandBaseTiledCameraCfg = _ShadowHandBaseTiledCameraCfg( + data_types=["rgb", "depth", "semantic_segmentation"] + ) + """Full modalities: RGB + depth + semantic segmentation (7 channels). Alias for default.""" + + rgb: _ShadowHandBaseTiledCameraCfg = _ShadowHandBaseTiledCameraCfg(data_types=["rgb"]) + """RGB only (3 CNN input channels).""" + + albedo: _ShadowHandBaseTiledCameraCfg = _ShadowHandBaseTiledCameraCfg(data_types=["albedo"]) + """Albedo (3 CNN input channels).""" + + simple_shading_constant_diffuse: _ShadowHandBaseTiledCameraCfg = _ShadowHandBaseTiledCameraCfg( + data_types=["simple_shading_constant_diffuse"] + ) + """Simple shading with constant diffuse (3 CNN input channels).""" + + simple_shading_diffuse_mdl: _ShadowHandBaseTiledCameraCfg = _ShadowHandBaseTiledCameraCfg( + data_types=["simple_shading_diffuse_mdl"] + ) + """Simple shading with diffuse MDL (3 CNN input channels).""" + + simple_shading_full_mdl: _ShadowHandBaseTiledCameraCfg = _ShadowHandBaseTiledCameraCfg( + data_types=["simple_shading_full_mdl"] + ) + """Simple shading with full MDL (3 CNN input channels).""" + + depth: _ShadowHandBaseTiledCameraCfg = _ShadowHandBaseTiledCameraCfg(data_types=["depth"]) + """Depth only (1 channel). + + .. warning:: + This preset is intended for **benchmarking only**. The keypoint-regression CNN + cannot be meaningfully trained from depth alone. Use it with + :class:`ShadowHandVisionBenchmarkEnvCfg` (``feature_extractor.enabled=False``) + to measure pure depth-rendering throughput, e.g.:: + + presets=depth # depth rendering, default renderer + presets=depth,newton_renderer # depth rendering with Newton renderer + presets=depth,ovrtx_renderer # depth rendering with OVRTX renderer + """ + + @configclass class ShadowHandVisionEnvCfg(ShadowHandEnvCfg): # scene scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=1225, env_spacing=2.0, replicate_physics=True) - # camera - tiled_camera: TiledCameraCfg = TiledCameraCfg( - prim_path="/World/envs/env_.*/Camera", - offset=TiledCameraCfg.OffsetCfg(pos=(0, -0.35, 1.0), rot=(0.0, 0.7071, 0.0, 0.7071), convention="world"), - data_types=["rgb", "depth", "semantic_segmentation"], - spawn=sim_utils.PinholeCameraCfg( - focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 20.0) - ), - width=120, - height=120, - ) - feature_extractor = FeatureExtractorCfg() + # camera — data-type and renderer backend selectable via CLI presets + tiled_camera: ShadowHandVisionTiledCameraCfg = ShadowHandVisionTiledCameraCfg() + feature_extractor: FeatureExtractorCfg = FeatureExtractorCfg() # env observation_space = 164 + 27 # state observation + vision CNN embedding - state_space = 187 + 27 # asymettric states + vision CNN embedding + state_space = 187 + 27 # asymmetric states + vision CNN embedding @configclass @@ -42,4 +123,23 @@ class ShadowHandVisionEnvPlayCfg(ShadowHandVisionEnvCfg): # scene scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=64, env_spacing=2.0, replicate_physics=True) # inference for CNN - feature_extractor = FeatureExtractorCfg(train=False, load_checkpoint=True) + feature_extractor: FeatureExtractorCfg = FeatureExtractorCfg(train=False, load_checkpoint=True) + + +@configclass +class ShadowHandVisionBenchmarkEnvCfg(ShadowHandVisionEnvCfg): + """Benchmark configuration with the feature extractor CNN disabled. + + The tiled camera renders frames each step as normal, but the CNN forward pass is + bypassed — zero embeddings are returned instead. This isolates rendering throughput + from CNN inference overhead when profiling. + + The renderer backend and camera data types can still be selected via ``presets``:: + + presets = newton_renderer # benchmark with Newton renderer + presets = ovrtx_renderer # benchmark with OVRTX renderer + presets = rgb # benchmark RGB rendering only + presets = depth, newton_renderer # benchmark depth rendering with Newton + """ + + feature_extractor: FeatureExtractorCfg = FeatureExtractorCfg(enabled=False) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/mdp/__init__.py index 984f1ffecd5c..1ce92a1952e8 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/mdp/__init__.py @@ -7,4 +7,4 @@ from isaaclab.utils.module import lazy_export -lazy_export(packages=["isaaclab.envs.mdp"]) +lazy_export() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/mdp/__init__.pyi index 3fafa5558790..8d3c2673c76d 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/mdp/__init__.pyi @@ -8,3 +8,4 @@ __all__ = [ ] from .rewards import joint_pos_target_l2 +from isaaclab.envs.mdp import * diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/__init__.py index e11218b79576..95af9b285800 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/__init__.py @@ -7,4 +7,4 @@ from isaaclab.utils.module import lazy_export -lazy_export(packages=["isaaclab.envs.mdp"]) +lazy_export() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/__init__.pyi index 98d3aad77928..94bf0f7f2c5c 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/__init__.pyi @@ -23,3 +23,4 @@ from .rewards import ( progress_reward, upright_posture_bonus, ) +from isaaclab.envs.mdp import * diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/__init__.py index d42488d31ed3..1a58aa96d858 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/__init__.py @@ -7,4 +7,4 @@ from isaaclab.utils.module import lazy_export -lazy_export(packages=["isaaclab.envs.mdp", "isaaclab_contrib.mdp"]) +lazy_export() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/__init__.pyi index d6426490ca35..d54142b0609d 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/__init__.pyi @@ -17,3 +17,5 @@ __all__ = [ from .commands import DroneUniformPoseCommand, DroneUniformPoseCommandCfg from .observations import base_roll_pitch, generated_drone_commands from .rewards import ang_vel_xyz_exp, distance_to_goal_exp, lin_vel_xyz_exp, yaw_aligned +from isaaclab.envs.mdp import * +from isaaclab_contrib.mdp import * diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/agents/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/agents/__init__.py new file mode 100644 index 000000000000..460a30569089 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/agents/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/fixed_base_upper_body_ik_g1_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/fixed_base_upper_body_ik_g1_env_cfg.py index cb9060cca9fe..a34748fa2ee9 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/fixed_base_upper_body_ik_g1_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/fixed_base_upper_body_ik_g1_env_cfg.py @@ -3,9 +3,7 @@ # # SPDX-License-Identifier: BSD-3-Clause - -from isaaclab_teleop.isaac_teleop_cfg import IsaacTeleopCfg -from isaaclab_teleop.xr_cfg import XrCfg +from isaaclab_teleop import IsaacTeleopCfg, XrCfg import isaaclab.envs.mdp as base_mdp import isaaclab.sim as sim_utils @@ -39,8 +37,10 @@ def _build_g1_upper_body_pipeline(): via TensorReorderer. Returns: - OutputCombiner with a single "action" output containing the flattened - 28D action tensor: [left_wrist(7), right_wrist(7), hand_joints(14)]. + Tuple of (OutputCombiner, list): the pipeline with a single "action" + output containing the flattened 28D action tensor + [left_wrist(7), right_wrist(7), hand_joints(14)], and the list of + retargeter instances [left_se3, right_se3] for tuning UI. """ from isaacteleop.retargeters import ( Se3AbsRetargeter, @@ -391,8 +391,10 @@ def __post_init__(self): anchor_pos=(0.0, 0.0, -0.45), anchor_rot=(0.0, 0.0, 0.0, 1.0), ) + self.isaac_teleop = IsaacTeleopCfg( - pipeline_builder=_build_g1_upper_body_pipeline, + pipeline_builder=lambda: _build_g1_upper_body_pipeline()[0], + # retargeters_to_tune=lambda: _build_g1_upper_body_pipeline()[1], sim_device=self.sim.device, xr_cfg=self.xr, ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/locomanipulation_g1_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/locomanipulation_g1_env_cfg.py index 18ba0cd8b754..1a2ef2826de4 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/locomanipulation_g1_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/locomanipulation_g1_env_cfg.py @@ -3,16 +3,7 @@ # # SPDX-License-Identifier: BSD-3-Clause -import logging - -try: - import isaacteleop # noqa: F401 -- pipeline builders need isaacteleop at runtime - from isaaclab_teleop import IsaacTeleopCfg, XrAnchorRotationMode, XrCfg - - _TELEOP_AVAILABLE = True -except ImportError: - _TELEOP_AVAILABLE = False - logging.getLogger(__name__).warning("isaaclab_teleop is not installed. XR teleoperation features will be disabled.") +from isaaclab_teleop import IsaacTeleopCfg, XrAnchorRotationMode, XrCfg import isaaclab.envs.mdp as base_mdp import isaaclab.sim as sim_utils @@ -438,17 +429,16 @@ def __post_init__(self): # Set the URDF path for the IK controller. Path resolution (Nucleus → local) happens at runtime. self.actions.upper_body_ik.controller.urdf_path = f"{ISAACLAB_NUCLEUS_DIR}/Controllers/LocomanipulationAssets/unitree_g1_kinematics_asset/g1_29dof_with_hand_only_kinematics.urdf" # noqa: E501 - if _TELEOP_AVAILABLE: - self.xr = XrCfg( - anchor_pos=(0.0, 0.0, -0.95), - anchor_rot=(0.0, 0.0, 0.0, 1.0), - ) - self.xr.anchor_prim_path = "/World/envs/env_0/Robot/pelvis" - self.xr.fixed_anchor_height = True - self.xr.anchor_rotation_mode = XrAnchorRotationMode.FOLLOW_PRIM_SMOOTHED - - self.isaac_teleop = IsaacTeleopCfg( - pipeline_builder=_build_g1_locomanipulation_pipeline, - sim_device=self.sim.device, - xr_cfg=self.xr, - ) + self.xr = XrCfg( + anchor_pos=(0.0, 0.0, -0.95), + anchor_rot=(0.0, 0.0, 0.0, 1.0), + ) + self.xr.anchor_prim_path = "/World/envs/env_0/Robot/pelvis" + self.xr.fixed_anchor_height = True + self.xr.anchor_rotation_mode = XrAnchorRotationMode.FOLLOW_PRIM_SMOOTHED + + self.isaac_teleop = IsaacTeleopCfg( + pipeline_builder=_build_g1_locomanipulation_pipeline, + sim_device=self.sim.device, + xr_cfg=self.xr, + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/__init__.py index cc0fe238f0fe..d5fea96a02fd 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/__init__.py @@ -8,4 +8,4 @@ from isaaclab.utils.module import lazy_export -lazy_export(packages=["isaaclab.envs.mdp"]) +lazy_export() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/__init__.pyi index c139c6ae98fe..7398059fa3ef 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/__init__.pyi @@ -13,3 +13,4 @@ __all__ = [ from .actions import AgileBasedLowerBodyAction from .observations import upper_body_last_action from .terminations import object_too_far_from_robot, task_done_pick_place_table_frame +from isaaclab.envs.mdp import * diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/anymal_d/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/anymal_d/__init__.py index 41b0398e5206..f5556701ee02 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/anymal_d/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/anymal_d/__init__.py @@ -18,9 +18,13 @@ kwargs={ "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalDFlatEnvCfg", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDFlatPPORunnerCfg", + "rsl_rl_recurrent_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDFlatPPORunnerRecurrentCfg", "rsl_rl_distillation_cfg_entry_point": ( f"{agents.__name__}.rsl_rl_distillation_cfg:AnymalDFlatDistillationRunnerCfg" ), + "rsl_rl_distillation_recurrent_cfg_entry_point": ( + f"{agents.__name__}.rsl_rl_distillation_cfg:AnymalDFlatDistillationRunnerRecurrentCfg" + ), "rsl_rl_with_symmetry_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDFlatPPORunnerWithSymmetryCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, @@ -33,9 +37,13 @@ kwargs={ "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalDFlatEnvCfg_PLAY", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDFlatPPORunnerCfg", + "rsl_rl_recurrent_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDFlatPPORunnerRecurrentCfg", "rsl_rl_distillation_cfg_entry_point": ( f"{agents.__name__}.rsl_rl_distillation_cfg:AnymalDFlatDistillationRunnerCfg" ), + "rsl_rl_distillation_recurrent_cfg_entry_point": ( + f"{agents.__name__}.rsl_rl_distillation_cfg:AnymalDFlatDistillationRunnerRecurrentCfg" + ), "rsl_rl_with_symmetry_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDFlatPPORunnerWithSymmetryCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", }, diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/anymal_d/agents/rsl_rl_distillation_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/anymal_d/agents/rsl_rl_distillation_cfg.py index ea3d5f521ac3..faf362709ed0 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/anymal_d/agents/rsl_rl_distillation_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/anymal_d/agents/rsl_rl_distillation_cfg.py @@ -8,7 +8,8 @@ from isaaclab_rl.rsl_rl import ( RslRlDistillationAlgorithmCfg, RslRlDistillationRunnerCfg, - RslRlDistillationStudentTeacherCfg, + RslRlMLPModelCfg, + RslRlRNNModelCfg, ) @@ -18,18 +19,43 @@ class AnymalDFlatDistillationRunnerCfg(RslRlDistillationRunnerCfg): max_iterations = 300 save_interval = 50 experiment_name = "anymal_d_flat" - obs_groups = {"policy": ["policy"], "teacher": ["policy"]} - policy = RslRlDistillationStudentTeacherCfg( - init_noise_std=0.1, - noise_std_type="scalar", - student_obs_normalization=False, - teacher_obs_normalization=False, - student_hidden_dims=[128, 128, 128], - teacher_hidden_dims=[128, 128, 128], + obs_groups = {"student": ["policy"], "teacher": ["policy"]} + student = RslRlMLPModelCfg( + hidden_dims=[128, 128, 128], activation="elu", + obs_normalization=False, + distribution_cfg=RslRlMLPModelCfg.GaussianDistributionCfg(init_std=0.1), + ) + teacher = RslRlMLPModelCfg( + hidden_dims=[128, 128, 128], + activation="elu", + obs_normalization=False, + distribution_cfg=RslRlMLPModelCfg.GaussianDistributionCfg(init_std=0.0), ) algorithm = RslRlDistillationAlgorithmCfg( num_learning_epochs=2, learning_rate=1.0e-3, gradient_length=15, ) + + +@configclass +class AnymalDFlatDistillationRunnerRecurrentCfg(AnymalDFlatDistillationRunnerCfg): + student = RslRlRNNModelCfg( + hidden_dims=[128, 128, 128], + activation="elu", + obs_normalization=False, + distribution_cfg=RslRlMLPModelCfg.GaussianDistributionCfg(init_std=0.1), + rnn_type="lstm", + rnn_hidden_dim=256, + rnn_num_layers=1, + ) + teacher = RslRlRNNModelCfg( + hidden_dims=[128, 128, 128], + activation="elu", + obs_normalization=False, + distribution_cfg=RslRlMLPModelCfg.GaussianDistributionCfg(init_std=0.0), + rnn_type="lstm", + rnn_hidden_dim=256, + rnn_num_layers=1, + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/anymal_d/agents/rsl_rl_ppo_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/anymal_d/agents/rsl_rl_ppo_cfg.py index 220efdd6e8c4..67690049ee6f 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/anymal_d/agents/rsl_rl_ppo_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/anymal_d/agents/rsl_rl_ppo_cfg.py @@ -5,7 +5,13 @@ from isaaclab.utils import configclass -from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlPpoActorCriticCfg, RslRlPpoAlgorithmCfg, RslRlSymmetryCfg +from isaaclab_rl.rsl_rl import ( + RslRlMLPModelCfg, + RslRlOnPolicyRunnerCfg, + RslRlPpoAlgorithmCfg, + RslRlRNNModelCfg, + RslRlSymmetryCfg, +) from isaaclab_tasks.manager_based.locomotion.velocity.mdp.symmetry import anymal @@ -16,13 +22,17 @@ class AnymalDRoughPPORunnerCfg(RslRlOnPolicyRunnerCfg): max_iterations = 1500 save_interval = 50 experiment_name = "anymal_d_rough" - policy = RslRlPpoActorCriticCfg( - init_noise_std=1.0, - actor_obs_normalization=False, - critic_obs_normalization=False, - actor_hidden_dims=[512, 256, 128], - critic_hidden_dims=[512, 256, 128], + obs_groups = {"actor": ["policy"], "critic": ["policy"]} + actor = RslRlMLPModelCfg( + hidden_dims=[512, 256, 128], activation="elu", + obs_normalization=False, + distribution_cfg=RslRlMLPModelCfg.GaussianDistributionCfg(init_std=1.0), + ) + critic = RslRlMLPModelCfg( + hidden_dims=[512, 256, 128], + activation="elu", + obs_normalization=False, ) algorithm = RslRlPpoAlgorithmCfg( value_loss_coef=1.0, @@ -47,8 +57,29 @@ def __post_init__(self): self.max_iterations = 300 self.experiment_name = "anymal_d_flat" - self.policy.actor_hidden_dims = [128, 128, 128] - self.policy.critic_hidden_dims = [128, 128, 128] + self.actor.hidden_dims = [128, 128, 128] + self.critic.hidden_dims = [128, 128, 128] + + +@configclass +class AnymalDFlatPPORunnerRecurrentCfg(AnymalDFlatPPORunnerCfg): + actor = RslRlRNNModelCfg( + hidden_dims=[128, 128, 128], + activation="elu", + obs_normalization=False, + distribution_cfg=RslRlMLPModelCfg.GaussianDistributionCfg(init_std=1.0), + rnn_type="lstm", + rnn_hidden_dim=256, + rnn_num_layers=1, + ) + critic = RslRlRNNModelCfg( + hidden_dims=[128, 128, 128], + activation="elu", + obs_normalization=False, + rnn_type="lstm", + rnn_hidden_dim=256, + rnn_num_layers=1, + ) @configclass diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/__init__.py index 66d6b0cbceca..8f9a146abdc8 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/__init__.py @@ -7,4 +7,4 @@ from isaaclab.utils.module import lazy_export -lazy_export(packages=["isaaclab.envs.mdp"]) +lazy_export() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/__init__.py index 8aa1f19834be..3feb2d711bc5 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/__init__.py @@ -7,4 +7,4 @@ from isaaclab.utils.module import lazy_export -lazy_export(packages=["isaaclab.envs.mdp"]) +lazy_export() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/__init__.pyi index a923dd8d48e2..40e046d8d22e 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/__init__.pyi @@ -34,3 +34,4 @@ from .rewards import ( multi_stage_open_drawer, open_drawer_bonus, ) +from isaaclab.envs.mdp import * diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.py index 9b1004376966..81af1c012910 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.py @@ -7,4 +7,4 @@ from isaaclab.utils.module import lazy_export -lazy_export(packages=["isaaclab.envs.mdp"]) +lazy_export() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.pyi index ed35195b924b..9f27fc7560c6 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.pyi @@ -31,3 +31,4 @@ from .rewards import ( keypoint_entity_error_exp, ) from .terminations import reset_when_gear_dropped, reset_when_gear_orientation_exceeds_threshold +from isaaclab.envs.mdp import * diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/__init__.py index e8b169e7df30..8c9e9617fce3 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/__init__.py @@ -61,60 +61,3 @@ "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:DexsuiteKukaAllegroPPORunnerCfg", }, ) - - -# # Vision-Based Environments -# gym.register( -# id="Isaac-Dexsuite-Kuka-Allegro-Lift-Single-Camera-v0", -# entry_point="isaaclab.envs:ManagerBasedRLEnv", -# disable_env_checker=True, -# kwargs={ -# "env_cfg_entry_point": ( -# f"{__name__}.dexsuite_kuka_allegro_vision_env_cfg:DexsuiteKukaAllegroLiftSingleCameraEnvCfg" -# ), -# "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_cnn_cfg:DexsuiteKukaAllegroPPOCNNRunnerCfg", -# "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cnn_cfg.yaml", -# }, -# ) - - -# gym.register( -# id="Isaac-Dexsuite-Kuka-Allegro-Lift-Single-Camera-Play-v0", -# entry_point="isaaclab.envs:ManagerBasedRLEnv", -# disable_env_checker=True, -# kwargs={ -# "env_cfg_entry_point": ( -# f"{__name__}.dexsuite_kuka_allegro_vision_env_cfg:DexsuiteKukaAllegroLiftSingleCameraEnvCfg_PLAY" -# ), -# "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cnn_cfg.yaml", -# "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_cnn_cfg:DexsuiteKukaAllegroPPOCNNRunnerCfg", -# }, -# ) - - -# gym.register( -# id="Isaac-Dexsuite-Kuka-Allegro-Lift-Duo-Camera-v0", -# entry_point="isaaclab.envs:ManagerBasedRLEnv", -# disable_env_checker=True, -# kwargs={ -# "env_cfg_entry_point": ( -# f"{__name__}.dexsuite_kuka_allegro_vision_env_cfg:DexsuiteKukaAllegroLiftDuoCameraEnvCfg" -# ), -# "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_duo_camera_cfg:DexsuiteKukaAllegroPPORunnerDuoCameraCfg", -# "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cnn_cfg.yaml", -# }, -# ) - - -# gym.register( -# id="Isaac-Dexsuite-Kuka-Allegro-Lift-Duo-Camera-Play-v0", -# entry_point="isaaclab.envs:ManagerBasedRLEnv", -# disable_env_checker=True, -# kwargs={ -# "env_cfg_entry_point": ( -# f"{__name__}.dexsuite_kuka_allegro_vision_env_cfg:DexsuiteKukaAllegroLiftDuoCameraEnvCfg_PLAY" -# ), -# "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cnn_cfg.yaml", -# "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_duo_camera_cfg:DexsuiteKukaAllegroPPORunnerDuoCameraCfg" -# }, -# ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/agents/rl_games_ppo_cfg.yaml b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/agents/rl_games_ppo_cfg.yaml deleted file mode 100644 index 3a9e96eaeb05..000000000000 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/agents/rl_games_ppo_cfg.yaml +++ /dev/null @@ -1,111 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -params: - seed: 42 - - # environment wrapper clipping - env: - clip_observations: 100.0 - clip_actions: 100.0 - obs_groups: - obs: ["policy", "proprio", "perception"] - states: ["policy", "proprio", "perception"] - concate_obs_groups: True - - algo: - name: a2c_continuous - - model: - name: continuous_a2c_logstd - - network: - name: actor_critic - separate: True - space: - continuous: - mu_activation: None - sigma_activation: None - mu_init: - name: default - sigma_init: - name: const_initializer - val: 0 - fixed_sigma: True - mlp: - units: [512, 256, 128] - activation: elu - d2rl: False - initializer: - name: default - regularizer: - name: None - - load_checkpoint: False # flag which sets whether to load the checkpoint - load_path: '' # path to the checkpoint to load - - config: - name: reorient - env_name: rlgpu - device: 'cuda:0' - device_name: 'cuda:0' - multi_gpu: False - ppo: True - mixed_precision: False - normalize_input: True - normalize_value: True - value_bootstrap: False - num_actors: -1 - reward_shaper: - scale_value: 0.01 - normalize_advantage: True - gamma: 0.99 - tau: 0.95 - learning_rate: 1e-3 - lr_schedule: adaptive - schedule_type: legacy - kl_threshold: 0.01 - score_to_win: 100000000 - max_epochs: 750000 - save_best_after: 100 - save_frequency: 50 - print_stats: True - grad_norm: 1.0 - entropy_coef: 0.001 - truncate_grads: True - e_clip: 0.2 - horizon_length: 36 - minibatch_size: 36864 - mini_epochs: 5 - critic_coef: 4 - clip_value: True - clip_actions: False - seq_len: 4 - bounds_loss_coef: 0.0001 - -pbt: - enabled: False - policy_idx: 0 # policy index in a population - num_policies: 8 # total number of policies in the population - directory: . - workspace: "pbt_workspace" # suffix of the workspace dir name inside train_dir - objective: episode.Curriculum/adr - - # PBT hyperparams - interval_steps: 50000000 - threshold_std: 0.1 - threshold_abs: 0.025 - mutation_rate: 0.25 - change_range: [1.1, 2.0] - mutation: - - agent.params.config.learning_rate: "mutate_float" - agent.params.config.grad_norm: "mutate_float" - agent.params.config.entropy_coef: "mutate_float" - agent.params.config.critic_coef: "mutate_float" - agent.params.config.bounds_loss_coef: "mutate_float" - agent.params.config.kl_threshold: "mutate_float" - agent.params.config.gamma: "mutate_discount" - agent.params.config.tau: "mutate_discount" diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/agents/rsl_rl_ppo_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/agents/rsl_rl_ppo_cfg.py index 2abc981d3060..d77f0e9fc57c 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/agents/rsl_rl_ppo_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/agents/rsl_rl_ppo_cfg.py @@ -8,34 +8,37 @@ from isaaclab.utils import configclass from isaaclab_rl.rsl_rl import ( - RslRlActorCriticCNNCfg, + RslRlCNNModelCfg, + RslRlMLPModelCfg, RslRlOnPolicyRunnerCfg, - RslRlPpoActorCriticCfg, RslRlPpoAlgorithmCfg, ) from isaaclab_tasks.utils import PresetCfg -STATE_POLICY_CFG = RslRlPpoActorCriticCfg( - init_noise_std=1.0, - actor_obs_normalization=True, - critic_obs_normalization=True, - actor_hidden_dims=[512, 256, 128], - critic_hidden_dims=[512, 256, 128], +STATE_POLICY_CFG = RslRlMLPModelCfg( + distribution_cfg=RslRlMLPModelCfg.GaussianDistributionCfg(init_std=1.0), + obs_normalization=True, + hidden_dims=[512, 256, 128], activation="elu", ) -CNN_POLICY_CFG = RslRlActorCriticCNNCfg( - init_noise_std=1.0, - actor_obs_normalization=True, - critic_obs_normalization=True, - actor_hidden_dims=[512, 256, 128], - critic_hidden_dims=[512, 256, 128], - actor_cnn_cfg=RslRlActorCriticCNNCfg.CNNCfg( - output_channels=[16, 32, 64], - kernel_size=[3, 3, 3], + +STATE_CRITIC_CFG = RslRlMLPModelCfg( + obs_normalization=True, + hidden_dims=[512, 256, 128], + activation="elu", +) + +CNN_POLICY_CFG = RslRlCNNModelCfg( + obs_normalization=True, + hidden_dims=[512, 256, 128], + distribution_cfg=RslRlCNNModelCfg.GaussianDistributionCfg(init_std=1.0), + cnn_cfg=RslRlCNNModelCfg.CNNCfg( + output_channels=[16, 32], + kernel_size=[3, 3], activation="elu", - max_pool=[True, True, True], + max_pool=[True, True], norm="batch", global_pool="avg", ), @@ -64,10 +67,10 @@ class DexsuiteKukaAllegroPPOBaseRunnerCfg(RslRlOnPolicyRunnerCfg): num_steps_per_env = 32 max_iterations = 15000 save_interval = 250 - empirical_normalization = True experiment_name = (MISSING,) # type: ignore obs_groups = (MISSING,) # type: ignore - policy = (MISSING,) # type: ignore + actor = (MISSING,) # type: ignore + critic = (MISSING,) # type: ignore algorithm = MISSING # type: ignore @@ -75,24 +78,27 @@ class DexsuiteKukaAllegroPPOBaseRunnerCfg(RslRlOnPolicyRunnerCfg): class DexsuiteKukaAllegroPPORunnerCfg(PresetCfg): default = DexsuiteKukaAllegroPPOBaseRunnerCfg().replace( experiment_name="dexsuite_kuka_allegro", - obs_groups={"policy": ["policy", "proprio", "perception"], "critic": ["policy", "proprio", "perception"]}, - policy=STATE_POLICY_CFG, + obs_groups={"actor": ["policy", "proprio", "perception"], "critic": ["policy", "proprio", "perception"]}, + actor=STATE_POLICY_CFG, + critic=STATE_CRITIC_CFG, algorithm=ALGO_CFG, ) single_camera = DexsuiteKukaAllegroPPOBaseRunnerCfg().replace( experiment_name="dexsuite_kuka_allegro_single_camera", - obs_groups={"policy": ["policy", "proprio", "base_image"], "critic": ["policy", "proprio", "perception"]}, - policy=CNN_POLICY_CFG, + obs_groups={"actor": ["policy", "proprio", "base_image"], "critic": ["policy", "proprio", "perception"]}, + actor=CNN_POLICY_CFG, + critic=STATE_CRITIC_CFG, algorithm=ALGO_CFG.replace(num_mini_batches=16), ) duo_camera = DexsuiteKukaAllegroPPOBaseRunnerCfg().replace( experiment_name="dexsuite_kuka_allegro_duo_camera", obs_groups={ - "policy": ["policy", "proprio", "base_image", "wrist_image"], + "actor": ["policy", "proprio", "base_image", "wrist_image"], "critic": ["policy", "proprio", "perception"], }, - policy=CNN_POLICY_CFG, + actor=CNN_POLICY_CFG, + critic=STATE_CRITIC_CFG, algorithm=ALGO_CFG.replace(num_mini_batches=16), ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/camera_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/camera_cfg.py index efda30050910..74b4a331adfe 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/camera_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/camera_cfg.py @@ -14,6 +14,7 @@ from isaaclab.utils.noise import UniformNoiseCfg as Unoise from isaaclab_tasks.utils import PresetCfg +from isaaclab_tasks.utils.presets import MultiBackendRendererCfg from ... import dexsuite_env_cfg as dexsuite from ... import mdp @@ -32,6 +33,7 @@ spawn=sim_utils.PinholeCameraCfg(clipping_range=(0.01, 2.5)), width=MISSING, height=MISSING, + renderer_cfg=MultiBackendRendererCfg(), ) WRIST_CAMERA_CFG = TiledCameraCfg( @@ -45,6 +47,7 @@ spawn=sim_utils.PinholeCameraCfg(clipping_range=(0.01, 2.5)), width=MISSING, height=MISSING, + renderer_cfg=MultiBackendRendererCfg(), ) @@ -61,6 +64,27 @@ class BaseTiledCameraCfg(PresetCfg): albedo64 = BASE_CAMERA_CFG.replace(data_types=["albedo"], width=64, height=64) albedo128 = BASE_CAMERA_CFG.replace(data_types=["albedo"], width=128, height=128) albedo256 = BASE_CAMERA_CFG.replace(data_types=["albedo"], width=256, height=256) + simple_shading_constant_diffuse64 = BASE_CAMERA_CFG.replace( + data_types=["simple_shading_constant_diffuse"], width=64, height=64 + ) + simple_shading_constant_diffuse128 = BASE_CAMERA_CFG.replace( + data_types=["simple_shading_constant_diffuse"], width=128, height=128 + ) + simple_shading_constant_diffuse256 = BASE_CAMERA_CFG.replace( + data_types=["simple_shading_constant_diffuse"], width=256, height=256 + ) + simple_shading_diffuse_mdl64 = BASE_CAMERA_CFG.replace( + data_types=["simple_shading_diffuse_mdl"], width=64, height=64 + ) + simple_shading_diffuse_mdl128 = BASE_CAMERA_CFG.replace( + data_types=["simple_shading_diffuse_mdl"], width=128, height=128 + ) + simple_shading_diffuse_mdl256 = BASE_CAMERA_CFG.replace( + data_types=["simple_shading_diffuse_mdl"], width=256, height=256 + ) + simple_shading_full_mdl64 = BASE_CAMERA_CFG.replace(data_types=["simple_shading_full_mdl"], width=64, height=64) + simple_shading_full_mdl128 = BASE_CAMERA_CFG.replace(data_types=["simple_shading_full_mdl"], width=128, height=128) + simple_shading_full_mdl256 = BASE_CAMERA_CFG.replace(data_types=["simple_shading_full_mdl"], width=256, height=256) default = rgb64 @@ -77,6 +101,27 @@ class WristTiledCameraCfg(PresetCfg): albedo64 = WRIST_CAMERA_CFG.replace(data_types=["albedo"], width=64, height=64) albedo128 = WRIST_CAMERA_CFG.replace(data_types=["albedo"], width=128, height=128) albedo256 = WRIST_CAMERA_CFG.replace(data_types=["albedo"], width=256, height=256) + simple_shading_constant_diffuse64 = BASE_CAMERA_CFG.replace( + data_types=["simple_shading_constant_diffuse"], width=64, height=64 + ) + simple_shading_constant_diffuse128 = BASE_CAMERA_CFG.replace( + data_types=["simple_shading_constant_diffuse"], width=128, height=128 + ) + simple_shading_constant_diffuse256 = BASE_CAMERA_CFG.replace( + data_types=["simple_shading_constant_diffuse"], width=256, height=256 + ) + simple_shading_diffuse_mdl64 = BASE_CAMERA_CFG.replace( + data_types=["simple_shading_diffuse_mdl"], width=64, height=64 + ) + simple_shading_diffuse_mdl128 = BASE_CAMERA_CFG.replace( + data_types=["simple_shading_diffuse_mdl"], width=128, height=128 + ) + simple_shading_diffuse_mdl256 = BASE_CAMERA_CFG.replace( + data_types=["simple_shading_diffuse_mdl"], width=256, height=256 + ) + simple_shading_full_mdl64 = BASE_CAMERA_CFG.replace(data_types=["simple_shading_full_mdl"], width=64, height=64) + simple_shading_full_mdl128 = BASE_CAMERA_CFG.replace(data_types=["simple_shading_full_mdl"], width=128, height=128) + simple_shading_full_mdl256 = BASE_CAMERA_CFG.replace(data_types=["simple_shading_full_mdl"], width=256, height=256) default = rgb64 diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/dexsuite_kuka_allegro_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/dexsuite_kuka_allegro_env_cfg.py index ce8ba8ef4376..c7ed7dc36bca 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/dexsuite_kuka_allegro_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/dexsuite_kuka_allegro_env_cfg.py @@ -3,6 +3,9 @@ # # SPDX-License-Identifier: BSD-3-Clause +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +from isaaclab_physx.physics import PhysxCfg + from isaaclab.assets import ArticulationCfg from isaaclab.managers import RewardTermCfg as RewTerm from isaaclab.managers import SceneEntityCfg @@ -24,6 +27,38 @@ ) FINGERTIP_LIST = ["index_link_3", "middle_link_3", "ring_link_3", "thumb_link_3"] +THUMB_SENSOR = "thumb_link_3_object_s" +FINGER_SENSORS = [f"{name}_object_s" for name in FINGERTIP_LIST if name != "thumb_link_3"] + + +@configclass +class KukaAllegroPhysicsCfg(PresetCfg): + default = PhysxCfg( + bounce_threshold_velocity=0.01, + gpu_max_rigid_patch_count=4 * 5 * 2**15, + gpu_found_lost_pairs_capacity=2**26, + gpu_found_lost_aggregate_pairs_capacity=2**29, + gpu_total_aggregate_pairs_capacity=2**25, + ) + newton = NewtonCfg( + solver_cfg=MJWarpSolverCfg( + solver="newton", + integrator="implicitfast", + njmax=300, + nconmax=70, + impratio=10.0, + cone="elliptic", + update_data_interval=2, + iterations=100, + ls_iterations=15, + ls_parallel=False, + use_mujoco_contacts=True, + ccd_iterations=5000, + ), + num_substeps=2, + debug_mode=False, + ) + physx = default @configclass @@ -62,23 +97,52 @@ class KukaAllegroRelJointPosActionCfg: @configclass class KukaAllegroReorientRewardCfg(dexsuite.RewardsCfg): - # bool awarding term if 2 finger tips are in contact with object, one of the contacting fingers has to be thumb. good_finger_contact = RewTerm( func=mdp.contacts, weight=0.5, - params={"threshold": 0.1}, + params={"threshold": 0.1, "thumb_name": THUMB_SENSOR, "finger_names": FINGER_SENSORS}, + ) + + contact_count = RewTerm( + func=mdp.contact_count, + weight=1.0, + params={ + "threshold": 0.01, + "sensor_names": FINGER_SENSORS + [THUMB_SENSOR], + }, ) def __post_init__(self: dexsuite.RewardsCfg): super().__post_init__() self.fingers_to_object.params["asset_cfg"] = SceneEntityCfg("robot", body_names=["palm_link", ".*_tip"]) + self.fingers_to_object.params["thumb_name"] = THUMB_SENSOR + self.fingers_to_object.params["finger_names"] = FINGER_SENSORS + self.position_tracking.params["thumb_name"] = THUMB_SENSOR + self.position_tracking.params["finger_names"] = FINGER_SENSORS + if self.orientation_tracking: + self.orientation_tracking.params["thumb_name"] = THUMB_SENSOR + self.orientation_tracking.params["finger_names"] = FINGER_SENSORS + self.success.params["thumb_name"] = THUMB_SENSOR + self.success.params["finger_names"] = FINGER_SENSORS @configclass class KukaAllegroObservationCfg(PresetCfg): - default = StateObservationCfg() + state = StateObservationCfg() single_camera = SingleCameraObservationsCfg() duo_camera = DuoCameraObservationsCfg() + default = state + + +@configclass +class KukaAllegroEventCfg(PresetCfg): + @configclass + class KukaAllegroPhysxEventCfg(dexsuite.StartupEventCfg, dexsuite.EventCfg): + pass + + default = KukaAllegroPhysxEventCfg() + newton = dexsuite.EventCfg() + physx = default @configclass @@ -86,10 +150,12 @@ class KukaAllegroMixinCfg: scene: KukaAllegroSceneCfg = KukaAllegroSceneCfg() rewards: KukaAllegroReorientRewardCfg = KukaAllegroReorientRewardCfg() observations: KukaAllegroObservationCfg = KukaAllegroObservationCfg() + events: KukaAllegroEventCfg = KukaAllegroEventCfg() actions: KukaAllegroRelJointPosActionCfg = KukaAllegroRelJointPosActionCfg() def __post_init__(self): super().__post_init__() + self.sim.physics = KukaAllegroPhysicsCfg() @configclass diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/dexsuite_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/dexsuite_env_cfg.py index a0687c2edaa9..70a54f1b7aad 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/dexsuite_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/dexsuite_env_cfg.py @@ -23,6 +23,8 @@ from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR from isaaclab.utils.noise import UniformNoiseCfg as Unoise +from isaaclab_tasks.utils import PresetCfg + from . import mdp from .adr_curriculum import CurriculumCfg @@ -35,6 +37,49 @@ ) +@configclass +class ObjectCfg(PresetCfg): + shapes = sim_utils.MultiAssetSpawnerCfg( + assets_cfg=[ + CuboidCfg(size=(0.05, 0.1, 0.1), physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CuboidCfg(size=(0.05, 0.05, 0.1), physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CuboidCfg(size=(0.025, 0.1, 0.1), physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CuboidCfg(size=(0.025, 0.05, 0.1), physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CuboidCfg(size=(0.025, 0.025, 0.1), physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CuboidCfg(size=(0.01, 0.1, 0.1), physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + SphereCfg(radius=0.05, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + SphereCfg(radius=0.025, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CapsuleCfg(radius=0.04, height=0.025, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CapsuleCfg(radius=0.04, height=0.01, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CapsuleCfg(radius=0.04, height=0.1, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CapsuleCfg(radius=0.025, height=0.1, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CapsuleCfg(radius=0.025, height=0.2, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + CapsuleCfg(radius=0.01, height=0.2, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + ConeCfg(radius=0.05, height=0.1, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + ConeCfg(radius=0.025, height=0.1, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), + ], + rigid_props=sim_utils.RigidBodyPropertiesCfg( + solver_position_iteration_count=16, + solver_velocity_iteration_count=0, + disable_gravity=False, + ), + collision_props=sim_utils.CollisionPropertiesCfg(), + mass_props=sim_utils.MassPropertiesCfg(mass=0.2), + ) + cube = sim_utils.CuboidCfg( + size=(0.05, 0.1, 0.1), + physics_material=RigidBodyMaterialCfg(static_friction=0.5), + rigid_props=sim_utils.RigidBodyPropertiesCfg( + solver_position_iteration_count=16, + solver_velocity_iteration_count=0, + disable_gravity=False, + ), + collision_props=sim_utils.CollisionPropertiesCfg(), + mass_props=sim_utils.MassPropertiesCfg(mass=0.2), + ) + default = shapes + + @configclass class SceneCfg(InteractiveSceneCfg): """Dexsuite Scene for multi-objects Lifting""" @@ -45,33 +90,7 @@ class SceneCfg(InteractiveSceneCfg): # object object: RigidObjectCfg = RigidObjectCfg( prim_path="{ENV_REGEX_NS}/Object", - spawn=sim_utils.MultiAssetSpawnerCfg( - assets_cfg=[ - CuboidCfg(size=(0.05, 0.1, 0.1), physics_material=RigidBodyMaterialCfg(static_friction=0.5)), - CuboidCfg(size=(0.05, 0.05, 0.1), physics_material=RigidBodyMaterialCfg(static_friction=0.5)), - CuboidCfg(size=(0.025, 0.1, 0.1), physics_material=RigidBodyMaterialCfg(static_friction=0.5)), - CuboidCfg(size=(0.025, 0.05, 0.1), physics_material=RigidBodyMaterialCfg(static_friction=0.5)), - CuboidCfg(size=(0.025, 0.025, 0.1), physics_material=RigidBodyMaterialCfg(static_friction=0.5)), - CuboidCfg(size=(0.01, 0.1, 0.1), physics_material=RigidBodyMaterialCfg(static_friction=0.5)), - SphereCfg(radius=0.05, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), - SphereCfg(radius=0.025, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), - CapsuleCfg(radius=0.04, height=0.025, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), - CapsuleCfg(radius=0.04, height=0.01, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), - CapsuleCfg(radius=0.04, height=0.1, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), - CapsuleCfg(radius=0.025, height=0.1, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), - CapsuleCfg(radius=0.025, height=0.2, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), - CapsuleCfg(radius=0.01, height=0.2, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), - ConeCfg(radius=0.05, height=0.1, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), - ConeCfg(radius=0.025, height=0.1, physics_material=RigidBodyMaterialCfg(static_friction=0.5)), - ], - rigid_props=sim_utils.RigidBodyPropertiesCfg( - solver_position_iteration_count=16, - solver_velocity_iteration_count=0, - disable_gravity=False, - ), - collision_props=sim_utils.CollisionPropertiesCfg(), - mass_props=sim_utils.MassPropertiesCfg(mass=0.2), - ), + spawn=ObjectCfg(), # type: ignore init_state=RigidObjectCfg.InitialStateCfg(pos=(-0.55, 0.1, 0.35)), ) @@ -198,8 +217,8 @@ def __post_init__(self): @configclass -class EventCfg: - """Configuration for randomization.""" +class StartupEventCfg: + """Startup-mode domain randomization (PhysX only — Newton does not support startup events).""" robot_physics_material = EventTerm( func=mdp.randomize_rigid_body_material, @@ -256,13 +275,20 @@ class EventCfg: }, ) - reset_table = EventTerm( - func=mdp.reset_root_state_uniform, + +@configclass +class EventCfg: + """Reset-mode events (shared by all physics backends).""" + + # Gravity scheduling is a deliberate curriculum trick — starting with no + # gravity (easy) and gradually introducing full gravity (hard) makes learning + # smoother and removes the need for a separate "Lift" reward. + variable_gravity = EventTerm( + func=mdp.randomize_physics_scene_gravity, mode="reset", params={ - "pose_range": {"x": [-0.05, 0.05], "y": [-0.05, 0.05], "z": [0.0, 0.0]}, - "velocity_range": {"x": [-0.0, 0.0], "y": [-0.0, 0.0], "z": [-0.0, 0.0]}, - "asset_cfg": SceneEntityCfg("table"), + "gravity_distribution_params": ([0.0, 0.0, 0.0], [0.0, 0.0, 0.0]), + "operation": "abs", }, ) @@ -312,20 +338,6 @@ class EventCfg: }, ) - # Note (Octi): This is a deliberate trick in Remake to accelerate learning. - # By scheduling gravity as a curriculum — starting with no gravity (easy) - # and gradually introducing full gravity (hard) — the agent learns more smoothly. - # This removes the need for a special "Lift" reward (often required to push the - # agent to counter gravity), which has bonus effect of simplifying reward composition overall. - variable_gravity = EventTerm( - func=mdp.randomize_physics_scene_gravity, - mode="reset", - params={ - "gravity_distribution_params": ([0.0, 0.0, 0.0], [0.0, 0.0, 0.0]), - "operation": "abs", - }, - ) - @configclass class ActionsCfg: @@ -410,7 +422,7 @@ class DexsuiteReorientEnvCfg(ManagerBasedEnvCfg): # MDP settings rewards: RewardsCfg = RewardsCfg() terminations: TerminationsCfg = TerminationsCfg() - events: EventCfg = EventCfg() + events: EventCfg = MISSING # type: ignore curriculum: CurriculumCfg | None = CurriculumCfg() def __post_init__(self): @@ -419,10 +431,10 @@ def __post_init__(self): self.decimation = 2 # 50 Hz # *single-goal setup - self.commands.object_pose.resampling_time_range = (10.0, 10.0) + self.commands.object_pose.resampling_time_range = (2.0, 3.0) self.commands.object_pose.position_only = False - self.episode_length_s = 4.0 - self.is_finite_horizon = True + self.episode_length_s = 6.0 + self.is_finite_horizon = False # simulation settings self.sim.dt = 1 / 120 @@ -435,10 +447,6 @@ def __post_init__(self): gpu_total_aggregate_pairs_capacity=2**25, ) - if self.curriculum is not None: - self.curriculum.adr.params["pos_tol"] = self.rewards.success.params["pos_std"] / 2 - self.curriculum.adr.params["rot_tol"] = self.rewards.success.params["rot_std"] / 2 - class DexsuiteLiftEnvCfg(DexsuiteReorientEnvCfg): """Dexsuite lift task definition""" @@ -449,7 +457,6 @@ def __post_init__(self): self.commands.object_pose.position_only = True if self.curriculum is not None: self.rewards.success.params["rot_std"] = None # make success reward not consider orientation - self.curriculum.adr.params["rot_tol"] = None # make adr not tracking orientation class DexsuiteReorientEnvCfg_PLAY(DexsuiteReorientEnvCfg): @@ -457,7 +464,6 @@ class DexsuiteReorientEnvCfg_PLAY(DexsuiteReorientEnvCfg): def __post_init__(self): super().__post_init__() - self.commands.object_pose.resampling_time_range = (2.0, 3.0) self.commands.object_pose.debug_vis = True self.curriculum.adr.params["init_difficulty"] = self.curriculum.adr.params["max_difficulty"] @@ -467,7 +473,6 @@ class DexsuiteLiftEnvCfg_PLAY(DexsuiteLiftEnvCfg): def __post_init__(self): super().__post_init__() - self.commands.object_pose.resampling_time_range = (2.0, 3.0) self.commands.object_pose.debug_vis = True self.commands.object_pose.position_only = True self.curriculum.adr.params["init_difficulty"] = self.curriculum.adr.params["max_difficulty"] diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/__init__.py index fe33decaeb6f..e14e0f6d52c5 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/__init__.py @@ -5,4 +5,4 @@ from isaaclab.utils.module import lazy_export -lazy_export(packages=["isaaclab.envs.mdp"]) +lazy_export() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/__init__.pyi index 92debc414a6d..4cb94f6e78ae 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/__init__.pyi @@ -13,9 +13,11 @@ __all__ = [ "object_pos_b", "object_quat_b", "vision_camera", + "time_left", "action_l2_clamped", "action_rate_l2_clamped", "contacts", + "contact_count", "object_ee_distance", "orientation_command_error_tanh", "position_command_error_tanh", @@ -33,14 +35,17 @@ from .observations import ( object_pos_b, object_quat_b, vision_camera, + time_left, ) from .rewards import ( action_l2_clamped, action_rate_l2_clamped, contacts, + contact_count, object_ee_distance, orientation_command_error_tanh, position_command_error_tanh, success_reward, ) from .terminations import abnormal_robot_state, out_of_bound +from isaaclab.envs.mdp import * diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/pose_commands.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/pose_commands.py index 626a1495c2a7..35549df614ab 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/pose_commands.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/pose_commands.py @@ -15,7 +15,6 @@ import warp as wp from isaaclab.managers import CommandTerm -from isaaclab.markers import VisualizationMarkers from isaaclab.utils.math import combine_frame_transforms, compute_pose_error, quat_from_euler_xyz, quat_unique if TYPE_CHECKING: @@ -66,7 +65,10 @@ def __init__(self, cfg: dex_cmd_cfgs.ObjectUniformPoseCommandCfg, env: ManagerBa # extract the robot and body index for which the command is generated self.robot: Articulation = env.scene[cfg.asset_name] self.object: RigidObject = env.scene[cfg.object_name] - self.success_vis_asset: RigidObject = env.scene[cfg.success_vis_asset_name] + if cfg.success_vis_asset_name in env.scene.keys(): + self.success_vis_asset: RigidObject = env.scene[cfg.success_vis_asset_name] + else: + self.success_vis_asset = None # create buffers # -- commands: (x, y, z, qw, qx, qy, qz) in root frame @@ -76,6 +78,7 @@ def __init__(self, cfg: dex_cmd_cfgs.ObjectUniformPoseCommandCfg, env: ManagerBa # -- metrics self.metrics["position_error"] = torch.zeros(self.num_envs, device=self.device) self.metrics["orientation_error"] = torch.zeros(self.num_envs, device=self.device) + from isaaclab.markers import VisualizationMarkers self.success_visualizer = VisualizationMarkers(self.cfg.success_visualizer_cfg) self.success_visualizer.set_visibility(True) @@ -124,9 +127,10 @@ def _update_metrics(self): success_id = self.metrics["position_error"] < 0.05 if not self.cfg.position_only: success_id &= self.metrics["orientation_error"] < 0.5 - self.success_visualizer.visualize( - wp.to_torch(self.success_vis_asset.data.root_pos_w), marker_indices=success_id.int() - ) + if self.success_vis_asset is not None: + self.success_visualizer.visualize( + wp.to_torch(self.success_vis_asset.data.root_pos_w), marker_indices=success_id.int() + ) def _resample_command(self, env_ids: Sequence[int]): # sample new pose targets @@ -148,12 +152,11 @@ def _update_command(self): pass def _set_debug_vis_impl(self, debug_vis: bool): - # create markers if necessary for the first tome if debug_vis: if not hasattr(self, "goal_visualizer"): - # -- goal pose + from isaaclab.markers import VisualizationMarkers + self.goal_visualizer = VisualizationMarkers(self.cfg.goal_pose_visualizer_cfg) - # -- current body pose self.curr_visualizer = VisualizationMarkers(self.cfg.curr_pose_visualizer_cfg) # set their visibility to true self.goal_visualizer.set_visibility(True) @@ -173,12 +176,15 @@ def _debug_vis_callback(self, event): # -- goal pose self.goal_visualizer.visualize(self.pose_command_w[:, :3], self.pose_command_w[:, 3:]) # -- current object pose - self.curr_visualizer.visualize(self.object.data.root_pos_w, self.object.data.root_quat_w) + obj_pos = wp.to_torch(self.object.data.root_pos_w) + obj_quat = wp.to_torch(self.object.data.root_quat_w) + self.curr_visualizer.visualize(obj_pos, obj_quat) else: - distance = torch.linalg.norm(self.pose_command_w[:, :3] - wp.to_torch(self.object.data.root_pos_w), dim=1) + obj_pos = wp.to_torch(self.object.data.root_pos_w) + distance = torch.linalg.norm(self.pose_command_w[:, :3] - obj_pos, dim=1) success_id = (distance < 0.05).int() # note: since marker indices for position is 1(far) and 2(near), we can simply shift the success_id by 1. # -- goal position self.goal_visualizer.visualize(self.pose_command_w[:, :3], marker_indices=success_id + 1) # -- current object position - self.curr_visualizer.visualize(self.object.data.root_pos_w, marker_indices=success_id + 1) + self.curr_visualizer.visualize(obj_pos, marker_indices=success_id + 1) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/curriculums.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/curriculums.py index c6b64556bf86..61f7c2390d73 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/curriculums.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/curriculums.py @@ -9,31 +9,25 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.envs import mdp -from isaaclab.managers import ManagerTermBase, SceneEntityCfg -from isaaclab.utils.math import combine_frame_transforms, compute_pose_error +from isaaclab.managers import ManagerTermBase if TYPE_CHECKING: - from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import ManagerBasedRLEnv def initial_final_interpolate_fn(env: ManagerBasedRLEnv, env_id, data, initial_value, final_value, difficulty_term_str): + """Interpolate between initial and final values scaled by the current difficulty fraction. + + Works on arbitrarily nested structures of lists/tuples; scalars (int/float) + are interpolated at the leaves. """ - Interpolate between initial value iv and final value fv, for any arbitrarily - nested structure of lists/tuples in 'data'. Scalars (int/float) are handled - at the leaves. - """ - # get the fraction scalar on the device difficulty_term: DifficultyScheduler = getattr(env.curriculum_manager.cfg, difficulty_term_str).func frac = difficulty_term.difficulty_frac if frac < 0.1: - # no-op during start, since the difficulty fraction near 0 is wasting of resource. return mdp.modify_env_param.NO_CHANGE - # convert iv/fv to tensors, but we'll peel them apart in recursion initial_value_tensor = torch.tensor(initial_value, device=env.device) final_value_tensor = torch.tensor(final_value, device=env.device) @@ -41,31 +35,24 @@ def initial_final_interpolate_fn(env: ManagerBasedRLEnv, env_id, data, initial_v def _recurse(iv_elem, fv_elem, data_elem, frac): - # If it's a sequence, rebuild the same type with each element recursed if isinstance(data_elem, Sequence) and not isinstance(data_elem, (str, bytes)): - # Note: we assume initial value element and final value element have the same structure as data return type(data_elem)(_recurse(iv_e, fv_e, d_e, frac) for iv_e, fv_e, d_e in zip(iv_elem, fv_elem, data_elem)) - # Otherwise it's a leaf scalar: do the interpolation new_val = frac * (fv_elem - iv_elem) + iv_elem if isinstance(data_elem, int): return int(new_val.item()) - else: - # cast floats or any numeric - return new_val.item() + return new_val.item() class DifficultyScheduler(ManagerTermBase): """Adaptive difficulty scheduler for curriculum learning. - Tracks per-environment difficulty levels and adjusts them based on task performance. Difficulty increases when - position/orientation errors fall below given tolerances, and decreases otherwise (unless `promotion_only` is set). - The normalized average difficulty across environments is exposed as `difficulty_frac` for use in curriculum - interpolation. - - Args: - cfg: Configuration object specifying scheduler parameters. - env: The manager-based RL environment. + Promotes difficulty when the ``success`` reward term's sticky + :attr:`succeeded` flag is ``True`` for an environment at episode end, + meaning the agent achieved the success condition at least once during the + episode. Demotes (unless ``promotion_only`` is set) otherwise. + The normalized average difficulty across environments is exposed as + :attr:`difficulty_frac` for use in curriculum interpolation. """ def __init__(self, cfg, env): @@ -84,36 +71,15 @@ def __call__( self, env: ManagerBasedRLEnv, env_ids: Sequence[int], - asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), - object_cfg: SceneEntityCfg = SceneEntityCfg("object"), - pos_tol: float = 0.1, - rot_tol: float | None = None, init_difficulty: int = 0, min_difficulty: int = 0, max_difficulty: int = 50, promotion_only: bool = False, ): - asset: Articulation = env.scene[asset_cfg.name] - object: RigidObject = env.scene[object_cfg.name] - command = env.command_manager.get_command("object_pose") - des_pos_w, des_quat_w = combine_frame_transforms( - wp.to_torch(asset.data.root_pos_w)[env_ids], - wp.to_torch(asset.data.root_quat_w)[env_ids], - command[env_ids, :3], - command[env_ids, 3:7], - ) - pos_err, rot_err = compute_pose_error( - des_pos_w, - des_quat_w, - wp.to_torch(object.data.root_pos_w)[env_ids], - wp.to_torch(object.data.root_quat_w)[env_ids], - ) - pos_dist = torch.linalg.norm(pos_err, dim=1) - rot_dist = torch.linalg.norm(rot_err, dim=1) - move_up = (pos_dist < pos_tol) & (rot_dist < rot_tol) if rot_tol else pos_dist < pos_tol + succeeded = env.reward_manager.get_term_cfg("success").func.succeeded[env_ids] demot = self.current_adr_difficulties[env_ids] if promotion_only else self.current_adr_difficulties[env_ids] - 1 self.current_adr_difficulties[env_ids] = torch.where( - move_up, + succeeded, self.current_adr_difficulties[env_ids] + 1, demot, ).clamp(min=min_difficulty, max=max_difficulty) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/rewards.py index 74dffaa435c9..2958352ded1f 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/rewards.py @@ -5,12 +5,13 @@ from __future__ import annotations +from collections.abc import Sequence from typing import TYPE_CHECKING import torch import warp as wp -from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import ManagerTermBase, SceneEntityCfg from isaaclab.utils import math as math_utils from isaaclab.utils.math import combine_frame_transforms, compute_pose_error @@ -33,100 +34,193 @@ def action_l2_clamped(env: ManagerBasedRLEnv) -> torch.Tensor: def object_ee_distance( env: ManagerBasedRLEnv, std: float, + thumb_name: str, + finger_names: list[str], + contact_threshold: float = 1.0, object_cfg: SceneEntityCfg = SceneEntityCfg("object"), asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), ) -> torch.Tensor: - """Reward reaching the object using a tanh-kernel on end-effector distance. - - The reward is close to 1 when the maximum distance between the object and any end-effector body is small. + """Reward reaching the object using a tanh-kernel on end-effector distance with contact bonus. + + The reward is close to 1 when the distance is small. The reward is scaled by contact: + - Full reward (1x) when good contact (thumb + finger) + - Reduced reward (0.1x) when no contact + + Args: + env: The environment instance. + std: Standard deviation for tanh kernel. + thumb_name: Name of the thumb contact sensor. + finger_names: Names of the finger contact sensors. + contact_threshold: Contact force magnitude threshold. + object_cfg: Configuration for the object. + asset_cfg: Configuration for the robot asset. """ asset: RigidObject = env.scene[asset_cfg.name] - object: RigidObject = env.scene[object_cfg.name] + obj: RigidObject = env.scene[object_cfg.name] asset_pos = wp.to_torch(asset.data.body_pos_w)[:, asset_cfg.body_ids] - object_pos = wp.to_torch(object.data.root_pos_w) - object_ee_distance = torch.linalg.norm(asset_pos - object_pos[:, None, :], dim=-1).max(dim=-1).values - return 1 - torch.tanh(object_ee_distance / std) - - -def contacts(env: ManagerBasedRLEnv, threshold: float) -> torch.Tensor: - """Penalize undesired contacts as the number of violations that are above a threshold.""" - - thumb_contact_sensor: ContactSensor = env.scene.sensors["thumb_link_3_object_s"] - index_contact_sensor: ContactSensor = env.scene.sensors["index_link_3_object_s"] - middle_contact_sensor: ContactSensor = env.scene.sensors["middle_link_3_object_s"] - ring_contact_sensor: ContactSensor = env.scene.sensors["ring_link_3_object_s"] - # check if contact force is above threshold - thumb_contact = wp.to_torch(thumb_contact_sensor.data.force_matrix_w).view(env.num_envs, 3) - index_contact = wp.to_torch(index_contact_sensor.data.force_matrix_w).view(env.num_envs, 3) - middle_contact = wp.to_torch(middle_contact_sensor.data.force_matrix_w).view(env.num_envs, 3) - ring_contact = wp.to_torch(ring_contact_sensor.data.force_matrix_w).view(env.num_envs, 3) - - thumb_contact_mag = torch.linalg.norm(thumb_contact, dim=-1) - index_contact_mag = torch.linalg.norm(index_contact, dim=-1) - middle_contact_mag = torch.linalg.norm(middle_contact, dim=-1) - ring_contact_mag = torch.linalg.norm(ring_contact, dim=-1) - good_contact_cond1 = (thumb_contact_mag > threshold) & ( - (index_contact_mag > threshold) | (middle_contact_mag > threshold) | (ring_contact_mag > threshold) - ) + object_pos = wp.to_torch(obj.data.root_pos_w) + distance = torch.linalg.norm(asset_pos - object_pos[:, None, :], dim=-1).max(dim=-1).values + contact_bonus = contacts(env, contact_threshold, thumb_name, finger_names).float().clamp(0.1, 1.0) + return (1 - torch.tanh(distance / std)) * contact_bonus - return good_contact_cond1 +def _contact_force_mag(sensor: ContactSensor, num_envs: int) -> torch.Tensor: + """Extract per-environment contact force magnitude from a sensor's force_matrix_w.""" + force = wp.to_torch(sensor.data.force_matrix_w).view(num_envs, 3) + return torch.linalg.norm(force, dim=-1) -def success_reward( - env: ManagerBasedRLEnv, - command_name: str, - asset_cfg: SceneEntityCfg, - align_asset_cfg: SceneEntityCfg, - pos_std: float, - rot_std: float | None = None, -) -> torch.Tensor: - """Reward success by comparing commanded pose to the object pose using tanh kernels on error.""" - asset: RigidObject = env.scene[asset_cfg.name] - object: RigidObject = env.scene[align_asset_cfg.name] - command = env.command_manager.get_command(command_name) - des_pos_w, des_quat_w = combine_frame_transforms( - wp.to_torch(asset.data.root_pos_w), wp.to_torch(asset.data.root_quat_w), command[:, :3], command[:, 3:7] - ) - pos_err, rot_err = compute_pose_error( - des_pos_w, des_quat_w, wp.to_torch(object.data.root_pos_w), wp.to_torch(object.data.root_quat_w) - ) - pos_dist = torch.linalg.norm(pos_err, dim=1) - if not rot_std: - # square is not necessary but this help to keep the final value between having rot_std or not roughly the same - return (1 - torch.tanh(pos_dist / pos_std)) ** 2 - rot_dist = torch.linalg.norm(rot_err, dim=1) - return (1 - torch.tanh(pos_dist / pos_std)) * (1 - torch.tanh(rot_dist / rot_std)) +def contacts(env: ManagerBasedRLEnv, threshold: float, thumb_name: str, finger_names: list[str]) -> torch.Tensor: + """Reward for good contact: thumb + at least one finger above threshold. + + Args: + env: The environment instance. + threshold: Contact force magnitude threshold. + thumb_name: Name of the thumb contact sensor in the scene. + finger_names: Names of the finger contact sensors in the scene. + + Returns: + Boolean tensor indicating good contact condition per environment. + """ + thumb_mag = _contact_force_mag(env.scene.sensors[thumb_name], env.num_envs) + + any_finger_contact = torch.zeros(env.num_envs, dtype=torch.bool, device=env.device) + for finger_name in finger_names: + finger_mag = _contact_force_mag(env.scene.sensors[finger_name], env.num_envs) + any_finger_contact = any_finger_contact | (finger_mag > threshold) + + return (thumb_mag > threshold) & any_finger_contact + + +def contact_count(env: ManagerBasedRLEnv, threshold: float, sensor_names: list[str]) -> torch.Tensor: + """Count the number of contact sensors with force above threshold. + + For each sensor that detects contact above the threshold, add 1 to the total. + This provides a reward proportional to the number of fingers in contact. + + Args: + env: The environment instance. + threshold: Contact force magnitude threshold. + sensor_names: Names of the contact sensors in the scene. + + Returns: + Tensor of shape (num_envs,) with the count of sensors in contact per environment. + """ + count = torch.zeros(env.num_envs, dtype=torch.float32, device=env.device) + + for sensor_name in sensor_names: + mag = _contact_force_mag(env.scene.sensors[sensor_name], env.num_envs) + count += (mag > threshold).float() + return count / len(sensor_names) + + +class success_reward(ManagerTermBase): + """Reward success by comparing commanded pose to the object pose using tanh kernels on error. + + The reward is gated by contact: only given when thumb + at least one finger are in contact. + + Maintains a sticky ``succeeded`` boolean tensor per environment that flips to ``True`` once + the success condition is met during an episode and resets to ``False`` on environment reset. + + Args: + cfg: Configuration object specifying term parameters. + env: The manager-based RL environment. + """ + + def __init__(self, cfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + self.succeeded = torch.zeros(env.num_envs, dtype=torch.bool, device=env.device) + + def reset(self, env_ids: Sequence[int] | None = None): + if env_ids is None: + self.succeeded[:] = False + else: + self.succeeded[env_ids] = False + + def __call__( + self, + env: ManagerBasedRLEnv, + command_name: str, + asset_cfg: SceneEntityCfg, + align_asset_cfg: SceneEntityCfg, + pos_std: float, + thumb_name: str, + finger_names: list[str], + contact_threshold: float = 0.01, + rot_std: float | None = None, + ) -> torch.Tensor: + asset: RigidObject = env.scene[asset_cfg.name] + obj: RigidObject = env.scene[align_asset_cfg.name] + command = env.command_manager.get_command(command_name) + des_pos_w, des_quat_w = combine_frame_transforms( + wp.to_torch(asset.data.root_pos_w), + wp.to_torch(asset.data.root_quat_w), + command[:, :3], + command[:, 3:7], + ) + pos_err, rot_err = compute_pose_error( + des_pos_w, + des_quat_w, + wp.to_torch(obj.data.root_pos_w), + wp.to_torch(obj.data.root_quat_w), + ) + pos_dist = torch.linalg.norm(pos_err, dim=1) + contact_mask = contacts(env, contact_threshold, thumb_name, finger_names) + + if rot_std: + rot_dist = torch.linalg.norm(rot_err, dim=1) + reward = (1 - torch.tanh(pos_dist / pos_std)) * (1 - torch.tanh(rot_dist / rot_std)) * contact_mask.float() + self.succeeded |= contact_mask & (pos_dist < pos_std) & (rot_dist < rot_std) + else: + reward = ((1 - torch.tanh(pos_dist / pos_std)) ** 2) * contact_mask.float() + self.succeeded |= contact_mask & (pos_dist < pos_std) + + return reward def position_command_error_tanh( - env: ManagerBasedRLEnv, std: float, command_name: str, asset_cfg: SceneEntityCfg, align_asset_cfg: SceneEntityCfg + env: ManagerBasedRLEnv, + std: float, + command_name: str, + asset_cfg: SceneEntityCfg, + align_asset_cfg: SceneEntityCfg, + thumb_name: str, + finger_names: list[str], + contact_threshold: float = 0.1, ) -> torch.Tensor: """Reward tracking of commanded position using tanh kernel, gated by contact presence.""" asset: RigidObject = env.scene[asset_cfg.name] - object: RigidObject = env.scene[align_asset_cfg.name] + obj: RigidObject = env.scene[align_asset_cfg.name] command = env.command_manager.get_command(command_name) - # obtain the desired and current positions des_pos_b = command[:, :3] des_pos_w, _ = combine_frame_transforms( - wp.to_torch(asset.data.root_pos_w), wp.to_torch(asset.data.root_quat_w), des_pos_b + wp.to_torch(asset.data.root_pos_w), + wp.to_torch(asset.data.root_quat_w), + des_pos_b, ) - distance = torch.linalg.norm(wp.to_torch(object.data.root_pos_w) - des_pos_w, dim=1) - return (1 - torch.tanh(distance / std)) * contacts(env, 1.0).float() + distance = torch.linalg.norm(wp.to_torch(obj.data.root_pos_w) - des_pos_w, dim=1) + return (1 - torch.tanh(distance / std)) * contacts(env, contact_threshold, thumb_name, finger_names).float() def orientation_command_error_tanh( - env: ManagerBasedRLEnv, std: float, command_name: str, asset_cfg: SceneEntityCfg, align_asset_cfg: SceneEntityCfg + env: ManagerBasedRLEnv, + std: float, + command_name: str, + asset_cfg: SceneEntityCfg, + align_asset_cfg: SceneEntityCfg, + thumb_name: str, + finger_names: list[str], + contact_threshold: float = 0.1, ) -> torch.Tensor: """Reward tracking of commanded orientation using tanh kernel, gated by contact presence.""" asset: RigidObject = env.scene[asset_cfg.name] - object: RigidObject = env.scene[align_asset_cfg.name] + obj: RigidObject = env.scene[align_asset_cfg.name] command = env.command_manager.get_command(command_name) - # obtain the desired and current orientations des_quat_b = command[:, 3:7] - des_quat_w = math_utils.quat_mul(wp.to_torch(asset.data.root_link_pose_w)[:, 3:7], des_quat_b) - quat_distance = math_utils.quat_error_magnitude(wp.to_torch(object.data.root_quat_w), des_quat_w) + root_state = wp.to_torch(asset.data.root_state_w) + des_quat_w = math_utils.quat_mul(root_state[:, 3:7], des_quat_b) + quat_distance = math_utils.quat_error_magnitude(wp.to_torch(obj.data.root_quat_w), des_quat_w) - return (1 - torch.tanh(quat_distance / std)) * contacts(env, 1.0).float() + return (1 - torch.tanh(quat_distance / std)) * contacts(env, contact_threshold, thumb_name, finger_names).float() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/terminations.py index 8d40665be119..8490847814d6 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/terminations.py @@ -31,22 +31,19 @@ class out_of_bound(ManagerTermBase): """ def __init__(self, cfg: TerminationTermCfg, env: ManagerBasedRLEnv): - """Initialize the termination term. - - Args: - cfg: The termination term configuration. - env: The environment instance. - """ super().__init__(cfg, env) - # Cache asset reference asset_cfg: SceneEntityCfg = cfg.params.get("asset_cfg", SceneEntityCfg("object")) self._object: RigidObject = env.scene[asset_cfg.name] - # Cache ranges tensor in_bound_range: dict[str, tuple[float, float]] = cfg.params.get("in_bound_range", {}) range_list = [in_bound_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]] - self._ranges = torch.tensor(range_list, device=env.device, dtype=torch.float32) + ranges = torch.tensor(range_list, device=env.device, dtype=torch.float32) + + # Pre-apply env_origins so we can compare directly against world-space positions. + origins = env.scene.env_origins # (N, 3) + self._lower = origins + ranges[:, 0] # (N, 3) + self._upper = origins + ranges[:, 1] # (N, 3) def __call__( self, @@ -54,19 +51,8 @@ def __call__( asset_cfg: SceneEntityCfg = SceneEntityCfg("object"), in_bound_range: dict[str, tuple[float, float]] = {}, ) -> torch.Tensor: - """Check if the object is out of bounds. - - Args: - env: The environment (unused, cached in __init__). - asset_cfg: The object configuration (unused, cached in __init__). - in_bound_range: The bound ranges (unused, cached in __init__). - - Returns: - Boolean tensor indicating which environments have objects out of bounds. - """ - object_pos_local = wp.to_torch(self._object.data.root_pos_w) - env.scene.env_origins - outside_bounds = ((object_pos_local < self._ranges[:, 0]) | (object_pos_local > self._ranges[:, 1])).any(dim=1) - return outside_bounds + pos_w = wp.to_torch(self._object.data.root_pos_w) + return ((pos_w < self._lower) | (pos_w > self._upper)).any(dim=1) def abnormal_robot_state(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/__init__.py index d9e133f51ee4..a69b931a8cfb 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/__init__.py @@ -7,4 +7,4 @@ from isaaclab.utils.module import lazy_export -lazy_export(packages=["isaaclab.envs.mdp"]) +lazy_export() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/__init__.pyi index 00531b8eed0f..fac53547437a 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/__init__.pyi @@ -21,3 +21,4 @@ from .events import reset_joints_within_limits_range from .observations import goal_quat_diff from .rewards import success_bonus, track_orientation_inv_l2, track_pos_l2 from .terminations import max_consecutive_success, object_away_from_goal, object_away_from_robot +from isaaclab.envs.mdp import * diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/__init__.py index fedbeb243fdd..1c3431f2637c 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/__init__.py @@ -7,4 +7,4 @@ from isaaclab.utils.module import lazy_export -lazy_export(packages=["isaaclab.envs.mdp"]) +lazy_export() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/__init__.pyi index 2f18398bf95a..a3487daa9a59 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/__init__.pyi @@ -14,3 +14,4 @@ __all__ = [ from .observations import object_position_in_robot_root_frame from .rewards import object_ee_distance, object_goal_distance, object_is_lifted from .terminations import object_reached_goal +from isaaclab.envs.mdp import * diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/terminations.py index 756cbae90492..0f1a392fe5a0 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/terminations.py @@ -14,6 +14,7 @@ from typing import TYPE_CHECKING import torch +import warp as wp from isaaclab.managers import SceneEntityCfg from isaaclab.utils.math import combine_frame_transforms @@ -46,9 +47,13 @@ def object_reached_goal( 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) + # Convert to torch for combine_frame_transforms (robot data may be Warp arrays under Newton) + root_pos_w = wp.to_torch(robot.data.root_pos_w) + root_quat_w = wp.to_torch(robot.data.root_quat_w) + des_pos_w, _ = combine_frame_transforms(root_pos_w, root_quat_w, des_pos_b) # distance of the end-effector to the object: (num_envs,) - distance = torch.linalg.norm(des_pos_w - object.data.root_pos_w[:, :3], dim=1) + object_pos_w = wp.to_torch(object.data.root_pos_w) + distance = torch.linalg.norm(des_pos_w - object_pos_w[:, :3], dim=1) # rewarded if the object is lifted above the threshold return distance < threshold diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/__init__.py index fedbeb243fdd..1c3431f2637c 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/__init__.py @@ -7,4 +7,4 @@ from isaaclab.utils.module import lazy_export -lazy_export(packages=["isaaclab.envs.mdp"]) +lazy_export() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/__init__.pyi index 5896be48b519..1421a52bc546 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/__init__.pyi @@ -24,3 +24,4 @@ from .observations import ( ) from .pick_place_events import reset_object_poses_nut_pour from .terminations import task_done_exhaust_pipe, task_done_nut_pour, task_done_pick_place +from isaaclab.envs.mdp import * diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/config/agibot/place_toy2box_rmp_rel_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/config/agibot/place_toy2box_rmp_rel_env_cfg.py index bf47a30f1fc8..4d70b1571a5f 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/config/agibot/place_toy2box_rmp_rel_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/config/agibot/place_toy2box_rmp_rel_env_cfg.py @@ -244,7 +244,6 @@ def __post_init__(self): controller=AGIBOT_RIGHT_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=self.use_relative_mode, ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/config/agibot/place_upright_mug_rmp_rel_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/config/agibot/place_upright_mug_rmp_rel_env_cfg.py index 29206a8f486e..ad9b4d9292ec 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/config/agibot/place_upright_mug_rmp_rel_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/config/agibot/place_upright_mug_rmp_rel_env_cfg.py @@ -197,7 +197,6 @@ def __post_init__(self): controller=AGIBOT_LEFT_ARM_RMPFLOW_CFG, scale=1.0, body_offset=RMPFlowActionCfg.OffsetCfg(pos=[0.0, 0.0, 0.0], rot=[0.0, -0.7071, 0.0, 0.7071]), - articulation_prim_expr="/World/envs/env_.*/Robot", use_relative_mode=self.use_relative_mode, ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/__init__.py index 074e37108c0b..5446e2c14ced 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/__init__.py @@ -7,4 +7,4 @@ from isaaclab.utils.module import lazy_export -lazy_export(packages=["isaaclab.envs.mdp"]) +lazy_export() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/__init__.pyi index 78ef67f21370..dcdd348b8bd6 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/__init__.pyi @@ -12,3 +12,4 @@ __all__ = [ from .observations import object_grasped, object_poses_in_base_frame from .terminations import object_a_is_into_b, object_placed_upright +from isaaclab.envs.mdp import * diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/__init__.py index 66d6b0cbceca..8f9a146abdc8 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/__init__.py @@ -7,4 +7,4 @@ from isaaclab.utils.module import lazy_export -lazy_export(packages=["isaaclab.envs.mdp"]) +lazy_export() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/__init__.pyi index 103991042109..90535ff4c7e0 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/__init__.pyi @@ -10,3 +10,4 @@ __all__ = [ ] from .rewards import orientation_command_error, position_command_error, position_command_error_tanh +from isaaclab.envs.mdp import * diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/reach_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/reach_env_cfg.py index 94fe3c911e4e..8d78c587aae5 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/reach_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/reach_env_cfg.py @@ -38,7 +38,7 @@ class ReachPhysicsCfg(PresetCfg): physx: PhysxCfg = PhysxCfg(bounce_threshold_velocity=0.2) newton: NewtonCfg = NewtonCfg( solver_cfg=MJWarpSolverCfg( - njmax=20, + njmax=30, nconmax=20, ls_iterations=20, cone="pyramidal", diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/galbot/stack_joint_pos_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/galbot/stack_joint_pos_env_cfg.py index 12b04bf012f9..8c8490ad1c7e 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/galbot/stack_joint_pos_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/galbot/stack_joint_pos_env_cfg.py @@ -4,18 +4,8 @@ # SPDX-License-Identifier: BSD-3-Clause -import logging - from isaaclab_physx.assets import SurfaceGripperCfg - -try: - import isaacteleop # noqa: F401 -- pipeline builders need isaacteleop at runtime - from isaaclab_teleop import IsaacTeleopCfg - - _TELEOP_AVAILABLE = True -except ImportError: - _TELEOP_AVAILABLE = False - logging.getLogger(__name__).warning("isaaclab_teleop is not installed. XR teleoperation features will be disabled.") +from isaaclab_teleop import IsaacTeleopCfg from isaaclab.assets import RigidObjectCfg from isaaclab.envs.mdp.actions.actions_cfg import SurfaceGripperBinaryActionCfg @@ -334,13 +324,11 @@ def __post_init__(self): ) # IsaacTeleop-based teleoperation pipeline (left hand) - if _TELEOP_AVAILABLE: - pipeline = _build_se3_abs_gripper_pipeline(hand_side="left") - self.isaac_teleop = IsaacTeleopCfg( - pipeline_builder=lambda: pipeline, - sim_device=self.sim.device, - xr_cfg=self.xr, - ) + self.isaac_teleop = IsaacTeleopCfg( + pipeline_builder=lambda: _build_se3_abs_gripper_pipeline(hand_side="left"), + sim_device=self.sim.device, + xr_cfg=self.xr, + ) @configclass @@ -377,10 +365,8 @@ def __post_init__(self): self.scene.ee_frame.target_frames[0].prim_path = "{ENV_REGEX_NS}/Robot/right_suction_cup_tcp_link" # IsaacTeleop-based teleoperation pipeline (right hand) - if _TELEOP_AVAILABLE: - pipeline = _build_se3_abs_gripper_pipeline(hand_side="right") - self.isaac_teleop = IsaacTeleopCfg( - pipeline_builder=lambda: pipeline, - sim_device=self.sim.device, - xr_cfg=self.xr, - ) + self.isaac_teleop = IsaacTeleopCfg( + pipeline_builder=lambda: _build_se3_abs_gripper_pipeline(hand_side="right"), + sim_device=self.sim.device, + xr_cfg=self.xr, + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/galbot/stack_rmp_rel_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/galbot/stack_rmp_rel_env_cfg.py index 236ad80dead1..1d52fd34d015 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/galbot/stack_rmp_rel_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/galbot/stack_rmp_rel_env_cfg.py @@ -52,7 +52,6 @@ def __post_init__(self): 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=self.use_relative_mode, ) @@ -86,7 +85,6 @@ def __post_init__(self): controller=GALBOT_RIGHT_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=self.use_relative_mode, ) # Set the simulation parameters diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/__init__.py index fedbeb243fdd..1c3431f2637c 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/__init__.py @@ -7,4 +7,4 @@ from isaaclab.utils.module import lazy_export -lazy_export(packages=["isaaclab.envs.mdp"]) +lazy_export() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/__init__.pyi index 9fdc0cc5ea1e..c7184d285121 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/__init__.pyi @@ -38,3 +38,4 @@ from .observations import ( object_stacked, ) from .terminations import cubes_stacked +from isaaclab.envs.mdp import * diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/__init__.py index 66d6b0cbceca..8f9a146abdc8 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/__init__.py @@ -7,4 +7,4 @@ from isaaclab.utils.module import lazy_export -lazy_export(packages=["isaaclab.envs.mdp"]) +lazy_export() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/__init__.pyi index cb4e77d8b3b4..af793c2c205a 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/__init__.pyi @@ -11,3 +11,4 @@ __all__ = [ from .pre_trained_policy_action_cfg import PreTrainedPolicyActionCfg from .rewards import heading_command_error_abs, position_command_error_tanh +from isaaclab.envs.mdp import * diff --git a/source/isaaclab_tasks/isaaclab_tasks/utils/hydra.py b/source/isaaclab_tasks/isaaclab_tasks/utils/hydra.py index 374088826bc3..065ca108720a 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/utils/hydra.py +++ b/source/isaaclab_tasks/isaaclab_tasks/utils/hydra.py @@ -411,7 +411,12 @@ def _path_reachable(sec: str, path: str) -> bool: def _apply_node(sec: str, path: str, node): """Replace a config node at the given section/path, handling root (path='').""" - node_dict = node.to_dict() if hasattr(node, "to_dict") else dict(node) + if node is None: + node_dict = None + elif hasattr(node, "to_dict"): + node_dict = node.to_dict() + else: + node_dict = dict(node) if path == "": cfgs[sec] = node hydra_cfg[sec] = node_dict diff --git a/source/isaaclab_tasks/isaaclab_tasks/utils/sim_launcher.py b/source/isaaclab_tasks/isaaclab_tasks/utils/sim_launcher.py index 1563aa1506c5..5e6a4c1068dc 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/utils/sim_launcher.py +++ b/source/isaaclab_tasks/isaaclab_tasks/utils/sim_launcher.py @@ -81,6 +81,9 @@ def _get_visualizer_types(launcher_args: argparse.Namespace | dict | None) -> se return set() if not visualizers: return set() + if isinstance(visualizers, str): + # CLI now uses comma-delimited syntax: --visualizer kit,newton,rerun + visualizers = [token.strip() for token in visualizers.split(",")] return {str(v).strip().lower() for v in visualizers if str(v).strip()} @@ -107,7 +110,7 @@ def compute_kit_requirements( Args: env_cfg: Resolved environment config (e.g. from :func:`resolve_task_config`). - launcher_args: Optional CLI args; if ``visualizer=kit`` is set, needs_kit is True. + launcher_args: Optional CLI args; if ``--visualizer`` includes ``kit``, needs_kit is True. Returns: (needs_kit, has_kit_cameras, visualizer_types) @@ -155,6 +158,21 @@ def launch_simulation( close_fn: Any = None if needs_kit: + # check if Isaac Sim is installed + import importlib.util + + if importlib.util.find_spec("omni.kit") is None: + logger.error( + "\n[ERROR] Isaac Sim is not installed or not found on PYTHONPATH.\n" + "\n" + " This environment requires Isaac Sim and Omniverse Kit.\n" + " PhysX backend and Kit visualizer currently requires Isaac Sim.\n" + "\n" + " To fix this, ensure Isaac Sim is installed and available in the current environment.\n" + "\n" + " See https://isaac-sim.github.io/IsaacLab/main/source/setup/installation for details.\n" + ) + raise SystemExit(1) from isaaclab.app import AppLauncher app_launcher = AppLauncher(launcher_args) @@ -165,7 +183,8 @@ def launch_simulation( # SimulationContext._get_cli_visualizer_types() can find it. from isaaclab.app.settings_manager import get_settings_manager - get_settings_manager().set_string("/isaaclab/visualizer", " ".join(sorted(visualizer_types))) + visualizer_str = " ".join(sorted(visualizer_types)) + get_settings_manager().set_string("/isaaclab/visualizer/types", visualizer_str) try: yield diff --git a/source/isaaclab_tasks/test/test_hydra.py b/source/isaaclab_tasks/test/test_hydra.py index 5c7c7ec7e605..39afda81a5c4 100644 --- a/source/isaaclab_tasks/test/test_hydra.py +++ b/source/isaaclab_tasks/test/test_hydra.py @@ -3,15 +3,6 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Launch Isaac Sim Simulator first.""" - -from isaaclab.app import AppLauncher - -# launch the simulator -app_launcher = AppLauncher(headless=True) -simulation_app = app_launcher.app - - """Self-contained tests for Hydra configuration utilities. These tests verify the REPLACE-only preset system without depending on @@ -656,6 +647,46 @@ def test_root_presetcfg_resolve_defaults(): assert resolved.sensor.renderer.backend == "rtx" +@configclass +class OptionalFeatureCfg: + buffer_size: int = 200 + export_path: str = "." + + +@configclass +class OptionalFeaturePresetCfg(PresetCfg): + default = None + enabled: OptionalFeatureCfg = OptionalFeatureCfg() + + +@configclass +class EnvWithOptionalFeatureCfg: + decimation: int = 4 + optional_feature: OptionalFeaturePresetCfg = OptionalFeaturePresetCfg() + + +def test_presetcfg_none_default_auto_applies(): + """PresetCfg with default=None auto-applies None without crashing.""" + env_cfg = EnvWithOptionalFeatureCfg() + agent_cfg = PresetCfgAgentCfg() + presets = {"env": collect_presets(env_cfg), "agent": collect_presets(agent_cfg)} + hydra_cfg = {"env": env_cfg.to_dict(), "agent": agent_cfg.to_dict()} + apply_overrides(env_cfg, agent_cfg, hydra_cfg, [], [], [], presets) + assert env_cfg.optional_feature is None + + +def test_presetcfg_none_default_cli_selects_enabled(): + """PresetCfg with default=None can be overridden to a real config via CLI.""" + env_cfg = EnvWithOptionalFeatureCfg() + agent_cfg = PresetCfgAgentCfg() + presets = {"env": collect_presets(env_cfg), "agent": collect_presets(agent_cfg)} + hydra_cfg = {"env": env_cfg.to_dict(), "agent": agent_cfg.to_dict()} + sel = [("env", "optional_feature", "enabled")] + apply_overrides(env_cfg, agent_cfg, hydra_cfg, [], sel, [], presets) + assert isinstance(env_cfg.optional_feature, OptionalFeatureCfg) + assert env_cfg.optional_feature.buffer_size == 200 + + def test_root_presetcfg_global_depth_resolves_nested(): """Global preset=depth on root PresetCfg also resolves nested sensor and renderer.""" env_cfg = RootPresetEnvCfg() diff --git a/source/isaaclab_tasks/test/test_lazy_export_stubs.py b/source/isaaclab_tasks/test/test_lazy_export_stubs.py new file mode 100644 index 000000000000..7a13dba6c6b0 --- /dev/null +++ b/source/isaaclab_tasks/test/test_lazy_export_stubs.py @@ -0,0 +1,82 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Verify that lazy_export() call sites use no arguments. + +Every ``__init__.py`` that calls ``lazy_export()`` should pass no arguments. +Fallback packages and wildcard re-exports are inferred from the ``.pyi`` +stub. Passing ``packages=`` is deprecated and indicates a stub that has +not been updated with the corresponding ``from pkg import *`` line. + +This test is purely static (AST-based) and requires no simulator. +""" + +import ast +import os +from pathlib import Path + +import pytest + +_SOURCE_ROOT = Path(__file__).resolve().parent.parent.parent + + +def _find_lazy_export_calls() -> list[tuple[Path, int, str]]: + """Return ``(file, lineno, source_line)`` for every ``lazy_export(...)`` with args.""" + results: list[tuple[Path, int, str]] = [] + for root, _dirs, files in os.walk(_SOURCE_ROOT): + for fname in files: + if fname != "__init__.py": + continue + path = Path(root) / fname + try: + source = path.read_text() + except OSError: + continue + if "lazy_export" not in source: + continue + + tree = ast.parse(source, filename=str(path)) + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + func = node.func + is_lazy_export = (isinstance(func, ast.Attribute) and func.attr == "lazy_export") or ( + isinstance(func, ast.Name) and func.id == "lazy_export" + ) + if not is_lazy_export: + continue + if node.args or node.keywords: + line = source.splitlines()[node.lineno - 1].strip() + results.append((path, node.lineno, line)) + + return sorted(results) + + +_VIOLATIONS = _find_lazy_export_calls() +_IDS = [f"{p.relative_to(_SOURCE_ROOT)}:{lineno}" for p, lineno, _ in _VIOLATIONS] + + +@pytest.mark.parametrize("violation", _VIOLATIONS or [None], ids=_IDS or ["no-violations"]) +def test_lazy_export_has_no_args(violation: tuple[Path, int, str] | None): + """lazy_export() must be called with no arguments.""" + if violation is None: + return + path, lineno, line = violation + pytest.fail( + f"{path.relative_to(_SOURCE_ROOT)}:{lineno}: {line}\n\n" + "lazy_export() should take no arguments. Move fallback packages into\n" + "the .pyi stub as 'from import *' and remove the packages= arg." + ) + + +def test_no_lazy_export_violations_found(): + """Canary: confirm we actually scanned files (guard against broken discovery).""" + init_count = sum( + 1 + for root, _dirs, files in os.walk(_SOURCE_ROOT) + for f in files + if f == "__init__.py" and "lazy_export" in (Path(root) / f).read_text(errors="ignore") + ) + assert init_count > 0, "No __init__.py files with lazy_export() found — discovery may be broken" diff --git a/source/isaaclab_tasks/test/test_shadow_hand_vision_presets.py b/source/isaaclab_tasks/test/test_shadow_hand_vision_presets.py new file mode 100644 index 000000000000..c4faf3b785b4 --- /dev/null +++ b/source/isaaclab_tasks/test/test_shadow_hand_vision_presets.py @@ -0,0 +1,440 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Tests for shadow hand vision environment preset combinations. + +Two test suites are provided: + +1. **Validation unit tests** — use lightweight ``types.SimpleNamespace`` mocks. + These exercise :func:`_validate_cfg` directly and do not require Isaac Sim. + +2. **Preset resolution tests** — verify that each named preset in + :class:`ShadowHandVisionTiledCameraCfg` and + :class:`~isaaclab_tasks.utils.renderer_cfg.RendererPresetCfg` resolves to the expected + concrete config class and data types, using the real config classes. + These require Isaac Sim to be launched so that the renderer cfg imports + are available. +""" + +"""Launch Isaac Sim Simulator first.""" + +from isaaclab.app import AppLauncher + +app_launcher = AppLauncher(headless=True, enable_cameras=True) +simulation_app = app_launcher.app + + +import copy # noqa: E402 +import types # noqa: E402 + +import pytest # noqa: E402 +import torch # noqa: E402 +from isaaclab_newton.renderers import NewtonWarpRendererCfg # noqa: E402 +from isaaclab_ov.renderers import OVRTXRendererCfg # noqa: E402 +from isaaclab_physx.renderers import IsaacRtxRendererCfg # noqa: E402 + +from isaaclab_tasks.direct.shadow_hand.shadow_hand_vision_env import ( # noqa: E402 + _WARP_SUPPORTED_DATA_TYPES, + ShadowHandVisionEnv, + _validate_cfg, +) +from isaaclab_tasks.direct.shadow_hand.shadow_hand_vision_env_cfg import ( # noqa: E402 + ShadowHandVisionBenchmarkEnvCfg, + ShadowHandVisionEnvCfg, +) +from isaaclab_tasks.utils.hydra import collect_presets, resolve_preset_defaults # noqa: E402 + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_cfg(renderer_type: str | None, data_types: list[str], feature_extractor_enabled: bool = True): + """Build a minimal mock cfg accepted by :func:`_validate_cfg`.""" + cfg = types.SimpleNamespace() + cfg.tiled_camera = types.SimpleNamespace( + renderer_cfg=types.SimpleNamespace(renderer_type=renderer_type), + data_types=data_types, + ) + cfg.feature_extractor = types.SimpleNamespace(enabled=feature_extractor_enabled) + return cfg + + +# --------------------------------------------------------------------------- +# Valid combinations — must not raise +# --------------------------------------------------------------------------- + +_VALID_COMBOS = [ + # renderer_type, data_types, feature_extractor_enabled + # ── Non-warp renderers accept every data type ── + (None, ["rgb"], True), + (None, ["rgb", "depth", "semantic_segmentation"], True), + (None, ["albedo"], True), + (None, ["simple_shading_constant_diffuse"], True), + (None, ["simple_shading_diffuse_mdl"], True), + (None, ["simple_shading_full_mdl"], True), + (None, ["depth"], False), # depth-only OK when CNN disabled + ("isaac_rtx", ["rgb"], True), + ("isaac_rtx", ["albedo"], True), + ("isaac_rtx", ["simple_shading_full_mdl"], True), + ("isaac_rtx", ["rgb", "depth", "semantic_segmentation"], True), + ("isaac_rtx", ["depth"], False), + pytest.param("ovrtx", ["rgb"], True, marks=pytest.mark.skip(reason="OVRTX testing disabled")), + pytest.param("ovrtx", ["albedo"], True, marks=pytest.mark.skip(reason="OVRTX testing disabled")), + pytest.param("ovrtx", ["depth"], False, marks=pytest.mark.skip(reason="OVRTX testing disabled")), + # ── Warp renderer: only rgb and depth are supported ── + ("newton_warp", ["rgb"], True), + ("newton_warp", ["depth"], False), # depth-only OK when CNN disabled + ("newton_warp", ["rgb", "depth"], True), # multiple supported types +] + + +@pytest.mark.parametrize("renderer_type,data_types,enabled", _VALID_COMBOS) +def test_valid_combinations_do_not_raise(renderer_type, data_types, enabled): + cfg = _make_cfg(renderer_type, data_types, enabled) + _validate_cfg(cfg) # must not raise + + +# --------------------------------------------------------------------------- +# Invalid combinations — must raise ValueError with a descriptive message +# --------------------------------------------------------------------------- + +_INVALID_COMBOS = [ + # renderer_type, data_types, enabled, substring expected in error message + # ── Warp does not support colour-space data types ── + ( + "newton_warp", + ["albedo"], + True, + "albedo", + ), + ( + "newton_warp", + ["simple_shading_constant_diffuse"], + True, + "simple_shading_constant_diffuse", + ), + ( + "newton_warp", + ["simple_shading_diffuse_mdl"], + True, + "simple_shading_diffuse_mdl", + ), + ( + "newton_warp", + ["simple_shading_full_mdl"], + True, + "simple_shading_full_mdl", + ), + ( + "newton_warp", + ["rgb", "depth", "semantic_segmentation"], + True, + "semantic_segmentation", + ), + # ── Depth-only with CNN enabled is not valid for training ── + ( + None, + ["depth"], + True, + "Depth-only", + ), + ( + "isaac_rtx", + ["depth"], + True, + "Depth-only", + ), + pytest.param( + "ovrtx", + ["depth"], + True, + "Depth-only", + marks=pytest.mark.skip(reason="OVRTX testing disabled"), + ), + ( + "newton_warp", + ["depth"], + True, + "Depth-only", # depth is warp-supported but CNN can't train on it + ), +] + + +@pytest.mark.parametrize("renderer_type,data_types,enabled,match", _INVALID_COMBOS) +def test_invalid_combinations_raise_value_error(renderer_type, data_types, enabled, match): + cfg = _make_cfg(renderer_type, data_types, enabled) + with pytest.raises(ValueError, match=match): + _validate_cfg(cfg) + + +# --------------------------------------------------------------------------- +# Warp supported data types constant +# --------------------------------------------------------------------------- + + +def test_warp_supported_data_types(): + """_WARP_SUPPORTED_DATA_TYPES must contain exactly rgb and depth.""" + assert {"rgb", "depth"} == _WARP_SUPPORTED_DATA_TYPES + + +# --------------------------------------------------------------------------- +# Preset resolution — camera data types +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def shadow_hand_vision_presets(): + """Collect all presets from ShadowHandVisionEnvCfg once for the module.""" + return collect_presets(ShadowHandVisionEnvCfg()) + + +_CAMERA_DATA_TYPE_PRESETS = [ + # preset_name, expected_data_types + ("default", ["rgb", "depth", "semantic_segmentation"]), + ("full", ["rgb", "depth", "semantic_segmentation"]), + ("rgb", ["rgb"]), + ("albedo", ["albedo"]), + ("simple_shading_constant_diffuse", ["simple_shading_constant_diffuse"]), + ("simple_shading_diffuse_mdl", ["simple_shading_diffuse_mdl"]), + ("simple_shading_full_mdl", ["simple_shading_full_mdl"]), + ("depth", ["depth"]), +] + + +@pytest.mark.parametrize("preset_name,expected_data_types", _CAMERA_DATA_TYPE_PRESETS) +def test_camera_preset_data_types(shadow_hand_vision_presets, preset_name, expected_data_types): + camera_presets = shadow_hand_vision_presets["tiled_camera"] + assert preset_name in camera_presets, f"Preset '{preset_name}' not found in tiled_camera presets" + resolved = camera_presets[preset_name] + assert resolved.data_types == expected_data_types, ( + f"Preset '{preset_name}': expected data_types={expected_data_types}, got {resolved.data_types}" + ) + + +@pytest.mark.parametrize("preset_name,_", _CAMERA_DATA_TYPE_PRESETS) +def test_camera_preset_cfg_is_valid(shadow_hand_vision_presets, preset_name, _): + """Every camera preset config must request at least one data type and have valid dimensions. + + Note: this is a config-level check only. It verifies that the preset is correctly wired up + (non-empty data_types, positive width/height) but does NOT run the renderer or inspect actual + pixel values. Verifying that rendered frames are non-empty requires a full integration test + that steps the simulation and checks the camera output tensors. + """ + resolved = shadow_hand_vision_presets["tiled_camera"][preset_name] + assert len(resolved.data_types) > 0, ( + f"Camera preset '{preset_name}' has an empty data_types list — nothing would be rendered." + ) + assert resolved.width > 0, f"Camera preset '{preset_name}' has non-positive width: {resolved.width}" + assert resolved.height > 0, f"Camera preset '{preset_name}' has non-positive height: {resolved.height}" + + +def test_all_camera_presets_present(shadow_hand_vision_presets): + """Every preset defined in ShadowHandVisionTiledCameraCfg is discoverable.""" + camera_presets = shadow_hand_vision_presets["tiled_camera"] + expected_names = {name for name, _ in _CAMERA_DATA_TYPE_PRESETS} + missing = expected_names - set(camera_presets.keys()) + assert not missing, f"Camera presets missing from collected presets: {missing}" + + +# --------------------------------------------------------------------------- +# Preset resolution — renderer +# --------------------------------------------------------------------------- + +_RENDERER_PRESETS = [ + # preset_name, expected_class + ("default", IsaacRtxRendererCfg), + ("isaacsim_rtx_renderer", IsaacRtxRendererCfg), + ("newton_renderer", NewtonWarpRendererCfg), + pytest.param("ovrtx_renderer", OVRTXRendererCfg, marks=pytest.mark.skip(reason="OVRTX testing disabled")), +] + + +@pytest.mark.parametrize("preset_name,expected_class", _RENDERER_PRESETS) +def test_renderer_preset_class(shadow_hand_vision_presets, preset_name, expected_class): + renderer_presets = shadow_hand_vision_presets["tiled_camera.renderer_cfg"] + assert preset_name in renderer_presets, f"Preset '{preset_name}' not found in renderer presets" + resolved = renderer_presets[preset_name] + assert isinstance(resolved, expected_class), ( + f"Renderer preset '{preset_name}': expected {expected_class.__name__}, got {type(resolved).__name__}" + ) + + +def test_warp_renderer_has_correct_renderer_type(shadow_hand_vision_presets): + """NewtonWarpRendererCfg must expose renderer_type='newton_warp' for validation to work.""" + warp_cfg = shadow_hand_vision_presets["tiled_camera.renderer_cfg"]["newton_renderer"] + assert warp_cfg.renderer_type == "newton_warp" + + +def test_all_renderer_presets_present(shadow_hand_vision_presets): + """Every preset in MultiBackendRendererCfg is discoverable.""" + renderer_presets = shadow_hand_vision_presets["tiled_camera.renderer_cfg"] + expected_names = {"default", "isaacsim_rtx_renderer", "newton_renderer", "ovrtx_renderer"} + missing = expected_names - set(renderer_presets.keys()) + assert not missing, f"Renderer presets missing from collected presets: {missing}" + + +# --------------------------------------------------------------------------- +# Cross-validation: every camera preset resolves to a valid warp combination +# when paired with the warp renderer preset +# --------------------------------------------------------------------------- + +_WARP_VALID_CAMERA_PRESETS = ["rgb", "depth"] +_WARP_INVALID_CAMERA_PRESETS = [ + "default", + "full", + "albedo", + "simple_shading_constant_diffuse", + "simple_shading_diffuse_mdl", + "simple_shading_full_mdl", +] + + +@pytest.mark.parametrize("camera_preset", _WARP_VALID_CAMERA_PRESETS) +def test_warp_with_valid_camera_preset(shadow_hand_vision_presets, camera_preset): + """Warp + {rgb, depth} camera presets must not raise (depth with CNN disabled).""" + camera_cfg = shadow_hand_vision_presets["tiled_camera"][camera_preset] + warp_cfg = shadow_hand_vision_presets["tiled_camera.renderer_cfg"]["newton_renderer"] + enabled = camera_cfg.data_types != ["depth"] # disable CNN for depth-only + cfg = _make_cfg(warp_cfg.renderer_type, camera_cfg.data_types, enabled) + _validate_cfg(cfg) # must not raise + + +@pytest.mark.parametrize("camera_preset", _WARP_INVALID_CAMERA_PRESETS) +def test_warp_with_invalid_camera_preset(shadow_hand_vision_presets, camera_preset): + """Warp + unsupported camera presets must raise ValueError.""" + camera_cfg = shadow_hand_vision_presets["tiled_camera"][camera_preset] + warp_cfg = shadow_hand_vision_presets["tiled_camera.renderer_cfg"]["newton_renderer"] + cfg = _make_cfg(warp_cfg.renderer_type, camera_cfg.data_types, True) + with pytest.raises(ValueError): + _validate_cfg(cfg) + + +# --------------------------------------------------------------------------- +# Integration: camera output tensors must contain non-zero pixel values +# --------------------------------------------------------------------------- + +_RENDER_CORRECTNESS_CASES = [ + # (renderer_preset, camera_preset, physics) — physics is "physx" or "newton" + # ── PhysX physics (default) + IsaacRTX: supports all data types ── + pytest.param(("isaacsim_rtx_renderer", "rgb", "physx"), id="physx-isaacsim_rtx-rgb"), + pytest.param(("isaacsim_rtx_renderer", "depth", "physx"), id="physx-isaacsim_rtx-depth"), + pytest.param(("isaacsim_rtx_renderer", "albedo", "physx"), id="physx-isaacsim_rtx-albedo"), + pytest.param( + ("isaacsim_rtx_renderer", "simple_shading_constant_diffuse", "physx"), + id="physx-isaacsim_rtx-simple_shading_constant_diffuse", + ), + pytest.param( + ("isaacsim_rtx_renderer", "simple_shading_diffuse_mdl", "physx"), + id="physx-isaacsim_rtx-simple_shading_diffuse_mdl", + ), + pytest.param( + ("isaacsim_rtx_renderer", "simple_shading_full_mdl", "physx"), + id="physx-isaacsim_rtx-simple_shading_full_mdl", + ), + # ── PhysX physics + Warp: only rgb and depth are supported ── + pytest.param(("newton_renderer", "rgb", "physx"), id="physx-warp-rgb"), + pytest.param(("newton_renderer", "depth", "physx"), id="physx-warp-depth"), + # ── Newton physics + Warp: Warp renderer is physics-backend agnostic ── + pytest.param(("newton_renderer", "rgb", "newton"), id="newton-warp-rgb"), + pytest.param(("newton_renderer", "depth", "newton"), id="newton-warp-depth"), + # ── Newton physics + IsaacRTX: known incompatibility — produces empty frames ── + # xfail(strict=True): if these ever pass the mark becomes a hard failure, prompting review. + pytest.param( + ("isaacsim_rtx_renderer", "rgb", "newton"), + id="newton-isaacsim_rtx-rgb", + marks=pytest.mark.xfail(strict=True, reason="Newton physics + IsaacRTX renderer produces empty frames"), + ), + pytest.param( + ("isaacsim_rtx_renderer", "depth", "newton"), + id="newton-isaacsim_rtx-depth", + marks=pytest.mark.xfail(strict=True, reason="Newton physics + IsaacRTX renderer produces empty frames"), + ), + pytest.param( + ("isaacsim_rtx_renderer", "albedo", "newton"), + id="newton-isaacsim_rtx-albedo", + marks=pytest.mark.xfail(strict=True, reason="Newton physics + IsaacRTX renderer produces empty frames"), + ), + pytest.param( + ("isaacsim_rtx_renderer", "simple_shading_constant_diffuse", "newton"), + id="newton-isaacsim_rtx-simple_shading_constant_diffuse", + marks=pytest.mark.xfail(strict=True, reason="Newton physics + IsaacRTX renderer produces empty frames"), + ), + pytest.param( + ("isaacsim_rtx_renderer", "simple_shading_diffuse_mdl", "newton"), + id="newton-isaacsim_rtx-simple_shading_diffuse_mdl", + marks=pytest.mark.xfail(strict=True, reason="Newton physics + IsaacRTX renderer produces empty frames"), + ), + pytest.param( + ("isaacsim_rtx_renderer", "simple_shading_full_mdl", "newton"), + id="newton-isaacsim_rtx-simple_shading_full_mdl", + marks=pytest.mark.xfail(strict=True, reason="Newton physics + IsaacRTX renderer produces empty frames"), + ), + # ── OVRTX: disabled ── + pytest.param( + ("ovrtx_renderer", "rgb", "physx"), + id="physx-ovrtx-rgb", + marks=pytest.mark.skip(reason="OVRTX testing disabled"), + ), +] + + +@pytest.fixture(params=_RENDER_CORRECTNESS_CASES) +def render_correctness_env(request, shadow_hand_vision_presets): + """Build an env with the specified renderer+camera+physics combination, step once, yield, close. + + Function-scoped so each parametrized case creates and closes its own env sequentially. + ``SimulationContext.clear_instance()`` (called by ``env.close()``) fully tears down the + singleton, allowing a new env with a different physics backend to be created next. + + The shared ``shadow_hand_vision_presets`` fixture is deepcopied before mutation so that + subsequent parametrized cases see clean preset configs. + """ + renderer_preset, camera_preset, physics = request.param + cfg = ShadowHandVisionBenchmarkEnvCfg() + # Wire in the requested camera and renderer presets. + camera_cfg = copy.deepcopy(shadow_hand_vision_presets["tiled_camera"][camera_preset]) + camera_cfg.renderer_cfg = copy.deepcopy(shadow_hand_vision_presets["tiled_camera.renderer_cfg"][renderer_preset]) + cfg.tiled_camera = camera_cfg + # Apply Newton presets before resolve_preset_defaults so they are not overwritten by defaults. + # Newton needs a specific solver config, a different robot USD, an articulation-based object, + # and a stripped-down event cfg (no PhysX-specific material randomization). + if physics == "newton": + cfg.sim.physics = copy.deepcopy(shadow_hand_vision_presets["sim.physics"]["newton"]) + cfg.robot_cfg = copy.deepcopy(shadow_hand_vision_presets["robot_cfg"]["newton"]) + cfg.object_cfg = copy.deepcopy(shadow_hand_vision_presets["object_cfg"]["newton"]) + if "events" in shadow_hand_vision_presets: + cfg.events = copy.deepcopy(shadow_hand_vision_presets["events"]["newton"]) + cfg = resolve_preset_defaults(cfg) + cfg.scene.num_envs = 4 + cfg.feature_extractor.write_image_to_file = False + env = ShadowHandVisionEnv(cfg) + env.reset() + actions = torch.zeros(cfg.scene.num_envs, env.action_space.shape[-1], device=env.device) + env.step(actions) + yield renderer_preset, camera_preset, physics, env + env.close() + + +def test_camera_renders_not_empty(render_correctness_env): + """Camera output must contain at least one non-zero pixel for every valid renderer+camera combo. + + Depth tensors may contain ``inf`` for background pixels (empty space). ``inf`` is replaced + with 0 before checking ``max()``; a non-zero max confirms the renderer produced geometry pixels. + + The ``newton-isaacsim_rtx-rgb`` case is marked ``xfail(strict=True)``: Newton physics + + IsaacRTX renderer is a known incompatibility that produces empty frames. If it ever starts + passing, the strict xfail will surface it as a regression for review. + """ + renderer_preset, camera_preset, physics, env = render_correctness_env + label = f"{physics}-{renderer_preset}+{camera_preset}" + camera_output = env._tiled_camera.data.output + assert len(camera_output) > 0, f"[{label}] Camera produced no output tensors at all." + for dt, tensor in camera_output.items(): + finite = torch.where(torch.isinf(tensor), torch.zeros_like(tensor), tensor) + assert finite.max() > 0, ( + f"[{label}] Camera output '{dt}' is all zeros or all inf " + f"after stepping. Tensor shape: {tensor.shape}, dtype: {tensor.dtype}." + ) diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/__init__.py b/source/isaaclab_visualizers/isaaclab_visualizers/__init__.py new file mode 100644 index 000000000000..186afae0b156 --- /dev/null +++ b/source/isaaclab_visualizers/isaaclab_visualizers/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Visualizer backends for Isaac Lab. + +Visualizers are loaded lazily by type (kit, newton, rerun, viser) via the factory in +isaaclab.visualizers. Import a specific backend only when needed: + + from isaaclab_visualizers.kit import KitVisualizer, KitVisualizerCfg + from isaaclab_visualizers.newton import NewtonVisualizer, NewtonVisualizerCfg + from isaaclab_visualizers.rerun import RerunVisualizer, RerunVisualizerCfg + from isaaclab_visualizers.viser import ViserVisualizer, ViserVisualizerCfg +""" diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/kit/__init__.py b/source/isaaclab_visualizers/isaaclab_visualizers/kit/__init__.py new file mode 100644 index 000000000000..5de5badade9c --- /dev/null +++ b/source/isaaclab_visualizers/isaaclab_visualizers/kit/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Kit visualizer backend (Isaac Sim viewport).""" + +from .kit_visualizer import KitVisualizer +from .kit_visualizer_cfg import KitVisualizerCfg + +__all__ = ["KitVisualizer", "KitVisualizerCfg"] diff --git a/source/isaaclab/isaaclab/visualizers/kit_visualizer.py b/source/isaaclab_visualizers/isaaclab_visualizers/kit/kit_visualizer.py similarity index 73% rename from source/isaaclab/isaaclab/visualizers/kit_visualizer.py rename to source/isaaclab_visualizers/isaaclab_visualizers/kit/kit_visualizer.py index dcbb6779a205..3505fb20eb79 100644 --- a/source/isaaclab/isaaclab/visualizers/kit_visualizer.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/kit/kit_visualizer.py @@ -13,19 +13,25 @@ from pxr import UsdGeom +from isaaclab.visualizers.base_visualizer import BaseVisualizer + from .kit_visualizer_cfg import KitVisualizerCfg -from .visualizer import Visualizer logger = logging.getLogger(__name__) if TYPE_CHECKING: - from isaaclab.sim.scene_data_providers import SceneDataProvider + from isaaclab.physics import BaseSceneDataProvider -class KitVisualizer(Visualizer): +class KitVisualizer(BaseVisualizer): """Kit visualizer using Isaac Sim viewport.""" def __init__(self, cfg: KitVisualizerCfg): + """Initialize Kit visualizer state. + + Args: + cfg: Kit visualizer configuration. + """ super().__init__(cfg) self.cfg: KitVisualizerCfg = cfg @@ -36,10 +42,16 @@ def __init__(self, cfg: KitVisualizerCfg): self._sim_time = 0.0 self._step_counter = 0 self._hidden_env_visibilities: dict[str, str] = {} + self._runtime_headless = bool(cfg.headless) # ---- Lifecycle ------------------------------------------------------------------------ - def initialize(self, scene_data_provider: SceneDataProvider) -> None: + def initialize(self, scene_data_provider: BaseSceneDataProvider) -> None: + """Initialize viewport resources and bind scene data provider. + + Args: + scene_data_provider: Scene data provider used by the visualizer. + """ if self._is_initialized: logger.debug("[KitVisualizer] initialize() called while already initialized.") return @@ -61,13 +73,28 @@ def initialize(self, scene_data_provider: SceneDataProvider) -> None: "[KitVisualizer] env_filter_ids filtering is cosmetic only (no perf gain) in OV; hiding other envs." ) self._apply_env_visibility(usd_stage, metadata) - cam_pos = self.cfg.camera_position - cam_target = self.cfg.camera_target - logger.info("[KitVisualizer] initialized | camera_pos=%s camera_target=%s", cam_pos, cam_target) + num_visualized_envs = len(self._env_ids) if self._env_ids is not None else int(metadata.get("num_envs", 0)) + self._log_initialization_table( + logger=logger, + title="KitVisualizer Configuration", + rows=[ + ("camera_position", self.cfg.camera_position), + ("camera_target", self.cfg.camera_target), + ("camera_source", self.cfg.camera_source), + ("num_visualized_envs", num_visualized_envs), + ("create_viewport", self.cfg.create_viewport), + ("headless", self._runtime_headless), + ], + ) self._is_initialized = True def step(self, dt: float) -> None: + """Advance visualizer/UI updates for one simulation step. + + Args: + dt: Simulation time-step in seconds. + """ if not self._is_initialized: return self._sim_time += dt @@ -79,6 +106,8 @@ def step(self, dt: float) -> None: app = omni.kit.app.get_app() if app is not None and app.is_running(): + # Keep app pumping for viewport/UI updates only. + # Simulation stepping is owned by SimulationContext. settings = get_settings_manager() settings.set_bool("/app/player/playSimulations", False) app.update() @@ -87,6 +116,7 @@ def step(self, dt: float) -> None: logger.debug("[KitVisualizer] App update skipped: %s", exc) def close(self) -> None: + """Close viewport resources and restore temporary state.""" if not self._is_initialized: return self._restore_env_visibility() @@ -99,6 +129,11 @@ def close(self) -> None: # ---- Capabilities --------------------------------------------------------------------- def is_running(self) -> bool: + """Return whether Kit app/runtime is still running. + + Returns: + ``True`` when the visualizer can continue stepping, otherwise ``False``. + """ if self._simulation_app is not None: return self._simulation_app.is_running() try: @@ -110,12 +145,22 @@ def is_running(self) -> bool: return False def is_training_paused(self) -> bool: - return False + """Return whether simulation play flag is paused in Kit settings.""" + try: + from isaaclab.app.settings_manager import get_settings_manager + + settings = get_settings_manager() + play_flag = settings.get("/app/player/playSimulations") + return play_flag is False + except Exception: + return False def supports_markers(self) -> bool: + """Kit viewport supports marker visualization through Omni UI rendering.""" return True def supports_live_plots(self) -> bool: + """Kit backend can host live plot widgets via viewport UI panels.""" return True def requires_forward_before_step(self) -> bool: @@ -129,6 +174,12 @@ def pumps_app_update(self) -> bool: def set_camera_view( self, eye: tuple[float, float, float] | list[float], target: tuple[float, float, float] | list[float] ) -> None: + """Set active viewport camera eye/target. + + Args: + eye: Camera eye position. + target: Camera look-at target. + """ if not self._is_initialized: logger.debug("[KitVisualizer] set_camera_view() ignored because visualizer is not initialized.") return @@ -137,6 +188,7 @@ def set_camera_view( # ---- Viewport + camera ---------------------------------------------------------------- def _ensure_simulation_app(self) -> None: + """Ensure a running Isaac Sim app is available and cache runtime mode.""" import omni.kit.app app = omni.kit.app.get_app() @@ -154,15 +206,27 @@ def _ensure_simulation_app(self) -> None: if sim_app is not None: self._simulation_app = sim_app - if self._simulation_app.config.get("headless", False): + self._runtime_headless = bool(self.cfg.headless or self._simulation_app.config.get("headless", False)) + if self._runtime_headless: logger.warning("[KitVisualizer] Running in headless mode. Viewport may not display.") except ImportError: pass def _setup_viewport(self, usd_stage) -> None: + """Create/resolve viewport and configure initial camera. + + Args: + usd_stage: USD stage used for camera prim setup. + """ import omni.kit.viewport.utility as vp_utils from omni.ui import DockPosition + if self._runtime_headless: + # In headless mode we keep the visualizer active but skip viewport/window setup. + self._viewport_window = None + self._viewport_api = None + return + if self.cfg.create_viewport and self.cfg.viewport_name: dock_position_name = self.cfg.dock_position.upper() dock_position_map = { @@ -185,17 +249,13 @@ def _setup_viewport(self, usd_stage) -> None: asyncio.ensure_future(self._dock_viewport_async(self.cfg.viewport_name, dock_pos)) self._create_and_assign_camera(usd_stage) else: - if self.cfg.viewport_name: - self._viewport_window = vp_utils.get_viewport_window_by_name(self.cfg.viewport_name) - if self._viewport_window is None: - logger.warning(f"[KitVisualizer] Viewport '{self.cfg.viewport_name}' not found. Using active.") - self._viewport_window = vp_utils.get_active_viewport_window() - else: - self._viewport_window = vp_utils.get_active_viewport_window() + self._viewport_window = vp_utils.get_active_viewport_window() + if self._viewport_window is None: + logger.warning("[KitVisualizer] No active viewport window found.") + self._viewport_api = None + return self._viewport_api = self._viewport_window.viewport_api - # TODO: Unify camera initialization with a renderer-level rendering_cfg/camera_cfg - # so visualizers can consume one canonical camera policy. if self.cfg.camera_source == "usd_path": if not self._set_active_camera_path(self.cfg.camera_usd_path): logger.warning( @@ -207,6 +267,7 @@ def _setup_viewport(self, usd_stage) -> None: self._set_viewport_camera(self.cfg.camera_position, self.cfg.camera_target) async def _dock_viewport_async(self, viewport_name: str, dock_position) -> None: + """Dock a created viewport window relative to main viewport.""" import omni.kit.app import omni.ui @@ -237,7 +298,7 @@ async def _dock_viewport_async(self, viewport_name: str, dock_position) -> None: viewport_window.focus() def _create_and_assign_camera(self, usd_stage) -> None: - # Create camera prim path based on viewport name (sanitize to ensure valid USD path). + """Create viewport camera prim (if needed) and set it active.""" camera_path = f"/World/Cameras/{self.cfg.viewport_name}_Camera".replace(" ", "_") camera_prim = usd_stage.GetPrimAtPath(camera_path) @@ -248,8 +309,11 @@ def _create_and_assign_camera(self, usd_stage) -> None: self._viewport_api.set_active_camera(camera_path) def _set_viewport_camera(self, position: tuple[float, float, float], target: tuple[float, float, float]) -> None: + """Apply eye/target camera view to the active viewport.""" import isaacsim.core.utils.viewports as isaacsim_viewports + if self._viewport_api is None: + return camera_path = self._viewport_api.get_active_camera() if not camera_path: camera_path = "/OmniverseKit_Persp" @@ -259,6 +323,11 @@ def _set_viewport_camera(self, position: tuple[float, float, float], target: tup ) def _set_active_camera_path(self, camera_path: str) -> bool: + """Set active camera path for viewport if the prim exists. + + Returns: + ``True`` if camera was set, otherwise ``False``. + """ if self._viewport_api is None: return False usd_stage = self._scene_data_provider.get_usd_stage() if self._scene_data_provider else None @@ -271,6 +340,7 @@ def _set_active_camera_path(self, camera_path: str) -> bool: return True def _apply_env_visibility(self, usd_stage, metadata: dict) -> None: + """Hide non-selected environments for cosmetic env filtering.""" if not self._env_ids: return num_envs = int(metadata.get("num_envs", 0)) @@ -294,6 +364,7 @@ def _apply_env_visibility(self, usd_stage, metadata: dict) -> None: attr.Set(UsdGeom.Tokens.invisible) def _restore_env_visibility(self) -> None: + """Restore environment visibilities modified by env filtering.""" if not self._hidden_env_visibilities: return usd_stage = self._scene_data_provider.get_usd_stage() if self._scene_data_provider else None diff --git a/source/isaaclab/isaaclab/visualizers/kit_visualizer_cfg.py b/source/isaaclab_visualizers/isaaclab_visualizers/kit/kit_visualizer_cfg.py similarity index 82% rename from source/isaaclab/isaaclab/visualizers/kit_visualizer_cfg.py rename to source/isaaclab_visualizers/isaaclab_visualizers/kit/kit_visualizer_cfg.py index 230242eed779..88112a6f20b4 100644 --- a/source/isaaclab/isaaclab/visualizers/kit_visualizer_cfg.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/kit/kit_visualizer_cfg.py @@ -6,8 +6,7 @@ """Configuration for Kit-based visualizer.""" from isaaclab.utils import configclass - -from .visualizer_cfg import VisualizerCfg +from isaaclab.visualizers.visualizer_cfg import VisualizerCfg @configclass @@ -20,9 +19,12 @@ class KitVisualizerCfg(VisualizerCfg): viewport_name: str | None = "Visualizer Viewport" """Viewport name to use. If None, uses active viewport.""" - create_viewport: bool = True + create_viewport: bool = False """Create new viewport with specified name and camera pose.""" + headless: bool = False + """Run without creating viewport windows when supported by the app.""" + dock_position: str = "SAME" """Dock position for new viewport. Options: 'LEFT', 'RIGHT', 'BOTTOM', 'SAME'.""" diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/newton/__init__.py b/source/isaaclab_visualizers/isaaclab_visualizers/newton/__init__.py new file mode 100644 index 000000000000..2f8287e138a2 --- /dev/null +++ b/source/isaaclab_visualizers/isaaclab_visualizers/newton/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Newton OpenGL visualizer backend.""" + +from .newton_visualizer import NewtonVisualizer +from .newton_visualizer_cfg import NewtonVisualizerCfg + +__all__ = ["NewtonVisualizer", "NewtonVisualizerCfg"] diff --git a/source/isaaclab/isaaclab/visualizers/newton_visualizer.py b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py similarity index 61% rename from source/isaaclab/isaaclab/visualizers/newton_visualizer.py rename to source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py index 8eb410d6706c..32f92b17ec8e 100644 --- a/source/isaaclab/isaaclab/visualizers/newton_visualizer.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py @@ -14,13 +14,14 @@ import warp as wp from newton.viewer import ViewerGL +from isaaclab.visualizers.base_visualizer import BaseVisualizer + from .newton_visualizer_cfg import NewtonVisualizerCfg -from .visualizer import Visualizer logger = logging.getLogger(__name__) if TYPE_CHECKING: - from isaaclab.sim.scene_data_providers import SceneDataProvider + from isaaclab.physics import BaseSceneDataProvider class NewtonViewerGL(ViewerGL): @@ -33,12 +34,21 @@ def __init__( update_frequency: int = 1, **kwargs, ): + """Initialize Newton viewer wrapper state. + + Args: + *args: Positional arguments forwarded to ``ViewerGL``. + metadata: Optional metadata shown in viewer panels. + update_frequency: Viewer refresh cadence in simulation frames. + **kwargs: Keyword arguments forwarded to ``ViewerGL``. + """ super().__init__(*args, **kwargs) self._paused_training = False self._paused_rendering = False self._metadata = metadata or {} self._fallback_draw_controls = False self._update_frequency = update_frequency + self._color_edit3_accepts_sequence: bool | None = None try: self.register_ui_callback(self._render_training_controls, position="side") @@ -46,12 +56,15 @@ def __init__( self._fallback_draw_controls = True def is_training_paused(self) -> bool: + """Return whether training is paused by viewer controls.""" return self._paused_training def is_rendering_paused(self) -> bool: + """Return whether rendering is paused by viewer controls.""" return self._paused_rendering def _render_training_controls(self, imgui): + """Render Isaac Lab-specific control widgets in the Newton viewer UI.""" imgui.separator() imgui.text("IsaacLab Controls") @@ -79,22 +92,13 @@ def _render_training_controls(self, imgui): ) def on_key_press(self, symbol, modifiers): + """Forward key presses unless UI is currently capturing input.""" if self.ui.is_capturing(): return - - try: - import pyglet # noqa: PLC0415 - except Exception: - return - - if symbol == pyglet.window.key.SPACE: - self._paused_rendering = not self._paused_rendering - self._paused = self._paused_rendering - return - super().on_key_press(symbol, modifiers) def _render_ui(self): + """Render default UI and fallback control window when callback hooks are unavailable.""" if not self._fallback_draw_controls: return super()._render_ui() @@ -111,12 +115,42 @@ def _render_ui(self): imgui.end() return None + def _coerce_color3(self, color) -> tuple[float, float, float]: + """Normalize color values from imgui/renderer into an RGB tuple.""" + if hasattr(color, "x") and hasattr(color, "y") and hasattr(color, "z"): + return (float(color.x), float(color.y), float(color.z)) + return (float(color[0]), float(color[1]), float(color[2])) + + def _color_edit3_compat(self, imgui, label: str, color): + """ + # Handle imgui.color_edit3 API differences between bindings. + # Some require a Sequence[float], others accept vector-like objects. + # This method tries both approaches, caching the one that works to avoid repeated exceptions. + # NOTE: This is a compatibility workaround, perhaps we can address the issue more directly. + """ + color_tuple = self._coerce_color3(color) + sequence_color = [color_tuple[0], color_tuple[1], color_tuple[2]] + if self._color_edit3_accepts_sequence is not False: + try: + changed, edited = imgui.color_edit3(label, sequence_color) + self._color_edit3_accepts_sequence = True + return changed, self._coerce_color3(edited) + except TypeError: + self._color_edit3_accepts_sequence = False + + try: + imvec4 = imgui.ImVec4(sequence_color[0], sequence_color[1], sequence_color[2], 1.0) + changed, edited = imgui.color_edit3(label, imvec4) + return changed, self._coerce_color3(edited) + except Exception as exc: + logger.debug("[NewtonVisualizer] color_edit3 failed for '%s': %s", label, exc) + return False, color_tuple + def _render_left_panel(self): """Override the left panel to remove the base pause checkbox.""" import newton as nt imgui = self.ui.imgui - nav_highlight_color = self.ui.get_theme_color(imgui.Col_.nav_cursor, (1.0, 1.0, 1.0, 1.0)) io = self.ui.io imgui.set_next_window_pos(imgui.ImVec2(10, 10)) @@ -170,9 +204,18 @@ def _render_left_panel(self): changed, self.renderer.draw_shadows = imgui.checkbox("Shadows", self.renderer.draw_shadows) changed, self.renderer.draw_wireframe = imgui.checkbox("Wireframe", self.renderer.draw_wireframe) - changed, self.renderer._light_color = imgui.color_edit3("Light Color", self.renderer._light_color) - changed, self.renderer.sky_upper = imgui.color_edit3("Upper Sky Color", self.renderer.sky_upper) - changed, self.renderer.sky_lower = imgui.color_edit3("Lower Sky Color", self.renderer.sky_lower) + try: + changed, self.renderer._light_color = self._color_edit3_compat( + imgui, "Light Color", self.renderer._light_color + ) + changed, self.renderer.sky_upper = self._color_edit3_compat( + imgui, "Upper Sky Color", self.renderer.sky_upper + ) + changed, self.renderer.sky_lower = self._color_edit3_compat( + imgui, "Lower Sky Color", self.renderer.sky_lower + ) + except Exception as exc: + logger.debug("[NewtonVisualizer] Rendering color controls failed: %s", exc) imgui.set_next_item_open(True, imgui.Cond_.appearing) if imgui.collapsing_header("Camera"): @@ -186,14 +229,10 @@ def _render_left_panel(self): imgui.text(f"Pitch: {self.camera.pitch:.1f}°") imgui.separator() - imgui.push_style_color(imgui.Col_.text, imgui.ImVec4(*nav_highlight_color)) - imgui.text("Controls:") - imgui.pop_style_color() imgui.text("WASD - Forward/Left/Back/Right") imgui.text("QE - Down/Up") imgui.text("Left Click - Look around") imgui.text("Scroll - Zoom") - imgui.text("Space - Pause/Resume Rendering") imgui.text("H - Toggle UI") imgui.text("ESC - Exit") @@ -201,10 +240,15 @@ def _render_left_panel(self): return -class NewtonVisualizer(Visualizer): +class NewtonVisualizer(BaseVisualizer): """Newton OpenGL visualizer for Isaac Lab.""" def __init__(self, cfg: NewtonVisualizerCfg): + """Initialize Newton visualizer state. + + Args: + cfg: Newton visualizer configuration. + """ super().__init__(cfg) self.cfg: NewtonVisualizerCfg = cfg self._viewer: NewtonViewerGL | None = None @@ -215,8 +259,14 @@ def __init__(self, cfg: NewtonVisualizerCfg): self._update_frequency = cfg.update_frequency self._scene_data_provider = None self._last_camera_pose: tuple[tuple[float, float, float], tuple[float, float, float]] | None = None + self._headless_no_viewer = False - def initialize(self, scene_data_provider: SceneDataProvider) -> None: + def initialize(self, scene_data_provider: BaseSceneDataProvider) -> None: + """Initialize viewer resources and bind scene data provider. + + Args: + scene_data_provider: Scene data provider used to fetch model/state data. + """ if self._is_initialized: logger.debug("[NewtonVisualizer] initialize() called while already initialized.") return @@ -236,49 +286,83 @@ def initialize(self, scene_data_provider: SceneDataProvider) -> None: self._model = scene_data_provider.get_newton_model() self._state = scene_data_provider.get_newton_state(self._env_ids) - self._viewer = NewtonViewerGL( - width=self.cfg.window_width, - height=self.cfg.window_height, - metadata=metadata, - update_frequency=self.cfg.update_frequency, - ) - - max_worlds = self.cfg.max_worlds - self._viewer.set_model(self._model, max_worlds=None if max_worlds in (None, 0) else max_worlds) - self._viewer.set_world_offsets((0.0, 0.0, 0.0)) - self._apply_camera_pose(self._resolve_initial_camera_pose()) - self._viewer.up_axis = 2 # Z-up - - self._viewer.scaling = 1.0 - self._viewer._paused = False - - self._viewer.show_joints = self.cfg.show_joints - self._viewer.show_contacts = self.cfg.show_contacts - self._viewer.show_springs = self.cfg.show_springs - self._viewer.show_com = self.cfg.show_com - - self._viewer.renderer.draw_shadows = self.cfg.enable_shadows - self._viewer.renderer.draw_sky = self.cfg.enable_sky - self._viewer.renderer.draw_wireframe = self.cfg.enable_wireframe - - self._viewer.renderer.sky_upper = self.cfg.sky_upper_color - self._viewer.renderer.sky_lower = self.cfg.sky_lower_color - self._viewer.renderer._light_color = self.cfg.light_color + try: + self._viewer = NewtonViewerGL( + width=self.cfg.window_width, + height=self.cfg.window_height, + headless=self.cfg.headless, + metadata=metadata, + update_frequency=self.cfg.update_frequency, + ) + except Exception as exc: + if not self.cfg.headless: + raise + self._viewer = None + self._headless_no_viewer = True + logger.info( + "[NewtonVisualizer] Headless fallback enabled (ViewerGL unavailable in this environment): %s", + exc, + ) - logger.info( - "[NewtonVisualizer] initialized | camera_pos=%s camera_target=%s", - self._viewer.camera.pos, - self._last_camera_pose[1] if self._last_camera_pose else self.cfg.camera_target, + if self._viewer is not None: + max_worlds = self.cfg.max_worlds + self._viewer.set_model(self._model, max_worlds=max_worlds) + self._viewer.set_world_offsets((0.0, 0.0, 0.0)) + self._apply_camera_pose(self._resolve_initial_camera_pose()) + self._viewer.up_axis = 2 # Z-up + + self._viewer.scaling = 1.0 + self._viewer._paused = False + + self._viewer.show_joints = self.cfg.show_joints + self._viewer.show_contacts = self.cfg.show_contacts + self._viewer.show_springs = self.cfg.show_springs + self._viewer.show_com = self.cfg.show_com + + self._viewer.renderer.draw_shadows = self.cfg.enable_shadows + self._viewer.renderer.draw_sky = self.cfg.enable_sky + self._viewer.renderer.draw_wireframe = self.cfg.enable_wireframe + + self._viewer.renderer.sky_upper = self.cfg.sky_upper_color + self._viewer.renderer.sky_lower = self.cfg.sky_lower_color + self._viewer.renderer._light_color = self.cfg.light_color + + num_visualized_envs = len(self._env_ids) if self._env_ids is not None else int(metadata.get("num_envs", 0)) + self._log_initialization_table( + logger=logger, + title="NewtonVisualizer Configuration", + rows=[ + ( + "camera_position", + tuple(float(x) for x in self._viewer.camera.pos) + if self._viewer is not None + else self.cfg.camera_position, + ), + ("camera_target", self._last_camera_pose[1] if self._last_camera_pose else self.cfg.camera_target), + ("camera_source", self.cfg.camera_source), + ("num_visualized_envs", num_visualized_envs), + ("headless", self.cfg.headless), + ], ) self._is_initialized = True def step(self, dt: float) -> None: - if not self._is_initialized or self._is_closed or self._viewer is None: + """Advance visualization by one simulation step. + + Args: + dt: Simulation time-step in seconds. + """ + if not self._is_initialized or self._is_closed: return self._sim_time += dt self._step_counter += 1 + if self._viewer is None: + if self._scene_data_provider is not None: + self._state = self._scene_data_provider.get_newton_state(self._env_ids) + return + if self.cfg.camera_source == "usd_path": self._update_camera_from_usd_path() @@ -313,10 +397,11 @@ def step(self, dt: float) -> None: self._viewer.end_frame() else: self._viewer._update() - except RuntimeError as exc: + except Exception as exc: logger.debug("[NewtonVisualizer] Viewer update failed: %s", exc) def close(self) -> None: + """Release viewer resources.""" if self._is_closed: return if self._viewer is not None: @@ -324,11 +409,25 @@ def close(self) -> None: self._is_closed = True def is_running(self) -> bool: - if not self._is_initialized or self._is_closed or self._viewer is None: + """Return whether the visualizer should continue stepping. + + Returns: + ``True`` while the visualizer is active, otherwise ``False``. + """ + if not self._is_initialized or self._is_closed: + return False + if self._headless_no_viewer and self._viewer is None: + return True + if self._viewer is None: return False return self._viewer.is_running() def _resolve_initial_camera_pose(self) -> tuple[tuple[float, float, float], tuple[float, float, float]]: + """Resolve initial camera pose from config or USD camera path. + + Returns: + Camera eye and target tuples. + """ if self.cfg.camera_source == "usd_path": pose = self._resolve_camera_pose_from_usd_path(self.cfg.camera_usd_path) if pose is not None: @@ -340,6 +439,11 @@ def _resolve_initial_camera_pose(self) -> tuple[tuple[float, float, float], tupl return self.cfg.camera_position, self.cfg.camera_target def _apply_camera_pose(self, pose: tuple[tuple[float, float, float], tuple[float, float, float]]) -> None: + """Apply camera eye/target pose to the Newton viewer. + + Args: + pose: Camera eye and target tuples. + """ if self._viewer is None: return cam_pos, cam_target = pose @@ -355,6 +459,7 @@ def _apply_camera_pose(self, pose: tuple[tuple[float, float, float], tuple[float self._last_camera_pose = (cam_pos, cam_target) def _update_camera_from_usd_path(self) -> None: + """Refresh camera pose from configured USD camera path when it changes.""" pose = self._resolve_camera_pose_from_usd_path(self.cfg.camera_usd_path) if pose is None: return @@ -363,17 +468,21 @@ def _update_camera_from_usd_path(self) -> None: self._apply_camera_pose(pose) def supports_markers(self) -> bool: + """Newton OpenGL viewer does not implement Isaac Lab marker primitives.""" return False def supports_live_plots(self) -> bool: + """Newton OpenGL viewer does not provide live-plot panels.""" return False def is_training_paused(self) -> bool: + """Return whether training is paused from viewer controls.""" if not self._is_initialized or self._viewer is None: return False return self._viewer.is_training_paused() def is_rendering_paused(self) -> bool: + """Return whether rendering is paused from viewer controls.""" if not self._is_initialized or self._viewer is None: return False return self._viewer.is_rendering_paused() diff --git a/source/isaaclab/isaaclab/visualizers/newton_visualizer_cfg.py b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer_cfg.py similarity index 85% rename from source/isaaclab/isaaclab/visualizers/newton_visualizer_cfg.py rename to source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer_cfg.py index 85900b085bb4..2449db73b5e3 100644 --- a/source/isaaclab/isaaclab/visualizers/newton_visualizer_cfg.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer_cfg.py @@ -6,8 +6,7 @@ """Configuration for Newton OpenGL Visualizer.""" from isaaclab.utils import configclass - -from .visualizer_cfg import VisualizerCfg +from isaaclab.visualizers.visualizer_cfg import VisualizerCfg @configclass @@ -23,8 +22,14 @@ class NewtonVisualizerCfg(VisualizerCfg): window_height: int = 1080 """Window height in pixels.""" - max_worlds: int | None = 0 - """Maximum number of worlds/environments rendered by the viewer (0/None = all).""" + headless: bool = False + """Run the Newton viewer without requiring a display server.""" + + max_worlds: int | None = None + """Maximum number of worlds/environments rendered by the viewer. + + Set to ``None`` to leave this option disabled. + """ update_frequency: int = 1 """Visualizer update frequency (updates every N frames).""" diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/rerun/__init__.py b/source/isaaclab_visualizers/isaaclab_visualizers/rerun/__init__.py new file mode 100644 index 000000000000..7fd7704ce009 --- /dev/null +++ b/source/isaaclab_visualizers/isaaclab_visualizers/rerun/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Rerun visualizer backend.""" + +from .rerun_visualizer import RerunVisualizer +from .rerun_visualizer_cfg import RerunVisualizerCfg + +__all__ = ["RerunVisualizer", "RerunVisualizerCfg"] diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/rerun/rerun_visualizer.py b/source/isaaclab_visualizers/isaaclab_visualizers/rerun/rerun_visualizer.py new file mode 100644 index 000000000000..013e1522c58a --- /dev/null +++ b/source/isaaclab_visualizers/isaaclab_visualizers/rerun/rerun_visualizer.py @@ -0,0 +1,326 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Rerun visualizer implementation for Isaac Lab.""" + +from __future__ import annotations + +import atexit +import logging +import socket +import webbrowser +from typing import TYPE_CHECKING +from urllib.parse import quote + +import rerun as rr +import rerun.blueprint as rrb +from newton.viewer import ViewerRerun + +from isaaclab.visualizers.base_visualizer import BaseVisualizer + +from .rerun_visualizer_cfg import RerunVisualizerCfg + +if TYPE_CHECKING: + from isaaclab.physics import BaseSceneDataProvider + +logger = logging.getLogger(__name__) + + +def _is_port_free(port: int, host: str = "127.0.0.1") -> bool: + """Return whether a TCP port can be bound on host.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + sock.bind((host, int(port))) + return True + except OSError: + return False + + +def _is_port_open(port: int, host: str = "127.0.0.1") -> bool: + """Return whether a TCP port is currently accepting connections.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(0.2) + return sock.connect_ex((host, int(port))) == 0 + + +def _normalize_host(addr: str) -> str: + """Normalize bind host to loopback-friendly address for client URLs.""" + if addr in ("0.0.0.0", "127.0.0.1", "localhost"): + return "127.0.0.1" + return addr + + +def _ensure_rerun_server(app_id: str, bind_address: str, grpc_port: int, web_port: int) -> tuple[str, bool]: + """Resolve rerun endpoint and whether viewer should start web/grpc server.""" + del app_id + connect_host = _normalize_host(bind_address) + expected_uri = f"rerun+http://{connect_host}:{int(grpc_port)}/proxy" + + if _is_port_open(grpc_port, host=connect_host): + # Reuse existing endpoint; do not create a new server here. + return expected_uri, False + + if not _is_port_free(web_port, host=connect_host): + raise RuntimeError(f"Rerun web port {web_port} is in use. Free the port or choose a different `web_port`.") + + # No existing gRPC server: NewtonViewerRerun should start and own it. + return expected_uri, True + + +def _open_rerun_web_viewer(host: str, web_port: int, connect_to: str) -> None: + """Open rerun web UI and prefill endpoint connection URL.""" + url = f"http://{host}:{int(web_port)}/?url={quote(connect_to, safe='')}" + try: + if not webbrowser.open_new_tab(url): + logger.info("[RerunVisualizer] Could not auto-open browser tab. Open manually: %s", url) + except Exception: + logger.info("[RerunVisualizer] Could not auto-open browser tab. Open manually: %s", url) + + +class NewtonViewerRerun(ViewerRerun): + """Wrapper around Newton's ViewerRerun with rendering pause controls.""" + + def __init__(self, *args, **kwargs): + """Initialize viewer wrapper and Isaac Lab pause state.""" + super().__init__(*args, **kwargs) + self._paused_rendering = False + + def is_rendering_paused(self) -> bool: + """Return whether rendering is paused by viewer controls.""" + return self._paused_rendering + + def _render_ui(self): + """Extend base UI with Isaac Lab rendering pause toggle.""" + super()._render_ui() + + if not self._has_imgui: + return + + imgui = self._imgui + if not imgui: + return + + if imgui.collapsing_header("IsaacLab Controls"): + if imgui.button("Pause Rendering" if not self._paused_rendering else "Resume Rendering"): + self._paused_rendering = not self._paused_rendering + + +class RerunVisualizer(BaseVisualizer): + """Rerun visualizer for Isaac Lab.""" + + def __init__(self, cfg: RerunVisualizerCfg): + """Initialize Rerun visualizer state. + + Args: + cfg: Rerun visualizer configuration. + """ + super().__init__(cfg) + self.cfg: RerunVisualizerCfg = cfg + self._viewer: NewtonViewerRerun | None = None + self._sim_time = 0.0 + self._step_counter = 0 + self._model = None + self._state = None + self._scene_data_provider = None + self._last_camera_pose: tuple[tuple[float, float, float], tuple[float, float, float]] | None = None + + def initialize(self, scene_data_provider: BaseSceneDataProvider) -> None: + """Initialize rerun viewer and bind scene data provider. + + Args: + scene_data_provider: Scene data provider used to fetch model/state data. + """ + if self._is_initialized: + return + if scene_data_provider is None: + raise RuntimeError("Rerun visualizer requires a scene_data_provider.") + + self._scene_data_provider = scene_data_provider + metadata = scene_data_provider.get_metadata() + self._env_ids = self._compute_visualized_env_ids() + if self._env_ids: + get_filtered_model = getattr(scene_data_provider, "get_newton_model_for_env_ids", None) + if callable(get_filtered_model): + self._model = get_filtered_model(self._env_ids) + else: + self._model = scene_data_provider.get_newton_model() + else: + self._model = scene_data_provider.get_newton_model() + self._state = scene_data_provider.get_newton_state(self._env_ids) + + grpc_port = int(self.cfg.grpc_port) + web_port = int(self.cfg.web_port) + bind_address = self.cfg.bind_address or "0.0.0.0" + rerun_address, start_server_in_viewer = _ensure_rerun_server( + app_id=self.cfg.app_id, + bind_address=bind_address, + grpc_port=grpc_port, + web_port=web_port, + ) + if not start_server_in_viewer: + logger.info("[RerunVisualizer] Reusing existing rerun server at %s.", rerun_address) + + viewer_address = None if start_server_in_viewer else rerun_address + self._viewer = NewtonViewerRerun( + app_id=self.cfg.app_id, + address=viewer_address, + serve_web_viewer=start_server_in_viewer, + web_port=web_port, + grpc_port=grpc_port, + keep_historical_data=self.cfg.keep_historical_data, + keep_scalar_history=self.cfg.keep_scalar_history, + record_to_rrd=self.cfg.record_to_rrd, + ) + if start_server_in_viewer: + rerun_address = getattr(self._viewer, "_grpc_server_uri", rerun_address) + if self.cfg.open_browser and not start_server_in_viewer: + _open_rerun_web_viewer(_normalize_host(bind_address), web_port, rerun_address) + self._viewer.set_model(self._model) + self._apply_camera_pose(self._resolve_initial_camera_pose()) + self._viewer.up_axis = 2 + self._viewer.scaling = 1.0 + self._viewer._paused = False + + num_visualized_envs = len(self._env_ids) if self._env_ids is not None else int(metadata.get("num_envs", 0)) + self._log_initialization_table( + logger=logger, + title="RerunVisualizer Configuration", + rows=[ + ("camera_position", self.cfg.camera_position), + ("camera_target", self.cfg.camera_target), + ("camera_source", self.cfg.camera_source), + ("num_visualized_envs", num_visualized_envs), + ("endpoint", f"http://{_normalize_host(bind_address)}:{web_port}"), + ("bind_address", bind_address), + ("grpc_port", grpc_port), + ("web_port", web_port), + ("open_browser", self.cfg.open_browser), + ("record_to_rrd", self.cfg.record_to_rrd or ""), + ], + ) + + self._is_initialized = True + atexit.register(self.close) + + def step(self, dt: float) -> None: + """Advance visualization by one simulation step. + + Args: + dt: Simulation time-step in seconds. + """ + if not self._is_initialized or self._is_closed or self._viewer is None: + return + + self._sim_time += dt + self._step_counter += 1 + + if self.cfg.camera_source == "usd_path": + self._update_camera_from_usd_path() + + self._state = self._scene_data_provider.get_newton_state(self._env_ids) + + if not self._viewer.is_paused(): + self._viewer.begin_frame(self._sim_time) + if self._state is not None: + body_q = getattr(self._state, "body_q", None) + if hasattr(body_q, "shape") and body_q.shape[0] == 0: + self._viewer.end_frame() + return + self._viewer.log_state(self._state) + self._viewer.end_frame() + + def close(self) -> None: + """Close viewer/session resources.""" + if self._is_closed: + return + + if self._viewer is not None: + try: + self._viewer.close() + except Exception as exc: + logger.warning("[RerunVisualizer] Failed while closing viewer: %s", exc) + finally: + self._viewer = None + + try: + rr.disconnect() + except Exception as exc: + logger.warning("[RerunVisualizer] Failed while disconnecting rerun: %s", exc) + self._is_closed = True + + def is_running(self) -> bool: + """Return whether the visualizer should continue stepping. + + Returns: + ``True`` while the visualizer is active, otherwise ``False``. + """ + if not self._is_initialized or self._is_closed: + return False + if self._viewer is None: + return False + return self._viewer.is_running() + + def _resolve_initial_camera_pose(self) -> tuple[tuple[float, float, float], tuple[float, float, float]]: + """Resolve initial camera pose from config or USD camera path.""" + if self.cfg.camera_source == "usd_path": + pose = self._resolve_camera_pose_from_usd_path(self.cfg.camera_usd_path) + if pose is not None: + return pose + return self.cfg.camera_position, self.cfg.camera_target + + def _apply_camera_pose(self, pose: tuple[tuple[float, float, float], tuple[float, float, float]]) -> None: + """Apply camera pose to rerun's 3D view controls. + + Args: + pose: Camera eye and target tuples. + """ + if self._viewer is None: + return + cam_pos, cam_target = pose + rr.send_blueprint( + rrb.Blueprint( + rrb.Spatial3DView( + name="3D View", + origin="/", + eye_controls=rrb.EyeControls3D( + position=cam_pos, + look_target=cam_target, + ), + ), + collapse_panels=True, + ) + ) + self._last_camera_pose = (cam_pos, cam_target) + + def _update_camera_from_usd_path(self) -> None: + """Refresh camera pose from configured USD camera path when it changes.""" + pose = self._resolve_camera_pose_from_usd_path(self.cfg.camera_usd_path) + if pose is None: + return + if self._last_camera_pose == pose: + return + self._apply_camera_pose(pose) + + def supports_markers(self) -> bool: + """Rerun backend currently does not expose Isaac Lab marker primitives.""" + return False + + def supports_live_plots(self) -> bool: + """Rerun backend currently does not expose Isaac Lab live-plot widgets.""" + return False + + def is_training_paused(self) -> bool: + """Return whether training is paused. + + Rerun viewer exposes rendering pause only. + """ + return False + + def is_rendering_paused(self) -> bool: + """Return whether rendering is paused from viewer controls.""" + if not self._is_initialized or self._viewer is None: + return False + return self._viewer.is_rendering_paused() diff --git a/source/isaaclab/isaaclab/visualizers/rerun_visualizer_cfg.py b/source/isaaclab_visualizers/isaaclab_visualizers/rerun/rerun_visualizer_cfg.py similarity index 69% rename from source/isaaclab/isaaclab/visualizers/rerun_visualizer_cfg.py rename to source/isaaclab_visualizers/isaaclab_visualizers/rerun/rerun_visualizer_cfg.py index a87b033afabc..5edd918929de 100644 --- a/source/isaaclab/isaaclab/visualizers/rerun_visualizer_cfg.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/rerun/rerun_visualizer_cfg.py @@ -8,8 +8,7 @@ from __future__ import annotations from isaaclab.utils import configclass - -from .visualizer_cfg import VisualizerCfg +from isaaclab.visualizers.visualizer_cfg import VisualizerCfg @configclass @@ -29,10 +28,16 @@ class RerunVisualizerCfg(VisualizerCfg): """Port of the rerun gRPC server (used when serving web viewer externally).""" bind_address: str | None = "0.0.0.0" - """If set, start a rerun server bound to this address (e.g. '0.0.0.0') and connect to it.""" + """Host used for endpoint formatting and reuse checks. + + Notes: + - If an existing rerun server is reachable on ``grpc_port``, it is reused. + - New server startup is managed by ``newton.viewer.ViewerRerun`` via the rerun Python SDK. + - Local browser links normalize common loopback/wildcard hosts to ``127.0.0.1``. + """ open_browser: bool = True - """Whether to auto-open a browser when serving the rerun web viewer.""" + """Whether to attempt opening the rerun web viewer URL in a browser.""" keep_historical_data: bool = False """Keep transform history for time scrubbing (False = constant memory for training).""" @@ -43,5 +48,8 @@ class RerunVisualizerCfg(VisualizerCfg): record_to_rrd: str | None = None """Path to save .rrd recording file. None = no recording.""" - max_worlds: int | None = 0 - """Maximum number of worlds/environments rendered by the viewer (0/None = all).""" + max_worlds: int | None = None + """Maximum number of worlds/environments rendered by the viewer. + + Set to ``None`` to leave this option disabled. + """ diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/viser/__init__.py b/source/isaaclab_visualizers/isaaclab_visualizers/viser/__init__.py new file mode 100644 index 000000000000..d38621a2398e --- /dev/null +++ b/source/isaaclab_visualizers/isaaclab_visualizers/viser/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Viser visualizer backend.""" + +from .viser_visualizer import ViserVisualizer +from .viser_visualizer_cfg import ViserVisualizerCfg + +__all__ = ["ViserVisualizer", "ViserVisualizerCfg"] diff --git a/source/isaaclab/isaaclab/visualizers/viser_visualizer.py b/source/isaaclab_visualizers/isaaclab_visualizers/viser/viser_visualizer.py similarity index 59% rename from source/isaaclab/isaaclab/visualizers/viser_visualizer.py rename to source/isaaclab_visualizers/isaaclab_visualizers/viser/viser_visualizer.py index b2e80af46a29..ca3fba2e672d 100644 --- a/source/isaaclab/isaaclab/visualizers/viser_visualizer.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/viser/viser_visualizer.py @@ -7,30 +7,28 @@ from __future__ import annotations +import contextlib +import io import logging import os +import webbrowser from pathlib import Path from typing import TYPE_CHECKING, Any from newton.viewer import ViewerViser +from isaaclab.visualizers.base_visualizer import BaseVisualizer + from .viser_visualizer_cfg import ViserVisualizerCfg -from .visualizer import Visualizer logger = logging.getLogger(__name__) if TYPE_CHECKING: - from isaaclab.sim.scene_data_providers import SceneDataProvider + from isaaclab.physics import BaseSceneDataProvider def _disable_viser_runtime_client_rebuild_if_bundled() -> None: - """Skip viser's runtime frontend rebuild when a bundled build is present. - - Newer viser versions may try to rebuild the client if source timestamps are newer - than the packaged build directory, which requires ``nodeenv`` at runtime. - For Isaac Lab usage, we prefer the prebuilt static assets shipped with the package. - """ - + """Skip viser's runtime frontend rebuild when a bundled build is present.""" try: import viser import viser._client_autobuild as client_autobuild @@ -45,6 +43,33 @@ def _disable_viser_runtime_client_rebuild_if_bundled() -> None: client_autobuild.ensure_client_is_built = lambda: None +@contextlib.contextmanager +def _suppress_viser_startup_logs(enabled: bool): + """Temporarily quiet noisy viser/websockets startup output.""" + if not enabled: + yield + return + + websockets_logger = logging.getLogger("websockets.server") + previous_level = websockets_logger.level + websockets_logger.setLevel(logging.WARNING) + try: + with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()): + yield + finally: + websockets_logger.setLevel(previous_level) + + +def _open_viser_web_viewer(port: int) -> None: + """Open the local viser web UI in a browser.""" + url = f"http://localhost:{int(port)}" + try: + if not webbrowser.open_new_tab(url): + logger.info("[ViserVisualizer] Could not auto-open browser tab. Open manually: %s", url) + except Exception: + logger.info("[ViserVisualizer] Could not auto-open browser tab. Open manually: %s", url) + + class NewtonViewerViser(ViewerViser): """Isaac Lab wrapper for Newton's ViewerViser.""" @@ -57,6 +82,16 @@ def __init__( record_to_viser: str | None = None, metadata: dict | None = None, ): + """Initialize Newton-backed viser viewer wrapper. + + Args: + port: HTTP port for viser server. + label: Optional viewer label. + verbose: Whether to keep verbose startup output enabled. + share: Whether to enable sharing/tunneling. + record_to_viser: Optional recording destination. + metadata: Optional metadata attached to the viewer. + """ _disable_viser_runtime_client_rebuild_if_bundled() super().__init__( port=port, @@ -68,25 +103,32 @@ def __init__( self._metadata = metadata or {} -class ViserVisualizer(Visualizer): +class ViserVisualizer(BaseVisualizer): """Viser web-based visualizer backed by Newton's ViewerViser.""" def __init__(self, cfg: ViserVisualizerCfg): - """Initialize the visualizer state from the resolved config.""" + """Initialize Viser visualizer state. + + Args: + cfg: Viser visualizer configuration. + """ super().__init__(cfg) self.cfg: ViserVisualizerCfg = cfg self._viewer: NewtonViewerViser | None = None self._model: Any | None = None self._state = None - self._is_initialized = False self._sim_time = 0.0 self._scene_data_provider = None self._active_record_path: str | None = None self._last_camera_pose: tuple[tuple[float, float, float], tuple[float, float, float]] | None = None self._pending_camera_pose: tuple[tuple[float, float, float], tuple[float, float, float]] | None = None - def initialize(self, scene_data_provider: SceneDataProvider) -> None: - """Create the viewer and bind it to the active scene data provider.""" + def initialize(self, scene_data_provider: BaseSceneDataProvider) -> None: + """Initialize viewer resources and bind scene data provider. + + Args: + scene_data_provider: Scene data provider used to fetch model/state data. + """ if self._is_initialized: logger.debug("[ViserVisualizer] initialize() called while already initialized.") return @@ -98,29 +140,38 @@ def initialize(self, scene_data_provider: SceneDataProvider) -> None: self._env_ids = self._compute_visualized_env_ids() if self._env_ids: get_filtered_model = getattr(scene_data_provider, "get_newton_model_for_env_ids", None) - if callable(get_filtered_model): - self._model = get_filtered_model(self._env_ids) - else: - self._model = scene_data_provider.get_newton_model() + self._model = ( + get_filtered_model(self._env_ids) + if callable(get_filtered_model) + else scene_data_provider.get_newton_model() + ) else: self._model = scene_data_provider.get_newton_model() self._state = scene_data_provider.get_newton_state(self._env_ids) - try: - self._active_record_path = self.cfg.record_to_viser - self._create_viewer(record_to_viser=self.cfg.record_to_viser, metadata=metadata) - logger.info( - "[ViserVisualizer] initialized | camera_pos=%s camera_target=%s", - self.cfg.camera_position, - self.cfg.camera_target, - ) - self._is_initialized = True - except Exception as exc: - logger.error(f"[ViserVisualizer] Failed to initialize viewer: {exc}") - raise + self._active_record_path = self.cfg.record_to_viser + self._create_viewer(record_to_viser=self.cfg.record_to_viser, metadata=metadata) + num_visualized_envs = len(self._env_ids) if self._env_ids is not None else int(metadata.get("num_envs", 0)) + self._log_initialization_table( + logger=logger, + title="ViserVisualizer Configuration", + rows=[ + ("camera_position", self.cfg.camera_position), + ("camera_target", self.cfg.camera_target), + ("camera_source", self.cfg.camera_source), + ("num_visualized_envs", num_visualized_envs), + ("port", self.cfg.port), + ("record_to_viser", self.cfg.record_to_viser or ""), + ], + ) + self._is_initialized = True def step(self, dt: float) -> None: - """Fetch the latest state and stream one frame to Viser.""" + """Advance visualization by one simulation step. + + Args: + dt: Simulation time-step in seconds. + """ if not self._is_initialized or self._viewer is None or self._scene_data_provider is None: return @@ -130,20 +181,18 @@ def step(self, dt: float) -> None: self._state = self._scene_data_provider.get_newton_state(self._env_ids) self._sim_time += dt - self._viewer.begin_frame(self._sim_time) self._viewer.log_state(self._state) self._viewer.end_frame() def close(self) -> None: - """Close the viewer and finalize any active recording.""" + """Close viewer resources and finalize optional recording.""" if not self._is_initialized: return - try: self._close_viewer(finalize_viser=bool(self.cfg.record_to_viser)) except Exception as exc: - logger.warning(f"[ViserVisualizer] Error during close: {exc}") + logger.warning("[ViserVisualizer] Error during close: %s", exc) self._viewer = None self._is_initialized = False @@ -152,60 +201,73 @@ def close(self) -> None: self._pending_camera_pose = None def is_running(self) -> bool: - """Return whether the underlying Viser server is still running.""" + """Return whether the visualizer should continue stepping. + + Returns: + ``True`` while the visualizer is active, otherwise ``False``. + """ + if not self._is_initialized or self._is_closed: + return False if self._viewer is None: return False return self._viewer.is_running() def is_training_paused(self) -> bool: - """Viser does not currently expose training pause controls.""" + """Return whether training is paused. + + Viser backend does not currently expose a training pause control. + """ return False def supports_markers(self) -> bool: - """Viser marker support is not implemented yet.""" + """Viser backend currently does not expose Isaac Lab marker primitives.""" return False def supports_live_plots(self) -> bool: - """Viser live plot support is not implemented yet.""" + """Viser backend currently does not expose Isaac Lab live-plot widgets.""" return False def _create_viewer(self, record_to_viser: str | None, metadata: dict | None = None) -> None: - """Create the Viser viewer and load the Newton model.""" + """Create Newton-backed Viser viewer and apply initial camera. + + Args: + record_to_viser: Optional output path for viser recording. + metadata: Optional metadata passed to viewer. + """ if self._model is None: raise RuntimeError("Viser visualizer requires a Newton model.") - self._viewer = NewtonViewerViser( - port=self.cfg.port, - label=self.cfg.label, - verbose=self.cfg.verbose, - share=self.cfg.share, - record_to_viser=record_to_viser, - metadata=metadata or {}, - ) + with _suppress_viser_startup_logs(enabled=not self.cfg.verbose): + self._viewer = NewtonViewerViser( + port=self.cfg.port, + label=self.cfg.label, + verbose=self.cfg.verbose, + share=self.cfg.share, + record_to_viser=record_to_viser, + metadata=metadata or {}, + ) max_worlds = self.cfg.max_worlds - self._viewer.set_model(self._model, max_worlds=None if max_worlds in (None, 0) else max_worlds) + self._viewer.set_model(self._model, max_worlds=max_worlds) + if self.cfg.open_browser: + _open_viser_web_viewer(self.cfg.port) self._set_viser_camera_view(self._resolve_initial_camera_pose()) self._sim_time = 0.0 def _close_viewer(self, finalize_viser: bool = False) -> None: - """Close the viewer and log the finalized recording, if any.""" + """Close viewer and log recording output when requested.""" if self._viewer is None: return self._viewer.close() if finalize_viser and self._active_record_path: if os.path.exists(self._active_record_path): size = os.path.getsize(self._active_record_path) - logger.info( - "[ViserVisualizer] Recording saved: %s (%s bytes)", - self._active_record_path, - size, - ) + logger.info("[ViserVisualizer] Recording saved: %s (%s bytes)", self._active_record_path, size) else: logger.warning("[ViserVisualizer] Recording file not found: %s", self._active_record_path) self._viewer = None def _resolve_initial_camera_pose(self) -> tuple[tuple[float, float, float], tuple[float, float, float]]: - """Resolve the startup camera pose from config or a USD camera path.""" + """Resolve initial camera pose from config or USD camera path.""" if self.cfg.camera_source == "usd_path": pose = self._resolve_camera_pose_from_usd_path(self.cfg.camera_usd_path) if pose is not None: @@ -217,10 +279,13 @@ def _resolve_initial_camera_pose(self) -> tuple[tuple[float, float, float], tupl return self.cfg.camera_position, self.cfg.camera_target def _try_apply_viser_camera_view(self, pose: tuple[tuple[float, float, float], tuple[float, float, float]]) -> bool: - """Best-effort apply a camera pose to any connected Viser clients.""" + """Try applying camera pose to active viser clients. + + Returns: + ``True`` if at least one client camera was updated, otherwise ``False``. + """ if self._viewer is None: return False - server = getattr(self._viewer, "_server", None) get_clients = getattr(server, "get_clients", None) if server is not None else None if not callable(get_clients): @@ -231,11 +296,7 @@ def _try_apply_viser_camera_view(self, pose: tuple[tuple[float, float, float], t except Exception: return False - if isinstance(clients, dict): - client_iterable = clients.values() - else: - client_iterable = clients - + client_iterable = clients.values() if isinstance(clients, dict) else clients cam_pos, cam_target = pose applied = False for client in client_iterable: @@ -254,7 +315,7 @@ def _try_apply_viser_camera_view(self, pose: tuple[tuple[float, float, float], t return applied def _set_viser_camera_view(self, pose: tuple[tuple[float, float, float], tuple[float, float, float]]) -> None: - """Apply the pose immediately or queue it until a client connects.""" + """Apply or defer camera pose update depending on client readiness.""" if self._try_apply_viser_camera_view(pose): self._last_camera_pose = pose self._pending_camera_pose = None @@ -262,7 +323,7 @@ def _set_viser_camera_view(self, pose: tuple[tuple[float, float, float], tuple[f self._pending_camera_pose = pose def _apply_pending_camera_pose(self) -> None: - """Retry applying any queued camera pose once clients are available.""" + """Apply deferred camera pose once client cameras are available.""" if self._pending_camera_pose is None: return if self._try_apply_viser_camera_view(self._pending_camera_pose): @@ -270,7 +331,7 @@ def _apply_pending_camera_pose(self) -> None: self._pending_camera_pose = None def _update_camera_from_usd_path(self) -> None: - """Mirror the configured USD camera path into the active Viser camera.""" + """Refresh camera pose from configured USD camera path when it changes.""" pose = self._resolve_camera_pose_from_usd_path(self.cfg.camera_usd_path) if pose is None: return diff --git a/source/isaaclab/isaaclab/visualizers/viser_visualizer_cfg.py b/source/isaaclab_visualizers/isaaclab_visualizers/viser/viser_visualizer_cfg.py similarity index 79% rename from source/isaaclab/isaaclab/visualizers/viser_visualizer_cfg.py rename to source/isaaclab_visualizers/isaaclab_visualizers/viser/viser_visualizer_cfg.py index d762afdc7bf1..c2400c7ee1e6 100644 --- a/source/isaaclab/isaaclab/visualizers/viser_visualizer_cfg.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/viser/viser_visualizer_cfg.py @@ -8,8 +8,7 @@ from __future__ import annotations from isaaclab.utils import configclass - -from .visualizer_cfg import VisualizerCfg +from isaaclab.visualizers.visualizer_cfg import VisualizerCfg @configclass @@ -22,6 +21,9 @@ class ViserVisualizerCfg(VisualizerCfg): port: int = 8080 """Port of the local viser web server.""" + open_browser: bool = True + """Whether to attempt opening the viser web viewer URL in a browser.""" + label: str | None = "Isaac Lab Simulation" """Optional label shown in the viewer page title.""" @@ -34,5 +36,8 @@ class ViserVisualizerCfg(VisualizerCfg): record_to_viser: str | None = None """Path to save a .viser recording file. None = no recording.""" - max_worlds: int | None = 0 - """Maximum number of worlds/environments rendered by the viewer (0/None = all).""" + max_worlds: int | None = None + """Maximum number of worlds/environments rendered by the viewer. + + Set to ``None`` to leave this option disabled. + """ diff --git a/source/isaaclab_visualizers/setup.py b/source/isaaclab_visualizers/setup.py new file mode 100644 index 000000000000..3deadef84a8c --- /dev/null +++ b/source/isaaclab_visualizers/setup.py @@ -0,0 +1,60 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Installation script for the 'isaaclab_visualizers' python package.""" + +from setuptools import setup + +# Base requirements shared across visualizer backends. +INSTALL_REQUIRES = [ + "isaaclab", + "numpy", +] + +EXTRAS_REQUIRE = { + "kit": [], + "newton": [ + "warp-lang", + "newton", + "PyOpenGL-accelerate", + "imgui-bundle>=1.92.5", + ], + "rerun": [ + "newton", + "rerun-sdk>=0.29.0", + ], + "viser": [ + "newton", + "viser>=1.0.16", + ], +} + +EXTRAS_REQUIRE["all"] = sorted({dep for group in EXTRAS_REQUIRE.values() for dep in group}) + +setup( + name="isaaclab_visualizers", + author="Isaac Lab Project Developers", + maintainer="Isaac Lab Project Developers", + url="https://github.com/isaac-sim/IsaacLab", + version="0.1.0", + description="Visualizer backends for Isaac Lab (Kit, Newton, Rerun, Viser).", + keywords=["robotics", "simulation", "visualization"], + license="BSD-3-Clause", + include_package_data=True, + package_data={"": ["*.pyi"]}, + python_requires=">=3.11", + install_requires=INSTALL_REQUIRES, + extras_require=EXTRAS_REQUIRE, + packages=["isaaclab_visualizers"], + classifiers=[ + "Natural Language :: English", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Isaac Sim :: 5.0.0", + "Isaac Sim :: 5.1.0", + "Isaac Sim :: 6.0.0", + ], + zip_safe=False, +) diff --git a/source/isaaclab_visualizers/test/__init__.py b/source/isaaclab_visualizers/test/__init__.py new file mode 100644 index 000000000000..460a30569089 --- /dev/null +++ b/source/isaaclab_visualizers/test/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause diff --git a/source/isaaclab_visualizers/test/test_visualizer_smoke_logs.py b/source/isaaclab_visualizers/test/test_visualizer_smoke_logs.py new file mode 100644 index 000000000000..e1e7ac84f534 --- /dev/null +++ b/source/isaaclab_visualizers/test/test_visualizer_smoke_logs.py @@ -0,0 +1,230 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Smoke test visualizer stepping and error logging.""" + +from isaaclab.app import AppLauncher + +# launch Kit app +simulation_app = AppLauncher(headless=True, enable_cameras=True).app + +import logging +import socket + +import pytest +import torch +from isaaclab_visualizers.kit import KitVisualizer, KitVisualizerCfg +from isaaclab_visualizers.newton import NewtonVisualizer, NewtonVisualizerCfg +from isaaclab_visualizers.rerun import RerunVisualizer, RerunVisualizerCfg +from isaaclab_visualizers.viser import ViserVisualizer, ViserVisualizerCfg + +import isaaclab.sim as sim_utils +from isaaclab.envs import DirectRLEnv, DirectRLEnvCfg +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sim import SimulationCfg, SimulationContext +from isaaclab.utils import configclass + +from isaaclab_tasks.manager_based.classic.cartpole.cartpole_env_cfg import ( + CartpolePhysicsCfg, + CartpoleSceneCfg, +) + +# Set to False to only fail on visualizer errors; when True, also fail on warnings. +ASSERT_VISUALIZER_WARNINGS = True + +_SMOKE_STEPS = 4 +_VIS_LOGGER_PREFIXES = ( + "isaaclab.visualizers", + "isaaclab_visualizers", + "isaaclab.sim.simulation_context", +) + + +def _find_free_tcp_port(host: str = "127.0.0.1") -> int: + """Ask OS for a currently free local TCP port.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind((host, 0)) + return int(sock.getsockname()[1]) + + +def _allocate_rerun_test_ports(host: str = "127.0.0.1") -> tuple[int, int]: + """Allocate distinct free ports for rerun web and gRPC endpoints.""" + grpc_port = _find_free_tcp_port(host) + web_port = _find_free_tcp_port(host) + while web_port == grpc_port: + web_port = _find_free_tcp_port(host) + return web_port, grpc_port + + +@configclass +class _SmokeEnvCfg(DirectRLEnvCfg): + decimation: int = 2 + action_space: int = 0 + observation_space: int = 0 + episode_length_s: float = 5.0 + sim: SimulationCfg = SimulationCfg(dt=0.005, render_interval=2, visualizer_cfgs=KitVisualizerCfg()) + scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=1, env_spacing=1.0) + + +class _SmokeEnv(DirectRLEnv): + def _pre_physics_step(self, actions): + return + + def _apply_action(self): + return + + def _get_observations(self): + return {} + + def _get_rewards(self): + return {} + + def _get_dones(self): + return torch.zeros(1, dtype=torch.bool), torch.zeros(1, dtype=torch.bool) + + +def _get_visualizer_cfg(visualizer_kind: str): + """Return (visualizer_cfg, expected_visualizer_cls) for the given visualizer kind.""" + if visualizer_kind == "newton": + __import__("newton") + return NewtonVisualizerCfg(headless=True), NewtonVisualizer + if visualizer_kind == "viser": + __import__("newton") + __import__("viser") + return ViserVisualizerCfg(open_browser=False), ViserVisualizer + if visualizer_kind == "rerun": + __import__("newton") + __import__("rerun") + web_port, grpc_port = _allocate_rerun_test_ports(host="127.0.0.1") + # Use dynamically allocated non-default ports in smoke tests to avoid collisions. + # TODO: Consider supporting cleanup/termination of stale rerun processes when ports are occupied. + return ( + RerunVisualizerCfg( + bind_address="127.0.0.1", + open_browser=False, + web_port=web_port, + grpc_port=grpc_port, + ), + RerunVisualizer, + ) + return KitVisualizerCfg(), KitVisualizer + + +def _get_physics_cfg(backend_kind: str): + """Return physics config and expected backend substring for the given backend kind. + + Uses cartpole preset instance so we work whether presets are class or instance attributes. + Fallback: build PhysxCfg/NewtonCfg in-test if preset does not expose that backend. + """ + if backend_kind == "physx": + __import__("isaaclab_physx") + preset = CartpolePhysicsCfg() + physics_cfg = getattr(preset, "physx", None) + if physics_cfg is None: + from isaaclab_physx.physics import PhysxCfg + + physics_cfg = PhysxCfg() + return physics_cfg, "physx" + if backend_kind == "newton": + __import__("newton") + __import__("isaaclab_newton") + preset = CartpolePhysicsCfg() + physics_cfg = getattr(preset, "newton", None) + if physics_cfg is None: + from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg + + physics_cfg = NewtonCfg( + solver_cfg=MJWarpSolverCfg( + njmax=5, + nconmax=3, + ls_iterations=10, + cone="pyramidal", + impratio=1, + ls_parallel=True, + integrator="implicitfast", + ), + num_substeps=1, + debug_mode=False, + use_cuda_graph=True, + ) + return physics_cfg, "newton" + raise ValueError(f"Unknown backend: {backend_kind!r}") + + +def _resolve_case(visualizer_kind: str, backend_kind: str): + """Resolve (env_cfg, expected_visualizer_cls, expected_backend_substring) for one smoke test. + + Uses cartpole scene for all combinations (works with both PhysX and Newton). + """ + scene_cfg = CartpoleSceneCfg(num_envs=1, env_spacing=1.0) + viz_cfg, expected_viz_cls = _get_visualizer_cfg(visualizer_kind) + physics_cfg, expected_backend = _get_physics_cfg(backend_kind) + + cfg = _SmokeEnvCfg() + cfg.scene = scene_cfg + cfg.sim = SimulationCfg( + dt=0.005, + render_interval=2, + visualizer_cfgs=viz_cfg, + physics=physics_cfg, + ) + return cfg, expected_viz_cls, expected_backend + + +def _run_smoke_test(cfg, expected_visualizer_cls, expected_backend: str, caplog) -> None: + """Run smoke steps and assert no visualizer errors; optionally no warnings (see ASSERT_VISUALIZER_WARNINGS).""" + env = None + try: + sim_utils.create_new_stage() + env = _SmokeEnv(cfg=cfg) + backend_name = env.sim.physics_manager.__name__.lower() + assert expected_backend in backend_name, ( + f"Expected physics backend containing {expected_backend!r}, got {backend_name!r}" + ) + env.sim.set_setting("/isaaclab/render/rtx_sensors", True) + env.sim._app_control_on_stop_handle = None # type: ignore[attr-defined] + + actions = torch.zeros((env.num_envs, 0), device=env.device) + with caplog.at_level(logging.WARNING): + env.reset() + assert env.sim.visualizers + assert isinstance(env.sim.visualizers[0], expected_visualizer_cls) + for _ in range(_SMOKE_STEPS): + env.step(action=actions) + + # Always fail on errors + error_logs = [ + r for r in caplog.records if r.levelno >= logging.ERROR and r.name.startswith(_VIS_LOGGER_PREFIXES) + ] + assert not error_logs, "Visualizer emitted error logs during smoke stepping: " + "; ".join( + f"{r.name}: {r.message}" for r in error_logs + ) + + # Optionally fail on warnings + if ASSERT_VISUALIZER_WARNINGS: + warning_logs = [ + r for r in caplog.records if r.levelno >= logging.WARNING and r.name.startswith(_VIS_LOGGER_PREFIXES) + ] + assert not warning_logs, "Visualizer emitted warning logs during smoke stepping: " + "; ".join( + f"{r.name}: {r.message}" for r in warning_logs + ) + finally: + if env is not None: + env.close() + else: + SimulationContext.clear_instance() + + +@pytest.mark.isaacsim_ci +@pytest.mark.parametrize("visualizer_kind", ["kit", "newton", "rerun", "viser"]) +@pytest.mark.parametrize("backend_kind", ["physx", "newton"]) +def test_visualizer_backend_smoke(visualizer_kind: str, backend_kind: str, caplog): + """Smoke test each (visualizer, backend) pair; assert no errors (optionally no warnings).""" + cfg, expected_viz_cls, expected_backend = _resolve_case(visualizer_kind, backend_kind) + _run_smoke_test(cfg, expected_viz_cls, expected_backend, caplog) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "--maxfail=1"]) diff --git a/tools/test_settings.py b/tools/test_settings.py index d0856150b24e..7b3bf0239299 100644 --- a/tools/test_settings.py +++ b/tools/test_settings.py @@ -53,6 +53,7 @@ "test_outdated_sensor.py": 500, "test_multi_tiled_camera.py": 1000, "test_multirotor.py": 1000, + "test_shadow_hand_vision_presets.py": 5000, } """A dictionary of tests and their timeouts in seconds.