diff --git a/.github/workflows/mlx-macos.yml b/.github/workflows/mlx-macos.yml index c80887736..7852c7657 100644 --- a/.github/workflows/mlx-macos.yml +++ b/.github/workflows/mlx-macos.yml @@ -39,11 +39,98 @@ jobs: run: uv venv --python 3.11 .venv - name: Install MLX smoke dependencies - run: uv pip install --python .venv/bin/python mlx pytest pytest-mock toml + run: | + uv pip install --python .venv/bin/python \ + -e source/isaaclab[macos-mlx,dev] \ + -e source/isaaclab_rl[dev] + + - name: Prepare benchmark artifact directory + run: mkdir -p logs/benchmarks/mlx + + - name: Run public install smoke + run: | + PYTHONPATH=.:source/isaaclab:source/isaaclab_rl .venv/bin/python - <<'PY' + import json + from pathlib import Path + + from isaaclab.backends import resolve_runtime_selection, set_runtime_selection + + set_runtime_selection(resolve_runtime_selection(compute_backend="mlx", sim_backend="mac-sim", device="cpu")) + + import isaaclab + import isaaclab.backends.mac_sim + import isaaclab.controllers + import isaaclab.envs.mdp.actions.actions_cfg + import isaaclab.envs.mdp.actions.rmpflow_actions_cfg + import isaaclab.sim + import isaaclab.sim.converters + import isaaclab.sim.schemas + import isaaclab.sim.spawners.from_files + import isaaclab_rl.sb3 + import isaaclab_rl.skrl + from isaaclab.backends.mac_sim import MacCartpoleEnv + + env = MacCartpoleEnv() + observations, _ = env.reset() + payload = { + "modules": [ + "isaaclab", + "isaaclab.backends.mac_sim", + "isaaclab.controllers", + "isaaclab.envs.mdp.actions.actions_cfg", + "isaaclab.envs.mdp.actions.rmpflow_actions_cfg", + "isaaclab.sim", + "isaaclab.sim.converters", + "isaaclab.sim.schemas", + "isaaclab.sim.spawners.from_files", + "isaaclab_rl.sb3", + "isaaclab_rl.skrl", + ], + "policy_shape": list(observations["policy"].shape), + "package": isaaclab.__name__, + } + out_path = Path("logs/benchmarks/mlx/install_import_smoke.json") + out_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(payload["policy_shape"]) + PY - name: Run MLX backend smoke tests run: | - PYTHONPATH=.:source/isaaclab .venv/bin/pytest \ + PYTHONPATH=.:source/isaaclab:source/isaaclab_rl .venv/bin/pytest \ scripts/tools/test/test_bootstrap_isaac_sources.py \ source/isaaclab/test/backends/test_runtime.py \ - source/isaaclab/test/backends/test_mac_cartpole.py -q + source/isaaclab_rl/test/test_import_safety.py \ + source/isaaclab/test/backends/test_portability_utils.py \ + source/isaaclab/test/backends/test_mac_cartpole.py \ + source/isaaclab/test/backends/test_mac_cartpole_showcase.py \ + source/isaaclab/test/backends/test_mac_cart_double_pendulum.py \ + source/isaaclab/test/backends/test_mac_quadcopter.py -q + + - name: Run MLX cart-double-pendulum smoke script + run: | + PYTHONPATH=.:source/isaaclab .venv/bin/python \ + scripts/reinforcement_learning/mlx/play_cart_double_pendulum.py \ + --num-envs 8 --episodes 1 --max-steps 256 --random-actions + + - name: Run MLX quadcopter smoke script + run: | + PYTHONPATH=.:source/isaaclab .venv/bin/python \ + scripts/reinforcement_learning/mlx/play_quadcopter.py \ + --num-envs 8 --episodes 1 --episode-length-s 0.5 --max-steps 256 --thrust-action 0.2 --no-random-actions + + - name: Run MLX benchmark smoke + run: | + PYTHONPATH=.:source/isaaclab .venv/bin/python \ + scripts/benchmarks/mlx/benchmark_mac_tasks.py \ + --tasks cartpole cart-double-pendulum quadcopter \ + --num-envs 64 \ + --steps 128 \ + --json-out logs/benchmarks/mlx/smoke.json + + - name: Upload MLX smoke artifacts + if: always() + uses: actions/upload-artifact@v4 + with: + name: mlx-macos-smoke-${{ github.run_id }} + path: logs/benchmarks/mlx + if-no-files-found: warn diff --git a/PORT_TO_MLX_TODO.md b/PORT_TO_MLX_TODO.md new file mode 100644 index 000000000..4d5d898e7 --- /dev/null +++ b/PORT_TO_MLX_TODO.md @@ -0,0 +1,528 @@ +# Port To MLX Task Board + +This file is the execution backlog for the full CUDA-to-MLX port of IsaacLab on Apple Silicon. + +It is intentionally written as a task source, not as a loose idea list. The goal is to keep shipping work on `main` +without pausing for replanning after every small success. + +## Working Rules + +- Always take the next unblocked task from the `Execution Order` section. +- Every merged task must update: + - this file + - `/Users/ilessio/.codex/skills/port-to-mlx/references/progress_log.md` +- Every compatibility task needs at least one of: + - import-safety test + - install smoke + - benchmark smoke + - replay/train smoke +- “Done” means: + - code landed on `main` + - validation ran + - follow-on tasks updated + +## Status Legend + +- `DONE`: landed and validated +- `ACTIVE`: current execution focus +- `READY`: unblocked and next in line +- `BLOCKED`: depends on earlier tasks +- `LATER`: intentionally deferred until the earlier substrate exists + +## Execution Order + +1. `MLX-IMPORT-001` through `MLX-IMPORT-008` +2. `MLX-PKG-001` through `MLX-PKG-004` +3. `MLX-SIM-001` through `MLX-SIM-010` +4. `MLX-TASK-001` through `MLX-TASK-012` +5. `MLX-KERNEL-001` through `MLX-KERNEL-008` +6. `MLX-SENSOR-001` through `MLX-SENSOR-006` +7. `MLX-RL-001` through `MLX-RL-006` +8. `MLX-CI-001` through `MLX-CI-006` +9. `MLX-ROS-001` through `MLX-ROS-003` + +## Current Wins + +- `DONE` Runtime backend seam for `torch-cuda|mlx` +- `DONE` Sim backend seam for `isaacsim|mac-sim` +- `DONE` Kernel backend seam for `warp|metal|cpu` +- `DONE` AppLauncher mac-sim bootstrap mode +- `DONE` MLX/mac-sim slices for cartpole, cart-double-pendulum, quadcopter +- `DONE` Focused backend tests and benchmark harness +- `DONE` Base package extras for `macos-mlx`, `cuda-isaacsim`, and `dev` +- `DONE` RL package base install no longer forces framework extras +- `DONE` `isaaclab_rl.sb3` and `isaaclab_rl.skrl` import-safe on MLX/mac path +- `DONE` `RmpFlowControllerCfg` split from the heavy controller implementation +- `DONE` Lightweight Nucleus constant module split from the heavy assets helper +- `DONE` `isaaclab.sim` config surfaces for schemas, converters, and file spawners import on `mlx + mac-sim` +- `DONE` `isaaclab.envs.mdp.actions` config surfaces import on `mlx + mac-sim` without runtime action modules +- `DONE` macOS install/import smoke now writes retained JSON artifacts and benchmark results in CI + +## Phase A: Import And Packaging Safety + +### MLX-IMPORT-001 + +- Status: `DONE` +- Title: Split controller config dataclasses away from heavy Isaac Sim runtime modules +- Validation: + - focused backend tests + - fresh import smoke + +### MLX-IMPORT-002 + +- Status: `DONE` +- Title: Split Nucleus constants away from `isaaclab.utils.assets` +- Validation: + - focused backend tests + - fresh import smoke + +### MLX-IMPORT-003 + +- Status: `DONE` +- Title: Audit `isaaclab.sim` public config/helper imports and isolate config-safe surfaces +- Scope: + - `isaaclab.sim.__init__` + - `isaaclab.sim.schemas` + - `isaaclab.sim.spawners.*_cfg` + - `isaaclab.sim.utils` modules that should be config-safe +- Acceptance: + - config-only imports do not pull `carb`, `omni`, `isaacsim`, or `pxr` unless explicitly upstream-only + - tests prove the new import-safe path +- Validation: + - focused backend suite + - torch-free local install/import smoke + - local benchmark smoke JSON output + +### MLX-IMPORT-004 + +- Status: `DONE` +- Depends on: `MLX-IMPORT-003` +- Title: Audit `isaaclab.envs.mdp.actions` config surfaces for config-only imports that still pull runtime-heavy modules +- Acceptance: + - action config modules import on `mlx + mac-sim` + - unsupported runtime paths fail when instantiated, not when imported +- Validation: + - focused backend suite + - torch-free local install/import smoke + +### MLX-IMPORT-005 + +- Status: `READY` +- Depends on: `MLX-IMPORT-003` +- Title: Audit `isaaclab.markers` config surfaces for any remaining heavy imports +- Acceptance: + - config modules use lightweight constants/helpers where possible + +### MLX-IMPORT-006 + +- Status: `READY` +- Depends on: `MLX-IMPORT-003` +- Title: Audit `isaaclab.devices.openxr` retargeter config/helpers for constant-only imports +- Acceptance: + - config constants come from lightweight modules + - heavy runtime helpers stay lazy + +### MLX-IMPORT-007 + +- Status: `DONE` +- Depends on: `MLX-IMPORT-003` +- Title: Add explicit import-safety test matrix for public bootstrap modules +- Acceptance: + - test covers `isaaclab`, `isaaclab.sim`, `isaaclab.controllers`, `isaaclab_rl` + - runs without Isaac Sim installed +- Validation: + - `source/isaaclab/test/backends/test_runtime.py` + - `source/isaaclab_rl/test/test_import_safety.py` + - torch-free local install/import smoke + +### MLX-IMPORT-008 + +- Status: `READY` +- Depends on: `MLX-IMPORT-007` +- Title: Publish a definitive “import-safe on mac” module surface +- Acceptance: + - documented in README + - kept in sync with tests + +### MLX-IMPORT-009 + +- Status: `ACTIVE` +- Depends on: `MLX-IMPORT-007` +- Title: Replace `isaaclab_tasks` eager recursive package registration with a lazy task registry +- Scope: + - `source/isaaclab_tasks/isaaclab_tasks/__init__.py` + - `source/isaaclab_tasks/isaaclab_tasks/utils/importer.py` +- Acceptance: + - `import isaaclab_tasks` does not recursively import every task package on the MLX/mac path + - task/config entry points resolve lazily from strings or a registry manifest + +### MLX-IMPORT-010 + +- Status: `READY` +- Depends on: `MLX-IMPORT-009` +- Title: Split direct-task CUDA/IsaacSim clusters from task package import surfaces +- Scope: + - AutoMate + - Factory + - FORGE + - first manager-based manipulation config clusters that import `carb` or `isaacsim` at module load +- Acceptance: + - task packages expose config/registration surfaces without importing `warp`, `carb`, `pxr`, or `isaacsim` + - runtime-only backends load on demand behind capability checks + +### MLX-PKG-001 + +- Status: `DONE` +- Title: Core package extras split for MLX vs CUDA users + +### MLX-PKG-002 + +- Status: `DONE` +- Title: RL package extras split so framework deps are optional + +### MLX-PKG-003 + +- Status: `READY` +- Title: Add clear support matrix for package extras and runtime combinations +- Acceptance: + - README section listing supported install combos + - one install smoke per documented combo + +### MLX-PKG-004 + +- Status: `READY` +- Depends on: `MLX-PKG-003` +- Title: Add fresh-env install smoke for documented public paths in CI +- Acceptance: + - MLX base install smoke + - MLX + RL base install smoke + - upstream CUDA path remains unaffected + +## Phase B: Shared mac-sim Substrate + +### MLX-SIM-001 + +- Status: `READY` +- Title: Extract shared articulated joint-state container from cartpole/cart-double-pendulum +- Acceptance: + - no duplicated joint state reset/write logic across those envs + +### MLX-SIM-002 + +- Status: `READY` +- Depends on: `MLX-SIM-001` +- Title: Extract shared root-state container from quadcopter path +- Acceptance: + - root pose/velocity read/write helpers centralized + +### MLX-SIM-003 + +- Status: `READY` +- Depends on: `MLX-SIM-001`, `MLX-SIM-002` +- Title: Create reusable mac-sim batched state primitives module +- Acceptance: + - cartpole/cart-double-pendulum/quadcopter use the shared substrate + +### MLX-SIM-004 + +- Status: `READY` +- Depends on: `MLX-SIM-003` +- Title: Add reusable environment origin/grid manager + +### MLX-SIM-005 + +- Status: `READY` +- Depends on: `MLX-SIM-003` +- Title: Add terrain representation for the first locomotion task + +### MLX-SIM-006 + +- Status: `BLOCKED` +- Depends on: `MLX-SIM-005` +- Title: Add contact approximation model sufficient for first locomotion bring-up + +### MLX-SIM-007 + +- Status: `BLOCKED` +- Depends on: `MLX-SIM-006` +- Title: Add contact-oriented reward/termination utilities + +### MLX-SIM-008 + +- Status: `READY` +- Depends on: `MLX-SIM-003` +- Title: Centralize reset sampling helpers and determinism controls + +### MLX-SIM-009 + +- Status: `READY` +- Depends on: `MLX-SIM-003` +- Title: Add capability reporting from concrete mac-sim adapters into benchmarks and diagnostics + +### MLX-SIM-010 + +- Status: `READY` +- Depends on: `MLX-SIM-003` +- Title: Add shared replay/rollout helpers for mac-native tasks + +## Phase C: Task Ports + +### MLX-TASK-001 + +- Status: `DONE` +- Title: Port cartpole + +### MLX-TASK-002 + +- Status: `DONE` +- Title: Port cart-double-pendulum + +### MLX-TASK-003 + +- Status: `DONE` +- Title: Port quadcopter + +### MLX-TASK-004 + +- Status: `BLOCKED` +- Depends on: `MLX-SIM-005`, `MLX-SIM-006`, `MLX-SIM-007` +- Title: Port first quadruped locomotion task +- Suggested target: + - pick the simplest existing velocity locomotion task with minimal sensor coupling + +### MLX-TASK-005 + +- Status: `BLOCKED` +- Depends on: `MLX-TASK-004` +- Title: Add quadruped replay smoke and benchmark + +### MLX-TASK-006 + +- Status: `BLOCKED` +- Depends on: `MLX-TASK-004` +- Title: Add quadruped training smoke on MLX + +### MLX-TASK-007 + +- Status: `LATER` +- Depends on: `MLX-TASK-004` +- Title: Port first humanoid locomotion task + +### MLX-TASK-008 + +- Status: `LATER` +- Depends on: `MLX-SIM-003` +- Title: Port first manipulation reach task + +### MLX-TASK-009 + +- Status: `LATER` +- Depends on: `MLX-TASK-008` +- Title: Port first manipulation lift task + +### MLX-TASK-010 + +- Status: `LATER` +- Depends on: `MLX-SENSOR-001` +- Title: Port first raycast-driven task + +### MLX-TASK-011 + +- Status: `READY` +- Title: Keep all current task slices benchmarked on every major substrate change + +### MLX-TASK-012 + +- Status: `READY` +- Title: Keep checkpoint/replay contracts stable across all mac-native tasks + +## Phase D: Kernel Replacement + +### MLX-KERNEL-001 + +- Status: `READY` +- Title: Inventory first real Warp/custom CUDA kernels needed by planned locomotion target + +### MLX-KERNEL-002 + +- Status: `READY` +- Depends on: `MLX-KERNEL-001` +- Title: Implement MLX-op replacements for non-hot helper kernels + +### MLX-KERNEL-003 + +- Status: `BLOCKED` +- Depends on: `MLX-KERNEL-001` +- Title: Implement Metal-backed replacements for locomotion hot loops + +### MLX-KERNEL-004 + +- Status: `READY` +- Depends on: `MLX-KERNEL-001` +- Title: Add per-kernel parity tests against upstream outputs where feasible + +### MLX-KERNEL-005 + +- Status: `READY` +- Title: Add benchmark reporting that detects accidental CPU fallback + +### MLX-KERNEL-006 + +- Status: `LATER` +- Title: Inventory raycast kernels for future sensor port + +### MLX-KERNEL-007 + +- Status: `LATER` +- Title: Inventory camera/tiled-camera reshape kernels + +### MLX-KERNEL-008 + +- Status: `LATER` +- Title: Create a shared kernel-compat layer instead of scattered replacements + +## Phase E: Sensors + +### MLX-SENSOR-001 + +- Status: `LATER` +- Title: Implement `mac-sensors` raycast substrate + +### MLX-SENSOR-002 + +- Status: `LATER` +- Depends on: `MLX-SENSOR-001` +- Title: Add raycast parity tests and benchmark + +### MLX-SENSOR-003 + +- Status: `LATER` +- Title: Define minimal depth output contract for task-usable cameras + +### MLX-SENSOR-004 + +- Status: `LATER` +- Depends on: `MLX-SENSOR-003` +- Title: Add basic camera/depth path for non-RTX tasks + +### MLX-SENSOR-005 + +- Status: `READY` +- Title: Ensure unsupported camera/RTX features fail explicitly via capability checks + +### MLX-SENSOR-006 + +- Status: `LATER` +- Title: Add benchmark coverage for sensor-heavy mac-native tasks + +## Phase F: RL And Training + +### MLX-RL-001 + +- Status: `READY` +- Title: Extract reusable PPO trainer substrate from cartpole-specific code + +### MLX-RL-002 + +- Status: `READY` +- Depends on: `MLX-RL-001` +- Title: Define shared MLX policy/checkpoint format for mac-native tasks + +### MLX-RL-003 + +- Status: `BLOCKED` +- Depends on: `MLX-TASK-004`, `MLX-RL-001` +- Title: Train first locomotion task on MLX + +### MLX-RL-004 + +- Status: `READY` +- Title: Add shared replay/eval scripts for all MLX task slices + +### MLX-RL-005 + +- Status: `LATER` +- Title: Define MLX-native wrapper surface instead of relying on torch-centric RL wrappers + +### MLX-RL-006 + +- Status: `LATER` +- Title: Add multi-task benchmark/training dashboard output + +## Phase G: CI And Release + +### MLX-CI-001 + +- Status: `DONE` +- Title: Add MLX macOS smoke workflow + +### MLX-CI-002 + +- Status: `READY` +- Title: Add benchmark smoke run and artifact upload to MLX macOS workflow + +### MLX-CI-003 + +- Status: `READY` +- Depends on: `MLX-CI-002` +- Title: Add import-safety lane that proves no Isaac Sim install is required + +### MLX-CI-004 + +- Status: `LATER` +- Title: Add nightly drift checks against selected upstream task semantics + +### MLX-CI-005 + +- Status: `READY` +- Title: Publish support matrix and benchmark expectations in README + +### MLX-CI-006 + +- Status: `LATER` +- Title: Archive benchmark trend JSON for M-series comparisons + +## Phase H: Follow-On Compatibility + +### MLX-ROS-001 + +- Status: `LATER` +- Title: Define planner compatibility seam to replace cuRobo progressively + +### MLX-ROS-002 + +- Status: `LATER` +- Title: Start plain ROS 2 process/message interoperability without CUDA assumptions + +### MLX-ROS-003 + +- Status: `LATER` +- Title: Document future CPU/Metal transport path for Isaac ROS compatibility + +## Continuous Work Queue + +This queue exists so work can continue without waiting for a new plan: + +- If `MLX-IMPORT-003` is incomplete, keep taking the next import blocker from `isaaclab.sim`. +- If `MLX-IMPORT-003` is complete, move directly to `MLX-IMPORT-004`. +- Once the import-safety pass stops yielding high-value wins, start `MLX-SIM-001`. +- Once `MLX-SIM-003` is complete, immediately start `MLX-TASK-004`. + +## Validation Commands + +```bash +PYTHONPATH=.:source/isaaclab:source/isaaclab_rl .venv/bin/pytest \ + source/isaaclab_rl/test/test_import_safety.py \ + source/isaaclab/test/backends/test_runtime.py \ + source/isaaclab/test/backends/test_portability_utils.py \ + source/isaaclab/test/backends/test_mac_cartpole.py \ + source/isaaclab/test/backends/test_mac_cartpole_showcase.py \ + source/isaaclab/test/backends/test_mac_cart_double_pendulum.py \ + source/isaaclab/test/backends/test_mac_quadcopter.py -q +``` + +```bash +PYTHONPATH=.:source/isaaclab .venv/bin/python \ + scripts/benchmarks/mlx/benchmark_mac_tasks.py \ + --tasks cartpole cart-double-pendulum quadcopter train-cartpole +``` diff --git a/README.md b/README.md index 0078cad45..06857708e 100644 --- a/README.md +++ b/README.md @@ -30,11 +30,15 @@ This fork exists to close that gap: What works today: - runtime/backend selection seam for `torch-cuda|mlx` and `isaacsim|mac-sim` -- lazy import boundaries so the macOS path fails explicitly instead of exploding on missing `omni.*` +- lazy import boundaries in `envs`, `sim`, `assets`, `sensors`, `managers`, `controllers`, `devices`, `scene`, and `markers` so the macOS path fails explicitly instead of exploding on missing `omni.*` - reproducible source bootstrap script for upstream repositories - a runnable `MLX + mac-sim` cartpole vertical slice +- cartpole showcase variants covering Box, Discrete, MultiDiscrete, Tuple, and Dict spaces +- a runnable `MLX + mac-sim` cart-double-pendulum MARL slice with dict actions/observations/rewards +- a runnable `MLX + mac-sim` quadcopter slice with root-state dynamics - MLX training, checkpoint save/load, and replay scripts -- smoke tests for the backend seam and mac-native cartpole flow +- portability guards for optional `torch`/`warp` utility imports on macOS +- smoke tests for the backend seam and mac-native task slices What this does not claim yet: @@ -80,7 +84,17 @@ Current public options: - `isaacsim`: upstream runtime adapter - `mac-sim`: Mac-native adapter path -The current `mac-sim` implementation is intentionally narrow. It only implements the cartpole slice required to prove the full MLX training loop can run without CUDA. +The current `mac-sim` implementation is intentionally narrow. It currently covers cartpole, cart-double-pendulum, and quadcopter slices needed to prove the MLX training/inference path can run without CUDA. + +### Kernel backend + +The kernel backend isolates Warp and future Metal custom-kernel paths: + +- `warp`: upstream Isaac Sim + Warp path +- `metal`: Apple Silicon path for MLX + Metal-backed kernel replacements +- `cpu`: correctness fallback for bring-up and unsupported kernels + +Today the public MLX tasks mostly use pure MLX ops. The explicit kernel selection seam exists now so future Metal kernels can land without rewriting task-facing APIs again. ## MLX Quick Start @@ -91,7 +105,13 @@ This path is for Apple Silicon macOS. ```bash cd IsaacLab uv venv --python 3.11 .venv -uv pip install --python .venv/bin/python mlx pytest pytest-mock toml +uv pip install --python .venv/bin/python -e source/isaaclab[macos-mlx,dev] +``` + +If you want the upstream Isaac Sim runtime on Linux/NVIDIA instead, install the CUDA-side extra: + +```bash +uv pip install --python .venv/bin/python -e source/isaaclab[cuda-isaacsim,dev] ``` ### 2. Train the MLX cartpole baseline @@ -115,13 +135,52 @@ PYTHONPATH=.:source/isaaclab .venv/bin/python \ --episodes 3 ``` -### 4. Run the focused backend test suite +Optional resume flow: + +```bash +PYTHONPATH=.:source/isaaclab .venv/bin/python \ + scripts/reinforcement_learning/mlx/train_cartpole.py \ + --resume-from logs/mlx/cartpole_policy.npz \ + --checkpoint logs/mlx/cartpole_policy_resumed.npz \ + --updates 50 +``` + +The torch-centric RL wrapper extension is now optional. If you need upstream SB3, rl-games, or RSL-RL wrappers, install them separately from the MLX path: + +```bash +uv pip install --python .venv/bin/python -e source/isaaclab_rl[sb3] +uv pip install --python .venv/bin/python -e source/isaaclab_rl[rsl-rl] +``` + +For the MLX/mac-sim task slices documented in this fork, `source/isaaclab_rl` is not required. + +### 4. Run cart-double-pendulum MARL smoke + +```bash +PYTHONPATH=.:source/isaaclab .venv/bin/python \ + scripts/reinforcement_learning/mlx/play_cart_double_pendulum.py \ + --num-envs 64 --episodes 3 --max-steps 10000 --random-actions +``` + +### 5. Run quadcopter smoke + +```bash +PYTHONPATH=.:source/isaaclab .venv/bin/python \ + scripts/reinforcement_learning/mlx/play_quadcopter.py \ + --num-envs 64 --episodes 3 --episode-length-s 0.5 --max-steps 10000 --thrust-action 0.2 --no-random-actions +``` + +### 6. Run the focused backend test suite ```bash PYTHONPATH=.:source/isaaclab .venv/bin/pytest \ scripts/tools/test/test_bootstrap_isaac_sources.py \ source/isaaclab/test/backends/test_runtime.py \ - source/isaaclab/test/backends/test_mac_cartpole.py -q + source/isaaclab/test/backends/test_portability_utils.py \ + source/isaaclab/test/backends/test_mac_cartpole.py \ + source/isaaclab/test/backends/test_mac_cartpole_showcase.py \ + source/isaaclab/test/backends/test_mac_cart_double_pendulum.py \ + source/isaaclab/test/backends/test_mac_quadcopter.py -q ``` ## Runtime Selection @@ -129,6 +188,7 @@ PYTHONPATH=.:source/isaaclab .venv/bin/pytest \ The backend seam is exposed through the app/runtime layer: - `--compute-backend torch-cuda|mlx` +- `--kernel-backend warp|metal|cpu` - `--sim-backend isaacsim|mac-sim` These flags are published through: @@ -139,21 +199,50 @@ Examples: ```bash # Upstream path -python some_script.py --compute-backend torch-cuda --sim-backend isaacsim +python some_script.py --compute-backend torch-cuda --kernel-backend warp --sim-backend isaacsim # macOS port path -python some_script.py --compute-backend mlx --sim-backend mac-sim +python some_script.py --compute-backend mlx --kernel-backend metal --sim-backend mac-sim +``` + +Important constraint: `AppLauncher` now accepts `--compute-backend mlx --sim-backend mac-sim` in bootstrap mode, but it does not launch Isaac Sim/Omniverse there. Use the dedicated MLX scripts for task execution. + +## M5 Benchmarking + +The fork now includes a dedicated MLX benchmark entrypoint for the current mac-native task set: + +- [`scripts/benchmarks/mlx/benchmark_mac_tasks.py`](scripts/benchmarks/mlx/benchmark_mac_tasks.py) + +Example: + +```bash +PYTHONPATH=.:source/isaaclab .venv/bin/python \ + scripts/benchmarks/mlx/benchmark_mac_tasks.py \ + --tasks cartpole cart-double-pendulum quadcopter train-cartpole \ + --num-envs 256 \ + --steps 512 \ + --train-updates 20 \ + --json-out logs/benchmarks/mlx/m5-baseline.json ``` -Important constraint: `AppLauncher` itself still only launches the upstream Isaac Sim runtime path. The MLX cartpole scripts run directly on the `mac-sim` backend and do not launch Isaac Sim. +The benchmark emits: + +- per-task `env_steps_per_s` for the current MLX/mac-sim env slices +- cartpole `train_frames_per_s` for the first MLX training loop +- runtime metadata including compute, kernel, sensor, and planner backend selection ## Implemented MLX Vertical Slice -The current vertical slice is cartpole: +Current implemented slices: -- environment config: [`source/isaaclab/isaaclab/backends/mac_sim/cartpole.py`](source/isaaclab/isaaclab/backends/mac_sim/cartpole.py) -- trainer entrypoint: [`scripts/reinforcement_learning/mlx/train_cartpole.py`](scripts/reinforcement_learning/mlx/train_cartpole.py) -- replay entrypoint: [`scripts/reinforcement_learning/mlx/play_cartpole.py`](scripts/reinforcement_learning/mlx/play_cartpole.py) +- cartpole environment and trainer in [`source/isaaclab/isaaclab/backends/mac_sim/cartpole.py`](source/isaaclab/isaaclab/backends/mac_sim/cartpole.py) +- cartpole showcase space variants in [`source/isaaclab/isaaclab/backends/mac_sim/showcase.py`](source/isaaclab/isaaclab/backends/mac_sim/showcase.py) +- cart-double-pendulum MARL environment in [`source/isaaclab/isaaclab/backends/mac_sim/cart_double_pendulum.py`](source/isaaclab/isaaclab/backends/mac_sim/cart_double_pendulum.py) +- quadcopter environment in [`source/isaaclab/isaaclab/backends/mac_sim/quadcopter.py`](source/isaaclab/isaaclab/backends/mac_sim/quadcopter.py) +- cartpole trainer entrypoint in [`scripts/reinforcement_learning/mlx/train_cartpole.py`](scripts/reinforcement_learning/mlx/train_cartpole.py) +- cartpole replay entrypoint in [`scripts/reinforcement_learning/mlx/play_cartpole.py`](scripts/reinforcement_learning/mlx/play_cartpole.py) +- cart-double-pendulum replay/smoke entrypoint in [`scripts/reinforcement_learning/mlx/play_cart_double_pendulum.py`](scripts/reinforcement_learning/mlx/play_cart_double_pendulum.py) +- quadcopter replay/smoke entrypoint in [`scripts/reinforcement_learning/mlx/play_quadcopter.py`](scripts/reinforcement_learning/mlx/play_quadcopter.py) The cartpole path preserves the important upstream task semantics: @@ -162,6 +251,8 @@ The cartpole path preserves the important upstream task semantics: - termination conditions remain cart out-of-bounds or pole angle exceeding `pi/2` - reset sampling keeps the pole-angle randomization behavior - observations returned after done/reset are post-reset observations, matching the upstream direct RL flow +- cart-double-pendulum preserves per-agent dict observations/rewards/dones for `cart` and `pendulum` +- quadcopter preserves a root-state-centric policy observation layout with vectorized thrust/moment control ## Bootstrapping Upstream Sources @@ -202,8 +293,8 @@ The most important fork changes are: Near-term priorities: -1. Expand `mac-sim` from cartpole-only logic into a more general articulation/scene layer. -2. Add a second and third task with compatible observation/action structure. +1. Expand `mac-sim` from the current task slices into a more general articulation/scene layer. +2. Add additional locomotion/manipulation tasks with compatible observation/action structure. 3. Push more task code through backend capability checks instead of import-time backend assumptions. 4. Define a stable checkpoint/config story across multiple MLX tasks. 5. Add Apple Silicon CI for the MLX backend slice. @@ -277,7 +368,7 @@ Common current pitfalls: - If `mlx` is missing, install it into the active `uv` environment. - If imports fail on `omni.*` or `isaacsim.*`, you are probably invoking an upstream Isaac Sim path rather than the `mac-sim` path. -- If you expect `AppLauncher` to boot the MLX cartpole slice, that is not wired yet. Use the dedicated MLX scripts instead. +- If you use `AppLauncher` with `--sim-backend mac-sim`, it runs in bootstrap mode and returns a placeholder app object. Use the dedicated MLX scripts for environment stepping/training. - If a task depends on cameras, Warp, cuRobo, or Omniverse-only features, it should currently be treated as unsupported on the macOS path. ## Support diff --git a/scripts/benchmarks/mlx/benchmark_mac_tasks.py b/scripts/benchmarks/mlx/benchmark_mac_tasks.py new file mode 100644 index 000000000..0628eb54b --- /dev/null +++ b/scripts/benchmarks/mlx/benchmark_mac_tasks.py @@ -0,0 +1,272 @@ +# 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 the MLX/mac-sim task slices on Apple Silicon.""" + +from __future__ import annotations + +import argparse +import json +import platform +import sys +import time +from pathlib import Path +from typing import Any + +import mlx.core as mx + +from isaaclab.backends import get_runtime_state +from isaaclab.backends.mac_sim import ( + MacCartDoublePendulumEnv, + MacCartDoublePendulumEnvCfg, + MacCartpoleEnv, + MacCartpoleEnvCfg, + MacCartpoleTrainCfg, + MacQuadcopterEnv, + MacQuadcopterEnvCfg, + train_cartpole_policy, +) + +TASK_CHOICES = ("cartpole", "cart-double-pendulum", "quadcopter", "train-cartpole") + + +def parse_args() -> argparse.Namespace: + """Parse benchmark arguments.""" + parser = argparse.ArgumentParser(description="Benchmark the MLX/mac-sim task slices.") + parser.add_argument("--tasks", nargs="+", choices=TASK_CHOICES, default=list(TASK_CHOICES)) + parser.add_argument("--num-envs", type=int, default=256) + parser.add_argument("--steps", type=int, default=512) + parser.add_argument("--train-updates", type=int, default=20) + parser.add_argument("--rollout-steps", type=int, default=64) + parser.add_argument("--epochs-per-update", type=int, default=4) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--quadcopter-thrust-action", type=float, default=0.2) + parser.add_argument("--artifact-dir", type=Path, default=Path("logs/benchmarks/mlx")) + parser.add_argument("--json-out", type=Path, default=None) + return parser.parse_args() + + +def _sync(values: list[mx.array]) -> None: + """Force MLX execution for stable timing.""" + mx.eval(*values) + + +def _make_benchmark_result( + name: str, + *, + num_envs: int, + steps: int, + elapsed_s: float, + runtime_state: dict[str, Any], + extra: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Build a normalized benchmark payload.""" + env_steps = num_envs * steps + result = { + "task": name, + "num_envs": num_envs, + "steps": steps, + "elapsed_s": elapsed_s, + "env_steps_per_s": env_steps / elapsed_s, + "mean_step_ms": (elapsed_s / steps) * 1000.0, + "runtime": runtime_state, + } + if extra: + result.update(extra) + return result + + +def _sim_backend_snapshot(env: Any) -> dict[str, Any]: + """Capture the concrete simulator adapter contract for the benchmarked env.""" + sim_backend = env.sim_backend + return { + "backend": sim_backend.name, + "capabilities": sim_backend.capabilities.__dict__, + "contract": { + "reset_signature": sim_backend.contract.reset_signature, + "step_signature": sim_backend.contract.step_signature, + "articulations": sim_backend.contract.articulations.__dict__, + }, + } + + +def benchmark_cartpole(num_envs: int, steps: int, seed: int) -> dict[str, Any]: + """Benchmark the cartpole MLX env step loop.""" + env = MacCartpoleEnv(MacCartpoleEnvCfg(num_envs=num_envs, seed=seed)) + observations, _ = env.reset() + actions = mx.zeros((num_envs, 1), dtype=mx.float32) + _sync([observations["policy"]]) + + start = time.perf_counter() + for _ in range(steps): + observations, rewards, terminated, truncated, _ = env.step(actions) + _sync([observations["policy"], rewards, terminated, truncated]) + elapsed_s = time.perf_counter() - start + + return _make_benchmark_result( + "cartpole", + num_envs=num_envs, + steps=steps, + elapsed_s=elapsed_s, + runtime_state=get_runtime_state(env.runtime), + extra={ + "observation_dim": env.cfg.observation_space, + "sim_adapter": _sim_backend_snapshot(env), + }, + ) + + +def benchmark_cart_double_pendulum(num_envs: int, steps: int, seed: int) -> dict[str, Any]: + """Benchmark the MARL cart-double-pendulum step loop.""" + env = MacCartDoublePendulumEnv(MacCartDoublePendulumEnvCfg(num_envs=num_envs, seed=seed)) + observations, _ = env.reset() + actions = { + "cart": mx.zeros((num_envs, 1), dtype=mx.float32), + "pendulum": mx.zeros((num_envs, 1), dtype=mx.float32), + } + _sync([observations["cart"], observations["pendulum"]]) + + start = time.perf_counter() + for _ in range(steps): + observations, rewards, terminated, truncated, _ = env.step(actions) + _sync( + [ + observations["cart"], + observations["pendulum"], + rewards["cart"], + rewards["pendulum"], + terminated["cart"], + truncated["cart"], + ] + ) + elapsed_s = time.perf_counter() - start + + return _make_benchmark_result( + "cart-double-pendulum", + num_envs=num_envs, + steps=steps, + elapsed_s=elapsed_s, + runtime_state=get_runtime_state(env.runtime), + extra={ + "cart_observation_dim": env.cfg.observation_spaces["cart"], + "pendulum_observation_dim": env.cfg.observation_spaces["pendulum"], + "sim_adapter": _sim_backend_snapshot(env), + }, + ) + + +def benchmark_quadcopter(num_envs: int, steps: int, seed: int, thrust_action: float) -> dict[str, Any]: + """Benchmark the quadcopter MLX env step loop.""" + env = MacQuadcopterEnv(MacQuadcopterEnvCfg(num_envs=num_envs, seed=seed)) + observations, _ = env.reset() + actions = mx.zeros((num_envs, 4), dtype=mx.float32) + actions[:, 0] = thrust_action + _sync([observations["policy"]]) + + start = time.perf_counter() + for _ in range(steps): + observations, rewards, terminated, truncated, _ = env.step(actions) + _sync([observations["policy"], rewards, terminated, truncated]) + elapsed_s = time.perf_counter() - start + + return _make_benchmark_result( + "quadcopter", + num_envs=num_envs, + steps=steps, + elapsed_s=elapsed_s, + runtime_state=get_runtime_state(env.runtime), + extra={ + "observation_dim": env.cfg.observation_space, + "thrust_action": thrust_action, + "sim_adapter": _sim_backend_snapshot(env), + }, + ) + + +def benchmark_train_cartpole( + num_envs: int, + updates: int, + rollout_steps: int, + epochs_per_update: int, + seed: int, + artifact_dir: Path, +) -> dict[str, Any]: + """Benchmark the cartpole MLX training loop.""" + artifact_dir.mkdir(parents=True, exist_ok=True) + checkpoint_path = artifact_dir / "cartpole_benchmark_policy.npz" + cfg = MacCartpoleTrainCfg( + env=MacCartpoleEnvCfg(num_envs=num_envs, seed=seed), + updates=updates, + rollout_steps=rollout_steps, + epochs_per_update=epochs_per_update, + checkpoint_path=str(checkpoint_path), + eval_interval=max(1, updates), + ) + + start = time.perf_counter() + result = train_cartpole_policy(cfg) + elapsed_s = time.perf_counter() - start + frames = num_envs * updates * rollout_steps + + benchmark = { + "task": "train-cartpole", + "num_envs": num_envs, + "updates": updates, + "rollout_steps": rollout_steps, + "epochs_per_update": epochs_per_update, + "elapsed_s": elapsed_s, + "train_frames": frames, + "train_frames_per_s": frames / elapsed_s, + "checkpoint_path": result["checkpoint_path"], + "completed_episodes": result["completed_episodes"], + "mean_recent_return": result["mean_recent_return"], + } + benchmark["runtime"] = get_runtime_state() + return benchmark + + +def main() -> int: + """Run the requested benchmarks and optionally write JSON output.""" + args = parse_args() + args.artifact_dir.mkdir(parents=True, exist_ok=True) + if args.json_out is None: + args.json_out = args.artifact_dir / "benchmark-results.json" + results = { + "platform": { + "system": platform.system(), + "machine": platform.machine(), + "platform": platform.platform(), + "python": sys.version.split()[0], + }, + "benchmarks": [], + } + + for task in args.tasks: + if task == "cartpole": + benchmark = benchmark_cartpole(args.num_envs, args.steps, args.seed) + elif task == "cart-double-pendulum": + benchmark = benchmark_cart_double_pendulum(args.num_envs, args.steps, args.seed) + elif task == "quadcopter": + benchmark = benchmark_quadcopter(args.num_envs, args.steps, args.seed, args.quadcopter_thrust_action) + else: + benchmark = benchmark_train_cartpole( + args.num_envs, + args.train_updates, + args.rollout_steps, + args.epochs_per_update, + args.seed, + args.artifact_dir, + ) + results["benchmarks"].append(benchmark) + + output = json.dumps(results, indent=2, sort_keys=True) + print(output) + args.json_out.parent.mkdir(parents=True, exist_ok=True) + args.json_out.write_text(output + "\n", encoding="utf-8") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/reinforcement_learning/mlx/play_cart_double_pendulum.py b/scripts/reinforcement_learning/mlx/play_cart_double_pendulum.py new file mode 100644 index 000000000..3f6bb25b8 --- /dev/null +++ b/scripts/reinforcement_learning/mlx/play_cart_double_pendulum.py @@ -0,0 +1,76 @@ +# 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 + +"""Run the mac-native cart-double-pendulum MARL slice with random or constant actions.""" + +from __future__ import annotations + +import argparse + +import mlx.core as mx + +from isaaclab.backends.mac_sim import MacCartDoublePendulumEnv, MacCartDoublePendulumEnvCfg + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Play the MLX/mac-sim cart-double-pendulum slice.") + parser.add_argument("--num-envs", type=int, default=64) + parser.add_argument("--episodes", type=int, default=3) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--max-steps", type=int, default=10000) + parser.add_argument("--random-actions", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--cart-action", type=float, default=0.0) + parser.add_argument("--pendulum-action", type=float, default=0.0) + return parser.parse_args() + + +def _actions(args: argparse.Namespace, num_envs: int) -> dict[str, mx.array]: + if args.random_actions: + return { + "cart": mx.random.uniform(low=-1.0, high=1.0, shape=(num_envs, 1)), + "pendulum": mx.random.uniform(low=-1.0, high=1.0, shape=(num_envs, 1)), + } + return { + "cart": mx.full((num_envs, 1), args.cart_action, dtype=mx.float32), + "pendulum": mx.full((num_envs, 1), args.pendulum_action, dtype=mx.float32), + } + + +def main() -> int: + args = parse_args() + cfg = MacCartDoublePendulumEnvCfg(num_envs=args.num_envs, seed=args.seed) + env = MacCartDoublePendulumEnv(cfg) + mx.random.seed(args.seed) + env.reset() + + completed: list[tuple[float, float, int]] = [] + for _ in range(args.max_steps): + _, _, _, _, extras = env.step(_actions(args, cfg.num_envs)) + if not extras.get("completed_returns"): + continue + + cart_returns = extras["completed_returns"]["cart"] + pendulum_returns = extras["completed_returns"]["pendulum"] + lengths = extras["completed_lengths"] + completed.extend(zip(cart_returns, pendulum_returns, lengths, strict=True)) + if len(completed) >= args.episodes: + break + + for index, (cart_return, pendulum_return, length) in enumerate(completed[: args.episodes], start=1): + print( + f"[mlx-cart-double-pendulum] episode={index} cart_return={cart_return:.4f} " + f"pendulum_return={pendulum_return:.4f} length={length}" + ) + + if len(completed) < args.episodes: + print( + "[mlx-cart-double-pendulum] completed fewer episodes than requested " + f"({len(completed)}/{args.episodes}) within max_steps={args.max_steps}" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/reinforcement_learning/mlx/play_quadcopter.py b/scripts/reinforcement_learning/mlx/play_quadcopter.py new file mode 100644 index 000000000..1a771da40 --- /dev/null +++ b/scripts/reinforcement_learning/mlx/play_quadcopter.py @@ -0,0 +1,72 @@ +# 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 + +"""Run the mac-native quadcopter slice with random or constant actions.""" + +from __future__ import annotations + +import argparse + +import mlx.core as mx + +from isaaclab.backends.mac_sim import MacQuadcopterEnv, MacQuadcopterEnvCfg + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Play the MLX/mac-sim quadcopter slice.") + parser.add_argument("--num-envs", type=int, default=64) + parser.add_argument("--episodes", type=int, default=3) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--episode-length-s", type=float, default=10.0) + parser.add_argument("--max-steps", type=int, default=10000) + parser.add_argument("--random-actions", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--thrust-action", type=float, default=0.2) + parser.add_argument("--roll-action", type=float, default=0.0) + parser.add_argument("--pitch-action", type=float, default=0.0) + parser.add_argument("--yaw-action", type=float, default=0.0) + return parser.parse_args() + + +def _actions(args: argparse.Namespace, num_envs: int) -> mx.array: + if args.random_actions: + return mx.random.uniform(low=-1.0, high=1.0, shape=(num_envs, 4)) + return mx.array( + [[args.thrust_action, args.roll_action, args.pitch_action, args.yaw_action]] * num_envs, + dtype=mx.float32, + ) + + +def main() -> int: + args = parse_args() + cfg = MacQuadcopterEnvCfg(num_envs=args.num_envs, seed=args.seed, episode_length_s=args.episode_length_s) + env = MacQuadcopterEnv(cfg) + mx.random.seed(args.seed) + env.reset() + + completed: list[tuple[int, float]] = [] + for _ in range(args.max_steps): + _, _, _, _, extras = env.step(_actions(args, cfg.num_envs)) + if not extras.get("completed_lengths"): + continue + + lengths = extras["completed_lengths"] + distance = float(extras["final_distance_to_goal"]) + completed.extend((int(length), distance) for length in lengths) + if len(completed) >= args.episodes: + break + + for index, (length, distance) in enumerate(completed[: args.episodes], start=1): + print(f"[mlx-quadcopter] episode={index} length={length} final_distance={distance:.4f}") + + if len(completed) < args.episodes: + print( + "[mlx-quadcopter] completed fewer episodes than requested " + f"({len(completed)}/{args.episodes}) within max_steps={args.max_steps}" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/reinforcement_learning/mlx/train_cartpole.py b/scripts/reinforcement_learning/mlx/train_cartpole.py index dfdeeb84d..ec3c2964c 100644 --- a/scripts/reinforcement_learning/mlx/train_cartpole.py +++ b/scripts/reinforcement_learning/mlx/train_cartpole.py @@ -8,6 +8,8 @@ from __future__ import annotations import argparse +import json +from pathlib import Path from isaaclab.backends.mac_sim import MacCartpoleEnvCfg, MacCartpoleTrainCfg, train_cartpole_policy @@ -19,27 +21,49 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--rollout-steps", type=int, default=64) parser.add_argument("--epochs-per-update", type=int, default=4) parser.add_argument("--learning-rate", type=float, default=3e-4) - parser.add_argument("--hidden-dim", type=int, default=128) + parser.add_argument("--hidden-dim", type=int, default=None) parser.add_argument("--checkpoint", type=str, default="logs/mlx/cartpole_policy.npz") + parser.add_argument("--resume-from", type=str, default=None) parser.add_argument("--eval-interval", type=int, default=10) parser.add_argument("--seed", type=int, default=42) return parser.parse_args() +def _resolve_hidden_dim(args: argparse.Namespace) -> int: + if args.hidden_dim is not None: + return args.hidden_dim + if args.resume_from is None: + return 128 + + metadata_path = Path(args.resume_from) + if metadata_path.suffix: + metadata_path = metadata_path.with_suffix(f"{metadata_path.suffix}.json") + else: + metadata_path = metadata_path.with_suffix(".json") + if not metadata_path.exists(): + return 128 + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + return int(metadata.get("hidden_dim", 128)) + + def main() -> int: args = parse_args() + hidden_dim = _resolve_hidden_dim(args) cfg = MacCartpoleTrainCfg( env=MacCartpoleEnvCfg(num_envs=args.num_envs, seed=args.seed), - hidden_dim=args.hidden_dim, + hidden_dim=hidden_dim, updates=args.updates, rollout_steps=args.rollout_steps, epochs_per_update=args.epochs_per_update, learning_rate=args.learning_rate, checkpoint_path=args.checkpoint, + resume_from=args.resume_from, eval_interval=args.eval_interval, ) result = train_cartpole_policy(cfg) print(f"[mlx-cartpole] checkpoint={result['checkpoint_path']}") + if result["resumed_from"] is not None: + print(f"[mlx-cartpole] resumed_from={result['resumed_from']}") print(f"[mlx-cartpole] completed_episodes={result['completed_episodes']}") print(f"[mlx-cartpole] mean_recent_return={result['mean_recent_return']:.4f}") return 0 diff --git a/source/isaaclab/isaaclab/app/app_launcher.py b/source/isaaclab/isaaclab/app/app_launcher.py index c845e3cc9..958a1754a 100644 --- a/source/isaaclab/isaaclab/app/app_launcher.py +++ b/source/isaaclab/isaaclab/app/app_launcher.py @@ -23,6 +23,8 @@ from isaaclab.backends import ( UnsupportedBackendError, + create_planner_backend, + create_sensor_backend, is_isaacsim_available, resolve_runtime_selection, set_runtime_selection, @@ -63,6 +65,16 @@ def __call__(self, parser, namespace, values, option_string=None): setattr(namespace, f"{self.dest}_explicit", True) +class _MacSimPlaceholderApp: + """Minimal placeholder app for the mac-sim bootstrap path.""" + + def is_running(self) -> bool: + return False + + def close(self) -> None: + return None + + class AppLauncher: """A utility class to launch Isaac Sim application based on command-line arguments and environment variables. @@ -147,6 +159,14 @@ def __init__(self, launcher_args: argparse.Namespace | dict | None = None, **kwa # Integrate env-vars and input keyword args into simulation app config self._config_resolution(launcher_args) + # The mac-sim path is intentionally bootstrap-only and does not launch Isaac Sim. + if self.sim_backend != "isaacsim": + self._app = _MacSimPlaceholderApp() + print( + "[INFO][AppLauncher]: Initialized mac-sim bootstrap mode without launching Isaac Sim." + ) + return + # Internal: Override SimulationApp._start_app method to apply patches after app has started. self.__patch_simulation_start_app(launcher_args) @@ -341,6 +361,16 @@ def add_app_launcher_args(parser: argparse.ArgumentParser) -> None: choices={"isaacsim", "mac-sim"}, help="Simulation backend. 'isaacsim' preserves the upstream runtime, 'mac-sim' selects the macOS port path.", ) + arg_group.add_argument( + "--kernel-backend", + type=str, + default=AppLauncher._APPLAUNCHER_CFG_INFO["kernel_backend"][1], + choices={"warp", "metal", "cpu"}, + help=( + "Kernel backend. Use 'warp' for the upstream Isaac Sim path, 'metal' for the macOS port path," + " and 'cpu' only for correctness bring-up." + ), + ) # Add the deprecated cpu flag to raise an error if it is used arg_group.add_argument("--cpu", action="store_true", help=argparse.SUPPRESS) arg_group.add_argument( @@ -426,6 +456,7 @@ def add_app_launcher_args(parser: argparse.ArgumentParser) -> None: "device": ([str], "cuda:0"), "compute_backend": ([str], "torch-cuda"), "sim_backend": ([str], "isaacsim"), + "kernel_backend": ([str], "warp"), "experience": ([str], ""), "rendering_mode": ([str], "balanced"), } @@ -535,6 +566,11 @@ def _config_resolution(self, launcher_args: dict): # Resolve compute/sim backends before Isaac Sim-specific setup. self._resolve_backend_settings(launcher_args) + # The mac-sim path does not launch SimulationApp/Kit and therefore skips Isaac Sim setup. + if self.sim_backend != "isaacsim": + self._sim_app_config = {} + return + # Handle experience file settings self._resolve_experience_file(launcher_args) @@ -752,22 +788,36 @@ def _resolve_backend_settings(self, launcher_args: dict): compute_backend=launcher_args.pop("compute_backend", None), sim_backend=launcher_args.pop("sim_backend", None), device=launcher_args.get("device", AppLauncher._APPLAUNCHER_CFG_INFO["device"][1]), + kernel_backend=launcher_args.pop("kernel_backend", None), ) set_runtime_selection(runtime) self.compute_backend = runtime.compute_backend self.sim_backend = runtime.sim_backend + self.kernel_backend = runtime.kernel_backend + self.sensor_backend = create_sensor_backend(runtime).name + self.planner_backend = create_planner_backend(runtime).name print( "[INFO][AppLauncher]: Using runtime backends:" - f" compute={self.compute_backend}, sim={self.sim_backend}" + f" compute={self.compute_backend}, sim={self.sim_backend}, kernel={self.kernel_backend}," + f" sensors={self.sensor_backend}, planners={self.planner_backend}" ) + if self.sim_backend == "mac-sim": + if self.compute_backend != "mlx": + raise UnsupportedBackendError( + "`AppLauncher` supports `--sim-backend mac-sim` only with `--compute-backend mlx`." + ) + if self.kernel_backend == "warp": + raise UnsupportedBackendError( + "`AppLauncher` does not support `--kernel-backend warp` with `--sim-backend mac-sim`." + " Use `metal` or `cpu` for the macOS port path." + ) + return + if self.sim_backend != "isaacsim": - raise UnsupportedBackendError( - "`AppLauncher` only supports `--sim-backend isaacsim` today." - " The `mac-sim` path is being introduced behind separate entrypoints." - ) + raise UnsupportedBackendError(f"Unsupported simulation backend for AppLauncher: {self.sim_backend}.") if self.compute_backend != "torch-cuda": raise UnsupportedBackendError( @@ -775,6 +825,12 @@ def _resolve_backend_settings(self, launcher_args: dict): " Use `--compute-backend mlx --sim-backend mac-sim` for the macOS port path." ) + if self.kernel_backend == "metal": + raise UnsupportedBackendError( + "`AppLauncher` does not support `--kernel-backend metal` with `--sim-backend isaacsim`." + " Use `warp` for the upstream runtime or `metal` with `mac-sim`." + ) + if not is_isaacsim_available(): raise ModuleNotFoundError( "Isaac Sim Python modules are not available for the selected runtime." diff --git a/source/isaaclab/isaaclab/assets/__init__.py b/source/isaaclab/isaaclab/assets/__init__.py index 41640e528..4d55965f5 100644 --- a/source/isaaclab/isaaclab/assets/__init__.py +++ b/source/isaaclab/isaaclab/assets/__init__.py @@ -3,45 +3,48 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Sub-package for different assets, such as rigid objects and articulations. - -An asset is a physical object that can be spawned in the simulation. The class handles both -the spawning of the asset into the USD stage as well as initialization of necessary physics -handles to interact with the asset. - -Upon construction of the asset instance, the prim corresponding to the asset is spawned into the -USD stage if the spawn configuration is not None. The spawn configuration is defined in the -:attr:`AssetBaseCfg.spawn` attribute. In case the configured :attr:`AssetBaseCfg.prim_path` is -an expression, then the prim is spawned at all the matching paths. Otherwise, a single prim is -spawned at the configured path. For more information on the spawn configuration, see the -:mod:`isaaclab.sim.spawners` module. - -The asset class also registers callbacks for the stage play/stop events. These are used to -construct the physics handles for the asset as the physics engine is only available when the -stage is playing. Additionally, the class registers a callback for debug visualization of the -asset. This can be enabled by setting the :attr:`AssetBaseCfg.debug_vis` attribute to True. - -The asset class follows the following naming convention for its methods: - -* **set_xxx()**: These are used to only set the buffers into the :attr:`data` instance. However, they - do not write the data into the simulator. The writing of data only happens when the - :meth:`write_data_to_sim` method is called. -* **write_xxx_to_sim()**: These are used to set the buffers into the :attr:`data` instance and write - the corresponding data into the simulator as well. -* **update(dt)**: These are used to update the buffers in the :attr:`data` instance. This should - be called after a simulation step is performed. - -The main reason to separate the ``set`` and ``write`` operations is to provide flexibility to the -user when they need to perform a post-processing operation of the buffers before applying them -into the simulator. A common example for this is dealing with explicit actuator models where the -specified joint targets are not directly applied to the simulator but are instead used to compute -the corresponding actuator torques. -""" - -from .articulation import Articulation, ArticulationCfg, ArticulationData -from .asset_base import AssetBase -from .asset_base_cfg import AssetBaseCfg -from .deformable_object import DeformableObject, DeformableObjectCfg, DeformableObjectData -from .rigid_object import RigidObject, RigidObjectCfg, RigidObjectData -from .rigid_object_collection import RigidObjectCollection, RigidObjectCollectionCfg, RigidObjectCollectionData -from .surface_gripper import SurfaceGripper, SurfaceGripperCfg +"""Sub-package for simulation assets, lazily loaded by backend capability.""" + +from __future__ import annotations + +import importlib + +from isaaclab.backends import UnsupportedBackendError, current_runtime + +_MODULE_EXPORTS = { + "Articulation": (".articulation", "Articulation"), + "ArticulationCfg": (".articulation", "ArticulationCfg"), + "ArticulationData": (".articulation", "ArticulationData"), + "AssetBase": (".asset_base", "AssetBase"), + "AssetBaseCfg": (".asset_base_cfg", "AssetBaseCfg"), + "DeformableObject": (".deformable_object", "DeformableObject"), + "DeformableObjectCfg": (".deformable_object", "DeformableObjectCfg"), + "DeformableObjectData": (".deformable_object", "DeformableObjectData"), + "RigidObject": (".rigid_object", "RigidObject"), + "RigidObjectCfg": (".rigid_object", "RigidObjectCfg"), + "RigidObjectData": (".rigid_object", "RigidObjectData"), + "RigidObjectCollection": (".rigid_object_collection", "RigidObjectCollection"), + "RigidObjectCollectionCfg": (".rigid_object_collection", "RigidObjectCollectionCfg"), + "RigidObjectCollectionData": (".rigid_object_collection", "RigidObjectCollectionData"), + "SurfaceGripper": (".surface_gripper", "SurfaceGripper"), + "SurfaceGripperCfg": (".surface_gripper", "SurfaceGripperCfg"), +} + +__all__ = [*_MODULE_EXPORTS.keys()] + + +def __getattr__(name: str): + target = _MODULE_EXPORTS.get(name) + if target is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + if current_runtime().sim_backend != "isaacsim": + raise UnsupportedBackendError( + f"`isaaclab.assets.{name}` currently requires `sim-backend=isaacsim`." + " The `mac-sim` path will expose asset interfaces as capabilities are implemented." + ) + + module = importlib.import_module(target[0], __name__) + value = getattr(module, target[1]) + globals()[name] = value + return value diff --git a/source/isaaclab/isaaclab/backends/__init__.py b/source/isaaclab/isaaclab/backends/__init__.py index 6ca809da7..b8353f8ca 100644 --- a/source/isaaclab/isaaclab/backends/__init__.py +++ b/source/isaaclab/isaaclab/backends/__init__.py @@ -8,19 +8,39 @@ from .runtime import ( ArticulationCapabilities, ENV_COMPUTE_BACKEND, + ENV_KERNEL_BACKEND, ENV_SIM_BACKEND, ComputeBackend, ComputeCapabilities, + CpuKernelBackend, IsaacSimBackend, + IsaacSimPlannerBackend, + IsaacSimSensorBackend, + KernelBackend, + KernelCapabilities, + MacPlannerBackend, + MacSensorBackend, + MlxComputeBackend, MacSimBackend, + MetalKernelBackend, + PlannerBackend, + PlannerCapabilities, RuntimeCapabilities, RuntimeSelection, + SensorBackend, + SensorCapabilities, SimBackend, SimBackendContract, SimCapabilities, + TorchCudaComputeBackend, UnsupportedBackendError, UnsupportedRuntimeFeatureError, + WarpKernelBackend, configure_torch_device, + create_compute_backend, + create_kernel_backend, + create_planner_backend, + create_sensor_backend, create_sim_backend, current_runtime, current_runtime_capabilities, @@ -28,6 +48,7 @@ is_apple_silicon, is_isaacsim_available, is_mlx_available, + normalize_kernel_backend, resolve_runtime_selection, set_runtime_selection, ) diff --git a/source/isaaclab/isaaclab/backends/mac_sim/__init__.py b/source/isaaclab/isaaclab/backends/mac_sim/__init__.py index af5155da9..e91746c7b 100644 --- a/source/isaaclab/isaaclab/backends/mac_sim/__init__.py +++ b/source/isaaclab/isaaclab/backends/mac_sim/__init__.py @@ -14,4 +14,32 @@ play_cartpole_policy, train_cartpole_policy, ) - +from .cart_double_pendulum import ( + MacCartDoublePendulumEnv, + MacCartDoublePendulumEnvCfg, + MacCartDoublePendulumSimBackend, +) +from .quadcopter import ( + MacQuadcopterEnv, + MacQuadcopterEnvCfg, + MacQuadcopterSimBackend, +) +from .showcase import ( + SHOWCASE_CFGS, + BoxBoxEnvCfg, + BoxDiscreteEnvCfg, + BoxMultiDiscreteEnvCfg, + DictBoxEnvCfg, + DictDiscreteEnvCfg, + DictMultiDiscreteEnvCfg, + DiscreteBoxEnvCfg, + DiscreteDiscreteEnvCfg, + DiscreteMultiDiscreteEnvCfg, + MacCartpoleShowcaseEnv, + MultiDiscreteBoxEnvCfg, + MultiDiscreteDiscreteEnvCfg, + MultiDiscreteMultiDiscreteEnvCfg, + TupleBoxEnvCfg, + TupleDiscreteEnvCfg, + TupleMultiDiscreteEnvCfg, +) diff --git a/source/isaaclab/isaaclab/backends/mac_sim/cart_double_pendulum.py b/source/isaaclab/isaaclab/backends/mac_sim/cart_double_pendulum.py new file mode 100644 index 000000000..34907ae22 --- /dev/null +++ b/source/isaaclab/isaaclab/backends/mac_sim/cart_double_pendulum.py @@ -0,0 +1,388 @@ +# 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 + +"""Mac-native cart-double-pendulum environment with MARL-style dict interfaces.""" + +from __future__ import annotations + +import math +from typing import Any + +import mlx.core as mx + +from isaaclab.backends.runtime import ( + ArticulationCapabilities, + MacSimBackend, + SimBackendContract, + SimCapabilities, + resolve_runtime_selection, + set_runtime_selection, +) +from isaaclab.utils.configclass import configclass + + +@configclass +class MacCartDoublePendulumEnvCfg: + """Configuration aligned with upstream cart-double-pendulum semantics where practical.""" + + num_envs: int = 256 + sim_dt: float = 1.0 / 120.0 + decimation: int = 2 + episode_length_s: float = 5.0 + possible_agents: tuple[str, str] = ("cart", "pendulum") + action_spaces: dict[str, int] = {"cart": 1, "pendulum": 1} + observation_spaces: dict[str, int] = {"cart": 4, "pendulum": 3} + state_space: int = -1 + + max_cart_pos: float = 3.0 + initial_pole_angle_range: tuple[float, float] = (-0.25, 0.25) + initial_pendulum_angle_range: tuple[float, float] = (-0.25, 0.25) + + cart_action_scale: float = 100.0 + pendulum_action_scale: float = 50.0 + + rew_scale_alive: float = 1.0 + rew_scale_terminated: float = -2.0 + rew_scale_cart_pos: float = 0.0 + rew_scale_cart_vel: float = -0.01 + rew_scale_pole_pos: float = -1.0 + rew_scale_pole_vel: float = -0.01 + rew_scale_pendulum_pos: float = -1.0 + rew_scale_pendulum_vel: float = -0.01 + + gravity: float = 9.81 + mass_cart: float = 1.0 + mass_pole: float = 0.1 + mass_pendulum: float = 0.08 + pole_half_length: float = 0.5 + pendulum_half_length: float = 0.45 + cart_force_scale: float = 1.0 + pendulum_torque_scale: float = 1.0 + pendulum_damping: float = 0.05 + + seed: int = 42 + + +def normalize_angle(angle: mx.array) -> mx.array: + return (angle + math.pi) % (2.0 * math.pi) - math.pi + + +def compute_rewards( + cfg: MacCartDoublePendulumEnvCfg, + cart_pos: mx.array, + cart_vel: mx.array, + pole_pos: mx.array, + pole_vel: mx.array, + pendulum_pos: mx.array, + pendulum_vel: mx.array, + reset_terminated: mx.array, +) -> dict[str, mx.array]: + terminated = reset_terminated.astype(mx.float32) + rew_alive = cfg.rew_scale_alive * (1.0 - terminated) + rew_termination = cfg.rew_scale_terminated * terminated + rew_pole_pos = cfg.rew_scale_pole_pos * mx.square(pole_pos) + rew_pendulum_pos = cfg.rew_scale_pendulum_pos * mx.square(pole_pos + pendulum_pos) + rew_cart_vel = cfg.rew_scale_cart_vel * mx.abs(cart_vel) + rew_pole_vel = cfg.rew_scale_pole_vel * mx.abs(pole_vel) + rew_pendulum_vel = cfg.rew_scale_pendulum_vel * mx.abs(pendulum_vel) + rew_cart_pos = cfg.rew_scale_cart_pos * mx.square(cart_pos) + return { + "cart": rew_alive + rew_termination + rew_cart_pos + rew_pole_pos + rew_cart_vel + rew_pole_vel, + "pendulum": rew_alive + rew_termination + rew_pendulum_pos + rew_pendulum_vel, + } + + +class MacCartDoublePendulumSimBackend(MacSimBackend): + """A batched double-pendulum simulator with the same joint-level contract as Isaac Sim tasks.""" + + capabilities = SimCapabilities( + batched_stepping=True, + articulated_rigid_bodies=True, + contacts=False, + proprioceptive_observations=True, + cameras=False, + planners=False, + ) + contract = SimBackendContract( + reset_signature="reset(soft: bool = False) -> None", + step_signature="step(render: bool = True, update_fabric: bool = False) -> None", + articulations=ArticulationCapabilities( + joint_state_io=True, + root_state_io=False, + effort_targets=True, + batched_views=True, + ), + ) + + def __init__(self, cfg: MacCartDoublePendulumEnvCfg): + self.cfg = cfg + self.num_envs = cfg.num_envs + self.cart_pos = mx.zeros((cfg.num_envs,), dtype=mx.float32) + self.cart_vel = mx.zeros((cfg.num_envs,), dtype=mx.float32) + self.pole_angle = mx.zeros((cfg.num_envs,), dtype=mx.float32) + self.pole_vel = mx.zeros((cfg.num_envs,), dtype=mx.float32) + self.pendulum_angle = mx.zeros((cfg.num_envs,), dtype=mx.float32) + self.pendulum_vel = mx.zeros((cfg.num_envs,), dtype=mx.float32) + self._cart_effort_target = mx.zeros((cfg.num_envs,), dtype=mx.float32) + self._pendulum_effort_target = mx.zeros((cfg.num_envs,), dtype=mx.float32) + self.reset() + + @property + def physics_dt(self) -> float: + return self.cfg.sim_dt + + def joint_state(self) -> tuple[mx.array, mx.array]: + return self.get_joint_state(None) + + def reset(self, *, soft: bool = False) -> None: + self.reset_envs(list(range(self.num_envs))) + + def reset_envs(self, env_ids: list[int]) -> None: + if len(env_ids) == 0: + return + ids = mx.array(env_ids) + self.cart_pos[ids] = 0.0 + self.cart_vel[ids] = 0.0 + self.pole_vel[ids] = 0.0 + self.pendulum_vel[ids] = 0.0 + self.pole_angle[ids] = mx.random.uniform( + low=self.cfg.initial_pole_angle_range[0] * math.pi, + high=self.cfg.initial_pole_angle_range[1] * math.pi, + shape=(len(env_ids),), + ) + self.pendulum_angle[ids] = mx.random.uniform( + low=self.cfg.initial_pendulum_angle_range[0] * math.pi, + high=self.cfg.initial_pendulum_angle_range[1] * math.pi, + shape=(len(env_ids),), + ) + self._cart_effort_target[ids] = 0.0 + self._pendulum_effort_target[ids] = 0.0 + + def step(self, *, render: bool = True, update_fabric: bool = False) -> None: + del render, update_fabric + cart_force = self._cart_effort_target * self.cfg.cart_force_scale + pendulum_torque = self._pendulum_effort_target * self.cfg.pendulum_torque_scale + + total_mass = self.cfg.mass_cart + self.cfg.mass_pole + self.cfg.mass_pendulum + costheta = mx.cos(self.pole_angle) + sintheta = mx.sin(self.pole_angle) + polemass_length = self.cfg.mass_pole * self.cfg.pole_half_length + temp = (cart_force + polemass_length * mx.square(self.pole_vel) * sintheta) / total_mass + pole_acc = ( + self.cfg.gravity * sintheta - costheta * temp + ) / ( + self.cfg.pole_half_length + * (4.0 / 3.0 - self.cfg.mass_pole * mx.square(costheta) / total_mass) + ) + cart_acc = temp - polemass_length * pole_acc * costheta / total_mass + + pend_abs_angle = self.pole_angle + self.pendulum_angle + inertia = self.cfg.mass_pendulum * self.cfg.pendulum_half_length * self.cfg.pendulum_half_length + 1e-6 + pendulum_acc = ( + -self.cfg.gravity * mx.sin(pend_abs_angle) / self.cfg.pendulum_half_length + + pendulum_torque / inertia + - self.cfg.pendulum_damping * self.pendulum_vel + + 0.2 * pole_acc + ) + + dt = self.physics_dt + self.cart_pos = self.cart_pos + dt * self.cart_vel + self.cart_vel = self.cart_vel + dt * cart_acc + self.pole_angle = self.pole_angle + dt * self.pole_vel + self.pole_vel = self.pole_vel + dt * pole_acc + self.pendulum_angle = self.pendulum_angle + dt * self.pendulum_vel + self.pendulum_vel = self.pendulum_vel + dt * pendulum_acc + + def get_joint_state(self, articulation: Any) -> tuple[mx.array, mx.array]: + del articulation + joint_pos = mx.stack([self.cart_pos, self.pole_angle, self.pendulum_angle], axis=-1) + joint_vel = mx.stack([self.cart_vel, self.pole_vel, self.pendulum_vel], axis=-1) + return joint_pos, joint_vel + + def set_joint_effort_target( + self, + articulation: Any, + efforts: Any, + *, + joint_ids: Any | None = None, + ) -> None: + del articulation + values = mx.array(efforts, dtype=mx.float32).reshape((-1,)) + if joint_ids is None: + self._cart_effort_target = values + return + joint_ids = list(joint_ids) + if 0 in joint_ids: + self._cart_effort_target = values + return + if 2 in joint_ids: + self._pendulum_effort_target = values + return + raise ValueError(f"Unsupported joint_ids for cart-double-pendulum mac-sim: {joint_ids}") + + def write_joint_state( + self, + articulation: Any, + joint_pos: Any, + joint_vel: Any, + *, + joint_acc: Any | None = None, + env_ids: Any | None = None, + ) -> None: + del articulation, joint_acc + joint_pos = mx.array(joint_pos, dtype=mx.float32) + joint_vel = mx.array(joint_vel, dtype=mx.float32) + if env_ids is None: + env_ids = list(range(self.num_envs)) + ids = mx.array(env_ids) + self.cart_pos[ids] = joint_pos[:, 0] + self.pole_angle[ids] = joint_pos[:, 1] + self.pendulum_angle[ids] = joint_pos[:, 2] + self.cart_vel[ids] = joint_vel[:, 0] + self.pole_vel[ids] = joint_vel[:, 1] + self.pendulum_vel[ids] = joint_vel[:, 2] + + def write_root_pose(self, articulation: Any, root_pose: Any, *, env_ids: Any | None = None) -> None: + del articulation, root_pose, env_ids + + def write_root_velocity( + self, + articulation: Any, + root_velocity: Any, + *, + env_ids: Any | None = None, + ) -> None: + del articulation, root_velocity, env_ids + + def state_dict(self) -> dict[str, Any]: + return { + "backend": self.name, + "num_envs": self.num_envs, + "cart_pos": self.cart_pos.tolist(), + "cart_vel": self.cart_vel.tolist(), + "pole_angle": self.pole_angle.tolist(), + "pole_vel": self.pole_vel.tolist(), + "pendulum_angle": self.pendulum_angle.tolist(), + "pendulum_vel": self.pendulum_vel.tolist(), + } + + +class MacCartDoublePendulumEnv: + """A vectorized MARL cart-double-pendulum env with per-agent dict interfaces.""" + + def __init__(self, cfg: MacCartDoublePendulumEnvCfg | None = None): + self.cfg = cfg or MacCartDoublePendulumEnvCfg() + mx.random.seed(self.cfg.seed) + runtime = set_runtime_selection(resolve_runtime_selection("mlx", "mac-sim", "cpu")) + self.runtime = runtime + self.device = runtime.device + self.num_envs = self.cfg.num_envs + self.step_dt = self.cfg.sim_dt * self.cfg.decimation + self.max_episode_length = math.ceil(self.cfg.episode_length_s / self.step_dt) + self.possible_agents = tuple(self.cfg.possible_agents) + self.actions: dict[str, mx.array] = { + "cart": mx.zeros((self.num_envs, 1), dtype=mx.float32), + "pendulum": mx.zeros((self.num_envs, 1), dtype=mx.float32), + } + + self.sim_backend = MacCartDoublePendulumSimBackend(self.cfg) + self.reset_terminated = mx.zeros((self.num_envs,), dtype=mx.bool_) + self.reset_time_outs = mx.zeros((self.num_envs,), dtype=mx.bool_) + self.reset_buf = mx.zeros((self.num_envs,), dtype=mx.bool_) + self.episode_length_buf = mx.zeros((self.num_envs,), dtype=mx.int32) + self.episode_return_cart = mx.zeros((self.num_envs,), dtype=mx.float32) + self.episode_return_pendulum = mx.zeros((self.num_envs,), dtype=mx.float32) + self.obs_buf = { + "cart": mx.zeros((self.num_envs, 4), dtype=mx.float32), + "pendulum": mx.zeros((self.num_envs, 3), dtype=mx.float32), + } + self.reset() + + def reset(self) -> tuple[dict[str, mx.array], dict[str, Any]]: + env_ids = list(range(self.num_envs)) + self._reset_idx(env_ids) + self.obs_buf = self._get_observations() + return self.obs_buf, {} + + def step( + self, actions: dict[str, Any] + ) -> tuple[dict[str, mx.array], dict[str, mx.array], dict[str, mx.array], dict[str, mx.array], dict[str, Any]]: + self._pre_physics_step(actions) + for _ in range(self.cfg.decimation): + self._apply_action() + self.sim_backend.step(render=False) + + self.episode_length_buf = self.episode_length_buf + 1 + joint_pos, joint_vel = self.sim_backend.joint_state() + cart_pos = joint_pos[:, 0] + cart_vel = joint_vel[:, 0] + pole_pos = normalize_angle(joint_pos[:, 1]) + pole_vel = joint_vel[:, 1] + pendulum_pos = normalize_angle(joint_pos[:, 2]) + pendulum_vel = joint_vel[:, 2] + + self.reset_terminated = (mx.abs(cart_pos) > self.cfg.max_cart_pos) | (mx.abs(pole_pos) > math.pi / 2.0) + self.reset_time_outs = self.episode_length_buf >= (self.max_episode_length - 1) + self.reset_buf = self.reset_terminated | self.reset_time_outs + + rewards = compute_rewards( + self.cfg, + cart_pos, + cart_vel, + pole_pos, + pole_vel, + pendulum_pos, + pendulum_vel, + self.reset_terminated, + ) + self.episode_return_cart = self.episode_return_cart + rewards["cart"] + self.episode_return_pendulum = self.episode_return_pendulum + rewards["pendulum"] + + reset_ids = [idx for idx, flag in enumerate(self.reset_buf.tolist()) if flag] + extras: dict[str, Any] = {} + if reset_ids: + extras = { + "completed_returns": { + "cart": [float(self.episode_return_cart[idx].item()) for idx in reset_ids], + "pendulum": [float(self.episode_return_pendulum[idx].item()) for idx in reset_ids], + }, + "completed_lengths": [int(self.episode_length_buf[idx].item()) for idx in reset_ids], + } + self._reset_idx(reset_ids) + + self.obs_buf = self._get_observations() + terminated = {agent: self.reset_terminated for agent in self.possible_agents} + truncated = {agent: self.reset_time_outs for agent in self.possible_agents} + return self.obs_buf, rewards, terminated, truncated, extras + + def _pre_physics_step(self, actions: dict[str, Any]) -> None: + cart = mx.array(actions["cart"], dtype=mx.float32).reshape((self.num_envs, 1)) + pendulum = mx.array(actions["pendulum"], dtype=mx.float32).reshape((self.num_envs, 1)) + self.actions["cart"] = mx.clip(cart, -1.0, 1.0) + self.actions["pendulum"] = mx.clip(pendulum, -1.0, 1.0) + + def _apply_action(self) -> None: + self.sim_backend.set_joint_effort_target( + None, self.actions["cart"][:, 0] * self.cfg.cart_action_scale, joint_ids=[0] + ) + self.sim_backend.set_joint_effort_target( + None, self.actions["pendulum"][:, 0] * self.cfg.pendulum_action_scale, joint_ids=[2] + ) + + def _get_observations(self) -> dict[str, mx.array]: + joint_pos, joint_vel = self.sim_backend.joint_state() + pole_pos = normalize_angle(joint_pos[:, 1]) + pendulum_pos = normalize_angle(joint_pos[:, 2]) + return { + "cart": mx.stack([joint_pos[:, 0], joint_vel[:, 0], pole_pos, joint_vel[:, 1]], axis=-1), + "pendulum": mx.stack([pole_pos + pendulum_pos, pendulum_pos, joint_vel[:, 2]], axis=-1), + } + + def _reset_idx(self, env_ids: list[int]) -> None: + self.sim_backend.reset_envs(env_ids) + ids = mx.array(env_ids) + self.episode_length_buf[ids] = 0 + self.episode_return_cart[ids] = 0.0 + self.episode_return_pendulum[ids] = 0.0 diff --git a/source/isaaclab/isaaclab/backends/mac_sim/cartpole.py b/source/isaaclab/isaaclab/backends/mac_sim/cartpole.py index 6c6884df0..cac268ac0 100644 --- a/source/isaaclab/isaaclab/backends/mac_sim/cartpole.py +++ b/source/isaaclab/isaaclab/backends/mac_sim/cartpole.py @@ -38,9 +38,9 @@ class MacCartpoleEnvCfg: episode_length_s: float = 5.0 action_scale: float = 100.0 - action_space: int = 1 - observation_space: int = 4 - state_space: int = 0 + action_space: Any = 1 + observation_space: Any = 4 + state_space: Any = 0 max_cart_pos: float = 3.0 initial_pole_angle_range: tuple[float, float] = (-0.25 * math.pi, 0.25 * math.pi) @@ -77,6 +77,7 @@ class MacCartpoleTrainCfg: entropy_coef: float = 0.01 checkpoint_path: str = "logs/mlx/cartpole_policy.npz" eval_interval: int = 10 + resume_from: str | None = None class MacCartpoleSimBackend(MacSimBackend): @@ -231,6 +232,8 @@ def __init__(self, cfg: MacCartpoleEnvCfg | None = None): self.physics_dt = self.cfg.sim_dt self.step_dt = self.cfg.sim_dt * self.cfg.decimation self.max_episode_length = math.ceil(self.cfg.episode_length_s / self.step_dt) + self.single_action_space = self.cfg.action_space + self.single_observation_space = {"policy": self.cfg.observation_space} self.sim_backend = MacCartpoleSimBackend(self.cfg) self.actions = mx.zeros((self.num_envs, 1), dtype=mx.float32) @@ -251,12 +254,10 @@ def reset(self) -> tuple[dict[str, mx.array], dict[str, Any]]: return self.obs_buf, {} def step(self, action: Any) -> tuple[dict[str, mx.array], mx.array, mx.array, mx.array, dict[str, Any]]: - action = mx.array(action, dtype=mx.float32).reshape((self.num_envs, self.cfg.action_space)) - action = mx.clip(action, -1.0, 1.0) - self.actions = self.cfg.action_scale * action + self._pre_physics_step(action) for _ in range(self.cfg.decimation): - self.sim_backend.set_joint_effort_target(None, self.actions[:, 0], joint_ids=[0]) + self._apply_action() self.sim_backend.step(render=False) self.episode_length_buf = self.episode_length_buf + 1 @@ -300,13 +301,24 @@ def step(self, action: Any) -> tuple[dict[str, mx.array], mx.array, mx.array, mx self.obs_buf = self._get_observations() return self.obs_buf, self.reward_buf, self.reset_terminated, self.reset_time_outs, extras + def _pre_physics_step(self, action: Any) -> None: + action = mx.array(action, dtype=mx.float32).reshape((self.num_envs, 1)) + self.actions = mx.clip(action, -1.0, 1.0) + + def _apply_action(self) -> None: + target = self.cfg.action_scale * self.actions[:, 0] + self.sim_backend.set_joint_effort_target(None, target, joint_ids=[0]) + def _reset_idx(self, env_ids: list[int]) -> None: self.sim_backend.reset_envs(env_ids) self.episode_length_buf[mx.array(env_ids)] = 0 self.episode_return_buf[mx.array(env_ids)] = 0.0 + def _joint_state(self) -> tuple[mx.array, mx.array]: + return self.sim_backend.joint_state() + def _get_observations(self) -> dict[str, mx.array]: - joint_pos, joint_vel = self.sim_backend.joint_state() + joint_pos, joint_vel = self._joint_state() obs = mx.stack( [ joint_pos[:, 1], @@ -430,6 +442,13 @@ def train_cartpole_policy(cfg: MacCartpoleTrainCfg) -> dict[str, Any]: model = MacCartpolePolicy(obs_dim=cfg.env.observation_space, hidden_dim=cfg.hidden_dim) optimizer = optim.Adam(learning_rate=cfg.learning_rate) loss_and_grad = nn.value_and_grad(model, _ppo_loss) + resumed_from = None + if cfg.resume_from: + resume_path = Path(cfg.resume_from) + if not resume_path.exists(): + raise FileNotFoundError(f"Checkpoint to resume from does not exist: {resume_path}") + model.load_weights(str(resume_path)) + resumed_from = str(resume_path) mx.eval(model.parameters()) obs = env.reset()[0]["policy"] @@ -518,6 +537,7 @@ def train_cartpole_policy(cfg: MacCartpoleTrainCfg) -> dict[str, Any]: return { "checkpoint_path": str(checkpoint_path), "metadata_path": str(metadata_path), + "resumed_from": resumed_from, "train_cfg": asdict(cfg), "completed_episodes": len(completed_returns), "mean_recent_return": sum(completed_returns[-10:]) / max(1, min(len(completed_returns), 10)), diff --git a/source/isaaclab/isaaclab/backends/mac_sim/quadcopter.py b/source/isaaclab/isaaclab/backends/mac_sim/quadcopter.py new file mode 100644 index 000000000..b9d3ed3f1 --- /dev/null +++ b/source/isaaclab/isaaclab/backends/mac_sim/quadcopter.py @@ -0,0 +1,375 @@ +# 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 + +"""Mac-native quadcopter environment with root-state dynamics.""" + +from __future__ import annotations + +import math +from typing import Any + +import mlx.core as mx + +from isaaclab.backends.runtime import ( + ArticulationCapabilities, + MacSimBackend, + SimBackendContract, + SimCapabilities, + resolve_runtime_selection, + set_runtime_selection, +) +from isaaclab.utils.configclass import configclass + + +@configclass +class MacQuadcopterEnvCfg: + """Configuration aligned with the upstream quadcopter task where practical.""" + + num_envs: int = 256 + sim_dt: float = 1.0 / 100.0 + decimation: int = 2 + episode_length_s: float = 10.0 + action_space: int = 4 + observation_space: int = 12 + state_space: int = 0 + + env_spacing: float = 2.5 + mass: float = 0.032 + gravity: float = 9.81 + thrust_to_weight: float = 1.9 + moment_scale: float = 0.01 + angular_damping: float = 0.08 + linear_damping_xy: float = 0.05 + linear_damping_z: float = 0.03 + lateral_accel_scale: float = 2.5 + + min_height: float = 0.1 + max_height: float = 2.0 + + lin_vel_reward_scale: float = -0.05 + ang_vel_reward_scale: float = -0.01 + distance_to_goal_reward_scale: float = 15.0 + + seed: int = 42 + + +def _grid_origins(num_envs: int, spacing: float) -> mx.array: + side = int(math.ceil(math.sqrt(num_envs))) + grid_x = mx.arange(side, dtype=mx.float32) + grid_y = mx.arange(side, dtype=mx.float32) + mesh_x = mx.repeat(grid_x.reshape((1, -1)), side, axis=0).reshape((-1,)) + mesh_y = mx.repeat(grid_y.reshape((-1, 1)), side, axis=1).reshape((-1,)) + centered_x = (mesh_x[:num_envs] - (side - 1) / 2.0) * spacing + centered_y = (mesh_y[:num_envs] - (side - 1) / 2.0) * spacing + zeros = mx.zeros((num_envs,), dtype=mx.float32) + return mx.stack([centered_x, centered_y, zeros], axis=-1) + + +def _quat_conjugate(quat: mx.array) -> mx.array: + return mx.concatenate([-quat[:, :3], quat[:, 3:4]], axis=-1) + + +def _quat_normalize(quat: mx.array, eps: float = 1e-8) -> mx.array: + norm = mx.linalg.norm(quat, axis=1, keepdims=True) + return quat / mx.maximum(norm, eps) + + +def _quat_multiply(lhs: mx.array, rhs: mx.array) -> mx.array: + x1, y1, z1, w1 = lhs[:, 0], lhs[:, 1], lhs[:, 2], lhs[:, 3] + x2, y2, z2, w2 = rhs[:, 0], rhs[:, 1], rhs[:, 2], rhs[:, 3] + return mx.stack( + [ + w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, + w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2, + w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2, + w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, + ], + axis=-1, + ) + + +def _quat_rotate(quat: mx.array, vector: mx.array) -> mx.array: + quat_xyz = quat[:, :3] + quat_w = quat[:, 3:4] + cross_1 = mx.linalg.cross(quat_xyz, vector, axis=1) + cross_2 = mx.linalg.cross(quat_xyz, cross_1, axis=1) + return vector + 2.0 * (quat_w * cross_1 + cross_2) + + +def _quat_from_angular_velocity(omega_body: mx.array, dt: float) -> mx.array: + delta = omega_body * dt + angle = mx.linalg.norm(delta, axis=1, keepdims=True) + half_angle = 0.5 * angle + axis = delta / mx.maximum(angle, 1e-8) + sin_half = mx.sin(half_angle) + cos_half = mx.cos(half_angle) + return _quat_normalize(mx.concatenate([axis * sin_half, cos_half], axis=-1)) + + +class MacQuadcopterSimBackend(MacSimBackend): + """A lightweight root-state quadcopter simulator for MLX/mac-sim.""" + + capabilities = SimCapabilities( + batched_stepping=True, + articulated_rigid_bodies=True, + contacts=False, + proprioceptive_observations=True, + cameras=False, + planners=False, + ) + contract = SimBackendContract( + reset_signature="reset(soft: bool = False) -> None", + step_signature="step(render: bool = True, update_fabric: bool = False) -> None", + articulations=ArticulationCapabilities( + joint_state_io=True, + root_state_io=True, + effort_targets=True, + batched_views=True, + ), + ) + + def __init__(self, cfg: MacQuadcopterEnvCfg): + self.cfg = cfg + self.num_envs = cfg.num_envs + self.env_origins = _grid_origins(cfg.num_envs, cfg.env_spacing) + self.root_pos_w = mx.zeros((cfg.num_envs, 3), dtype=mx.float32) + self.root_lin_vel_b = mx.zeros((cfg.num_envs, 3), dtype=mx.float32) + self.root_ang_vel_b = mx.zeros((cfg.num_envs, 3), dtype=mx.float32) + self.root_quat_w = mx.tile(mx.array([[0.0, 0.0, 0.0, 1.0]], dtype=mx.float32), (cfg.num_envs, 1)) + self.projected_gravity_b = mx.tile(mx.array([[0.0, 0.0, -1.0]], dtype=mx.float32), (cfg.num_envs, 1)) + self._thrust = mx.zeros((cfg.num_envs,), dtype=mx.float32) + self._moments = mx.zeros((cfg.num_envs, 3), dtype=mx.float32) + self.reset() + + @property + def physics_dt(self) -> float: + return self.cfg.sim_dt + + def reset(self, *, soft: bool = False) -> None: + self.reset_envs(list(range(self.num_envs))) + + def reset_envs(self, env_ids: list[int]) -> None: + if len(env_ids) == 0: + return + ids = mx.array(env_ids) + self.root_pos_w[ids] = self.env_origins[ids] + mx.array([0.0, 0.0, 1.0], dtype=mx.float32) + self.root_lin_vel_b[ids] = 0.0 + self.root_ang_vel_b[ids] = 0.0 + self.root_quat_w[ids] = mx.array([0.0, 0.0, 0.0, 1.0], dtype=mx.float32) + self.projected_gravity_b[ids] = mx.array([0.0, 0.0, -1.0], dtype=mx.float32) + self._thrust[ids] = self.cfg.mass * self.cfg.gravity + self._moments[ids] = 0.0 + + def step(self, *, render: bool = True, update_fabric: bool = False) -> None: + del render, update_fabric + dt = self.physics_dt + thrust_accel_z = self._thrust / self.cfg.mass - self.cfg.gravity + lateral_accel = self.cfg.lateral_accel_scale * self._moments[:, :2] + + accel = mx.stack([lateral_accel[:, 0], lateral_accel[:, 1], thrust_accel_z], axis=-1) + damping = mx.array( + [self.cfg.linear_damping_xy, self.cfg.linear_damping_xy, self.cfg.linear_damping_z], dtype=mx.float32 + ) + self.root_lin_vel_b = self.root_lin_vel_b + dt * (accel - damping * self.root_lin_vel_b) + self.root_pos_w = self.root_pos_w + dt * self.root_lin_vel_b + + ang_acc = self._moments - self.cfg.angular_damping * self.root_ang_vel_b + self.root_ang_vel_b = self.root_ang_vel_b + dt * ang_acc + delta_quat = _quat_from_angular_velocity(self.root_ang_vel_b, dt) + self.root_quat_w = _quat_normalize(_quat_multiply(self.root_quat_w, delta_quat)) + gravity_w = mx.tile(mx.array([[0.0, 0.0, -1.0]], dtype=mx.float32), (self.num_envs, 1)) + self.projected_gravity_b = _quat_rotate(_quat_conjugate(self.root_quat_w), gravity_w) + + def set_thrust_and_moment(self, thrust: Any, moment: Any) -> None: + self._thrust = mx.array(thrust, dtype=mx.float32).reshape((self.num_envs,)) + self._moments = mx.array(moment, dtype=mx.float32).reshape((self.num_envs, 3)) + + def get_joint_state(self, articulation: Any) -> tuple[mx.array, mx.array]: + del articulation + joint_pos = mx.zeros((self.num_envs, 4), dtype=mx.float32) + joint_vel = mx.zeros((self.num_envs, 4), dtype=mx.float32) + return joint_pos, joint_vel + + def set_joint_effort_target( + self, + articulation: Any, + efforts: Any, + *, + joint_ids: Any | None = None, + ) -> None: + del articulation, efforts, joint_ids + + def write_joint_state( + self, + articulation: Any, + joint_pos: Any, + joint_vel: Any, + *, + joint_acc: Any | None = None, + env_ids: Any | None = None, + ) -> None: + del articulation, joint_pos, joint_vel, joint_acc, env_ids + + def write_root_pose(self, articulation: Any, root_pose: Any, *, env_ids: Any | None = None) -> None: + del articulation + pose = mx.array(root_pose, dtype=mx.float32) + if env_ids is None: + env_ids = list(range(self.num_envs)) + ids = mx.array(env_ids) + self.root_pos_w[ids] = pose[:, :3] + if pose.shape[1] >= 7: + self.root_quat_w[ids] = pose[:, 3:7] + + def write_root_velocity( + self, + articulation: Any, + root_velocity: Any, + *, + env_ids: Any | None = None, + ) -> None: + del articulation + velocity = mx.array(root_velocity, dtype=mx.float32) + if env_ids is None: + env_ids = list(range(self.num_envs)) + ids = mx.array(env_ids) + self.root_lin_vel_b[ids] = velocity[:, :3] + if velocity.shape[1] >= 6: + self.root_ang_vel_b[ids] = velocity[:, 3:6] + + def state_dict(self) -> dict[str, Any]: + return { + "backend": self.name, + "num_envs": self.num_envs, + "root_pos_w": self.root_pos_w.tolist(), + "root_lin_vel_b": self.root_lin_vel_b.tolist(), + "root_ang_vel_b": self.root_ang_vel_b.tolist(), + } + + +class MacQuadcopterEnv: + """A vectorized quadcopter env that mirrors upstream DirectRL task structure.""" + + def __init__(self, cfg: MacQuadcopterEnvCfg | None = None): + self.cfg = cfg or MacQuadcopterEnvCfg() + mx.random.seed(self.cfg.seed) + runtime = set_runtime_selection(resolve_runtime_selection("mlx", "mac-sim", "cpu")) + self.runtime = runtime + self.device = runtime.device + self.num_envs = self.cfg.num_envs + self.step_dt = self.cfg.sim_dt * self.cfg.decimation + self.max_episode_length = math.ceil(self.cfg.episode_length_s / self.step_dt) + + self.sim_backend = MacQuadcopterSimBackend(self.cfg) + self._actions = mx.zeros((self.num_envs, 4), dtype=mx.float32) + self._thrust = mx.zeros((self.num_envs,), dtype=mx.float32) + self._moment = mx.zeros((self.num_envs, 3), dtype=mx.float32) + self._desired_pos_w = mx.zeros((self.num_envs, 3), dtype=mx.float32) + + self.reset_terminated = mx.zeros((self.num_envs,), dtype=mx.bool_) + self.reset_time_outs = mx.zeros((self.num_envs,), dtype=mx.bool_) + self.reset_buf = mx.zeros((self.num_envs,), dtype=mx.bool_) + self.episode_length_buf = mx.zeros((self.num_envs,), dtype=mx.int32) + self._episode_sums = { + "lin_vel": mx.zeros((self.num_envs,), dtype=mx.float32), + "ang_vel": mx.zeros((self.num_envs,), dtype=mx.float32), + "distance_to_goal": mx.zeros((self.num_envs,), dtype=mx.float32), + } + self.obs_buf = {"policy": mx.zeros((self.num_envs, 12), dtype=mx.float32)} + self.reset() + + def reset(self) -> tuple[dict[str, mx.array], dict[str, Any]]: + env_ids = list(range(self.num_envs)) + self._reset_idx(env_ids) + self.obs_buf = self._get_observations() + return self.obs_buf, {} + + def step(self, actions: Any) -> tuple[dict[str, mx.array], mx.array, mx.array, mx.array, dict[str, Any]]: + self._pre_physics_step(actions) + for _ in range(self.cfg.decimation): + self._apply_action() + self.sim_backend.step(render=False) + + self.episode_length_buf = self.episode_length_buf + 1 + rewards = self._get_rewards() + self.reset_terminated, self.reset_time_outs = self._get_dones() + self.reset_buf = self.reset_terminated | self.reset_time_outs + + reset_ids = [idx for idx, flag in enumerate(self.reset_buf.tolist()) if flag] + extras: dict[str, Any] = {} + if reset_ids: + final_distance = mx.linalg.norm( + self._desired_pos_w[mx.array(reset_ids)] - self.sim_backend.root_pos_w[mx.array(reset_ids)], axis=1 + ) + extras = { + "completed_lengths": [int(self.episode_length_buf[idx].item()) for idx in reset_ids], + "final_distance_to_goal": float(mx.mean(final_distance).item()), + } + self._reset_idx(reset_ids) + + self.obs_buf = self._get_observations() + return self.obs_buf, rewards, self.reset_terminated, self.reset_time_outs, extras + + def _pre_physics_step(self, actions: Any) -> None: + self._actions = mx.clip(mx.array(actions, dtype=mx.float32).reshape((self.num_envs, 4)), -1.0, 1.0) + hover_thrust = self.cfg.thrust_to_weight * self.cfg.mass * self.cfg.gravity + self._thrust = hover_thrust * (self._actions[:, 0] + 1.0) / 2.0 + self._moment = self.cfg.moment_scale * self._actions[:, 1:] + + def _apply_action(self) -> None: + self.sim_backend.set_thrust_and_moment(self._thrust, self._moment) + + def _get_observations(self) -> dict[str, mx.array]: + desired_pos_w = self._desired_pos_w - self.sim_backend.root_pos_w + desired_pos_b = _quat_rotate(_quat_conjugate(self.sim_backend.root_quat_w), desired_pos_w) + obs = mx.concatenate( + [ + self.sim_backend.root_lin_vel_b, + self.sim_backend.root_ang_vel_b, + self.sim_backend.projected_gravity_b, + desired_pos_b, + ], + axis=-1, + ) + return {"policy": obs} + + def _get_rewards(self) -> mx.array: + lin_vel = mx.sum(mx.square(self.sim_backend.root_lin_vel_b), axis=1) + ang_vel = mx.sum(mx.square(self.sim_backend.root_ang_vel_b), axis=1) + distance_to_goal = mx.linalg.norm(self._desired_pos_w - self.sim_backend.root_pos_w, axis=1) + distance_to_goal_mapped = 1.0 - mx.tanh(distance_to_goal / 0.8) + + rewards = { + "lin_vel": lin_vel * self.cfg.lin_vel_reward_scale * self.step_dt, + "ang_vel": ang_vel * self.cfg.ang_vel_reward_scale * self.step_dt, + "distance_to_goal": distance_to_goal_mapped * self.cfg.distance_to_goal_reward_scale * self.step_dt, + } + for key, value in rewards.items(): + self._episode_sums[key] = self._episode_sums[key] + value + return rewards["lin_vel"] + rewards["ang_vel"] + rewards["distance_to_goal"] + + def _get_dones(self) -> tuple[mx.array, mx.array]: + time_out = self.episode_length_buf >= (self.max_episode_length - 1) + died = ( + (self.sim_backend.root_pos_w[:, 2] < self.cfg.min_height) + | (self.sim_backend.root_pos_w[:, 2] > self.cfg.max_height) + ) + return died, time_out + + def _reset_idx(self, env_ids: list[int]) -> None: + if not env_ids: + return + self.sim_backend.reset_envs(env_ids) + ids = mx.array(env_ids) + + self._actions[ids] = 0.0 + self.episode_length_buf[ids] = 0 + for key in self._episode_sums: + self._episode_sums[key][ids] = 0.0 + + desired = mx.zeros((len(env_ids), 3), dtype=mx.float32) + desired[:, :2] = mx.random.uniform(low=-2.0, high=2.0, shape=(len(env_ids), 2)) + desired[:, :2] = desired[:, :2] + self.sim_backend.env_origins[ids, :2] + desired[:, 2] = mx.random.uniform(low=0.5, high=1.5, shape=(len(env_ids),)) + self._desired_pos_w[ids] = desired diff --git a/source/isaaclab/isaaclab/backends/mac_sim/showcase.py b/source/isaaclab/isaaclab/backends/mac_sim/showcase.py new file mode 100644 index 000000000..9fa87e424 --- /dev/null +++ b/source/isaaclab/isaaclab/backends/mac_sim/showcase.py @@ -0,0 +1,231 @@ +# 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 + +"""Cartpole showcase environments for exercising non-Box action and observation spaces on macOS.""" + +from __future__ import annotations + +from typing import Any + +import gymnasium as gym +from gymnasium import spaces +import mlx.core as mx + +from isaaclab.backends.mac_sim.cartpole import MacCartpoleEnv, MacCartpoleEnvCfg +from isaaclab.utils.configclass import configclass + + +def _box_obs() -> spaces.Box: + return spaces.Box(low=float("-inf"), high=float("inf"), shape=(4,)) + + +def _joint_box() -> spaces.Box: + return spaces.Box(low=float("-inf"), high=float("inf"), shape=(2,)) + + +def _box_action() -> spaces.Box: + return spaces.Box(low=-1.0, high=1.0, shape=(1,)) + + +def _discrete_obs() -> spaces.Discrete: + return spaces.Discrete(16) + + +def _discrete_action() -> spaces.Discrete: + return spaces.Discrete(3) + + +def _multidiscrete_obs() -> spaces.MultiDiscrete: + return spaces.MultiDiscrete([2, 2, 2, 2]) + + +def _multidiscrete_action() -> spaces.MultiDiscrete: + return spaces.MultiDiscrete([3, 2]) + + +def _dict_obs() -> spaces.Dict: + return spaces.Dict({"joint-positions": _joint_box(), "joint-velocities": _joint_box()}) + + +def _tuple_obs() -> spaces.Tuple: + return spaces.Tuple((_joint_box(), _joint_box())) + + +@configclass +class BoxBoxEnvCfg(MacCartpoleEnvCfg): + observation_space: Any = _box_obs() + action_space: Any = _box_action() + + +@configclass +class BoxDiscreteEnvCfg(MacCartpoleEnvCfg): + observation_space: Any = _box_obs() + action_space: Any = _discrete_action() + + +@configclass +class BoxMultiDiscreteEnvCfg(MacCartpoleEnvCfg): + observation_space: Any = _box_obs() + action_space: Any = _multidiscrete_action() + + +@configclass +class DiscreteBoxEnvCfg(MacCartpoleEnvCfg): + observation_space: Any = _discrete_obs() + action_space: Any = _box_action() + + +@configclass +class DiscreteDiscreteEnvCfg(MacCartpoleEnvCfg): + observation_space: Any = _discrete_obs() + action_space: Any = _discrete_action() + + +@configclass +class DiscreteMultiDiscreteEnvCfg(MacCartpoleEnvCfg): + observation_space: Any = _discrete_obs() + action_space: Any = _multidiscrete_action() + + +@configclass +class MultiDiscreteBoxEnvCfg(MacCartpoleEnvCfg): + observation_space: Any = _multidiscrete_obs() + action_space: Any = _box_action() + + +@configclass +class MultiDiscreteDiscreteEnvCfg(MacCartpoleEnvCfg): + observation_space: Any = _multidiscrete_obs() + action_space: Any = _discrete_action() + + +@configclass +class MultiDiscreteMultiDiscreteEnvCfg(MacCartpoleEnvCfg): + observation_space: Any = _multidiscrete_obs() + action_space: Any = _multidiscrete_action() + + +@configclass +class DictBoxEnvCfg(MacCartpoleEnvCfg): + observation_space: Any = _dict_obs() + action_space: Any = _box_action() + + +@configclass +class DictDiscreteEnvCfg(MacCartpoleEnvCfg): + observation_space: Any = _dict_obs() + action_space: Any = _discrete_action() + + +@configclass +class DictMultiDiscreteEnvCfg(MacCartpoleEnvCfg): + observation_space: Any = _dict_obs() + action_space: Any = _multidiscrete_action() + + +@configclass +class TupleBoxEnvCfg(MacCartpoleEnvCfg): + observation_space: Any = _tuple_obs() + action_space: Any = _box_action() + + +@configclass +class TupleDiscreteEnvCfg(MacCartpoleEnvCfg): + observation_space: Any = _tuple_obs() + action_space: Any = _discrete_action() + + +@configclass +class TupleMultiDiscreteEnvCfg(MacCartpoleEnvCfg): + observation_space: Any = _tuple_obs() + action_space: Any = _multidiscrete_action() + + +SHOWCASE_CFGS = ( + BoxBoxEnvCfg, + BoxDiscreteEnvCfg, + BoxMultiDiscreteEnvCfg, + DiscreteBoxEnvCfg, + DiscreteDiscreteEnvCfg, + DiscreteMultiDiscreteEnvCfg, + MultiDiscreteBoxEnvCfg, + MultiDiscreteDiscreteEnvCfg, + MultiDiscreteMultiDiscreteEnvCfg, + DictBoxEnvCfg, + DictDiscreteEnvCfg, + DictMultiDiscreteEnvCfg, + TupleBoxEnvCfg, + TupleDiscreteEnvCfg, + TupleMultiDiscreteEnvCfg, +) + + +def _sign_bits(joint_pos: mx.array, joint_vel: mx.array) -> mx.array: + return mx.stack( + [ + (joint_pos[:, 1] >= 0).astype(mx.int32), + (joint_pos[:, 0] >= 0).astype(mx.int32), + (joint_vel[:, 1] >= 0).astype(mx.int32), + (joint_vel[:, 0] >= 0).astype(mx.int32), + ], + axis=-1, + ) + + +class MacCartpoleShowcaseEnv(MacCartpoleEnv): + """Mac-native cartpole showcase mirroring the upstream action/observation space variations.""" + + cfg: MacCartpoleEnvCfg + + def _pre_physics_step(self, action: Any) -> None: + if isinstance(self.single_action_space, gym.spaces.Box): + actions = mx.array(action, dtype=mx.float32).reshape((self.num_envs, 1)) + self.actions = mx.clip(actions, -1.0, 1.0) + return + if isinstance(self.single_action_space, gym.spaces.Discrete): + self.actions = mx.array([int(item) for item in action], dtype=mx.int32).reshape((self.num_envs,)) + return + if isinstance(self.single_action_space, gym.spaces.MultiDiscrete): + normalized = [[int(value) for value in item] for item in action] + self.actions = mx.array(normalized, dtype=mx.int32).reshape((self.num_envs, 2)) + return + raise NotImplementedError(f"Action space {type(self.single_action_space)} not implemented") + + def _apply_action(self) -> None: + if isinstance(self.single_action_space, gym.spaces.Box): + target = self.cfg.action_scale * self.actions[:, 0] + elif isinstance(self.single_action_space, gym.spaces.Discrete): + target = mx.zeros((self.num_envs,), dtype=mx.float32) + target = mx.where(self.actions == 1, -self.cfg.action_scale, target) + target = mx.where(self.actions == 2, self.cfg.action_scale, target) + elif isinstance(self.single_action_space, gym.spaces.MultiDiscrete): + magnitude = mx.zeros((self.num_envs,), dtype=mx.float32) + magnitude = mx.where(self.actions[:, 0] == 1, self.cfg.action_scale / 2.0, magnitude) + magnitude = mx.where(self.actions[:, 0] == 2, self.cfg.action_scale, magnitude) + target = mx.where(self.actions[:, 1] == 0, -magnitude, magnitude) + else: + raise NotImplementedError(f"Action space {type(self.single_action_space)} not implemented") + + self.sim_backend.set_joint_effort_target(None, target, joint_ids=[0]) + + def _get_observations(self) -> dict[str, Any]: + joint_pos, joint_vel = self._joint_state() + policy_space = self.single_observation_space["policy"] + + if isinstance(policy_space, gym.spaces.Box): + obs = mx.stack([joint_pos[:, 1], joint_vel[:, 1], joint_pos[:, 0], joint_vel[:, 0]], axis=-1) + elif isinstance(policy_space, gym.spaces.Discrete): + bits = _sign_bits(joint_pos, joint_vel) + obs = bits[:, 0] * 8 + bits[:, 1] * 4 + bits[:, 2] * 2 + bits[:, 3] + elif isinstance(policy_space, gym.spaces.MultiDiscrete): + obs = _sign_bits(joint_pos, joint_vel) + elif isinstance(policy_space, gym.spaces.Tuple): + obs = (joint_pos, joint_vel) + elif isinstance(policy_space, gym.spaces.Dict): + obs = {"joint-positions": joint_pos, "joint-velocities": joint_vel} + else: + raise NotImplementedError(f"Observation space {type(policy_space)} not implemented") + + return {"policy": obs} diff --git a/source/isaaclab/isaaclab/backends/runtime.py b/source/isaaclab/isaaclab/backends/runtime.py index bfe84fc4d..b9b4b294b 100644 --- a/source/isaaclab/isaaclab/backends/runtime.py +++ b/source/isaaclab/isaaclab/backends/runtime.py @@ -10,6 +10,7 @@ import builtins import importlib.util import os +import pickle import platform from abc import ABC, abstractmethod from dataclasses import dataclass @@ -18,9 +19,13 @@ ComputeBackendName = Literal["torch-cuda", "mlx"] SimBackendName = Literal["isaacsim", "mac-sim"] +KernelBackendName = Literal["warp", "metal", "cpu"] +SensorBackendName = Literal["isaacsim-sensors", "mac-sensors", "cpu"] +PlannerBackendName = Literal["isaacsim-planners", "mac-planners", "none"] ENV_COMPUTE_BACKEND = "ISAACLAB_COMPUTE_BACKEND" ENV_SIM_BACKEND = "ISAACLAB_SIM_BACKEND" +ENV_KERNEL_BACKEND = "ISAACLAB_KERNEL_BACKEND" class UnsupportedBackendError(RuntimeError): @@ -53,12 +58,47 @@ class SimCapabilities: planners: bool = False +@dataclass(frozen=True) +class KernelCapabilities: + """Capabilities exposed by a kernel backend.""" + + custom_kernels: bool = False + raycast: bool = False + fabric_transforms: bool = False + mesh_queries: bool = False + cpu_fallback: bool = False + + +@dataclass(frozen=True) +class SensorCapabilities: + """Capabilities exposed by a sensor backend.""" + + proprioception: bool = True + raycast: bool = False + cameras: bool = False + depth: bool = False + segmentation: bool = False + rgb: bool = False + + +@dataclass(frozen=True) +class PlannerCapabilities: + """Capabilities exposed by a planner backend.""" + + motion_generation: bool = False + inverse_kinematics: bool = False + batched_planning: bool = False + + @dataclass(frozen=True) class RuntimeCapabilities: """Capability set for the selected compute/sim pair.""" compute: ComputeCapabilities sim: SimCapabilities + kernel: KernelCapabilities + sensor: SensorCapabilities + planner: PlannerCapabilities @dataclass(frozen=True) @@ -67,6 +107,7 @@ class RuntimeSelection: compute_backend: ComputeBackendName sim_backend: SimBackendName + kernel_backend: KernelBackendName device: str @@ -112,6 +153,142 @@ def load_checkpoint(self, path: str) -> Any: """Load a backend-native checkpoint.""" +class TorchCudaComputeBackend(ComputeBackend): + """Torch/CUDA compute adapter for the upstream runtime.""" + + name: ComputeBackendName = "torch-cuda" + capabilities = ComputeCapabilities( + autograd=True, + checkpoint_io=True, + custom_kernels=True, + torch_interop=True, + ) + + def configure_device(self, device: str) -> None: + if not device.startswith("cuda"): + return + import torch + + torch.cuda.set_device(device) + + def seed(self, seed: int) -> int: + import torch + + torch.manual_seed(seed) + if hasattr(torch, "cuda"): + torch.cuda.manual_seed_all(seed) + return seed + + def save_checkpoint(self, path: str, payload: Any) -> None: + import torch + + torch.save(payload, path) + + def load_checkpoint(self, path: str) -> Any: + import torch + + return torch.load(path, map_location="cpu") + + +class MlxComputeBackend(ComputeBackend): + """MLX compute adapter for Apple Silicon runtime.""" + + name: ComputeBackendName = "mlx" + capabilities = ComputeCapabilities( + autograd=True, + checkpoint_io=True, + custom_kernels=True, + torch_interop=False, + ) + + def configure_device(self, device: str) -> None: + # MLX manages device placement implicitly through the runtime. + del device + + def seed(self, seed: int) -> int: + if not is_mlx_available(): + raise UnsupportedRuntimeFeatureError("MLX backend selected but `mlx` is not importable.") + import mlx.core as mx + + mx.random.seed(seed) + return seed + + def save_checkpoint(self, path: str, payload: Any) -> None: + # Use pickle as a backend-agnostic transport for now. Callers control payload structure. + with open(path, "wb") as file: + pickle.dump(payload, file) + + def load_checkpoint(self, path: str) -> Any: + with open(path, "rb") as file: + return pickle.load(file) + + +class KernelBackend(ABC): + """Abstract backend seam for Warp, Metal, and CPU kernel implementations.""" + + name: KernelBackendName + capabilities: KernelCapabilities + + def require_feature(self, feature: str) -> None: + if getattr(self.capabilities, feature, False): + return + raise UnsupportedRuntimeFeatureError( + f"Kernel backend '{self.name}' does not implement required feature '{feature}'." + ) + + @abstractmethod + def state_dict(self) -> dict[str, Any]: + """Return serializable backend state.""" + + +class WarpKernelBackend(KernelBackend): + """Kernel backend for the upstream Warp/CUDA path.""" + + name: KernelBackendName = "warp" + capabilities = KernelCapabilities( + custom_kernels=True, + raycast=True, + fabric_transforms=True, + mesh_queries=True, + cpu_fallback=False, + ) + + def state_dict(self) -> dict[str, Any]: + return {"backend": self.name, "capabilities": self.capabilities.__dict__} + + +class MetalKernelBackend(KernelBackend): + """Kernel backend for MLX plus custom Metal kernels on Apple Silicon.""" + + name: KernelBackendName = "metal" + capabilities = KernelCapabilities( + custom_kernels=True, + raycast=True, + fabric_transforms=False, + mesh_queries=True, + cpu_fallback=True, + ) + + def state_dict(self) -> dict[str, Any]: + return {"backend": self.name, "capabilities": self.capabilities.__dict__} + + +class CpuKernelBackend(KernelBackend): + """Fallback kernel backend used during correctness bring-up.""" + + name: KernelBackendName = "cpu" + capabilities = KernelCapabilities( + custom_kernels=False, + raycast=False, + fabric_transforms=False, + mesh_queries=False, + cpu_fallback=True, + ) + + def state_dict(self) -> dict[str, Any]: + return {"backend": self.name, "capabilities": self.capabilities.__dict__} + + class SimBackend(ABC): """Abstract simulation backend seam for Isaac Sim and future macOS simulators.""" @@ -324,6 +501,135 @@ def state_dict(self) -> dict[str, Any]: return {"attached": False, "backend": self.name} +class SensorBackend(ABC): + """Abstract backend seam for sensor implementations.""" + + name: SensorBackendName + capabilities: SensorCapabilities + + def require_feature(self, feature: str) -> None: + if getattr(self.capabilities, feature, False): + return + raise UnsupportedRuntimeFeatureError( + f"Sensor backend '{self.name}' does not implement required feature '{feature}'." + ) + + @abstractmethod + def state_dict(self) -> dict[str, Any]: + """Return serializable backend state.""" + + +class IsaacSimSensorBackend(SensorBackend): + """Sensor backend for the upstream Isaac Sim runtime.""" + + name: SensorBackendName = "isaacsim-sensors" + capabilities = SensorCapabilities( + proprioception=True, + raycast=True, + cameras=True, + depth=True, + segmentation=True, + rgb=True, + ) + + def state_dict(self) -> dict[str, Any]: + return {"backend": self.name, "capabilities": self.capabilities.__dict__} + + +class MacSensorBackend(SensorBackend): + """Sensor backend for the mac-native simulator path.""" + + name: SensorBackendName = "mac-sensors" + capabilities = SensorCapabilities( + proprioception=True, + raycast=False, + cameras=False, + depth=False, + segmentation=False, + rgb=False, + ) + + def state_dict(self) -> dict[str, Any]: + return {"backend": self.name, "capabilities": self.capabilities.__dict__} + + +class CpuSensorBackend(SensorBackend): + """Fallback sensor backend for CPU bring-up paths.""" + + name: SensorBackendName = "cpu" + capabilities = SensorCapabilities( + proprioception=True, + raycast=False, + cameras=False, + depth=False, + segmentation=False, + rgb=False, + ) + + def state_dict(self) -> dict[str, Any]: + return {"backend": self.name, "capabilities": self.capabilities.__dict__} + + +class PlannerBackend(ABC): + """Abstract backend seam for planners and motion generation.""" + + name: PlannerBackendName + capabilities: PlannerCapabilities + + def require_feature(self, feature: str) -> None: + if getattr(self.capabilities, feature, False): + return + raise UnsupportedRuntimeFeatureError( + f"Planner backend '{self.name}' does not implement required feature '{feature}'." + ) + + @abstractmethod + def state_dict(self) -> dict[str, Any]: + """Return serializable backend state.""" + + +class IsaacSimPlannerBackend(PlannerBackend): + """Planner backend for the upstream Isaac Sim path.""" + + name: PlannerBackendName = "isaacsim-planners" + capabilities = PlannerCapabilities( + motion_generation=True, + inverse_kinematics=True, + batched_planning=True, + ) + + def state_dict(self) -> dict[str, Any]: + return {"backend": self.name, "capabilities": self.capabilities.__dict__} + + +class MacPlannerBackend(PlannerBackend): + """Placeholder planner backend for the mac-native path.""" + + name: PlannerBackendName = "mac-planners" + capabilities = PlannerCapabilities( + motion_generation=False, + inverse_kinematics=False, + batched_planning=False, + ) + + def state_dict(self) -> dict[str, Any]: + return {"backend": self.name, "capabilities": self.capabilities.__dict__} + + +class NullPlannerBackend(PlannerBackend): + """Planner backend used when no planner path exists yet.""" + + name: PlannerBackendName = "none" + capabilities = PlannerCapabilities( + motion_generation=False, + inverse_kinematics=False, + batched_planning=False, + ) + + def state_dict(self) -> dict[str, Any]: + return {"backend": self.name, "capabilities": self.capabilities.__dict__} + + def is_mlx_available() -> bool: """Return True when MLX can be imported in the current interpreter.""" return importlib.util.find_spec("mlx") is not None @@ -363,6 +669,20 @@ def normalize_sim_backend(name: str | None) -> SimBackendName | None: raise ValueError(f"Unsupported sim backend: {name!r}. Expected one of: 'isaacsim', 'mac-sim'.") +def normalize_kernel_backend(name: str | None) -> KernelBackendName | None: + """Normalize public kernel backend aliases.""" + if name is None: + return None + normalized = name.strip().lower() + if normalized == "warp": + return "warp" + if normalized in {"metal", "mlx-metal"}: + return "metal" + if normalized == "cpu": + return "cpu" + raise ValueError(f"Unsupported kernel backend: {name!r}. Expected one of: 'warp', 'metal', 'cpu'.") + + def _default_compute_backend(device: str, sim_backend: SimBackendName) -> ComputeBackendName: if sim_backend == "mac-sim": if is_apple_silicon() or is_mlx_available(): @@ -372,21 +692,51 @@ def _default_compute_backend(device: str, sim_backend: SimBackendName) -> Comput return "torch-cuda" +def _default_kernel_backend(sim_backend: SimBackendName) -> KernelBackendName: + if sim_backend == "mac-sim": + if is_apple_silicon() or is_mlx_available(): + return "metal" + return "cpu" + return "warp" + + +def _sensor_backend_name(runtime: RuntimeSelection) -> SensorBackendName: + if runtime.sim_backend == "isaacsim": + return "isaacsim-sensors" + if runtime.kernel_backend == "cpu": + return "cpu" + return "mac-sensors" + + +def _planner_backend_name(runtime: RuntimeSelection) -> PlannerBackendName: + if runtime.sim_backend == "isaacsim": + return "isaacsim-planners" + if runtime.compute_backend == "mlx": + return "mac-planners" + return "none" + + def resolve_runtime_selection( compute_backend: str | None = None, sim_backend: str | None = None, device: str | None = None, + kernel_backend: str | None = None, ) -> RuntimeSelection: """Resolve compute/sim backends from explicit inputs plus environment defaults.""" - device = device or "cuda:0" normalized_sim = normalize_sim_backend(sim_backend or os.environ.get(ENV_SIM_BACKEND)) or "isaacsim" + device = device or ("cpu" if normalized_sim == "mac-sim" else "cuda:0") normalized_compute = ( normalize_compute_backend(compute_backend or os.environ.get(ENV_COMPUTE_BACKEND)) or _default_compute_backend(device, normalized_sim) ) + normalized_kernel = ( + normalize_kernel_backend(kernel_backend or os.environ.get(ENV_KERNEL_BACKEND)) + or _default_kernel_backend(normalized_sim) + ) return RuntimeSelection( compute_backend=normalized_compute, sim_backend=normalized_sim, + kernel_backend=normalized_kernel, device=device, ) @@ -395,13 +745,15 @@ def set_runtime_selection(runtime: RuntimeSelection) -> RuntimeSelection: """Persist the runtime selection in environment variables and builtins.""" os.environ[ENV_COMPUTE_BACKEND] = runtime.compute_backend os.environ[ENV_SIM_BACKEND] = runtime.sim_backend + os.environ[ENV_KERNEL_BACKEND] = runtime.kernel_backend builtins.ISAACLAB_COMPUTE_BACKEND = runtime.compute_backend builtins.ISAACLAB_SIM_BACKEND = runtime.sim_backend + builtins.ISAACLAB_KERNEL_BACKEND = runtime.kernel_backend builtins.ISAACLAB_RUNTIME_DEVICE = runtime.device return runtime -def current_runtime(default_device: str = "cuda:0") -> RuntimeSelection: +def current_runtime(default_device: str | None = None) -> RuntimeSelection: """Return the currently selected runtime from process state.""" return resolve_runtime_selection(device=getattr(builtins, "ISAACLAB_RUNTIME_DEVICE", default_device)) @@ -436,7 +788,36 @@ def current_runtime_capabilities(runtime: RuntimeSelection | None = None) -> Run else: sim = SimCapabilities() - return RuntimeCapabilities(compute=compute, sim=sim) + if runtime.kernel_backend == "warp": + kernel = WarpKernelBackend.capabilities + elif runtime.kernel_backend == "metal": + kernel = MetalKernelBackend.capabilities + else: + kernel = CpuKernelBackend.capabilities + + sensor_name = _sensor_backend_name(runtime) + if sensor_name == "isaacsim-sensors": + sensor = IsaacSimSensorBackend.capabilities + elif sensor_name == "mac-sensors": + sensor = MacSensorBackend.capabilities + else: + sensor = CpuSensorBackend.capabilities + + planner_name = _planner_backend_name(runtime) + if planner_name == "isaacsim-planners": + planner = IsaacSimPlannerBackend.capabilities + elif planner_name == "mac-planners": + planner = MacPlannerBackend.capabilities + else: + planner = NullPlannerBackend.capabilities + + return RuntimeCapabilities( + compute=compute, + sim=sim, + kernel=kernel, + sensor=sensor, + planner=planner, + ) def create_sim_backend(runtime: RuntimeSelection | None = None, *, simulation_context: Any | None = None) -> SimBackend: @@ -447,6 +828,46 @@ def create_sim_backend(runtime: RuntimeSelection | None = None, *, simulation_co return MacSimBackend() +def create_kernel_backend(runtime: RuntimeSelection | None = None) -> KernelBackend: + """Instantiate the kernel adapter for the active runtime.""" + runtime = runtime or current_runtime() + if runtime.kernel_backend == "warp": + return WarpKernelBackend() + if runtime.kernel_backend == "metal": + return MetalKernelBackend() + return CpuKernelBackend() + + +def create_sensor_backend(runtime: RuntimeSelection | None = None) -> SensorBackend: + """Instantiate the sensor adapter for the active runtime.""" + runtime = runtime or current_runtime() + sensor_backend = _sensor_backend_name(runtime) + if sensor_backend == "isaacsim-sensors": + return IsaacSimSensorBackend() + if sensor_backend == "mac-sensors": + return MacSensorBackend() + return CpuSensorBackend() + + +def create_planner_backend(runtime: RuntimeSelection | None = None) -> PlannerBackend: + """Instantiate the planner adapter for the active runtime.""" + runtime = runtime or current_runtime() + planner_backend = _planner_backend_name(runtime) + if planner_backend == "isaacsim-planners": + return IsaacSimPlannerBackend() + if planner_backend == "mac-planners": + return MacPlannerBackend() + return NullPlannerBackend() + + +def create_compute_backend(runtime: RuntimeSelection | None = None) -> ComputeBackend: + """Instantiate the compute adapter for the active runtime.""" + runtime = runtime or current_runtime() + if runtime.compute_backend == "torch-cuda": + return TorchCudaComputeBackend() + return MlxComputeBackend() + + def get_runtime_state(runtime: RuntimeSelection | None = None) -> dict[str, Any]: """Return runtime selection plus capability metadata for diagnostics and tests.""" runtime = runtime or current_runtime() @@ -454,10 +875,16 @@ def get_runtime_state(runtime: RuntimeSelection | None = None) -> dict[str, Any] return { "compute_backend": runtime.compute_backend, "sim_backend": runtime.sim_backend, + "kernel_backend": runtime.kernel_backend, + "sensor_backend": _sensor_backend_name(runtime), + "planner_backend": _planner_backend_name(runtime), "device": runtime.device, "capabilities": { "compute": capabilities.compute.__dict__, "sim": capabilities.sim.__dict__, + "kernel": capabilities.kernel.__dict__, + "sensor": capabilities.sensor.__dict__, + "planner": capabilities.planner.__dict__, }, } @@ -467,6 +894,7 @@ def require_runtime_backends( *, compute_backend: ComputeBackendName | None = None, sim_backend: SimBackendName | None = None, + kernel_backend: KernelBackendName | None = None, ) -> RuntimeSelection: """Validate the active runtime against required backends.""" runtime = current_runtime() @@ -479,6 +907,11 @@ def require_runtime_backends( raise UnsupportedBackendError( f"{caller} requires sim backend '{sim_backend}', but the active backend is '{runtime.sim_backend}'." ) + if kernel_backend is not None and runtime.kernel_backend != kernel_backend: + raise UnsupportedBackendError( + f"{caller} requires kernel backend '{kernel_backend}', but the active backend is" + f" '{runtime.kernel_backend}'." + ) return runtime diff --git a/source/isaaclab/isaaclab/controllers/__init__.py b/source/isaaclab/isaaclab/controllers/__init__.py index 3a055c508..6379c4943 100644 --- a/source/isaaclab/isaaclab/controllers/__init__.py +++ b/source/isaaclab/isaaclab/controllers/__init__.py @@ -3,15 +3,45 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Sub-package for different controllers and motion-generators. - -Controllers or motion generators are responsible for closed-loop tracking of a given command. The -controller can be a simple PID controller or a more complex controller such as impedance control -or inverse kinematics control. The controller is responsible for generating the desired joint-level -commands to be sent to the robot. -""" - -from .differential_ik import DifferentialIKController -from .differential_ik_cfg import DifferentialIKControllerCfg -from .operational_space import OperationalSpaceController -from .operational_space_cfg import OperationalSpaceControllerCfg +"""Sub-package for controllers and motion generators, lazily loaded by backend capability.""" + +from __future__ import annotations + +import importlib + +from isaaclab.backends import UnsupportedBackendError, current_runtime + +_MODULE_EXPORTS = { + "DifferentialIKControllerCfg": (".differential_ik_cfg", "DifferentialIKControllerCfg"), + "OperationalSpaceControllerCfg": (".operational_space_cfg", "OperationalSpaceControllerCfg"), +} +_ISAACSIM_ONLY_EXPORTS = { + "DifferentialIKController": (".differential_ik", "DifferentialIKController"), + "OperationalSpaceController": (".operational_space", "OperationalSpaceController"), +} + +__all__ = [*_MODULE_EXPORTS.keys(), *_ISAACSIM_ONLY_EXPORTS.keys()] + + +def __getattr__(name: str): + target = _MODULE_EXPORTS.get(name) + if target is not None: + module = importlib.import_module(target[0], __name__) + value = getattr(module, target[1]) + globals()[name] = value + return value + + if current_runtime().sim_backend != "isaacsim": + raise UnsupportedBackendError( + f"`isaaclab.controllers.{name}` currently requires `sim-backend=isaacsim`." + " Controller interfaces for `mac-sim` are exposed progressively via backend capabilities." + ) + + target = _ISAACSIM_ONLY_EXPORTS.get(name) + if target is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + module = importlib.import_module(target[0], __name__) + value = getattr(module, target[1]) + globals()[name] = value + return value diff --git a/source/isaaclab/isaaclab/controllers/config/rmp_flow.py b/source/isaaclab/isaaclab/controllers/config/rmp_flow.py index 55c6d8e1f..434f8f4ce 100644 --- a/source/isaaclab/isaaclab/controllers/config/rmp_flow.py +++ b/source/isaaclab/isaaclab/controllers/config/rmp_flow.py @@ -5,10 +5,13 @@ import os -from isaacsim.core.utils.extensions import get_extension_path_from_name +from isaaclab.controllers.rmp_flow_cfg import RmpFlowControllerCfg +from isaaclab.utils.nucleus import ISAACLAB_NUCLEUS_DIR -from isaaclab.controllers.rmp_flow import RmpFlowControllerCfg -from isaaclab.utils.assets import ISAACLAB_NUCLEUS_DIR +try: + from isaacsim.core.utils.extensions import get_extension_path_from_name +except Exception: + get_extension_path_from_name = None # Directory on Nucleus Server for RMP-Flow assets (URDFs, collision models, etc.) ISAACLAB_NUCLEUS_RMPFLOW_DIR = os.path.join(ISAACLAB_NUCLEUS_DIR, "Controllers", "RmpFlowAssets") @@ -16,10 +19,9 @@ # Note: RMP-Flow config files for supported robots are stored in the motion_generation extension # We need to move it here for doc building purposes. try: - _RMP_CONFIG_DIR = os.path.join( - get_extension_path_from_name("isaacsim.robot_motion.motion_generation"), - "motion_policy_configs", - ) + if get_extension_path_from_name is None: + raise ModuleNotFoundError("Isaac Sim extensions are unavailable.") + _RMP_CONFIG_DIR = os.path.join(get_extension_path_from_name("isaacsim.robot_motion.motion_generation"), "motion_policy_configs") except Exception: _RMP_CONFIG_DIR = "" diff --git a/source/isaaclab/isaaclab/controllers/differential_ik_cfg.py b/source/isaaclab/isaaclab/controllers/differential_ik_cfg.py index 315a76275..9c8e4f142 100644 --- a/source/isaaclab/isaaclab/controllers/differential_ik_cfg.py +++ b/source/isaaclab/isaaclab/controllers/differential_ik_cfg.py @@ -3,19 +3,19 @@ # # SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations + from dataclasses import MISSING from typing import Literal -from isaaclab.utils import configclass - -from .differential_ik import DifferentialIKController +from isaaclab.utils.configclass import configclass @configclass class DifferentialIKControllerCfg: """Configuration for differential inverse kinematics controller.""" - class_type: type = DifferentialIKController + class_type: type | str = "isaaclab.controllers.differential_ik:DifferentialIKController" """The associated controller class.""" command_type: Literal["position", "pose"] = MISSING diff --git a/source/isaaclab/isaaclab/controllers/operational_space_cfg.py b/source/isaaclab/isaaclab/controllers/operational_space_cfg.py index d2fc3575b..b955b3755 100644 --- a/source/isaaclab/isaaclab/controllers/operational_space_cfg.py +++ b/source/isaaclab/isaaclab/controllers/operational_space_cfg.py @@ -3,19 +3,19 @@ # # SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations + from collections.abc import Sequence from dataclasses import MISSING -from isaaclab.utils import configclass - -from .operational_space import OperationalSpaceController +from isaaclab.utils.configclass import configclass @configclass class OperationalSpaceControllerCfg: """Configuration for operational-space controller.""" - class_type: type = OperationalSpaceController + class_type: type | str = "isaaclab.controllers.operational_space:OperationalSpaceController" """The associated controller class.""" target_types: Sequence[str] = MISSING diff --git a/source/isaaclab/isaaclab/controllers/rmp_flow.py b/source/isaaclab/isaaclab/controllers/rmp_flow.py index f2766d702..32df60c19 100644 --- a/source/isaaclab/isaaclab/controllers/rmp_flow.py +++ b/source/isaaclab/isaaclab/controllers/rmp_flow.py @@ -3,8 +3,6 @@ # # SPDX-License-Identifier: BSD-3-Clause -from dataclasses import MISSING - import torch from isaacsim.core.prims import SingleArticulation @@ -19,28 +17,8 @@ from isaacsim.robot_motion.motion_generation.lula.motion_policies import RmpFlow, RmpFlowSmoothed import isaaclab.sim as sim_utils -from isaaclab.utils import configclass from isaaclab.utils.assets import retrieve_file_path - - -@configclass -class RmpFlowControllerCfg: - """Configuration for RMP-Flow controller (provided through LULA library).""" - - name: str = "rmp_flow" - """Name of the controller. Supported: "rmp_flow", "rmp_flow_smoothed". Defaults to "rmp_flow".""" - config_file: str = MISSING - """Path to the configuration file for the controller.""" - urdf_file: str = MISSING - """Path to the URDF model of the robot.""" - collision_file: str = MISSING - """Path to collision model description of the robot.""" - frame_name: str = MISSING - """Name of the robot frame for task space (must be present in the URDF).""" - evaluations_per_frame: float = MISSING - """Number of substeps during Euler integration inside LULA world model.""" - ignore_robot_state_updates: bool = False - """If true, then state of the world model inside controller is rolled out. Defaults to False.""" +from .rmp_flow_cfg import RmpFlowControllerCfg class RmpFlowController: diff --git a/source/isaaclab/isaaclab/controllers/rmp_flow_cfg.py b/source/isaaclab/isaaclab/controllers/rmp_flow_cfg.py new file mode 100644 index 000000000..41bd90747 --- /dev/null +++ b/source/isaaclab/isaaclab/controllers/rmp_flow_cfg.py @@ -0,0 +1,32 @@ +# 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 + +"""Lightweight configuration types for the Isaac Sim RMPFlow controller.""" + +from __future__ import annotations + +from dataclasses import MISSING + +from isaaclab.utils.configclass import configclass + + +@configclass +class RmpFlowControllerCfg: + """Configuration for RMP-Flow controller (provided through LULA library).""" + + name: str = "rmp_flow" + """Name of the controller. Supported: "rmp_flow", "rmp_flow_smoothed". Defaults to "rmp_flow".""" + config_file: str = MISSING + """Path to the configuration file for the controller.""" + urdf_file: str = MISSING + """Path to the URDF model of the robot.""" + collision_file: str = MISSING + """Path to collision model description of the robot.""" + frame_name: str = MISSING + """Name of the robot frame for task space (must be present in the URDF).""" + evaluations_per_frame: float = MISSING + """Number of substeps during Euler integration inside LULA world model.""" + ignore_robot_state_updates: bool = False + """If true, then state of the world model inside controller is rolled out. Defaults to False.""" diff --git a/source/isaaclab/isaaclab/devices/__init__.py b/source/isaaclab/isaaclab/devices/__init__.py index b2605d39c..16a864c51 100644 --- a/source/isaaclab/isaaclab/devices/__init__.py +++ b/source/isaaclab/isaaclab/devices/__init__.py @@ -3,28 +3,56 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Sub-package providing interfaces to different teleoperation devices. - -Currently, the following categories of devices are supported: - -* **Keyboard**: Standard keyboard with WASD and arrow keys. -* **Spacemouse**: 3D mouse with 6 degrees of freedom. -* **Gamepad**: Gamepad with 2D two joysticks and buttons. Example: Xbox controller. -* **OpenXR**: Uses hand tracking of index/thumb tip avg to drive the target pose. Gripping is done with pinching. -* **Haply**: Haptic device (Inverse3 + VerseGrip) with position, orientation tracking and force feedback. - -All device interfaces inherit from the :class:`DeviceBase` class, which provides a -common interface for all devices. The device interface reads the input data when -the :meth:`DeviceBase.advance` method is called. It also provides the function :meth:`DeviceBase.add_callback` -to add user-defined callback functions to be called when a particular input is pressed from -the peripheral device. -""" - -from .device_base import DeviceBase, DeviceCfg, DevicesCfg -from .gamepad import Se2Gamepad, Se2GamepadCfg, Se3Gamepad, Se3GamepadCfg -from .haply import HaplyDevice, HaplyDeviceCfg -from .keyboard import Se2Keyboard, Se2KeyboardCfg, Se3Keyboard, Se3KeyboardCfg -from .openxr import ManusVive, ManusViveCfg, OpenXRDevice, OpenXRDeviceCfg -from .retargeter_base import RetargeterBase, RetargeterCfg -from .spacemouse import Se2SpaceMouse, Se2SpaceMouseCfg, Se3SpaceMouse, Se3SpaceMouseCfg -from .teleop_device_factory import create_teleop_device +"""Sub-package for teleoperation device interfaces, lazily loaded by backend capability.""" + +from __future__ import annotations + +import importlib + +from isaaclab.backends import UnsupportedBackendError, current_runtime + +_MODULE_EXPORTS = { + "DeviceBase": (".device_base", "DeviceBase"), + "DeviceCfg": (".device_base", "DeviceCfg"), + "DevicesCfg": (".device_base", "DevicesCfg"), + "Se2Gamepad": (".gamepad", "Se2Gamepad"), + "Se2GamepadCfg": (".gamepad", "Se2GamepadCfg"), + "Se3Gamepad": (".gamepad", "Se3Gamepad"), + "Se3GamepadCfg": (".gamepad", "Se3GamepadCfg"), + "HaplyDevice": (".haply", "HaplyDevice"), + "HaplyDeviceCfg": (".haply", "HaplyDeviceCfg"), + "Se2Keyboard": (".keyboard", "Se2Keyboard"), + "Se2KeyboardCfg": (".keyboard", "Se2KeyboardCfg"), + "Se3Keyboard": (".keyboard", "Se3Keyboard"), + "Se3KeyboardCfg": (".keyboard", "Se3KeyboardCfg"), + "ManusVive": (".openxr", "ManusVive"), + "ManusViveCfg": (".openxr", "ManusViveCfg"), + "OpenXRDevice": (".openxr", "OpenXRDevice"), + "OpenXRDeviceCfg": (".openxr", "OpenXRDeviceCfg"), + "RetargeterBase": (".retargeter_base", "RetargeterBase"), + "RetargeterCfg": (".retargeter_base", "RetargeterCfg"), + "Se2SpaceMouse": (".spacemouse", "Se2SpaceMouse"), + "Se2SpaceMouseCfg": (".spacemouse", "Se2SpaceMouseCfg"), + "Se3SpaceMouse": (".spacemouse", "Se3SpaceMouse"), + "Se3SpaceMouseCfg": (".spacemouse", "Se3SpaceMouseCfg"), + "create_teleop_device": (".teleop_device_factory", "create_teleop_device"), +} + +__all__ = [*_MODULE_EXPORTS.keys()] + + +def __getattr__(name: str): + target = _MODULE_EXPORTS.get(name) + if target is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + if current_runtime().sim_backend != "isaacsim": + raise UnsupportedBackendError( + f"`isaaclab.devices.{name}` currently requires `sim-backend=isaacsim`." + " Device interfaces for `mac-sim` are exposed progressively via backend capabilities." + ) + + module = importlib.import_module(target[0], __name__) + value = getattr(module, target[1]) + globals()[name] = value + return value diff --git a/source/isaaclab/isaaclab/envs/mdp/__init__.py b/source/isaaclab/isaaclab/envs/mdp/__init__.py index bdb507f3e..ee15f3e85 100644 --- a/source/isaaclab/isaaclab/envs/mdp/__init__.py +++ b/source/isaaclab/isaaclab/envs/mdp/__init__.py @@ -15,11 +15,65 @@ """ -from .actions import * # noqa: F401, F403 -from .commands import * # noqa: F401, F403 -from .curriculums import * # noqa: F401, F403 -from .events import * # noqa: F401, F403 -from .observations import * # noqa: F401, F403 -from .recorders import * # noqa: F401, F403 -from .rewards import * # noqa: F401, F403 -from .terminations import * # noqa: F401, F403 +from __future__ import annotations + +import importlib + +from isaaclab.backends import UnsupportedBackendError, current_runtime + +_SAFE_MODULES = (".actions",) +_ISAACSIM_MODULES = ( + ".commands", + ".curriculums", + ".events", + ".observations", + ".recorders", + ".rewards", + ".terminations", +) +_OPTIONAL_IMPORT_PREFIXES = ("pink", "omni", "isaacsim", "warp", "torch", "carb", "pxr") + +__all__ = ["actions", "commands", "curriculums", "events", "observations", "recorders", "rewards", "terminations"] + + +def _import_optional(module_name: str): + try: + return importlib.import_module(module_name, __name__) + except ModuleNotFoundError as exc: + if exc.name and exc.name.startswith(_OPTIONAL_IMPORT_PREFIXES): + return None + raise + + +def __getattr__(name: str): + if name in __all__: + if name != "actions" and current_runtime().sim_backend != "isaacsim": + raise UnsupportedBackendError( + f"`isaaclab.envs.mdp.{name}` currently requires `sim-backend=isaacsim`." + " The `mac-sim` bootstrap path currently exposes action configuration surfaces first." + ) + module = importlib.import_module(f".{name}", __name__) + globals()[name] = module + return module + + for module_name in _SAFE_MODULES: + module = _import_optional(module_name) + if module is not None and hasattr(module, name): + value = getattr(module, name) + globals()[name] = value + return value + + if current_runtime().sim_backend != "isaacsim": + raise UnsupportedBackendError( + f"`isaaclab.envs.mdp.{name}` currently requires `sim-backend=isaacsim`." + " The `mac-sim` bootstrap path currently exposes action configuration surfaces first." + ) + + for module_name in _ISAACSIM_MODULES: + module = _import_optional(module_name) + if module is not None and hasattr(module, name): + value = getattr(module, name) + globals()[name] = value + return value + + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/__init__.py b/source/isaaclab/isaaclab/envs/mdp/actions/__init__.py index 56b3ae25f..4c3093b31 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/__init__.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/__init__.py @@ -5,9 +5,58 @@ """Various action terms that can be used in the environment.""" -from .actions_cfg import * -from .binary_joint_actions import * -from .joint_actions import * -from .joint_actions_to_limits import * -from .non_holonomic_actions import * -from .surface_gripper_actions import * +from __future__ import annotations + +import importlib + +from isaaclab.backends import UnsupportedBackendError, current_runtime + +_SAFE_MODULES = ( + ".actions_cfg", + ".rmpflow_actions_cfg", + ".pink_actions_cfg", +) +_ISAACSIM_MODULES = ( + ".binary_joint_actions", + ".joint_actions", + ".joint_actions_to_limits", + ".non_holonomic_actions", + ".surface_gripper_actions", + ".task_space_actions", + ".rmpflow_task_space_actions", + ".pink_task_space_actions", +) +_OPTIONAL_IMPORT_PREFIXES = ("pink", "omni", "isaacsim", "warp", "torch", "carb", "pxr") + + +def _import_optional(module_name: str): + try: + return importlib.import_module(module_name, __name__) + except ModuleNotFoundError as exc: + if exc.name and exc.name.startswith(_OPTIONAL_IMPORT_PREFIXES): + return None + raise + + +def __getattr__(name: str): + for module_name in _SAFE_MODULES: + module = _import_optional(module_name) + if module is not None and hasattr(module, name): + value = getattr(module, name) + globals()[name] = value + return value + + if current_runtime().sim_backend != "isaacsim": + raise UnsupportedBackendError( + f"`isaaclab.envs.mdp.actions.{name}` currently requires `sim-backend=isaacsim`." + " Action configuration objects remain import-safe in the `mac-sim` bootstrap path." + ) + + for module_name in _ISAACSIM_MODULES: + module = _import_optional(module_name) + if module is not None and hasattr(module, name): + value = getattr(module, name) + globals()[name] = value + return value + + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py b/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py index bf8748bdf..d02ddea29 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py @@ -3,20 +3,18 @@ # # SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations + from dataclasses import MISSING +from typing import TYPE_CHECKING -from isaaclab.controllers import DifferentialIKControllerCfg, OperationalSpaceControllerCfg -from isaaclab.managers.action_manager import ActionTerm, ActionTermCfg -from isaaclab.utils import configclass +from isaaclab.controllers.differential_ik_cfg import DifferentialIKControllerCfg +from isaaclab.controllers.operational_space_cfg import OperationalSpaceControllerCfg +from isaaclab.managers.manager_term_cfg import ActionTermCfg +from isaaclab.utils.configclass import configclass -from . import ( - binary_joint_actions, - joint_actions, - joint_actions_to_limits, - non_holonomic_actions, - surface_gripper_actions, - task_space_actions, -) +if TYPE_CHECKING: + from isaaclab.managers.action_manager import ActionTerm ## # Joint actions. @@ -47,7 +45,7 @@ class JointPositionActionCfg(JointActionCfg): See :class:`JointPositionAction` for more details. """ - class_type: type[ActionTerm] = joint_actions.JointPositionAction + class_type: type[ActionTerm] | str = "isaaclab.envs.mdp.actions.joint_actions:JointPositionAction" use_default_offset: bool = True """Whether to use default joint positions configured in the articulation asset as offset. @@ -65,7 +63,7 @@ class RelativeJointPositionActionCfg(JointActionCfg): See :class:`RelativeJointPositionAction` for more details. """ - class_type: type[ActionTerm] = joint_actions.RelativeJointPositionAction + class_type: type[ActionTerm] | str = "isaaclab.envs.mdp.actions.joint_actions:RelativeJointPositionAction" use_zero_offset: bool = True """Whether to ignore the offset defined in articulation asset. Defaults to True. @@ -81,7 +79,7 @@ class JointVelocityActionCfg(JointActionCfg): See :class:`JointVelocityAction` for more details. """ - class_type: type[ActionTerm] = joint_actions.JointVelocityAction + class_type: type[ActionTerm] | str = "isaaclab.envs.mdp.actions.joint_actions:JointVelocityAction" use_default_offset: bool = True """Whether to use default joint velocities configured in the articulation asset as offset. @@ -98,7 +96,7 @@ class JointEffortActionCfg(JointActionCfg): See :class:`JointEffortAction` for more details. """ - class_type: type[ActionTerm] = joint_actions.JointEffortAction + class_type: type[ActionTerm] | str = "isaaclab.envs.mdp.actions.joint_actions:JointEffortAction" ## @@ -113,7 +111,7 @@ class JointPositionToLimitsActionCfg(ActionTermCfg): See :class:`JointPositionToLimitsAction` for more details. """ - class_type: type[ActionTerm] = joint_actions_to_limits.JointPositionToLimitsAction + class_type: type[ActionTerm] | str = "isaaclab.envs.mdp.actions.joint_actions_to_limits:JointPositionToLimitsAction" joint_names: list[str] = MISSING """List of joint names or regex expressions that the action will be mapped to.""" @@ -142,7 +140,9 @@ class EMAJointPositionToLimitsActionCfg(JointPositionToLimitsActionCfg): See :class:`EMAJointPositionToLimitsAction` for more details. """ - class_type: type[ActionTerm] = joint_actions_to_limits.EMAJointPositionToLimitsAction + class_type: type[ActionTerm] | str = ( + "isaaclab.envs.mdp.actions.joint_actions_to_limits:EMAJointPositionToLimitsAction" + ) alpha: float | dict[str, float] = 1.0 """The weight for the moving average (float or dict of regex expressions). Defaults to 1.0. @@ -178,7 +178,7 @@ class BinaryJointPositionActionCfg(BinaryJointActionCfg): See :class:`BinaryJointPositionAction` for more details. """ - class_type: type[ActionTerm] = binary_joint_actions.BinaryJointPositionAction + class_type: type[ActionTerm] | str = "isaaclab.envs.mdp.actions.binary_joint_actions:BinaryJointPositionAction" @configclass @@ -188,7 +188,7 @@ class BinaryJointVelocityActionCfg(BinaryJointActionCfg): See :class:`BinaryJointVelocityAction` for more details. """ - class_type: type[ActionTerm] = binary_joint_actions.BinaryJointVelocityAction + class_type: type[ActionTerm] | str = "isaaclab.envs.mdp.actions.binary_joint_actions:BinaryJointVelocityAction" @configclass @@ -223,7 +223,9 @@ class AbsBinaryJointPositionActionCfg(ActionTermCfg): positive_threshold: bool = True """Whether to use positive (Open actions > Close actions) threshold. Defaults to True.""" - class_type: type[ActionTerm] = binary_joint_actions.AbsBinaryJointPositionAction + class_type: type[ActionTerm] | str = ( + "isaaclab.envs.mdp.actions.binary_joint_actions:AbsBinaryJointPositionAction" + ) ## @@ -238,7 +240,7 @@ class NonHolonomicActionCfg(ActionTermCfg): See :class:`NonHolonomicAction` for more details. """ - class_type: type[ActionTerm] = non_holonomic_actions.NonHolonomicAction + class_type: type[ActionTerm] | str = "isaaclab.envs.mdp.actions.non_holonomic_actions:NonHolonomicAction" body_name: str = MISSING """Name of the body which has the dummy mechanism connected to.""" @@ -281,7 +283,9 @@ class OffsetCfg: rot: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0) """Quaternion rotation ``(w, x, y, z)`` w.r.t. the parent frame. Defaults to (1.0, 0.0, 0.0, 0.0).""" - class_type: type[ActionTerm] = task_space_actions.DifferentialInverseKinematicsAction + class_type: type[ActionTerm] | str = ( + "isaaclab.envs.mdp.actions.task_space_actions:DifferentialInverseKinematicsAction" + ) joint_names: list[str] = MISSING """List of joint names or regex expressions that the action will be mapped to.""" @@ -317,7 +321,9 @@ class OffsetCfg: rot: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0) """Quaternion rotation ``(w, x, y, z)`` w.r.t. the parent frame. Defaults to (1.0, 0.0, 0.0, 0.0).""" - class_type: type[ActionTerm] = task_space_actions.OperationalSpaceControllerAction + class_type: type[ActionTerm] | str = ( + "isaaclab.envs.mdp.actions.task_space_actions:OperationalSpaceControllerAction" + ) joint_names: list[str] = MISSING """List of joint names or regex expressions that the action will be mapped to.""" @@ -376,4 +382,6 @@ class SurfaceGripperBinaryActionCfg(ActionTermCfg): close_command: float = 1.0 """The command value to close the gripper. Defaults to 1.0.""" - class_type: type[ActionTerm] = surface_gripper_actions.SurfaceGripperBinaryAction + class_type: type[ActionTerm] | str = ( + "isaaclab.envs.mdp.actions.surface_gripper_actions:SurfaceGripperBinaryAction" + ) 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 dfd7a72dc..bbc809208 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_actions_cfg.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_actions_cfg.py @@ -3,14 +3,18 @@ # # SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations + from dataclasses import MISSING +from typing import TYPE_CHECKING -from isaaclab.controllers.rmp_flow import RmpFlowControllerCfg -from isaaclab.managers.action_manager import ActionTerm, ActionTermCfg -from isaaclab.utils import configclass +from isaaclab.controllers.rmp_flow_cfg import RmpFlowControllerCfg +from isaaclab.managers.manager_term_cfg import ActionTermCfg +from isaaclab.utils.configclass import configclass -from . import rmpflow_task_space_actions +if TYPE_CHECKING: + from isaaclab.managers.action_manager import ActionTerm @configclass @@ -29,7 +33,7 @@ class OffsetCfg: """Translation w.r.t. the parent frame. Defaults to (0.0, 0.0, 0.0).""" rot: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0) - class_type: type[ActionTerm] = rmpflow_task_space_actions.RMPFlowAction + class_type: type[ActionTerm] | str = "isaaclab.envs.mdp.actions.rmpflow_task_space_actions:RMPFlowAction" joint_names: list[str] = MISSING """List of joint names or regex expressions that the action will be mapped to.""" diff --git a/source/isaaclab/isaaclab/managers/__init__.py b/source/isaaclab/isaaclab/managers/__init__.py index 4a8266801..274a75e32 100644 --- a/source/isaaclab/isaaclab/managers/__init__.py +++ b/source/isaaclab/isaaclab/managers/__init__.py @@ -3,32 +3,58 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Sub-module for environment managers. - -The managers are used to handle various aspects of the environment such as randomization events, curriculum, -and observations. Each manager implements a specific functionality for the environment. The managers are -designed to be modular and can be easily extended to support new functionality. -""" - -from .action_manager import ActionManager, ActionTerm -from .command_manager import CommandManager, CommandTerm -from .curriculum_manager import CurriculumManager -from .event_manager import EventManager -from .manager_base import ManagerBase, ManagerTermBase -from .manager_term_cfg import ( - ActionTermCfg, - CommandTermCfg, - CurriculumTermCfg, - EventTermCfg, - ManagerTermBaseCfg, - ObservationGroupCfg, - ObservationTermCfg, - RecorderTermCfg, - RewardTermCfg, - TerminationTermCfg, -) -from .observation_manager import ObservationManager -from .recorder_manager import DatasetExportMode, RecorderManager, RecorderManagerBaseCfg, RecorderTerm -from .reward_manager import RewardManager -from .scene_entity_cfg import SceneEntityCfg -from .termination_manager import TerminationManager +"""Sub-module for environment managers, lazily loaded by backend capability.""" + +from __future__ import annotations + +import importlib + +from isaaclab.backends import UnsupportedBackendError, current_runtime + +_MODULE_EXPORTS = { + "ActionManager": (".action_manager", "ActionManager"), + "ActionTerm": (".action_manager", "ActionTerm"), + "CommandManager": (".command_manager", "CommandManager"), + "CommandTerm": (".command_manager", "CommandTerm"), + "CurriculumManager": (".curriculum_manager", "CurriculumManager"), + "EventManager": (".event_manager", "EventManager"), + "ManagerBase": (".manager_base", "ManagerBase"), + "ManagerTermBase": (".manager_base", "ManagerTermBase"), + "ActionTermCfg": (".manager_term_cfg", "ActionTermCfg"), + "CommandTermCfg": (".manager_term_cfg", "CommandTermCfg"), + "CurriculumTermCfg": (".manager_term_cfg", "CurriculumTermCfg"), + "EventTermCfg": (".manager_term_cfg", "EventTermCfg"), + "ManagerTermBaseCfg": (".manager_term_cfg", "ManagerTermBaseCfg"), + "ObservationGroupCfg": (".manager_term_cfg", "ObservationGroupCfg"), + "ObservationTermCfg": (".manager_term_cfg", "ObservationTermCfg"), + "RecorderTermCfg": (".manager_term_cfg", "RecorderTermCfg"), + "RewardTermCfg": (".manager_term_cfg", "RewardTermCfg"), + "TerminationTermCfg": (".manager_term_cfg", "TerminationTermCfg"), + "ObservationManager": (".observation_manager", "ObservationManager"), + "DatasetExportMode": (".recorder_manager", "DatasetExportMode"), + "RecorderManager": (".recorder_manager", "RecorderManager"), + "RecorderManagerBaseCfg": (".recorder_manager", "RecorderManagerBaseCfg"), + "RecorderTerm": (".recorder_manager", "RecorderTerm"), + "RewardManager": (".reward_manager", "RewardManager"), + "SceneEntityCfg": (".scene_entity_cfg", "SceneEntityCfg"), + "TerminationManager": (".termination_manager", "TerminationManager"), +} + +__all__ = [*_MODULE_EXPORTS.keys()] + + +def __getattr__(name: str): + target = _MODULE_EXPORTS.get(name) + if target is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + if current_runtime().sim_backend != "isaacsim": + raise UnsupportedBackendError( + f"`isaaclab.managers.{name}` currently requires `sim-backend=isaacsim`." + " Manager interfaces for `mac-sim` are exposed progressively via backend capabilities." + ) + + module = importlib.import_module(target[0], __name__) + value = getattr(module, target[1]) + globals()[name] = value + return value diff --git a/source/isaaclab/isaaclab/managers/action_manager.py b/source/isaaclab/isaaclab/managers/action_manager.py index 6d138451f..3493a84c3 100644 --- a/source/isaaclab/isaaclab/managers/action_manager.py +++ b/source/isaaclab/isaaclab/managers/action_manager.py @@ -20,6 +20,7 @@ import omni.kit.app from isaaclab.envs.utils.io_descriptors import GenericActionIODescriptor +from isaaclab.utils.string import string_to_callable from .manager_base import ManagerBase, ManagerTermBase from .manager_term_cfg import ActionTermCfg @@ -445,6 +446,8 @@ def _prepare_terms(self): f"Configuration for the term '{term_name}' is not of type ActionTermCfg." f" Received: '{type(term_cfg)}'." ) + if isinstance(term_cfg.class_type, str): + term_cfg.class_type = string_to_callable(term_cfg.class_type) # create the action term term = term_cfg.class_type(term_cfg, self._env) # sanity check if term is valid type diff --git a/source/isaaclab/isaaclab/managers/manager_term_cfg.py b/source/isaaclab/isaaclab/managers/manager_term_cfg.py index de7c23aa2..1ffff5336 100644 --- a/source/isaaclab/isaaclab/managers/manager_term_cfg.py +++ b/source/isaaclab/isaaclab/managers/manager_term_cfg.py @@ -11,15 +11,15 @@ from dataclasses import MISSING from typing import TYPE_CHECKING, Any -import torch - -from isaaclab.utils import configclass -from isaaclab.utils.modifiers import ModifierCfg -from isaaclab.utils.noise import NoiseCfg, NoiseModelCfg +from isaaclab.utils.configclass import configclass from .scene_entity_cfg import SceneEntityCfg if TYPE_CHECKING: + import torch + from isaaclab.utils.modifiers import ModifierCfg + from isaaclab.utils.noise import NoiseCfg, NoiseModelCfg + from .action_manager import ActionTerm from .command_manager import CommandTerm from .manager_base import ManagerTermBase diff --git a/source/isaaclab/isaaclab/managers/scene_entity_cfg.py b/source/isaaclab/isaaclab/managers/scene_entity_cfg.py index f3b01f7ee..abd2a72ae 100644 --- a/source/isaaclab/isaaclab/managers/scene_entity_cfg.py +++ b/source/isaaclab/isaaclab/managers/scene_entity_cfg.py @@ -10,7 +10,7 @@ from dataclasses import MISSING from typing import TYPE_CHECKING -from isaaclab.utils import configclass +from isaaclab.utils.configclass import configclass if TYPE_CHECKING: from isaaclab.assets import Articulation, RigidObject, RigidObjectCollection diff --git a/source/isaaclab/isaaclab/markers/__init__.py b/source/isaaclab/isaaclab/markers/__init__.py index eb7b69761..b15c616e1 100644 --- a/source/isaaclab/isaaclab/markers/__init__.py +++ b/source/isaaclab/isaaclab/markers/__init__.py @@ -3,23 +3,47 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Sub-package for marker utilities to simplify creation of UI elements in the GUI. - -Currently, the sub-package provides the following classes: - -* :class:`VisualizationMarkers` for creating a group of markers using `UsdGeom.PointInstancer - `_. - - -.. note:: - - For some simple use-cases, it may be sufficient to use the debug drawing utilities from Isaac Sim. - The debug drawing API is available in the `isaacsim.util.debug_drawing`_ module. It allows drawing of - points and splines efficiently on the UI. - - .. _isaacsim.util.debug_drawing: https://docs.omniverse.nvidia.com/app_isaacsim/app_isaacsim/ext_omni_isaac_debug_drawing.html - -""" - -from .config import * # noqa: F401, F403 -from .visualization_markers import VisualizationMarkers, VisualizationMarkersCfg +"""Sub-package for marker utilities, lazily loaded by backend capability.""" + +from __future__ import annotations + +import importlib + +from isaaclab.backends import UnsupportedBackendError, current_runtime + +_CORE_EXPORTS = { + "VisualizationMarkers": (".visualization_markers", "VisualizationMarkers"), + "VisualizationMarkersCfg": (".visualization_markers", "VisualizationMarkersCfg"), +} +_SEARCH_MODULES = (".config",) + +__all__ = [*_CORE_EXPORTS.keys()] + + +def __getattr__(name: str): + target = _CORE_EXPORTS.get(name) + if target is not None: + if current_runtime().sim_backend != "isaacsim": + raise UnsupportedBackendError( + f"`isaaclab.markers.{name}` currently requires `sim-backend=isaacsim`." + " Marker interfaces for `mac-sim` are exposed progressively via backend capabilities." + ) + module = importlib.import_module(target[0], __name__) + value = getattr(module, target[1]) + globals()[name] = value + return value + + if current_runtime().sim_backend != "isaacsim": + raise UnsupportedBackendError( + f"`isaaclab.markers.{name}` currently requires `sim-backend=isaacsim`." + " Marker interfaces for `mac-sim` are exposed progressively via backend capabilities." + ) + + for module_name in _SEARCH_MODULES: + module = importlib.import_module(module_name, __name__) + if hasattr(module, name): + value = getattr(module, name) + globals()[name] = value + return value + + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/source/isaaclab/isaaclab/markers/config/__init__.py b/source/isaaclab/isaaclab/markers/config/__init__.py index 54d300512..cd6db7739 100644 --- a/source/isaaclab/isaaclab/markers/config/__init__.py +++ b/source/isaaclab/isaaclab/markers/config/__init__.py @@ -5,7 +5,7 @@ import isaaclab.sim as sim_utils from isaaclab.markers.visualization_markers import VisualizationMarkersCfg -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.nucleus import ISAAC_NUCLEUS_DIR ## # Sensors. diff --git a/source/isaaclab/isaaclab/scene/__init__.py b/source/isaaclab/isaaclab/scene/__init__.py index 4b614ea4b..c375bd4ad 100644 --- a/source/isaaclab/isaaclab/scene/__init__.py +++ b/source/isaaclab/isaaclab/scene/__init__.py @@ -3,27 +3,34 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Sub-package containing an interactive scene definition. - -A scene is a collection of entities (e.g., terrain, articulations, sensors, lights, etc.) that can be added to the -simulation. However, only a subset of these entities are of direct interest for the user to interact with. -For example, the user may want to interact with a robot in the scene, but not with the terrain or the lights. -For this reason, we integrate the different entities into a single class called :class:`InteractiveScene`. - -The interactive scene performs the following tasks: - -1. It parses the configuration class :class:`InteractiveSceneCfg` to create the scene. This configuration class is - inherited by the user to add entities to the scene. -2. It clones the entities based on the number of environments specified by the user. -3. It clubs the entities into different groups based on their type (e.g., articulations, sensors, etc.). -4. It provides a set of methods to unify the common operations on the entities in the scene (e.g., resetting internal - buffers, writing buffers to simulation and updating buffers from simulation). - -The interactive scene can be passed around to different modules in the framework to perform different tasks. -For instance, computing the observations based on the state of the scene, or randomizing the scene, or applying -actions to the scene. All these are handled by different "managers" in the framework. Please refer to the -:mod:`isaaclab.managers` sub-package for more details. -""" - -from .interactive_scene import InteractiveScene -from .interactive_scene_cfg import InteractiveSceneCfg +"""Sub-package for interactive scene definitions, lazily loaded by backend capability.""" + +from __future__ import annotations + +import importlib + +from isaaclab.backends import UnsupportedBackendError, current_runtime + +_MODULE_EXPORTS = { + "InteractiveScene": (".interactive_scene", "InteractiveScene"), + "InteractiveSceneCfg": (".interactive_scene_cfg", "InteractiveSceneCfg"), +} + +__all__ = [*_MODULE_EXPORTS.keys()] + + +def __getattr__(name: str): + target = _MODULE_EXPORTS.get(name) + if target is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + if current_runtime().sim_backend != "isaacsim": + raise UnsupportedBackendError( + f"`isaaclab.scene.{name}` currently requires `sim-backend=isaacsim`." + " Scene interfaces for `mac-sim` are exposed progressively via backend capabilities." + ) + + module = importlib.import_module(target[0], __name__) + value = getattr(module, target[1]) + globals()[name] = value + return value diff --git a/source/isaaclab/isaaclab/sensors/__init__.py b/source/isaaclab/isaaclab/sensors/__init__.py index f0a5719e5..e18280d19 100644 --- a/source/isaaclab/isaaclab/sensors/__init__.py +++ b/source/isaaclab/isaaclab/sensors/__init__.py @@ -3,42 +3,61 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Sub-package containing various sensor classes implementations. - -This subpackage contains the sensor classes that are compatible with Isaac Sim. We include both -USD-based and custom sensors: - -* **USD-prim sensors**: Available in Omniverse and require creating a USD prim for them. - For instance, RTX ray tracing camera and lidar sensors. -* **USD-schema sensors**: Available in Omniverse and require creating a USD schema on an existing prim. - For instance, contact sensors and frame transformers. -* **Custom sensors**: Implemented in Python and do not require creating any USD prim or schema. - For instance, warp-based ray-casters. - -Due to the above categorization, the prim paths passed to the sensor's configuration class -are interpreted differently based on the sensor type. The following table summarizes the -interpretation of the prim paths for different sensor types: - -+---------------------+---------------------------+---------------------------------------------------------------+ -| Sensor Type | Example Prim Path | Pre-check | -+=====================+===========================+===============================================================+ -| Camera | /World/robot/base/camera | Leaf is available, and it will spawn a USD camera | -+---------------------+---------------------------+---------------------------------------------------------------+ -| Contact Sensor | /World/robot/feet_* | Leaf is available and checks if the schema exists | -+---------------------+---------------------------+---------------------------------------------------------------+ -| Ray Caster | /World/robot/base | Leaf exists and is a physics body (Articulation / Rigid Body) | -+---------------------+---------------------------+---------------------------------------------------------------+ -| Frame Transformer | /World/robot/base | Leaf exists and is a physics body (Articulation / Rigid Body) | -+---------------------+---------------------------+---------------------------------------------------------------+ -| Imu | /World/robot/base | Leaf exists and is a physics body (Rigid Body) | -+---------------------+---------------------------+---------------------------------------------------------------+ - -""" - -from .camera import * # noqa: F401, F403 -from .contact_sensor import * # noqa: F401, F403 -from .frame_transformer import * # noqa: F401 -from .imu import * # noqa: F401, F403 -from .ray_caster import * # noqa: F401, F403 -from .sensor_base import SensorBase # noqa: F401 -from .sensor_base_cfg import SensorBaseCfg # noqa: F401 +"""Sub-package containing sensor classes, lazily loaded by backend capability.""" + +from __future__ import annotations + +import importlib + +from isaaclab.backends import UnsupportedBackendError, current_runtime + +_MODULE_EXPORTS = { + "SensorBase": (".sensor_base", "SensorBase"), + "SensorBaseCfg": (".sensor_base_cfg", "SensorBaseCfg"), + "Camera": (".camera", "Camera"), + "CameraCfg": (".camera", "CameraCfg"), + "CameraData": (".camera", "CameraData"), + "TiledCamera": (".camera", "TiledCamera"), + "TiledCameraCfg": (".camera", "TiledCameraCfg"), + "ContactSensor": (".contact_sensor", "ContactSensor"), + "ContactSensorCfg": (".contact_sensor", "ContactSensorCfg"), + "ContactSensorData": (".contact_sensor", "ContactSensorData"), + "FrameTransformer": (".frame_transformer", "FrameTransformer"), + "FrameTransformerCfg": (".frame_transformer", "FrameTransformerCfg"), + "FrameTransformerData": (".frame_transformer", "FrameTransformerData"), + "OffsetCfg": (".frame_transformer", "OffsetCfg"), + "Imu": (".imu", "Imu"), + "ImuCfg": (".imu", "ImuCfg"), + "ImuData": (".imu", "ImuData"), + "patterns": (".ray_caster", "patterns"), + "RayCaster": (".ray_caster", "RayCaster"), + "RayCasterCfg": (".ray_caster", "RayCasterCfg"), + "RayCasterData": (".ray_caster", "RayCasterData"), + "RayCasterCamera": (".ray_caster", "RayCasterCamera"), + "RayCasterCameraCfg": (".ray_caster", "RayCasterCameraCfg"), + "MultiMeshRayCaster": (".ray_caster", "MultiMeshRayCaster"), + "MultiMeshRayCasterCfg": (".ray_caster", "MultiMeshRayCasterCfg"), + "MultiMeshRayCasterData": (".ray_caster", "MultiMeshRayCasterData"), + "MultiMeshRayCasterCamera": (".ray_caster", "MultiMeshRayCasterCamera"), + "MultiMeshRayCasterCameraCfg": (".ray_caster", "MultiMeshRayCasterCameraCfg"), + "MultiMeshRayCasterCameraData": (".ray_caster", "MultiMeshRayCasterCameraData"), +} + +__all__ = [*_MODULE_EXPORTS.keys()] + + +def __getattr__(name: str): + target = _MODULE_EXPORTS.get(name) + if target is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + if current_runtime().sim_backend != "isaacsim": + raise UnsupportedBackendError( + f"`isaaclab.sensors.{name}` currently requires `sim-backend=isaacsim`." + " Sensor interfaces on `mac-sim` are exposed progressively via backend capabilities." + ) + + module = importlib.import_module(target[0], __name__) + value = getattr(module, target[1]) + globals()[name] = value + return value diff --git a/source/isaaclab/isaaclab/sim/__init__.py b/source/isaaclab/isaaclab/sim/__init__.py index b2da1fac4..ae66e823a 100644 --- a/source/isaaclab/isaaclab/sim/__init__.py +++ b/source/isaaclab/isaaclab/sim/__init__.py @@ -16,6 +16,11 @@ "RenderCfg": (".simulation_cfg", "RenderCfg"), "SimulationCfg": (".simulation_cfg", "SimulationCfg"), } +_SAFE_MODULES = { + "converters": ".converters", + "schemas": ".schemas", + "spawners": ".spawners", +} _ISAACSIM_EXPORTS = { "SimulationContext": (".simulation_context", "SimulationContext"), "build_simulation_context": (".simulation_context", "build_simulation_context"), @@ -28,7 +33,7 @@ ".views", ) -__all__ = [*_SAFE_EXPORTS.keys(), *_ISAACSIM_EXPORTS.keys()] +__all__ = [*_SAFE_EXPORTS.keys(), *_SAFE_MODULES.keys(), *_ISAACSIM_EXPORTS.keys()] def __getattr__(name: str): @@ -39,6 +44,12 @@ def __getattr__(name: str): globals()[name] = value return value + safe_module = _SAFE_MODULES.get(name) + if safe_module is not None: + module = importlib.import_module(safe_module, __name__) + globals()[name] = module + return module + if current_runtime().sim_backend != "isaacsim": raise UnsupportedBackendError( f"`isaaclab.sim.{name}` currently requires `sim-backend=isaacsim`." diff --git a/source/isaaclab/isaaclab/sim/converters/__init__.py b/source/isaaclab/isaaclab/sim/converters/__init__.py index 7503c53bd..33a0c88b2 100644 --- a/source/isaaclab/isaaclab/sim/converters/__init__.py +++ b/source/isaaclab/isaaclab/sim/converters/__init__.py @@ -16,11 +16,47 @@ """ -from .asset_converter_base import AssetConverterBase -from .asset_converter_base_cfg import AssetConverterBaseCfg -from .mesh_converter import MeshConverter -from .mesh_converter_cfg import MeshConverterCfg -from .mjcf_converter import MjcfConverter -from .mjcf_converter_cfg import MjcfConverterCfg -from .urdf_converter import UrdfConverter -from .urdf_converter_cfg import UrdfConverterCfg +from __future__ import annotations + +import importlib + +from isaaclab.backends import UnsupportedBackendError, current_runtime + +_SAFE_EXPORTS = { + "AssetConverterBaseCfg": (".asset_converter_base_cfg", "AssetConverterBaseCfg"), + "MeshConverterCfg": (".mesh_converter_cfg", "MeshConverterCfg"), + "MjcfConverterCfg": (".mjcf_converter_cfg", "MjcfConverterCfg"), + "UrdfConverterCfg": (".urdf_converter_cfg", "UrdfConverterCfg"), +} +_ISAACSIM_EXPORTS = { + "AssetConverterBase": (".asset_converter_base", "AssetConverterBase"), + "MeshConverter": (".mesh_converter", "MeshConverter"), + "MjcfConverter": (".mjcf_converter", "MjcfConverter"), + "UrdfConverter": (".urdf_converter", "UrdfConverter"), +} + +__all__ = [*_SAFE_EXPORTS.keys(), *_ISAACSIM_EXPORTS.keys()] + + +def __getattr__(name: str): + target = _SAFE_EXPORTS.get(name) + if target is not None: + module = importlib.import_module(target[0], __name__) + value = getattr(module, target[1]) + globals()[name] = value + return value + + if current_runtime().sim_backend != "isaacsim": + raise UnsupportedBackendError( + f"`isaaclab.sim.converters.{name}` currently requires `sim-backend=isaacsim`." + " Converter configuration objects remain import-safe in the `mac-sim` bootstrap path." + ) + + target = _ISAACSIM_EXPORTS.get(name) + if target is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + module = importlib.import_module(target[0], __name__) + value = getattr(module, target[1]) + globals()[name] = value + return value diff --git a/source/isaaclab/isaaclab/sim/converters/asset_converter_base_cfg.py b/source/isaaclab/isaaclab/sim/converters/asset_converter_base_cfg.py index 79bb8d17d..6dd98325d 100644 --- a/source/isaaclab/isaaclab/sim/converters/asset_converter_base_cfg.py +++ b/source/isaaclab/isaaclab/sim/converters/asset_converter_base_cfg.py @@ -5,7 +5,7 @@ from dataclasses import MISSING -from isaaclab.utils import configclass +from isaaclab.utils.configclass import configclass @configclass diff --git a/source/isaaclab/isaaclab/sim/converters/mesh_converter_cfg.py b/source/isaaclab/isaaclab/sim/converters/mesh_converter_cfg.py index 3d231ac1e..91a6605fc 100644 --- a/source/isaaclab/isaaclab/sim/converters/mesh_converter_cfg.py +++ b/source/isaaclab/isaaclab/sim/converters/mesh_converter_cfg.py @@ -5,7 +5,7 @@ from isaaclab.sim.converters.asset_converter_base_cfg import AssetConverterBaseCfg from isaaclab.sim.schemas import schemas_cfg -from isaaclab.utils import configclass +from isaaclab.utils.configclass import configclass @configclass diff --git a/source/isaaclab/isaaclab/sim/converters/mjcf_converter_cfg.py b/source/isaaclab/isaaclab/sim/converters/mjcf_converter_cfg.py index 7cbd83e3e..4b3491870 100644 --- a/source/isaaclab/isaaclab/sim/converters/mjcf_converter_cfg.py +++ b/source/isaaclab/isaaclab/sim/converters/mjcf_converter_cfg.py @@ -6,7 +6,7 @@ from dataclasses import MISSING from isaaclab.sim.converters.asset_converter_base_cfg import AssetConverterBaseCfg -from isaaclab.utils import configclass +from isaaclab.utils.configclass import configclass @configclass diff --git a/source/isaaclab/isaaclab/sim/converters/urdf_converter_cfg.py b/source/isaaclab/isaaclab/sim/converters/urdf_converter_cfg.py index c04ede240..5501b8e49 100644 --- a/source/isaaclab/isaaclab/sim/converters/urdf_converter_cfg.py +++ b/source/isaaclab/isaaclab/sim/converters/urdf_converter_cfg.py @@ -7,7 +7,7 @@ from typing import Literal from isaaclab.sim.converters.asset_converter_base_cfg import AssetConverterBaseCfg -from isaaclab.utils import configclass +from isaaclab.utils.configclass import configclass @configclass diff --git a/source/isaaclab/isaaclab/sim/schemas/__init__.py b/source/isaaclab/isaaclab/sim/schemas/__init__.py index c8402fdb1..b4e5cbf9a 100644 --- a/source/isaaclab/isaaclab/sim/schemas/__init__.py +++ b/source/isaaclab/isaaclab/sim/schemas/__init__.py @@ -32,96 +32,75 @@ """ -from .schemas import ( - MESH_APPROXIMATION_TOKENS, - PHYSX_MESH_COLLISION_CFGS, - USD_MESH_COLLISION_CFGS, - activate_contact_sensors, - define_articulation_root_properties, - define_collision_properties, - define_deformable_body_properties, - define_mass_properties, - define_mesh_collision_properties, - define_rigid_body_properties, - modify_articulation_root_properties, - modify_collision_properties, - modify_deformable_body_properties, - modify_fixed_tendon_properties, - modify_joint_drive_properties, - modify_mass_properties, - modify_mesh_collision_properties, - modify_rigid_body_properties, - modify_spatial_tendon_properties, -) -from .schemas_cfg import ( - ArticulationRootPropertiesCfg, - BoundingCubePropertiesCfg, - BoundingSpherePropertiesCfg, - CollisionPropertiesCfg, - ConvexDecompositionPropertiesCfg, - ConvexHullPropertiesCfg, - DeformableBodyPropertiesCfg, - FixedTendonPropertiesCfg, - JointDrivePropertiesCfg, - MassPropertiesCfg, - MeshCollisionPropertiesCfg, - RigidBodyPropertiesCfg, - SDFMeshPropertiesCfg, - SpatialTendonPropertiesCfg, - TriangleMeshPropertiesCfg, - TriangleMeshSimplificationPropertiesCfg, -) - -__all__ = [ - # articulation root - "ArticulationRootPropertiesCfg", - "define_articulation_root_properties", - "modify_articulation_root_properties", - # rigid bodies - "RigidBodyPropertiesCfg", - "define_rigid_body_properties", - "modify_rigid_body_properties", - "activate_contact_sensors", - # colliders - "CollisionPropertiesCfg", - "define_collision_properties", - "modify_collision_properties", - # deformables - "DeformableBodyPropertiesCfg", - "define_deformable_body_properties", - "modify_deformable_body_properties", - # joints - "JointDrivePropertiesCfg", - "modify_joint_drive_properties", - # mass - "MassPropertiesCfg", - "define_mass_properties", - "modify_mass_properties", - # mesh colliders - "MeshCollisionPropertiesCfg", - "define_mesh_collision_properties", - "modify_mesh_collision_properties", - # bounding cube - "BoundingCubePropertiesCfg", - # bounding sphere - "BoundingSpherePropertiesCfg", - # convex decomposition - "ConvexDecompositionPropertiesCfg", - # convex hull - "ConvexHullPropertiesCfg", - # sdf mesh - "SDFMeshPropertiesCfg", - # triangle mesh - "TriangleMeshPropertiesCfg", - # triangle mesh simplification - "TriangleMeshSimplificationPropertiesCfg", - # tendons - "FixedTendonPropertiesCfg", - "SpatialTendonPropertiesCfg", - "modify_fixed_tendon_properties", - "modify_spatial_tendon_properties", - # Constants for configs that use PhysX vs USD API - "PHYSX_MESH_COLLISION_CFGS", - "USD_MESH_COLLISION_CFGS", - "MESH_APPROXIMATION_TOKENS", -] +from __future__ import annotations + +import importlib + +from isaaclab.backends import UnsupportedBackendError, current_runtime + +_SAFE_EXPORTS = { + "ArticulationRootPropertiesCfg": (".schemas_cfg", "ArticulationRootPropertiesCfg"), + "BoundingCubePropertiesCfg": (".schemas_cfg", "BoundingCubePropertiesCfg"), + "BoundingSpherePropertiesCfg": (".schemas_cfg", "BoundingSpherePropertiesCfg"), + "CollisionPropertiesCfg": (".schemas_cfg", "CollisionPropertiesCfg"), + "ConvexDecompositionPropertiesCfg": (".schemas_cfg", "ConvexDecompositionPropertiesCfg"), + "ConvexHullPropertiesCfg": (".schemas_cfg", "ConvexHullPropertiesCfg"), + "DeformableBodyPropertiesCfg": (".schemas_cfg", "DeformableBodyPropertiesCfg"), + "FixedTendonPropertiesCfg": (".schemas_cfg", "FixedTendonPropertiesCfg"), + "JointDrivePropertiesCfg": (".schemas_cfg", "JointDrivePropertiesCfg"), + "MassPropertiesCfg": (".schemas_cfg", "MassPropertiesCfg"), + "MeshCollisionPropertiesCfg": (".schemas_cfg", "MeshCollisionPropertiesCfg"), + "RigidBodyPropertiesCfg": (".schemas_cfg", "RigidBodyPropertiesCfg"), + "SDFMeshPropertiesCfg": (".schemas_cfg", "SDFMeshPropertiesCfg"), + "SpatialTendonPropertiesCfg": (".schemas_cfg", "SpatialTendonPropertiesCfg"), + "TriangleMeshPropertiesCfg": (".schemas_cfg", "TriangleMeshPropertiesCfg"), + "TriangleMeshSimplificationPropertiesCfg": (".schemas_cfg", "TriangleMeshSimplificationPropertiesCfg"), +} +_ISAACSIM_EXPORTS = { + "MESH_APPROXIMATION_TOKENS": (".schemas", "MESH_APPROXIMATION_TOKENS"), + "PHYSX_MESH_COLLISION_CFGS": (".schemas", "PHYSX_MESH_COLLISION_CFGS"), + "USD_MESH_COLLISION_CFGS": (".schemas", "USD_MESH_COLLISION_CFGS"), + "activate_contact_sensors": (".schemas", "activate_contact_sensors"), + "define_articulation_root_properties": (".schemas", "define_articulation_root_properties"), + "define_collision_properties": (".schemas", "define_collision_properties"), + "define_deformable_body_properties": (".schemas", "define_deformable_body_properties"), + "define_mass_properties": (".schemas", "define_mass_properties"), + "define_mesh_collision_properties": (".schemas", "define_mesh_collision_properties"), + "define_rigid_body_properties": (".schemas", "define_rigid_body_properties"), + "modify_articulation_root_properties": (".schemas", "modify_articulation_root_properties"), + "modify_collision_properties": (".schemas", "modify_collision_properties"), + "modify_deformable_body_properties": (".schemas", "modify_deformable_body_properties"), + "modify_fixed_tendon_properties": (".schemas", "modify_fixed_tendon_properties"), + "modify_joint_drive_properties": (".schemas", "modify_joint_drive_properties"), + "modify_mass_properties": (".schemas", "modify_mass_properties"), + "modify_mesh_collision_properties": (".schemas", "modify_mesh_collision_properties"), + "modify_rigid_body_properties": (".schemas", "modify_rigid_body_properties"), + "modify_spatial_tendon_properties": (".schemas", "modify_spatial_tendon_properties"), +} + +__all__ = [*_SAFE_EXPORTS.keys(), *_ISAACSIM_EXPORTS.keys()] + + +def __getattr__(name: str): + target = _SAFE_EXPORTS.get(name) + if target is not None: + module = importlib.import_module(target[0], __name__) + value = getattr(module, target[1]) + globals()[name] = value + return value + + if current_runtime().sim_backend != "isaacsim": + raise UnsupportedBackendError( + f"`isaaclab.sim.schemas.{name}` currently requires `sim-backend=isaacsim`." + " Schema configuration objects are available in the `mac-sim` bootstrap path, but USD/PhysX schema" + " mutation utilities remain Isaac Sim only." + ) + + target = _ISAACSIM_EXPORTS.get(name) + if target is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + module = importlib.import_module(target[0], __name__) + value = getattr(module, target[1]) + globals()[name] = value + return value diff --git a/source/isaaclab/isaaclab/sim/schemas/schemas.py b/source/isaaclab/isaaclab/sim/schemas/schemas.py index 91eb3b2cd..2f9bb9fa1 100644 --- a/source/isaaclab/isaaclab/sim/schemas/schemas.py +++ b/source/isaaclab/isaaclab/sim/schemas/schemas.py @@ -16,6 +16,7 @@ from pxr import PhysxSchema, Usd, UsdPhysics from isaaclab.sim.utils.stage import get_current_stage +from isaaclab.utils.string import string_to_callable from ..utils import ( apply_nested, @@ -29,6 +30,13 @@ logger = logging.getLogger(__name__) +def _resolve_schema_api(api_func: Callable | str) -> Callable: + """Resolve deferred schema callables stored in config objects.""" + if callable(api_func): + return api_func + return string_to_callable(api_func) + + """ Constants. """ @@ -1027,9 +1035,9 @@ def extract_mesh_collision_api_and_attrs( if use_usd_api: # Use USD API for corresponding attributes # For mesh collision approximation attribute, we set it explicitly in `modify_mesh_collision_properties`` - api_func = cfg.usd_func + api_func = _resolve_schema_api(cfg.usd_func) elif use_phsyx_api: - api_func = cfg.physx_func + api_func = _resolve_schema_api(cfg.physx_func) else: raise ValueError("Either USD or PhysX API should be used for modifying mesh collision attributes!") diff --git a/source/isaaclab/isaaclab/sim/schemas/schemas_cfg.py b/source/isaaclab/isaaclab/sim/schemas/schemas_cfg.py index 446d7faa1..b5a559741 100644 --- a/source/isaaclab/isaaclab/sim/schemas/schemas_cfg.py +++ b/source/isaaclab/isaaclab/sim/schemas/schemas_cfg.py @@ -6,9 +6,14 @@ from dataclasses import MISSING from typing import Literal -from pxr import PhysxSchema, UsdPhysics +from isaaclab.utils.configclass import configclass -from isaaclab.utils import configclass +_USD_MESH_COLLISION_API = "pxr.UsdPhysics:MeshCollisionAPI" +_PHYSX_CONVEX_DECOMPOSITION_API = "pxr.PhysxSchema:PhysxConvexDecompositionCollisionAPI" +_PHYSX_CONVEX_HULL_API = "pxr.PhysxSchema:PhysxConvexHullCollisionAPI" +_PHYSX_TRIANGLE_MESH_API = "pxr.PhysxSchema:PhysxTriangleMeshCollisionAPI" +_PHYSX_TRIANGLE_MESH_SIMPLIFICATION_API = "pxr.PhysxSchema:PhysxTriangleMeshSimplificationCollisionAPI" +_PHYSX_SDF_MESH_API = "pxr.PhysxSchema:PhysxSDFMeshCollisionAPI" @configclass @@ -441,14 +446,14 @@ class MeshCollisionPropertiesCfg: the properties and leave the rest as-is. """ - usd_func: callable = MISSING + usd_func: callable | str = MISSING """USD API function for modifying mesh collision properties. Refer to `original USD Documentation `_ for more information. """ - physx_func: callable = MISSING + physx_func: callable | str = MISSING """PhysX API function for modifying mesh collision properties. Refer to `original PhysX Documentation `_ @@ -463,7 +468,7 @@ class MeshCollisionPropertiesCfg: @configclass class BoundingCubePropertiesCfg(MeshCollisionPropertiesCfg): - usd_func: callable = UsdPhysics.MeshCollisionAPI + usd_func: callable | str = _USD_MESH_COLLISION_API """Original USD Documentation: https://docs.omniverse.nvidia.com/kit/docs/omni_usd_schema_physics/latest/class_usd_physics_mesh_collision_a_p_i.html """ @@ -476,7 +481,7 @@ class BoundingCubePropertiesCfg(MeshCollisionPropertiesCfg): @configclass class BoundingSpherePropertiesCfg(MeshCollisionPropertiesCfg): - usd_func: callable = UsdPhysics.MeshCollisionAPI + usd_func: callable | str = _USD_MESH_COLLISION_API """Original USD Documentation: https://docs.omniverse.nvidia.com/kit/docs/omni_usd_schema_physics/latest/class_usd_physics_mesh_collision_a_p_i.html """ @@ -489,12 +494,12 @@ class BoundingSpherePropertiesCfg(MeshCollisionPropertiesCfg): @configclass class ConvexDecompositionPropertiesCfg(MeshCollisionPropertiesCfg): - usd_func: callable = UsdPhysics.MeshCollisionAPI + usd_func: callable | str = _USD_MESH_COLLISION_API """Original USD Documentation: https://docs.omniverse.nvidia.com/kit/docs/omni_usd_schema_physics/latest/class_usd_physics_mesh_collision_a_p_i.html """ - physx_func: callable = PhysxSchema.PhysxConvexDecompositionCollisionAPI + physx_func: callable | str = _PHYSX_CONVEX_DECOMPOSITION_API """Original PhysX Documentation: https://docs.omniverse.nvidia.com/kit/docs/omni_usd_schema_physics/latest/class_physx_schema_physx_convex_decomposition_collision_a_p_i.html """ @@ -538,12 +543,12 @@ class ConvexDecompositionPropertiesCfg(MeshCollisionPropertiesCfg): @configclass class ConvexHullPropertiesCfg(MeshCollisionPropertiesCfg): - usd_func: callable = UsdPhysics.MeshCollisionAPI + usd_func: callable | str = _USD_MESH_COLLISION_API """Original USD Documentation: https://docs.omniverse.nvidia.com/kit/docs/omni_usd_schema_physics/latest/class_usd_physics_mesh_collision_a_p_i.html """ - physx_func: callable = PhysxSchema.PhysxConvexHullCollisionAPI + physx_func: callable | str = _PHYSX_CONVEX_HULL_API """Original PhysX Documentation: https://docs.omniverse.nvidia.com/kit/docs/omni_usd_schema_physics/latest/class_physx_schema_physx_convex_hull_collision_a_p_i.html """ @@ -567,7 +572,7 @@ class ConvexHullPropertiesCfg(MeshCollisionPropertiesCfg): @configclass class TriangleMeshPropertiesCfg(MeshCollisionPropertiesCfg): - physx_func: callable = PhysxSchema.PhysxTriangleMeshCollisionAPI + physx_func: callable | str = _PHYSX_TRIANGLE_MESH_API """Triangle mesh is only supported by PhysX API. Original PhysX Documentation: @@ -589,12 +594,12 @@ class TriangleMeshPropertiesCfg(MeshCollisionPropertiesCfg): @configclass class TriangleMeshSimplificationPropertiesCfg(MeshCollisionPropertiesCfg): - usd_func: callable = UsdPhysics.MeshCollisionAPI + usd_func: callable | str = _USD_MESH_COLLISION_API """Original USD Documentation: https://docs.omniverse.nvidia.com/kit/docs/omni_usd_schema_physics/latest/class_usd_physics_mesh_collision_a_p_i.html """ - physx_func: callable = PhysxSchema.PhysxTriangleMeshSimplificationCollisionAPI + physx_func: callable | str = _PHYSX_TRIANGLE_MESH_SIMPLIFICATION_API """Original PhysX Documentation: https://docs.omniverse.nvidia.com/kit/docs/omni_usd_schema_physics/latest/class_physx_schema_physx_triangle_mesh_simplification_collision_a_p_i.html """ @@ -619,7 +624,7 @@ class TriangleMeshSimplificationPropertiesCfg(MeshCollisionPropertiesCfg): @configclass class SDFMeshPropertiesCfg(MeshCollisionPropertiesCfg): - physx_func: callable = PhysxSchema.PhysxSDFMeshCollisionAPI + physx_func: callable | str = _PHYSX_SDF_MESH_API """SDF mesh is only supported by PhysX API. Original PhysX documentation: diff --git a/source/isaaclab/isaaclab/sim/spawners/from_files/__init__.py b/source/isaaclab/isaaclab/sim/spawners/from_files/__init__.py index a95ac491b..b5b4282c2 100644 --- a/source/isaaclab/isaaclab/sim/spawners/from_files/__init__.py +++ b/source/isaaclab/isaaclab/sim/spawners/from_files/__init__.py @@ -13,11 +13,52 @@ """ -from .from_files import ( - spawn_from_mjcf, - spawn_from_urdf, - spawn_from_usd, - spawn_from_usd_with_compliant_contact_material, - spawn_ground_plane, -) -from .from_files_cfg import GroundPlaneCfg, MjcfFileCfg, UrdfFileCfg, UsdFileCfg, UsdFileWithCompliantContactCfg +from __future__ import annotations + +import importlib + +from isaaclab.backends import UnsupportedBackendError, current_runtime + +_SAFE_EXPORTS = { + "GroundPlaneCfg": (".from_files_cfg", "GroundPlaneCfg"), + "MjcfFileCfg": (".from_files_cfg", "MjcfFileCfg"), + "UrdfFileCfg": (".from_files_cfg", "UrdfFileCfg"), + "UsdFileCfg": (".from_files_cfg", "UsdFileCfg"), + "UsdFileWithCompliantContactCfg": (".from_files_cfg", "UsdFileWithCompliantContactCfg"), +} +_ISAACSIM_EXPORTS = { + "spawn_from_mjcf": (".from_files", "spawn_from_mjcf"), + "spawn_from_urdf": (".from_files", "spawn_from_urdf"), + "spawn_from_usd": (".from_files", "spawn_from_usd"), + "spawn_from_usd_with_compliant_contact_material": ( + ".from_files", + "spawn_from_usd_with_compliant_contact_material", + ), + "spawn_ground_plane": (".from_files", "spawn_ground_plane"), +} + +__all__ = [*_SAFE_EXPORTS.keys(), *_ISAACSIM_EXPORTS.keys()] + + +def __getattr__(name: str): + target = _SAFE_EXPORTS.get(name) + if target is not None: + module = importlib.import_module(target[0], __name__) + value = getattr(module, target[1]) + globals()[name] = value + return value + + if current_runtime().sim_backend != "isaacsim": + raise UnsupportedBackendError( + f"`isaaclab.sim.spawners.from_files.{name}` currently requires `sim-backend=isaacsim`." + " File-based spawner configuration objects remain import-safe in the `mac-sim` bootstrap path." + ) + + target = _ISAACSIM_EXPORTS.get(name) + if target is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + module = importlib.import_module(target[0], __name__) + value = getattr(module, target[1]) + globals()[name] = value + return value diff --git a/source/isaaclab/isaaclab/sim/spawners/from_files/from_files_cfg.py b/source/isaaclab/isaaclab/sim/spawners/from_files/from_files_cfg.py index ad9f55859..590329a18 100644 --- a/source/isaaclab/isaaclab/sim/spawners/from_files/from_files_cfg.py +++ b/source/isaaclab/isaaclab/sim/spawners/from_files/from_files_cfg.py @@ -8,13 +8,41 @@ from collections.abc import Callable from dataclasses import MISSING +import isaaclab.sim.spawners.materials as materials from isaaclab.sim import converters, schemas -from isaaclab.sim.spawners import materials from isaaclab.sim.spawners.spawner_cfg import DeformableObjectSpawnerCfg, RigidObjectSpawnerCfg, SpawnerCfg -from isaaclab.utils import configclass -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.configclass import configclass +from isaaclab.utils.nucleus import ISAAC_NUCLEUS_DIR -from . import from_files + +def _spawn_from_usd(*args, **kwargs): + from .from_files import spawn_from_usd + + return spawn_from_usd(*args, **kwargs) + + +def _spawn_from_urdf(*args, **kwargs): + from .from_files import spawn_from_urdf + + return spawn_from_urdf(*args, **kwargs) + + +def _spawn_from_mjcf(*args, **kwargs): + from .from_files import spawn_from_mjcf + + return spawn_from_mjcf(*args, **kwargs) + + +def _spawn_from_usd_with_compliant_contact_material(*args, **kwargs): + from .from_files import spawn_from_usd_with_compliant_contact_material + + return spawn_from_usd_with_compliant_contact_material(*args, **kwargs) + + +def _spawn_ground_plane(*args, **kwargs): + from .from_files import spawn_ground_plane + + return spawn_ground_plane(*args, **kwargs) @configclass @@ -96,7 +124,7 @@ class UsdFileCfg(FileCfg): This is done by calling the respective function with the specified properties. """ - func: Callable = from_files.spawn_from_usd + func: Callable = _spawn_from_usd usd_path: str = MISSING """Path to the USD file to spawn asset from.""" @@ -129,7 +157,7 @@ class UrdfFileCfg(FileCfg, converters.UrdfConverterCfg): """ - func: Callable = from_files.spawn_from_urdf + func: Callable = _spawn_from_urdf @configclass @@ -151,7 +179,7 @@ class MjcfFileCfg(FileCfg, converters.MjcfConverterCfg): """ - func: Callable = from_files.spawn_from_mjcf + func: Callable = _spawn_from_mjcf """ @@ -169,7 +197,7 @@ class UsdFileWithCompliantContactCfg(UsdFileCfg): material application. """ - func: Callable = from_files.spawn_from_usd_with_compliant_contact_material + func: Callable = _spawn_from_usd_with_compliant_contact_material compliant_contact_stiffness: float | None = None """Stiffness of the compliant contact. Defaults to None. @@ -201,7 +229,7 @@ class GroundPlaneCfg(SpawnerCfg): This uses the USD for the standard grid-world ground plane from Isaac Sim by default. """ - func: Callable = from_files.spawn_ground_plane + func: Callable = _spawn_ground_plane usd_path: str = f"{ISAAC_NUCLEUS_DIR}/Environments/Grid/default_environment.usd" """Path to the USD file to spawn asset from. Defaults to the grid-world ground plane.""" diff --git a/source/isaaclab/isaaclab/sim/spawners/materials/__init__.py b/source/isaaclab/isaaclab/sim/spawners/materials/__init__.py index 45fa55def..52691bf99 100644 --- a/source/isaaclab/isaaclab/sim/spawners/materials/__init__.py +++ b/source/isaaclab/isaaclab/sim/spawners/materials/__init__.py @@ -9,20 +9,46 @@ import importlib -_SEARCH_MODULES = ( - ".physics_materials", - ".physics_materials_cfg", - ".visual_materials", - ".visual_materials_cfg", -) +from isaaclab.backends import UnsupportedBackendError, current_runtime + +_SAFE_EXPORTS = { + "PhysicsMaterialCfg": (".physics_materials_cfg", "PhysicsMaterialCfg"), + "RigidBodyMaterialCfg": (".physics_materials_cfg", "RigidBodyMaterialCfg"), + "DeformableBodyMaterialCfg": (".physics_materials_cfg", "DeformableBodyMaterialCfg"), + "VisualMaterialCfg": (".visual_materials_cfg", "VisualMaterialCfg"), + "PreviewSurfaceCfg": (".visual_materials_cfg", "PreviewSurfaceCfg"), + "MdlFileCfg": (".visual_materials_cfg", "MdlFileCfg"), + "GlassMdlCfg": (".visual_materials_cfg", "GlassMdlCfg"), +} +_ISAACSIM_EXPORTS = { + "spawn_rigid_body_material": (".physics_materials", "spawn_rigid_body_material"), + "spawn_deformable_body_material": (".physics_materials", "spawn_deformable_body_material"), + "spawn_preview_surface": (".visual_materials", "spawn_preview_surface"), + "spawn_from_mdl_file": (".visual_materials", "spawn_from_mdl_file"), +} + +__all__ = [*_SAFE_EXPORTS.keys(), *_ISAACSIM_EXPORTS.keys()] def __getattr__(name: str): - for module_name in _SEARCH_MODULES: - module = importlib.import_module(module_name, __name__) - if hasattr(module, name): - value = getattr(module, name) - globals()[name] = value - return value + target = _SAFE_EXPORTS.get(name) + if target is not None: + module = importlib.import_module(target[0], __name__) + value = getattr(module, target[1]) + globals()[name] = value + return value + + if current_runtime().sim_backend != "isaacsim": + raise UnsupportedBackendError( + f"`isaaclab.sim.spawners.materials.{name}` currently requires `sim-backend=isaacsim`." + " Material configuration objects are available in the `mac-sim` bootstrap path." + ) + + target = _ISAACSIM_EXPORTS.get(name) + if target is not None: + module = importlib.import_module(target[0], __name__) + value = getattr(module, target[1]) + globals()[name] = value + return value raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/source/isaaclab/isaaclab/sim/spawners/materials/visual_materials_cfg.py b/source/isaaclab/isaaclab/sim/spawners/materials/visual_materials_cfg.py index a0c2874b0..629bab324 100644 --- a/source/isaaclab/isaaclab/sim/spawners/materials/visual_materials_cfg.py +++ b/source/isaaclab/isaaclab/sim/spawners/materials/visual_materials_cfg.py @@ -6,9 +6,19 @@ from collections.abc import Callable from dataclasses import MISSING -from isaaclab.utils import configclass +from isaaclab.utils.configclass import configclass -from . import visual_materials + +def _spawn_preview_surface(*args, **kwargs): + from .visual_materials import spawn_preview_surface + + return spawn_preview_surface(*args, **kwargs) + + +def _spawn_from_mdl_file(*args, **kwargs): + from .visual_materials import spawn_from_mdl_file + + return spawn_from_mdl_file(*args, **kwargs) @configclass @@ -26,7 +36,7 @@ class PreviewSurfaceCfg(VisualMaterialCfg): See :meth:`spawn_preview_surface` for more information. """ - func: Callable = visual_materials.spawn_preview_surface + func: Callable = _spawn_preview_surface diffuse_color: tuple[float, float, float] = (0.18, 0.18, 0.18) """The RGB diffusion color. This is the base color of the surface. Defaults to a dark gray.""" @@ -51,7 +61,7 @@ class MdlFileCfg(VisualMaterialCfg): See :meth:`spawn_from_mdl_file` for more information. """ - func: Callable = visual_materials.spawn_from_mdl_file + func: Callable = _spawn_from_mdl_file mdl_path: str = MISSING """The path to the MDL material. @@ -93,7 +103,7 @@ class GlassMdlCfg(VisualMaterialCfg): The default values are taken from the glass material in the NVIDIA Nucleus. """ - func: Callable = visual_materials.spawn_from_mdl_file + func: Callable = _spawn_from_mdl_file mdl_path: str = "OmniGlass.mdl" """The path to the MDL material. Defaults to the glass material in the NVIDIA Nucleus.""" diff --git a/source/isaaclab/isaaclab/sim/spawners/spawner_cfg.py b/source/isaaclab/isaaclab/sim/spawners/spawner_cfg.py index 2dea8db8f..e8ce5f5d0 100644 --- a/source/isaaclab/isaaclab/sim/spawners/spawner_cfg.py +++ b/source/isaaclab/isaaclab/sim/spawners/spawner_cfg.py @@ -8,10 +8,8 @@ from collections.abc import Callable from dataclasses import MISSING -from pxr import Usd - from isaaclab.sim import schemas -from isaaclab.utils import configclass +from isaaclab.utils.configclass import configclass @configclass @@ -28,7 +26,7 @@ class SpawnerCfg: parameter. """ - func: Callable[..., Usd.Prim] = MISSING + func: Callable[..., object] = MISSING """Function to use for spawning the asset. The function takes in the prim path (or expression) to spawn the asset at, the configuration instance diff --git a/source/isaaclab/isaaclab/utils/__init__.py b/source/isaaclab/isaaclab/utils/__init__.py index 75448242e..a48e7dab6 100644 --- a/source/isaaclab/isaaclab/utils/__init__.py +++ b/source/isaaclab/isaaclab/utils/__init__.py @@ -27,6 +27,7 @@ ) __all__ = [*_EXPLICIT_EXPORTS.keys()] +_OPTIONAL_IMPORT_PREFIXES = ("omni", "isaacsim", "warp", "torch", "carb", "pxr") def __getattr__(name: str): @@ -38,7 +39,12 @@ def __getattr__(name: str): return value for module_name in _SEARCH_MODULES: - module = importlib.import_module(module_name, __name__) + try: + module = importlib.import_module(module_name, __name__) + except ModuleNotFoundError as exc: + if exc.name and exc.name.startswith(_OPTIONAL_IMPORT_PREFIXES): + continue + raise if hasattr(module, name): value = getattr(module, name) globals()[name] = value diff --git a/source/isaaclab/isaaclab/utils/array.py b/source/isaaclab/isaaclab/utils/array.py index d15fbc275..40cfcb8ea 100644 --- a/source/isaaclab/isaaclab/utils/array.py +++ b/source/isaaclab/isaaclab/utils/array.py @@ -11,10 +11,18 @@ from typing import Union import numpy as np -import torch -import warp as wp -TensorData = Union[np.ndarray, torch.Tensor, wp.array] # noqa: UP007 +try: + import torch +except ModuleNotFoundError: # pragma: no cover - exercised in mac-sim tests. + torch = None + +try: + import warp as wp +except ModuleNotFoundError: # pragma: no cover - exercised in mac-sim tests. + wp = None + +TensorData = Union[np.ndarray] # noqa: UP007 """Type definition for a tensor data. Union of numpy, torch, and warp arrays. @@ -22,8 +30,6 @@ TENSOR_TYPES = { "numpy": np.ndarray, - "torch": torch.Tensor, - "warp": wp.array, } """A dictionary containing the types for each backend. @@ -32,9 +38,7 @@ """ TENSOR_TYPE_CONVERSIONS = { - "numpy": {wp.array: lambda x: x.numpy(), torch.Tensor: lambda x: x.detach().cpu().numpy()}, - "torch": {wp.array: lambda x: wp.torch.to_torch(x), np.ndarray: lambda x: torch.from_numpy(x)}, - "warp": {np.array: lambda x: wp.array(x), torch.Tensor: lambda x: wp.torch.from_torch(x)}, + "numpy": {}, } """A nested dictionary containing the conversion functions for each backend. @@ -42,10 +46,26 @@ inner dictionary are the source backend (``np.ndarray``, ``torch.Tensor``, ``wp.array``). """ +if wp is not None: + TensorData = Union[np.ndarray, wp.array] # noqa: UP007 + TENSOR_TYPES["warp"] = wp.array + TENSOR_TYPE_CONVERSIONS["numpy"][wp.array] = lambda x: x.numpy() + TENSOR_TYPE_CONVERSIONS["warp"] = {np.ndarray: lambda x: wp.array(x)} + +if torch is not None: + TensorData = Union[np.ndarray, torch.Tensor] # noqa: UP007 + TENSOR_TYPES["torch"] = torch.Tensor + TENSOR_TYPE_CONVERSIONS["numpy"][torch.Tensor] = lambda x: x.detach().cpu().numpy() + TENSOR_TYPE_CONVERSIONS["torch"] = {np.ndarray: lambda x: torch.from_numpy(x)} + if wp is not None: + TensorData = Union[np.ndarray, torch.Tensor, wp.array] # noqa: UP007 + TENSOR_TYPE_CONVERSIONS["torch"][wp.array] = lambda x: wp.torch.to_torch(x) + TENSOR_TYPE_CONVERSIONS["warp"][torch.Tensor] = lambda x: wp.torch.from_torch(x) + def convert_to_torch( array: TensorData, - dtype: torch.dtype = None, + dtype: torch.dtype | None = None, device: torch.device | str | None = None, ) -> torch.Tensor: """Converts a given array into a torch tensor. @@ -69,6 +89,9 @@ def convert_to_torch( Returns: The converted array as torch tensor. """ + if torch is None: + raise ModuleNotFoundError("PyTorch is required for `convert_to_torch`, but it is not installed.") + # Convert array to tensor # if the datatype is not currently supported by torch we need to improvise # supported types are: https://pytorch.org/docs/stable/tensors.html @@ -79,7 +102,7 @@ def convert_to_torch( array = array.astype(np.int32) # need to deal with object arrays (np.void) separately tensor = torch.from_numpy(array) - elif isinstance(array, wp.array): + elif wp is not None and isinstance(array, wp.array): if array.dtype == wp.uint32: array = array.view(wp.int32) tensor = wp.to_torch(array) diff --git a/source/isaaclab/isaaclab/utils/assets.py b/source/isaaclab/isaaclab/utils/assets.py index 22c514158..a31df8adc 100644 --- a/source/isaaclab/isaaclab/utils/assets.py +++ b/source/isaaclab/isaaclab/utils/assets.py @@ -24,21 +24,11 @@ import carb import omni.client +from .nucleus import ISAACLAB_NUCLEUS_DIR, ISAAC_NUCLEUS_DIR, NVIDIA_NUCLEUS_DIR, NUCLEUS_ASSET_ROOT_DIR + # import logger logger = logging.getLogger(__name__) -NUCLEUS_ASSET_ROOT_DIR = carb.settings.get_settings().get("/persistent/isaac/asset_root/cloud") -"""Path to the root directory on the Nucleus Server.""" - -NVIDIA_NUCLEUS_DIR = f"{NUCLEUS_ASSET_ROOT_DIR}/NVIDIA" -"""Path to the root directory on the NVIDIA Nucleus Server.""" - -ISAAC_NUCLEUS_DIR = f"{NUCLEUS_ASSET_ROOT_DIR}/Isaac" -"""Path to the ``Isaac`` directory on the NVIDIA Nucleus Server.""" - -ISAACLAB_NUCLEUS_DIR = f"{ISAAC_NUCLEUS_DIR}/IsaacLab" -"""Path to the ``Isaac/IsaacLab`` directory on the NVIDIA Nucleus Server.""" - def check_file_path(path: str) -> Literal[0, 1, 2]: """Checks if a file exists on the Nucleus Server or locally. diff --git a/source/isaaclab/isaaclab/utils/nucleus.py b/source/isaaclab/isaaclab/utils/nucleus.py new file mode 100644 index 000000000..79f3b846c --- /dev/null +++ b/source/isaaclab/isaaclab/utils/nucleus.py @@ -0,0 +1,51 @@ +# 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 + +"""Lightweight Nucleus path constants that can be imported without Isaac Sim.""" + +from __future__ import annotations + +import os + +NUCLEUS_ASSET_ROOT_ENV_VAR = "ISAACLAB_NUCLEUS_ASSET_ROOT" +"""Environment variable that can override the root Nucleus asset path.""" + + +def _resolve_nucleus_asset_root() -> str: + """Resolve the root Nucleus asset directory without requiring Isaac Sim modules.""" + override = os.environ.get(NUCLEUS_ASSET_ROOT_ENV_VAR) + if override: + return override.rstrip("/") + + try: + import carb + except Exception: + return "" + + value = carb.settings.get_settings().get("/persistent/isaac/asset_root/cloud") + if not value: + return "" + return str(value).rstrip("/") + + +def _join_nucleus_path(root: str, *parts: str) -> str: + """Join a root path with path segments while preserving import safety.""" + clean_parts = [part.strip("/") for part in parts if part] + if not root: + return "/".join(clean_parts) + return "/".join([root.rstrip("/"), *clean_parts]) + + +NUCLEUS_ASSET_ROOT_DIR = _resolve_nucleus_asset_root() +"""Path to the root directory on the Nucleus Server, if available.""" + +NVIDIA_NUCLEUS_DIR = _join_nucleus_path(NUCLEUS_ASSET_ROOT_DIR, "NVIDIA") +"""Path to the root directory on the NVIDIA Nucleus Server.""" + +ISAAC_NUCLEUS_DIR = _join_nucleus_path(NUCLEUS_ASSET_ROOT_DIR, "Isaac") +"""Path to the ``Isaac`` directory on the NVIDIA Nucleus Server.""" + +ISAACLAB_NUCLEUS_DIR = _join_nucleus_path(ISAAC_NUCLEUS_DIR, "IsaacLab") +"""Path to the ``Isaac/IsaacLab`` directory on the NVIDIA Nucleus Server.""" diff --git a/source/isaaclab/isaaclab/utils/seed.py b/source/isaaclab/isaaclab/utils/seed.py index 6b2a8ff97..c757069f4 100644 --- a/source/isaaclab/isaaclab/utils/seed.py +++ b/source/isaaclab/isaaclab/utils/seed.py @@ -7,8 +7,15 @@ import random import numpy as np -import torch -import warp as wp + +try: + import torch +except ModuleNotFoundError: # pragma: no cover - exercised in mac-sim tests. + torch = None +try: + import warp as wp +except ModuleNotFoundError: # pragma: no cover - exercised in mac-sim tests. + wp = None def configure_seed(seed: int | None, torch_deterministic: bool = False) -> int: @@ -26,20 +33,26 @@ def configure_seed(seed: int | None, torch_deterministic: bool = False) -> int: random.seed(seed) np.random.seed(seed) - torch.manual_seed(seed) os.environ["PYTHONHASHSEED"] = str(seed) - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - wp.rand_init(seed) + if torch is not None: + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + if wp is not None: + wp.rand_init(seed) if torch_deterministic: # refer to https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility - os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" - torch.backends.cudnn.benchmark = False - torch.backends.cudnn.deterministic = True - torch.use_deterministic_algorithms(True) + if torch is not None and torch.cuda.is_available(): + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + if torch is not None: + torch.use_deterministic_algorithms(True) else: - torch.backends.cudnn.benchmark = True - torch.backends.cudnn.deterministic = False + if torch is not None and torch.cuda.is_available(): + torch.backends.cudnn.benchmark = True + torch.backends.cudnn.deterministic = False return seed diff --git a/source/isaaclab/setup.py b/source/isaaclab/setup.py index 369bf5bdf..56d31f82e 100644 --- a/source/isaaclab/setup.py +++ b/source/isaaclab/setup.py @@ -8,7 +8,7 @@ import os import toml -from setuptools import setup +from setuptools import find_packages, setup # Obtain the extension data from the extension.toml file EXTENSION_PATH = os.path.dirname(os.path.realpath(__file__)) @@ -33,20 +33,29 @@ # image processing "transformers==4.57.6", "einops", # needed for transformers, doesn't always auto-install - "warp-lang", # make sure this is consistent with isaac sim version "pillow==11.3.0", # livestream "starlette==0.49.1", - # testing - "pytest", - "pytest-mock", - "junitparser", - "flatdict==4.0.0", - "flaky", "packaging", ] +EXTRAS_REQUIRE = { + "cuda-isaacsim": [ + "warp-lang ; platform_system != 'Darwin'", + ], + "macos-mlx": [ + "mlx ; platform_system == 'Darwin' and platform_machine in 'arm64,aarch64'", + ], + "dev": [ + "pytest", + "pytest-mock", + "junitparser", + "flatdict==4.0.0", + "flaky", + ], +} + # Append Linux x86_64 and ARM64 deps via PEP 508 markers SUPPORTED_ARCHS_ARM = "platform_machine in 'x86_64,AMD64,aarch64,arm64'" SUPPORTED_ARCHS = "platform_machine in 'x86_64,AMD64'" @@ -74,7 +83,8 @@ python_requires=">=3.10", install_requires=INSTALL_REQUIRES, dependency_links=PYTORCH_INDEX_URL, - packages=["isaaclab"], + extras_require=EXTRAS_REQUIRE, + packages=find_packages(include=["isaaclab", "isaaclab.*"]), classifiers=[ "Natural Language :: English", "Programming Language :: Python :: 3.10", diff --git a/source/isaaclab/test/backends/test_mac_cart_double_pendulum.py b/source/isaaclab/test/backends/test_mac_cart_double_pendulum.py new file mode 100644 index 000000000..3a9fa473b --- /dev/null +++ b/source/isaaclab/test/backends/test_mac_cart_double_pendulum.py @@ -0,0 +1,66 @@ +# Copyright (c) 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 tests for mac-native cart-double-pendulum MARL env.""" + +from __future__ import annotations + +import pytest + +mx = pytest.importorskip("mlx.core") + +from isaaclab.backends.mac_sim import MacCartDoublePendulumEnv, MacCartDoublePendulumEnvCfg # noqa: E402 + + +def _random_actions(num_envs: int) -> dict[str, list[list[float]]]: + return { + "cart": [[0.5] for _ in range(num_envs)], + "pendulum": [[-0.25] for _ in range(num_envs)], + } + + +def test_mac_cart_double_pendulum_reset_and_step_shapes(): + """The MARL env should return per-agent observation/reward/done dictionaries.""" + cfg = MacCartDoublePendulumEnvCfg(num_envs=8, seed=11, episode_length_s=0.3) + env = MacCartDoublePendulumEnv(cfg) + + obs, extras = env.reset() + + assert extras == {} + assert set(obs.keys()) == {"cart", "pendulum"} + assert obs["cart"].shape == (8, 4) + assert obs["pendulum"].shape == (8, 3) + + next_obs, rewards, terminated, truncated, step_extras = env.step(_random_actions(cfg.num_envs)) + + assert next_obs["cart"].shape == (8, 4) + assert next_obs["pendulum"].shape == (8, 3) + assert set(rewards.keys()) == {"cart", "pendulum"} + assert rewards["cart"].shape == (8,) + assert rewards["pendulum"].shape == (8,) + assert set(terminated.keys()) == {"cart", "pendulum"} + assert set(truncated.keys()) == {"cart", "pendulum"} + assert terminated["cart"].shape == (8,) + assert truncated["cart"].shape == (8,) + assert isinstance(step_extras, dict) + + +def test_mac_cart_double_pendulum_termination_propagates_to_all_agents(): + """Out-of-bounds termination should mark both MARL agents and return reset extras.""" + cfg = MacCartDoublePendulumEnvCfg(num_envs=4, seed=5, episode_length_s=0.3) + env = MacCartDoublePendulumEnv(cfg) + + joint_pos, joint_vel = env.sim_backend.joint_state() + joint_pos = mx.array(joint_pos) + joint_vel = mx.array(joint_vel) + joint_pos[:, 0] = cfg.max_cart_pos + 1.0 + env.sim_backend.write_joint_state(None, joint_pos, joint_vel) + + _, _, terminated, _, extras = env.step(_random_actions(cfg.num_envs)) + + assert bool(mx.all(terminated["cart"]).item()) + assert bool(mx.all(terminated["pendulum"]).item()) + assert "completed_returns" in extras + assert set(extras["completed_returns"].keys()) == {"cart", "pendulum"} diff --git a/source/isaaclab/test/backends/test_mac_cartpole.py b/source/isaaclab/test/backends/test_mac_cartpole.py index 45334c5ff..eafa63298 100644 --- a/source/isaaclab/test/backends/test_mac_cartpole.py +++ b/source/isaaclab/test/backends/test_mac_cartpole.py @@ -71,3 +71,42 @@ def test_train_and_play_cartpole_smoke(tmp_path: Path): assert len(episode_returns) == 1 assert isinstance(episode_returns[0], float) + + +def test_train_cartpole_resume_from_checkpoint(tmp_path: Path): + """Cartpole training should support warm-starting from a previous checkpoint.""" + first_checkpoint = tmp_path / "cartpole_policy_first.npz" + resumed_checkpoint = tmp_path / "cartpole_policy_resumed.npz" + + base_cfg = MacCartpoleTrainCfg( + env=MacCartpoleEnvCfg(num_envs=8, seed=29, episode_length_s=0.3), + hidden_dim=32, + updates=1, + rollout_steps=8, + epochs_per_update=1, + checkpoint_path=str(first_checkpoint), + ) + first_result = train_cartpole_policy(base_cfg) + assert first_result["resumed_from"] is None + assert first_checkpoint.exists() + + resumed_cfg = MacCartpoleTrainCfg( + env=MacCartpoleEnvCfg(num_envs=8, seed=29, episode_length_s=0.3), + hidden_dim=32, + updates=1, + rollout_steps=8, + epochs_per_update=1, + checkpoint_path=str(resumed_checkpoint), + resume_from=str(first_checkpoint), + ) + resumed_result = train_cartpole_policy(resumed_cfg) + + assert resumed_result["resumed_from"] == str(first_checkpoint) + assert resumed_checkpoint.exists() + + episode_returns = play_cartpole_policy( + str(resumed_checkpoint), + env_cfg=MacCartpoleEnvCfg(num_envs=1, seed=29, episode_length_s=0.3), + episodes=1, + ) + assert len(episode_returns) == 1 diff --git a/source/isaaclab/test/backends/test_mac_cartpole_showcase.py b/source/isaaclab/test/backends/test_mac_cartpole_showcase.py new file mode 100644 index 000000000..8fdd0c9a4 --- /dev/null +++ b/source/isaaclab/test/backends/test_mac_cartpole_showcase.py @@ -0,0 +1,70 @@ +# Copyright (c) 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 tests for the mac-native cartpole showcase spaces.""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("mlx.core") +gym = pytest.importorskip("gymnasium") + +from isaaclab.backends.mac_sim import SHOWCASE_CFGS, MacCartpoleShowcaseEnv # noqa: E402 + + +def _batched_action(space: gym.Space, num_envs: int): + if isinstance(space, gym.spaces.Box): + return [space.sample().tolist() for _ in range(num_envs)] + if isinstance(space, gym.spaces.Discrete): + return [space.sample() for _ in range(num_envs)] + if isinstance(space, gym.spaces.MultiDiscrete): + return [space.sample().tolist() for _ in range(num_envs)] + raise TypeError(f"Unsupported action space for test generation: {type(space)}") + + +def _assert_policy_obs_matches_space(obs, space: gym.Space, num_envs: int): + if isinstance(space, gym.spaces.Box): + assert obs.shape == (num_envs, *space.shape) + return + if isinstance(space, gym.spaces.Discrete): + assert obs.shape == (num_envs,) + return + if isinstance(space, gym.spaces.MultiDiscrete): + assert obs.shape == (num_envs, len(space.nvec)) + return + if isinstance(space, gym.spaces.Tuple): + assert isinstance(obs, tuple) + assert len(obs) == len(space.spaces) + for item, item_space in zip(obs, space.spaces, strict=True): + assert item.shape == (num_envs, *item_space.shape) + return + if isinstance(space, gym.spaces.Dict): + assert isinstance(obs, dict) + assert set(obs.keys()) == set(space.spaces.keys()) + for key, item_space in space.spaces.items(): + assert obs[key].shape == (num_envs, *item_space.shape) + return + raise TypeError(f"Unsupported observation space for assertions: {type(space)}") + + +@pytest.mark.parametrize("cfg_type", SHOWCASE_CFGS, ids=lambda cfg_type: cfg_type.__name__) +def test_mac_cartpole_showcase_step_shapes(cfg_type): + """Every showcase config should reset and step through the mac-native cartpole backend.""" + cfg = cfg_type(num_envs=8, seed=17, episode_length_s=0.3) + env = MacCartpoleShowcaseEnv(cfg) + + obs, extras = env.reset() + + assert extras == {} + _assert_policy_obs_matches_space(obs["policy"], cfg.observation_space, cfg.num_envs) + + next_obs, reward, terminated, truncated, step_extras = env.step(_batched_action(cfg.action_space, cfg.num_envs)) + + _assert_policy_obs_matches_space(next_obs["policy"], cfg.observation_space, cfg.num_envs) + assert reward.shape == (cfg.num_envs,) + assert terminated.shape == (cfg.num_envs,) + assert truncated.shape == (cfg.num_envs,) + assert isinstance(step_extras, dict) diff --git a/source/isaaclab/test/backends/test_mac_quadcopter.py b/source/isaaclab/test/backends/test_mac_quadcopter.py new file mode 100644 index 000000000..d8a0eeb0f --- /dev/null +++ b/source/isaaclab/test/backends/test_mac_quadcopter.py @@ -0,0 +1,78 @@ +# Copyright (c) 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 tests for the mac-native quadcopter environment.""" + +from __future__ import annotations + +import pytest + +mx = pytest.importorskip("mlx.core") + +from isaaclab.backends.mac_sim import MacQuadcopterEnv, MacQuadcopterEnvCfg # noqa: E402 + + +def test_mac_quadcopter_reset_and_step_shapes(): + """The quadcopter env should expose IsaacLab-style vectorized tensors.""" + cfg = MacQuadcopterEnvCfg(num_envs=8, seed=17, episode_length_s=0.5) + env = MacQuadcopterEnv(cfg) + + obs, extras = env.reset() + assert extras == {} + assert obs["policy"].shape == (8, 12) + assert env.runtime.compute_backend == "mlx" + assert env.runtime.sim_backend == "mac-sim" + + actions = mx.zeros((cfg.num_envs, 4), dtype=mx.float32) + next_obs, reward, terminated, truncated, step_extras = env.step(actions) + + assert next_obs["policy"].shape == (8, 12) + assert reward.shape == (8,) + assert terminated.shape == (8,) + assert truncated.shape == (8,) + assert isinstance(step_extras, dict) + + +def test_mac_quadcopter_termination_on_height_violation(): + """A root-height violation should terminate and emit reset extras.""" + cfg = MacQuadcopterEnvCfg(num_envs=4, seed=9, episode_length_s=0.5) + env = MacQuadcopterEnv(cfg) + + env.sim_backend.root_pos_w[:, 2] = cfg.min_height - 0.05 + actions = mx.zeros((cfg.num_envs, 4), dtype=mx.float32) + _, _, terminated, _, extras = env.step(actions) + + assert bool(mx.all(terminated).item()) + assert "completed_lengths" in extras + assert "final_distance_to_goal" in extras + assert len(extras["completed_lengths"]) == cfg.num_envs + + +def test_mac_quadcopter_root_state_io(): + """The backend root pose/velocity write APIs should update state buffers.""" + cfg = MacQuadcopterEnvCfg(num_envs=3, seed=3, episode_length_s=0.5) + env = MacQuadcopterEnv(cfg) + + pose = mx.array( + [ + [0.1, 0.2, 0.9, 0.0, 0.0, 0.0, 1.0], + [0.3, -0.1, 1.1, 0.0, 0.0, 0.0, 1.0], + ], + dtype=mx.float32, + ) + velocity = mx.array( + [ + [0.4, 0.0, -0.2, 0.01, 0.02, 0.03], + [0.0, -0.4, 0.1, -0.02, 0.01, 0.0], + ], + dtype=mx.float32, + ) + + env.sim_backend.write_root_pose(None, pose, env_ids=[0, 2]) + env.sim_backend.write_root_velocity(None, velocity, env_ids=[0, 2]) + + assert mx.allclose(env.sim_backend.root_pos_w[[0, 2]], pose[:, :3]).item() + assert mx.allclose(env.sim_backend.root_lin_vel_b[[0, 2]], velocity[:, :3]).item() + assert mx.allclose(env.sim_backend.root_ang_vel_b[[0, 2]], velocity[:, 3:6]).item() diff --git a/source/isaaclab/test/backends/test_portability_utils.py b/source/isaaclab/test/backends/test_portability_utils.py new file mode 100644 index 000000000..48297c1e5 --- /dev/null +++ b/source/isaaclab/test/backends/test_portability_utils.py @@ -0,0 +1,113 @@ +# Copyright (c) 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 + +"""Portability tests for utility modules used by the MLX/mac path.""" + +from __future__ import annotations + +import os +import subprocess +import sys +import importlib + +import numpy as np +import pytest + +from isaaclab.utils.seed import configure_seed + +array_utils = importlib.import_module("isaaclab.utils.array") + + +def test_convert_to_torch_handles_missing_torch_dependency(): + """convert_to_torch should fail explicitly when PyTorch is unavailable.""" + array = np.array([1.0, 2.0, 3.0], dtype=np.float32) + if array_utils.torch is None: + with pytest.raises(ModuleNotFoundError, match="PyTorch is required"): + array_utils.convert_to_torch(array) + else: + tensor = array_utils.convert_to_torch(array) + assert tensor.shape == (3,) + + +def test_configure_seed_runs_without_cuda(): + """Seed utility should work even when no CUDA device is present.""" + value = configure_seed(123, torch_deterministic=False) + assert value == 123 + + +def test_utils_import_without_warp_module(): + """Importing core utility modules should not fail when Warp is unavailable.""" + code = """ +import importlib +import sys + +sys.modules['warp'] = None +array_mod = importlib.import_module('isaaclab.utils.array') +seed_mod = importlib.import_module('isaaclab.utils.seed') +assert 'warp' not in array_mod.TENSOR_TYPES +assert hasattr(seed_mod, 'configure_seed') +print('ok') +""" + result = subprocess.run( + [sys.executable, "-c", code], + env=os.environ.copy(), + check=False, + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + assert "ok" in result.stdout + + +def test_seed_module_import_without_torch_or_warp(): + """Seed utilities should still import when torch/warp are both unavailable.""" + code = """ +import importlib +import sys + +sys.modules['torch'] = None +sys.modules['warp'] = None +seed_mod = importlib.import_module('isaaclab.utils.seed') +value = seed_mod.configure_seed(7, torch_deterministic=False) +assert value == 7 +print('ok') +""" + result = subprocess.run( + [sys.executable, "-c", code], + env=os.environ.copy(), + check=False, + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + assert "ok" in result.stdout + + +def test_utils_dynamic_export_skips_optional_dependency_modules(): + """isaaclab.utils dynamic exports should not crash when optional deps are absent.""" + code = """ +import importlib +import sys + +import numpy as np + +sys.modules['torch'] = None +sys.modules['warp'] = None +utils_mod = importlib.import_module('isaaclab.utils') +convert_to_torch = getattr(utils_mod, 'convert_to_torch') +try: + convert_to_torch(np.array([1.0], dtype=np.float32)) +except ModuleNotFoundError: + print('ok') +""" + result = subprocess.run( + [sys.executable, "-c", code], + env=os.environ.copy(), + check=False, + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + assert "ok" in result.stdout diff --git a/source/isaaclab/test/backends/test_runtime.py b/source/isaaclab/test/backends/test_runtime.py index 440dc4cd8..24116ca15 100644 --- a/source/isaaclab/test/backends/test_runtime.py +++ b/source/isaaclab/test/backends/test_runtime.py @@ -16,17 +16,31 @@ from isaaclab.app import AppLauncher from isaaclab.backends import ( + CpuKernelBackend, IsaacSimBackend, + IsaacSimPlannerBackend, + IsaacSimSensorBackend, + MlxComputeBackend, + MacPlannerBackend, + MacSensorBackend, MacSimBackend, + MetalKernelBackend, + TorchCudaComputeBackend, ENV_COMPUTE_BACKEND, + ENV_KERNEL_BACKEND, ENV_SIM_BACKEND, UnsupportedBackendError, UnsupportedRuntimeFeatureError, configure_torch_device, + create_compute_backend, + create_kernel_backend, + create_planner_backend, + create_sensor_backend, create_sim_backend, get_runtime_state, resolve_runtime_selection, set_runtime_selection, + WarpKernelBackend, ) @@ -34,18 +48,21 @@ def clear_runtime_env(monkeypatch: pytest.MonkeyPatch): """Keep runtime selection isolated per test.""" monkeypatch.delenv(ENV_COMPUTE_BACKEND, raising=False) + monkeypatch.delenv(ENV_KERNEL_BACKEND, raising=False) monkeypatch.delenv(ENV_SIM_BACKEND, raising=False) def test_resolve_runtime_selection_from_environment(monkeypatch: pytest.MonkeyPatch): """Environment variables should seed the runtime selection.""" monkeypatch.setenv(ENV_COMPUTE_BACKEND, "mlx") + monkeypatch.setenv(ENV_KERNEL_BACKEND, "metal") monkeypatch.setenv(ENV_SIM_BACKEND, "mac-sim") runtime = resolve_runtime_selection(device="cpu") assert runtime.compute_backend == "mlx" assert runtime.sim_backend == "mac-sim" + assert runtime.kernel_backend == "metal" assert runtime.device == "cpu" @@ -56,6 +73,9 @@ def test_set_runtime_selection_persists_process_state(): assert state["compute_backend"] == "mlx" assert state["sim_backend"] == "mac-sim" + assert state["kernel_backend"] == "metal" + assert state["sensor_backend"] == "mac-sensors" + assert state["planner_backend"] == "mac-planners" assert state["device"] == "cpu" @@ -89,9 +109,56 @@ def test_app_launcher_arg_parser_exposes_backend_flags(): AppLauncher.add_app_launcher_args(parser) assert "--compute-backend" in parser._option_string_actions + assert "--kernel-backend" in parser._option_string_actions assert "--sim-backend" in parser._option_string_actions +def test_app_launcher_macsim_bootstrap_mode(): + """AppLauncher should support a mac-sim bootstrap path without requiring Isaac Sim modules.""" + launcher = AppLauncher( + { + "compute_backend": "mlx", + "sim_backend": "mac-sim", + "device": "cpu", + "headless": True, + } + ) + + assert launcher.sim_backend == "mac-sim" + assert launcher.compute_backend == "mlx" + assert launcher.kernel_backend == "metal" + assert launcher.sensor_backend == "mac-sensors" + assert launcher.planner_backend == "mac-planners" + assert launcher.app.is_running() is False + + +def test_app_launcher_rejects_non_mlx_compute_for_macsim(): + """AppLauncher should reject unsupported compute backends for mac-sim.""" + with pytest.raises(UnsupportedBackendError, match="mac-sim"): + AppLauncher( + { + "compute_backend": "torch-cuda", + "sim_backend": "mac-sim", + "device": "cuda:0", + "headless": True, + } + ) + + +def test_app_launcher_rejects_warp_kernel_for_macsim(): + """AppLauncher should reject Warp kernels on the mac-sim path.""" + with pytest.raises(UnsupportedBackendError, match="kernel-backend warp"): + AppLauncher( + { + "compute_backend": "mlx", + "sim_backend": "mac-sim", + "kernel_backend": "warp", + "device": "cpu", + "headless": True, + } + ) + + def test_create_sim_backend_returns_isaacsim_adapter(): """Isaac Sim runtime should produce the Isaac Sim adapter.""" backend = create_sim_backend(resolve_runtime_selection("torch-cuda", "isaacsim", "cuda:0")) @@ -108,6 +175,154 @@ def test_create_sim_backend_returns_macsim_adapter(): assert backend.contract.articulations.effort_targets is False +def test_create_kernel_backend_returns_warp_adapter(): + """Isaac Sim runtime should default to the Warp kernel adapter.""" + backend = create_kernel_backend(resolve_runtime_selection("torch-cuda", "isaacsim", "cuda:0")) + + assert isinstance(backend, WarpKernelBackend) + assert backend.capabilities.custom_kernels is True + + +def test_create_kernel_backend_returns_metal_adapter(): + """mac-sim runtime should default to the Metal kernel adapter.""" + backend = create_kernel_backend(resolve_runtime_selection("mlx", "mac-sim", "cpu")) + + assert isinstance(backend, MetalKernelBackend) + assert backend.capabilities.raycast is True + + +def test_create_sensor_backend_follows_runtime(): + """Sensor backends should follow the selected simulation runtime.""" + upstream = create_sensor_backend(resolve_runtime_selection("torch-cuda", "isaacsim", "cuda:0")) + mac = create_sensor_backend(resolve_runtime_selection("mlx", "mac-sim", "cpu")) + + assert isinstance(upstream, IsaacSimSensorBackend) + assert isinstance(mac, MacSensorBackend) + + +def test_create_planner_backend_follows_runtime(): + """Planner backends should follow the selected simulation runtime.""" + upstream = create_planner_backend(resolve_runtime_selection("torch-cuda", "isaacsim", "cuda:0")) + mac = create_planner_backend(resolve_runtime_selection("mlx", "mac-sim", "cpu")) + + assert isinstance(upstream, IsaacSimPlannerBackend) + assert isinstance(mac, MacPlannerBackend) + + +def test_create_cpu_kernel_backend(): + """Explicit CPU kernel selection should return the CPU backend.""" + backend = create_kernel_backend(resolve_runtime_selection("mlx", "mac-sim", "cpu", kernel_backend="cpu")) + + assert isinstance(backend, CpuKernelBackend) + assert backend.capabilities.cpu_fallback is True + + +def test_create_compute_backend_returns_torch_adapter(): + """Torch runtime should produce the torch compute adapter.""" + backend = create_compute_backend(resolve_runtime_selection("torch-cuda", "isaacsim", "cuda:0")) + + assert isinstance(backend, TorchCudaComputeBackend) + assert backend.capabilities.torch_interop is True + + +def test_create_compute_backend_returns_mlx_adapter(): + """MLX runtime should produce the MLX compute adapter.""" + backend = create_compute_backend(resolve_runtime_selection("mlx", "mac-sim", "cpu")) + + assert isinstance(backend, MlxComputeBackend) + assert backend.capabilities.torch_interop is False + + +def test_mac_runtime_entrypoints_import_without_isaacsim(): + """The public mac entrypoints should import without requiring Isaac Sim modules.""" + pytest.importorskip("mlx.core") + set_runtime_selection(resolve_runtime_selection(compute_backend="mlx", sim_backend="mac-sim", device="cpu")) + + importlib.import_module("isaaclab") + importlib.import_module("isaaclab.backends.runtime") + importlib.import_module("isaaclab.backends.mac_sim") + importlib.import_module("isaaclab.controllers") + importlib.import_module("isaaclab.controllers.differential_ik_cfg") + importlib.import_module("isaaclab.controllers.operational_space_cfg") + importlib.import_module("isaaclab.sim") + importlib.import_module("isaaclab.sim.schemas") + importlib.import_module("isaaclab.sim.converters") + importlib.import_module("isaaclab.sim.spawners.from_files") + importlib.import_module("isaaclab.utils.nucleus") + importlib.import_module("isaaclab.envs.mdp") + importlib.import_module("isaaclab.envs.mdp.actions") + importlib.import_module("isaaclab.envs.mdp.actions.actions_cfg") + importlib.import_module("isaaclab.envs.mdp.actions.rmpflow_actions_cfg") + importlib.import_module("isaaclab.controllers.rmp_flow_cfg") + importlib.import_module("isaaclab.controllers.config.rmp_flow") + + from isaaclab.sim.spawners.from_files import GroundPlaneCfg + + cfg = GroundPlaneCfg() + assert isinstance(cfg, GroundPlaneCfg) + assert cfg.usd_path.endswith("default_environment.usd") + + +def test_torch_compute_backend_routes_device_seed_and_checkpoint(monkeypatch: pytest.MonkeyPatch, tmp_path): + """Torch compute adapter should forward device, seed, and checkpoint I/O to torch.""" + calls: list[tuple[str, object]] = [] + payload = {"value": 123} + + class FakeTorch: + class cuda: + @staticmethod + def set_device(device): + calls.append(("set_device", device)) + + @staticmethod + def manual_seed_all(seed): + calls.append(("manual_seed_all", seed)) + + @staticmethod + def manual_seed(seed): + calls.append(("manual_seed", seed)) + + @staticmethod + def save(data, path): + calls.append(("save", path)) + path.write_text(str(data), encoding="utf-8") + + @staticmethod + def load(path, map_location=None): + calls.append(("load", (path, map_location))) + return {"loaded": path.read_text(encoding="utf-8")} + + monkeypatch.setitem(sys.modules, "torch", FakeTorch) + backend = TorchCudaComputeBackend() + checkpoint = tmp_path / "torch_ckpt.txt" + + backend.configure_device("cuda:1") + backend.seed(9) + backend.save_checkpoint(checkpoint, payload) + loaded = backend.load_checkpoint(checkpoint) + + assert calls[0] == ("set_device", "cuda:1") + assert ("manual_seed", 9) in calls + assert ("manual_seed_all", 9) in calls + assert ("save", checkpoint) in calls + assert ("load", (checkpoint, "cpu")) in calls + assert loaded == {"loaded": str(payload)} + + +def test_mlx_compute_backend_seed_and_checkpoint(tmp_path): + """MLX compute adapter should seed the backend and round-trip checkpoints.""" + pytest.importorskip("mlx.core") + backend = MlxComputeBackend() + checkpoint = tmp_path / "mlx_ckpt.pkl" + payload = {"episode": 4, "reward": 1.5} + + backend.seed(17) + backend.save_checkpoint(str(checkpoint), payload) + loaded = backend.load_checkpoint(str(checkpoint)) + + assert loaded == payload + + def test_isaacsim_backend_proxies_simulation_and_articulation_calls(): """IsaacSim backend should mirror the upstream simulation/articulation surface used by cartpole.""" calls: list[tuple[str, object]] = [] @@ -185,3 +400,66 @@ def test_mac_sim_can_load_simulation_cfg_but_not_simulation_context(): assert sim.SimulationCfg.__name__ == "SimulationCfg" with pytest.raises(UnsupportedBackendError, match="sim-backend=isaacsim"): _ = sim.SimulationContext + + +def test_mac_sim_asset_exports_fail_with_clear_backend_error(): + """Import-time asset access on mac-sim should raise a backend error instead of importing Omniverse modules.""" + set_runtime_selection(resolve_runtime_selection(compute_backend="mlx", sim_backend="mac-sim", device="cpu")) + assets = importlib.import_module("isaaclab.assets") + + with pytest.raises(UnsupportedBackendError, match="sim-backend=isaacsim"): + _ = assets.Articulation + + +def test_mac_sim_sensor_exports_fail_with_clear_backend_error(): + """Import-time sensor access on mac-sim should raise a backend error instead of importing Omniverse modules.""" + set_runtime_selection(resolve_runtime_selection(compute_backend="mlx", sim_backend="mac-sim", device="cpu")) + sensors = importlib.import_module("isaaclab.sensors") + + with pytest.raises(UnsupportedBackendError, match="sim-backend=isaacsim"): + _ = sensors.Camera + + +def test_mac_sim_manager_exports_fail_with_clear_backend_error(): + """Import-time manager access on mac-sim should raise a backend error instead of importing Omniverse modules.""" + set_runtime_selection(resolve_runtime_selection(compute_backend="mlx", sim_backend="mac-sim", device="cpu")) + managers = importlib.import_module("isaaclab.managers") + + with pytest.raises(UnsupportedBackendError, match="sim-backend=isaacsim"): + _ = managers.ActionManager + + +def test_mac_sim_controller_exports_fail_with_clear_backend_error(): + """Import-time controller access on mac-sim should raise a backend error instead of importing Omniverse modules.""" + set_runtime_selection(resolve_runtime_selection(compute_backend="mlx", sim_backend="mac-sim", device="cpu")) + controllers = importlib.import_module("isaaclab.controllers") + + with pytest.raises(UnsupportedBackendError, match="sim-backend=isaacsim"): + _ = controllers.DifferentialIKController + + +def test_mac_sim_device_exports_fail_with_clear_backend_error(): + """Import-time device access on mac-sim should raise a backend error instead of importing Omniverse modules.""" + set_runtime_selection(resolve_runtime_selection(compute_backend="mlx", sim_backend="mac-sim", device="cpu")) + devices = importlib.import_module("isaaclab.devices") + + with pytest.raises(UnsupportedBackendError, match="sim-backend=isaacsim"): + _ = devices.Se2Keyboard + + +def test_mac_sim_scene_exports_fail_with_clear_backend_error(): + """Import-time scene access on mac-sim should raise a backend error instead of importing Omniverse modules.""" + set_runtime_selection(resolve_runtime_selection(compute_backend="mlx", sim_backend="mac-sim", device="cpu")) + scene = importlib.import_module("isaaclab.scene") + + with pytest.raises(UnsupportedBackendError, match="sim-backend=isaacsim"): + _ = scene.InteractiveScene + + +def test_mac_sim_marker_exports_fail_with_clear_backend_error(): + """Import-time marker access on mac-sim should raise a backend error instead of importing Omniverse modules.""" + set_runtime_selection(resolve_runtime_selection(compute_backend="mlx", sim_backend="mac-sim", device="cpu")) + markers = importlib.import_module("isaaclab.markers") + + with pytest.raises(UnsupportedBackendError, match="sim-backend=isaacsim"): + _ = markers.VisualizationMarkers diff --git a/source/isaaclab_rl/isaaclab_rl/sb3.py b/source/isaaclab_rl/isaaclab_rl/sb3.py index 2177bc625..d653e3943 100644 --- a/source/isaaclab_rl/isaaclab_rl/sb3.py +++ b/source/isaaclab_rl/isaaclab_rl/sb3.py @@ -23,13 +23,24 @@ import gymnasium as gym import numpy as np -import torch -import torch.nn as nn # noqa: F401 -from stable_baselines3.common.preprocessing import is_image_space, is_image_space_channels_first -from stable_baselines3.common.utils import constant_fn -from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvObs, VecEnvStepReturn -from isaaclab.envs import DirectRLEnv, ManagerBasedRLEnv +from isaaclab.backends import UnsupportedBackendError + +try: + import torch +except ModuleNotFoundError: + torch = None + +try: + from stable_baselines3.common.preprocessing import is_image_space, is_image_space_channels_first + from stable_baselines3.common.utils import constant_fn + from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvObs, VecEnvStepReturn + _SB3_IMPORT_ERROR = None +except ModuleNotFoundError as exc: + VecEnv = object + VecEnvObs = Any + VecEnvStepReturn = Any + _SB3_IMPORT_ERROR = exc # remove SB3 warnings because PPO with bigger net actually benefits from GPU warnings.filterwarnings("ignore", message="You are trying to run PPO on the GPU") @@ -53,13 +64,16 @@ def process_sb3_cfg(cfg: dict, num_envs: int) -> dict: https://github.com/DLR-RM/rl-baselines3-zoo/blob/0e5eb145faefa33e7d79c7f8c179788574b20da5/utils/exp_manager.py#L358 """ + _require_sb3() + torch_module = _require_torch() + def update_dict(hyperparams: dict[str, Any], depth: int) -> dict[str, Any]: for key, value in hyperparams.items(): if isinstance(value, dict): update_dict(value, depth + 1) if isinstance(value, str): if value.startswith("nn."): - hyperparams[key] = getattr(nn, value[3:]) + hyperparams[key] = getattr(torch_module.nn, value[3:]) if depth == 0: if key in ["learning_rate", "clip_range", "clip_range_vf"]: if isinstance(value, str): @@ -135,7 +149,7 @@ class Sb3VecEnvWrapper(VecEnv): """ - def __init__(self, env: ManagerBasedRLEnv | DirectRLEnv, fast_variant: bool = True): + def __init__(self, env: Any, fast_variant: bool = True): """Initialize the wrapper. Args: @@ -145,8 +159,11 @@ def __init__(self, env: ManagerBasedRLEnv | DirectRLEnv, fast_variant: bool = Tr Raises: ValueError: When the environment is not an instance of :class:`ManagerBasedRLEnv` or :class:`DirectRLEnv`. """ + _require_sb3() + _require_torch() + supported_env_types = _get_supported_env_types() # check that input is valid - if not isinstance(env.unwrapped, ManagerBasedRLEnv) and not isinstance(env.unwrapped, DirectRLEnv): + if not isinstance(env.unwrapped, supported_env_types): raise ValueError( "The environment must be inherited from ManagerBasedRLEnv or DirectRLEnv. Environment type:" f" {type(env)}" @@ -182,7 +199,7 @@ def class_name(cls) -> str: return cls.__name__ @property - def unwrapped(self) -> ManagerBasedRLEnv | DirectRLEnv: + def unwrapped(self) -> Any: """Returns the base environment of the wrapper. This will be the bare :class:`gymnasium.Env` environment, underneath all layers of wrappers. @@ -217,12 +234,13 @@ def reset(self) -> VecEnvObs: # noqa: D102 return self._process_obs(obs_dict) def step_async(self, actions): # noqa: D102 + torch_module = _require_torch() # convert input to numpy array - if not isinstance(actions, torch.Tensor): + if not isinstance(actions, torch_module.Tensor): actions = np.asarray(actions) - actions = torch.from_numpy(actions).to(device=self.sim_device, dtype=torch.float32) + actions = torch_module.from_numpy(actions).to(device=self.sim_device, dtype=torch_module.float32) else: - actions = actions.to(device=self.sim_device, dtype=torch.float32) + actions = actions.to(device=self.sim_device, dtype=torch_module.float32) # convert to tensor self._async_actions = actions @@ -258,6 +276,7 @@ def close(self): # noqa: D102 self.env.close() def get_attr(self, attr_name, indices=None): # noqa: D102 + torch_module = _require_torch() # resolve indices if indices is None: indices = slice(None) @@ -267,7 +286,7 @@ def get_attr(self, attr_name, indices=None): # noqa: D102 # obtain attribute value attr_val = getattr(self.env, attr_name) # return the value - if not isinstance(attr_val, torch.Tensor): + if not isinstance(attr_val, torch_module.Tensor): return [attr_val] * num_indices else: return attr_val[indices].detach().cpu().numpy() @@ -297,11 +316,13 @@ def get_images(self): # noqa: D102 """ def _process_spaces(self): + _require_sb3() + torch_module = _require_torch() # process observation space observation_space = self.unwrapped.single_observation_space["policy"] if isinstance(observation_space, gym.spaces.Dict): for obs_key, obs_space in observation_space.spaces.items(): - processors: list[callable[[torch.Tensor], Any]] = [] + processors: list[callable[[Any], Any]] = [] # assume normalized, if not, it won't pass is_image_space, which check [0-255]. # for scale like image space that has right shape but not scaled, we will scale it later if is_image_space(obs_space, check_channels=True, normalized_image=True): @@ -314,7 +335,7 @@ def _process_spaces(self): ) # sb3 will handle normalization and transpose, but sb3 expects uint8 images if obs_space.dtype != np.uint8: - processors.append(lambda obs: obs.to(torch.uint8)) + processors.append(lambda obs: obs.to(torch_module.uint8)) observation_space.spaces[obs_key] = gym.spaces.Box(0, 255, obs_space.shape, np.uint8) else: # sb3 will NOT handle the normalization, while sb3 will transpose, its transpose applies to all @@ -322,14 +343,14 @@ def _process_spaces(self): # sb3 transpose it in numpy with cpu. if not is_image_space_channels_first(obs_space): - def tranp(img: torch.Tensor) -> torch.Tensor: + def tranp(img): return img.permute(2, 0, 1) if len(img.shape) == 3 else img.permute(0, 3, 1, 2) processors.append(tranp) h, w, c = obs_space.shape observation_space.spaces[obs_key] = gym.spaces.Box(-1.0, 1.0, (c, h, w), obs_space.dtype) - def chained_processor(obs: torch.Tensor, procs=processors) -> Any: + def chained_processor(obs: Any, procs=processors) -> Any: for proc in procs: obs = proc(obs) return obs @@ -350,6 +371,7 @@ def chained_processor(obs: torch.Tensor, procs=processors) -> Any: def _process_obs(self, obs_dict: torch.Tensor | dict[str, torch.Tensor]) -> np.ndarray | dict[str, np.ndarray]: """Convert observations into NumPy data type.""" + torch_module = _require_torch() # Sb3 doesn't support asymmetric observation spaces, so we only use "policy" obs = obs_dict["policy"] # note: ManagerBasedRLEnv uses torch backend (by default). @@ -358,7 +380,7 @@ def _process_obs(self, obs_dict: torch.Tensor | dict[str, torch.Tensor]) -> np.n if key in self.observation_processors: obs[key] = self.observation_processors[key](value) obs[key] = obs[key].detach().cpu().numpy() - elif isinstance(obs, torch.Tensor): + elif isinstance(obs, torch_module.Tensor): obs = obs.detach().cpu().numpy() else: raise NotImplementedError(f"Unsupported data type: {type(obs)}") @@ -431,3 +453,35 @@ def _process_extras( infos[idx]["terminal_observation"] = None # return list of dictionaries return infos + + +def _require_sb3() -> None: + """Raise a clear error when Stable-Baselines3 is unavailable.""" + if _SB3_IMPORT_ERROR is None: + return + raise ModuleNotFoundError( + "`isaaclab_rl.sb3` requires Stable-Baselines3. Install the package with the `sb3` extra," + " for example `uv pip install -e source/isaaclab_rl[sb3]`." + ) from _SB3_IMPORT_ERROR + + +def _require_torch(): + """Raise a clear error when PyTorch is unavailable.""" + if torch is not None: + return torch + raise ModuleNotFoundError( + "`isaaclab_rl.sb3` requires PyTorch. Install the package with a torch-enabled extra," + " for example `uv pip install -e source/isaaclab_rl[sb3]`." + ) + + +def _get_supported_env_types() -> tuple[type, ...]: + """Resolve Isaac Lab RL env base types lazily for mac-safe imports.""" + try: + from isaaclab.envs import DirectRLEnv, ManagerBasedRLEnv + except Exception as exc: + raise UnsupportedBackendError( + "`isaaclab_rl.sb3` requires Isaac Lab RL environment base classes that are only available on the" + " upstream Isaac Sim path today." + ) from exc + return (ManagerBasedRLEnv, DirectRLEnv) diff --git a/source/isaaclab_rl/isaaclab_rl/skrl.py b/source/isaaclab_rl/isaaclab_rl/skrl.py index 5aba12152..5bb53f167 100644 --- a/source/isaaclab_rl/isaaclab_rl/skrl.py +++ b/source/isaaclab_rl/isaaclab_rl/skrl.py @@ -27,9 +27,9 @@ # needed to import for type hinting: Agent | list[Agent] from __future__ import annotations -from typing import Literal +from typing import Any, Literal -from isaaclab.envs import DirectMARLEnv, DirectRLEnv, ManagerBasedRLEnv +from isaaclab.backends import UnsupportedBackendError """ Vectorized environment wrapper. @@ -37,7 +37,7 @@ def SkrlVecEnvWrapper( - env: ManagerBasedRLEnv | DirectRLEnv | DirectMARLEnv, + env: Any, ml_framework: Literal["torch", "jax", "jax-numpy"] = "torch", wrapper: Literal["auto", "isaaclab", "isaaclab-single-agent", "isaaclab-multi-agent"] = "isaaclab", ): @@ -62,11 +62,8 @@ def SkrlVecEnvWrapper( https://skrl.readthedocs.io/en/latest/api/envs/wrapping.html """ # check that input is valid - if ( - not isinstance(env.unwrapped, ManagerBasedRLEnv) - and not isinstance(env.unwrapped, DirectRLEnv) - and not isinstance(env.unwrapped, DirectMARLEnv) - ): + supported_env_types = _get_supported_env_types() + if not isinstance(env.unwrapped, supported_env_types): raise ValueError( "The environment must be inherited from ManagerBasedRLEnv, DirectRLEnv or DirectMARLEnv. Environment type:" f" {type(env)}" @@ -84,3 +81,15 @@ def SkrlVecEnvWrapper( # wrap and return the environment return wrap_env(env, wrapper) + + +def _get_supported_env_types() -> tuple[type, ...]: + """Resolve Isaac Lab RL env base types lazily for mac-safe imports.""" + try: + from isaaclab.envs import DirectMARLEnv, DirectRLEnv, ManagerBasedRLEnv + except Exception as exc: + raise UnsupportedBackendError( + "`isaaclab_rl.skrl` requires Isaac Lab RL environment base classes that are only available on the" + " upstream Isaac Sim path today." + ) from exc + return (ManagerBasedRLEnv, DirectRLEnv, DirectMARLEnv) diff --git a/source/isaaclab_rl/setup.py b/source/isaaclab_rl/setup.py index f1cf55d32..b3166abb6 100644 --- a/source/isaaclab_rl/setup.py +++ b/source/isaaclab_rl/setup.py @@ -9,7 +9,7 @@ import os import toml -from setuptools import setup +from setuptools import find_packages, setup # Obtain the extension data from the extension.toml file EXTENSION_PATH = os.path.dirname(os.path.realpath(__file__)) @@ -20,8 +20,6 @@ INSTALL_REQUIRES = [ # generic "numpy<2", - "torch>=2.7", - "torchvision>=0.14.1", # ensure compatibility with torch 1.13.1 "protobuf>=4.25.8,!=5.26.0", # configuration management "hydra-core", @@ -37,16 +35,27 @@ ] PYTORCH_INDEX_URL = ["https://download.pytorch.org/whl/cu128"] +TORCH_DEPENDENCIES = [ + "torch>=2.7", + "torchvision>=0.14.1", # ensure compatibility with torch 1.13.1 +] # Extra dependencies for RL agents EXTRAS_REQUIRE = { - "sb3": ["stable-baselines3>=2.6", "tqdm", "rich"], # tqdm/rich for progress bar + "torch": TORCH_DEPENDENCIES, + "sb3": [*TORCH_DEPENDENCIES, "stable-baselines3>=2.6", "tqdm", "rich"], # tqdm/rich for progress bar "skrl": ["skrl>=1.4.3"], "rl-games": [ + *TORCH_DEPENDENCIES, "rl-games @ git+https://github.com/isaac-sim/rl_games.git@python3.11", "gym", ], # rl-games still needs gym :( - "rsl-rl": ["rsl-rl-lib==5.0.1", "onnxscript>=0.5"], # linux aarch 64 requires manual onnxscript installation + "rsl-rl": [ + *TORCH_DEPENDENCIES, + "rsl-rl-lib==5.0.1", + "onnxscript>=0.5", + ], # linux aarch 64 requires manual onnxscript installation + "dev": ["pytest", "pytest-mock"], } # Add the names with hyphens as aliases for convenience EXTRAS_REQUIRE["rl_games"] = EXTRAS_REQUIRE["rl-games"] @@ -71,7 +80,7 @@ install_requires=INSTALL_REQUIRES, dependency_links=PYTORCH_INDEX_URL, extras_require=EXTRAS_REQUIRE, - packages=["isaaclab_rl"], + packages=find_packages(include=["isaaclab_rl", "isaaclab_rl.*"]), classifiers=[ "Natural Language :: English", "Programming Language :: Python :: 3.10", diff --git a/source/isaaclab_rl/test/test_import_safety.py b/source/isaaclab_rl/test/test_import_safety.py new file mode 100644 index 000000000..5099c0ad0 --- /dev/null +++ b/source/isaaclab_rl/test/test_import_safety.py @@ -0,0 +1,21 @@ +# 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 + +"""Import-safety tests for the RL extension on the MLX/mac path.""" + +from __future__ import annotations + +import importlib + +from isaaclab.backends import resolve_runtime_selection, set_runtime_selection + + +def test_isaaclab_rl_modules_import_on_mlx_mac_path(): + """The RL extension root and lightweight wrappers should import without Isaac Sim.""" + set_runtime_selection(resolve_runtime_selection(compute_backend="mlx", sim_backend="mac-sim", device="cpu")) + + importlib.import_module("isaaclab_rl") + importlib.import_module("isaaclab_rl.skrl") + importlib.import_module("isaaclab_rl.sb3") diff --git a/source/isaaclab_tasks/setup.py b/source/isaaclab_tasks/setup.py index 455c56689..843f4e424 100644 --- a/source/isaaclab_tasks/setup.py +++ b/source/isaaclab_tasks/setup.py @@ -8,7 +8,7 @@ import os import toml -from setuptools import setup +from setuptools import find_packages, setup # Obtain the extension data from the extension.toml file EXTENSION_PATH = os.path.dirname(os.path.realpath(__file__)) @@ -42,7 +42,7 @@ python_requires=">=3.10", install_requires=INSTALL_REQUIRES, dependency_links=PYTORCH_INDEX_URL, - packages=["isaaclab_tasks"], + packages=find_packages(include=["isaaclab_tasks", "isaaclab_tasks.*"]), classifiers=[ "Natural Language :: English", "Programming Language :: Python :: 3.10",