diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 1811ca6fd5eb..17e5dcd01062 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -123,6 +123,7 @@ Guidelines for modifications: * Jiwen Cai * Johnson Sun * Juana Du +* Kai Pei * Kaixi Bao * Kourosh Darvish * Kousheek Chakraborty diff --git a/docs/source/_static/css/environment-browser.js b/docs/source/_static/css/environment-browser.js index b40b0f503ec9..e2aef9ecfdf1 100644 --- a/docs/source/_static/css/environment-browser.js +++ b/docs/source/_static/css/environment-browser.js @@ -91,6 +91,7 @@ ["IsaacContrib-Lift-Cube-Franka-IK-Abs", "", "", "", "", "tasks/manipulation/franka_lift.jpg"], ["IsaacContrib-Lift-Cube-Franka-IK-Rel", "", "", "", "", "tasks/manipulation/franka_lift.jpg"], ["IsaacContrib-Lift-Cube-OpenArm", "rl_games,rsl_rl", "", "", "", "tasks/manipulation/openarm_uni_lift.jpg"], + ["IsaacContrib-Multitask-Manipulation", "rsl_rl", "", "", "", "tasks/manipulation/multitask_manipulation.jpg"], ["IsaacContrib-Navigation-3DObstacles-ARL-Robot-1", "rl_games,rsl_rl,skrl", "", "", "", "tasks/drone_arl/arl_robot_1_navigation.jpg"], ["IsaacContrib-Navigation-Flat-AnymalC", "rsl_rl,skrl", "", "", "", "tasks/navigation/anymal_c_nav.jpg"], ["IsaacContrib-NutPour-GR1T2-Pink-IK-Abs", "", "", "", ""], diff --git a/docs/source/features/imitation-learning/humanoids_imitation.rst b/docs/source/features/imitation-learning/humanoids_imitation.rst index 11c6d77efb32..e0f03adec0a8 100644 --- a/docs/source/features/imitation-learning/humanoids_imitation.rst +++ b/docs/source/features/imitation-learning/humanoids_imitation.rst @@ -774,12 +774,11 @@ Integrating 3D Gaussian Splatting into SDG This section extends :ref:`locomanipulation SDG ` -by replacing the synthetic background with a 3D Gaussian Splatting (NuRec) scene. As in the -base pipeline, the workflow takes a manipulation dataset as input and produces a combined -navigation and manipulation dataset as an HDF5 file — but here the robot navigates and -manipulates objects inside a neurally-rendered environment, and an ego-centric camera -captures the result, producing more realistic training data than a purely synthetic scene. -NVIDIA Isaac Sim renders 3DGS models stored as USD assets; see +by replacing the synthetic background with a 3D Gaussian Splatting (NuRec) scene. The workflow +takes a manipulation dataset as input and produces a combined navigation and manipulation +dataset as an HDF5 file. The navigation replay happens in a neurally-reconstructed environment, +and an ego-centric camera captures the result, producing more realistic training data. +For more info on NuRec rendering in Isaac Sim, see `Neural Volume Rendering `__ for details. @@ -805,21 +804,17 @@ Log in to Hugging Face: hf auth login --token -Download the required USDZ stage files and occupancy maps: +Download the required USDZ stage file and occupancy maps: .. code:: bash hf download nvidia/PhysicalAI-Robotics-NuRec \ - hand_hold-voyager-babyboom/stage_volume.usdz \ - hand_hold-voyager-babyboom/stage_particle.usdz \ + hand_hold-voyager-babyboom/particle_sh_optimized.usdz \ hand_hold-voyager-babyboom/occupancy_map.png \ hand_hold-voyager-babyboom/occupancy_map.yaml \ --repo-type dataset \ --local-dir -The sample includes both a volume-based USD (``stage_volume.usdz``) and a particle-field USD -(``stage_particle.usdz``). Either can be used as the background asset. - Asset requirements """""""""""""""""" @@ -831,11 +826,6 @@ compatible with the SDG pipeline: object placement. - An occupancy map is required for path planning. - - If your scene was reconstructed using the `Stereo Workflow `__, - the occupancy map is generated via ``nvblox``. - - If your background includes a mesh, use the `Occupancy Map Generator `__ - to create a map via physical simulation. - Generating the dataset """""""""""""""""""""" @@ -843,13 +833,7 @@ Before proceeding, ensure you have generated a manipulation dataset or downloade dataset provided in the :ref:`Generate the manipulation dataset ` section. -Once you have gathered: - -- A manipulation dataset -- A background USD asset -- A matched occupancy map - -you can run the generation command. At runtime, the script adds a ground plane at ``z=0`` to +At runtime, the script adds a ground plane at ``z=0`` to the scene. It then proceeds through four stages: 1. **Pick**: The robot picks up an object at the start location by replaying the manipulation @@ -873,18 +857,20 @@ Run the generation command: --num_runs 1 \ --lift_step 60 \ --navigate_step 130 \ - --output_file /generated_dataset_g1_locomanipulation_sdg_gaussian_background.hdf5 \ - --visualizer kit \ - --background_usd_path /stage_particle.usdz \ - --background_occupancy_yaml_file /occupancy_map.yaml \ - --randomize_placement \ + --output_file /generated_dataset_g1_locomanipulation_sdg_nurec.hdf5 \ + --background_usd_path /hand_hold-voyager-babyboom/particle_sh_optimized.usdz \ + --background_occupancy_yaml_file /hand_hold-voyager-babyboom/occupancy_map.yaml \ --high_res_video +To inspect the run in the GUI, add ``--viz kit``; adjust the perspective camera, or switch +the viewport to the robot camera view for clearer visualization. + The key parameters are: -- ``--background_usd_path``: Path to the 3D Gaussian background USD asset. -- ``--background_occupancy_yaml_file``: Path to the occupancy map file. -- ``--high_res_video``: Capture the ego-centric camera at 960×540 instead of the default +- ``--background_usd_path``: Path to the 3D Gaussian NuRec USD asset. +- ``--background_occupancy_yaml_file``: Path to the NuRec occupancy map file. +- Placement randomization is enabled automatically when a NuRec scene is used. +- ``--high_res_video``: Capture the ego-centric camera at 512×320 instead of the default 256×160. When the run completes successfully, an HDF5 dataset is generated containing camera @@ -892,23 +878,15 @@ observations. You can convert the ego-centric camera view to MP4: .. code:: bash - uv run python scripts/tools/hdf5_to_mp4.py \ - --input_file /generated_dataset_g1_locomanipulation_sdg_gaussian_background.hdf5 \ + uv run --extra mimic python scripts/tools/hdf5_to_mp4.py \ + --input_file /generated_dataset_g1_locomanipulation_sdg_nurec.hdf5 \ --output_dir / \ --input_keys robot_pov_cam \ - --video_width 960 \ - --video_height 540 + --video_width 512 \ + --video_height 320 Set ``--video_width`` and ``--video_height`` to match the resolution captured during -generation: 960×540 with ``--high_res_video``, or 256×160 without it. - -To play the generated MP4 video on Ubuntu, install the following multimedia packages: - -.. code:: bash - - sudo apt update - sudo apt install libavcodec-extra gstreamer1.0-libav gstreamer1.0-plugins-ugly - +generation: 512×320 with ``--high_res_video``, or 256×160 without it. .. figure:: https://download.isaacsim.omniverse.nvidia.com/isaaclab/images/locomanipulation_sdg_gaussian_background_2x.webp :width: 100% diff --git a/docs/source/refs/issues.rst b/docs/source/refs/issues.rst index 09f574ffe9d9..8d17d43e06cc 100644 --- a/docs/source/refs/issues.rst +++ b/docs/source/refs/issues.rst @@ -118,22 +118,13 @@ OpenUSD provider that contains the fix or to Isaac Sim 6.1 or later. .. _known-issues-closed-loop-newton: -Closed-loop articulations are not available on Newton (e.g. Agility Digit) +Closed-loop articulations are not validated on Kamino (e.g. Agility Digit) ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -**Affects:** ``physics=newton_mjwarp``, ``physics=newton_kamino``. +**Affects:** ``physics=newton_kamino``. Robots whose USD encodes a closed kinematic loop — such as the achilles rod and toe push-rods -on the Agility Digit — are not currently validated on the Newton backends. The Digit-based -contrib tasks are PhysX-only and do not expose a Newton preset at all: - -* ``IsaacContrib-Velocity-Flat-Digit`` -* ``IsaacContrib-Velocity-Rough-Digit`` -* ``IsaacContrib-Tracking-LocoManip-Digit`` - -Passing ``presets=newton_mjwarp`` to these tasks is rejected, because the preset never -selected a Newton backend on them; it only stripped a center-of-mass randomization from a -PhysX run. Use the default PhysX configuration for Digit-based environments. +on the Agility Digit — are not validated on ``newton_kamino``. Renderers diff --git a/scripts/benchmarks/benchmark_renderer.py b/scripts/benchmarks/benchmark_renderer.py new file mode 100644 index 000000000000..4cb93e3d032c --- /dev/null +++ b/scripts/benchmarks/benchmark_renderer.py @@ -0,0 +1,360 @@ +# 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 + +# +# Command to run: +# uv run --no-sync python scripts/benchmarks/benchmark_renderer.py [PROFILE] +# + +import argparse +import fnmatch +import json +import os +import re +import shutil +import site +import statistics +import subprocess +import sys +from pathlib import Path + +PROFILES = [ + { + "name": "ovrtx_constant_diffuse_oldpipe", + "preset": "ovrtx_renderer,simple_shading_constant_diffuse", + "settings": {"min-pipe": False}, + }, + { + "name": "ovrtx_constant_diffuse_newpipe", + "preset": "ovrtx_renderer,simple_shading_constant_diffuse", + "settings": {"min-pipe": True}, + }, + { + "name": "ovrtx_diffuse_mdl_oldpipe", + "preset": "ovrtx_renderer,simple_shading_diffuse_mdl", + "settings": {"min-pipe": False}, + }, + { + "name": "ovrtx_diffuse_mdl_newpipe", + "preset": "ovrtx_renderer,simple_shading_diffuse_mdl", + "settings": {"min-pipe": True}, + }, + { + "name": "ovrtx_full_mdl_oldpipe", + "preset": "ovrtx_renderer,simple_shading_full_mdl", + "settings": {"min-pipe": False}, + }, + { + "name": "ovrtx_full_mdl_newpipe", + "preset": "ovrtx_renderer,simple_shading_full_mdl", + "settings": {"min-pipe": True}, + }, + {"name": "newton_lbvh_lbvh", "preset": "newton_renderer,rgb", "settings": {"tlas": "lbvh", "blas": "lbvh"}}, + {"name": "newton_lbvh_sah", "preset": "newton_renderer,rgb", "settings": {"tlas": "lbvh", "blas": "sah"}}, + {"name": "newton_lbvh_cubql", "preset": "newton_renderer,rgb", "settings": {"tlas": "lbvh", "blas": "cubql"}}, + {"name": "newton_sah_lbvh", "preset": "newton_renderer,rgb", "settings": {"tlas": "sah", "blas": "lbvh"}}, + {"name": "newton_sah_sah", "preset": "newton_renderer,rgb", "settings": {"tlas": "sah", "blas": "sah"}}, + {"name": "newton_sah_cubql", "preset": "newton_renderer,rgb", "settings": {"tlas": "sah", "blas": "cubql"}}, +] + +TASK_NAME = "Isaac-RenderBenchmark-Franka-Cabinet" +FRAME_PADDING = 5 + +# Resolved from this file rather than the working directory, so the script runs from anywhere. +SCRIPT_DIR = Path(__file__).resolve().parent +RUNTIME_SCRIPT = SCRIPT_DIR / "runtime.py" +OUTPUT_PATH = str(SCRIPT_DIR.parent.parent / "benchmarks") + +OVRTX_RENDERER = "ovrtx_renderer" +NEWTON_RENDERER = "newton_renderer" + +RENDER_SCOPE = "IsaacLab::Renderer::render" +"""Backend-agnostic timer name around ``BaseRenderer.render``, enabled by ``ISAACLAB_RENDER_PROFILE``. + +See :data:`isaaclab.renderers.render_context.RENDER_PROFILE_SCOPE`. It brackets the render alone, +excluding the scene-state sync before it and the output readback after it. ``wp.ScopedTimer`` prints +one ``" took X.XX ms"`` line per call, which :func:`parse_log` regexes out of the run's log. +""" + +RENDER_SCOPE_PATTERN = re.compile(rf"{re.escape(RENDER_SCOPE)} took ([\d.]+) ms") + +log_stream = sys.stdout +"""Destination for progress and diagnostics. ``--json`` points it at stderr so stdout holds only JSON.""" + + +def log(message: str = "") -> None: + """Write one human-readable line to :data:`log_stream`.""" + print(message, file=log_stream) + + +def pixels_per_second(median_ms: float, num_envs: int, resolution: int) -> float: + """Rendered pixels per second [px/s] implied by a median frame time. + + Args: + median_ms: Median time to render one frame [ms]. + num_envs: Number of environments rendered per frame. + resolution: Width and height of each environment's tile [px]. + + Returns: + Throughput in pixels per second [px/s]. + """ + return num_envs * resolution * resolution * 1000.0 / median_ms + + +def build_record(profile: dict, results: dict | None, num_envs: int, resolution: int) -> dict: + """Summarize one profile's outcome for reporting. + + Args: + profile: Profile entry from :data:`PROFILES`. + results: Timing statistics from :func:`parse_log`, or ``None`` if the run failed. + num_envs: Number of environments the profile rendered. + resolution: Width and height of each environment's tile [px]. + + Returns: + A record carrying the profile's identity plus either its timings or the log to inspect. + """ + record = {"name": profile["name"], "preset": profile["preset"], "settings": profile["settings"]} + if not results: + return record | {"status": "failed", "log": os.path.join(OUTPUT_PATH, profile["name"] + ".log")} + return ( + record + | { + "status": "ok", + "size": results["size"], + "pixels_per_second": pixels_per_second(results["median"], num_envs, resolution), + } + | {f"{key}_ms": results[key] for key in ("median", "mean", "min", "max", "stdev")} + ) + + +def parse_log(filename: str, num_frames: int): + """Summarize per-frame render times [ms] from a run's captured log. + + Every backend is measured the same way: the wall time of :data:`RENDER_SCOPE`, printed once per + render by ``wp.ScopedTimer`` when ``ISAACLAB_RENDER_PROFILE`` is set. The timer synchronizes the + device on both ends, so it covers completed rather than merely submitted work — including for + the RTX backends, whose Vulkan render is consumed by warp extraction kernels inside the scope. + + Args: + filename: Path to the captured run log. + num_frames: Number of frames to measure, after skipping :data:`FRAME_PADDING` warm-up frames. + + Returns: + Timing statistics, or ``None`` if the log holds no usable frames. + """ + with open(filename) as file: + frames = [float(match.group(1)) for line in file if (match := RENDER_SCOPE_PATTERN.search(line))] + + if not frames: + log(f"No '{RENDER_SCOPE}' timings in {filename}; was ISAACLAB_RENDER_PROFILE set for this run?") + return None + out = frames[FRAME_PADDING : FRAME_PADDING + num_frames] + + if out: + return { + "size": len(out), + "median": statistics.median(out), + "mean": statistics.mean(out), + "min": min(out), + "max": max(out), + "stdev": statistics.stdev(out) if len(out) > 1 else 0, + } + return None + + +def run_profile(profile: dict, args: argparse.Namespace): + """Run one entry of :data:`PROFILES` and summarize its render times from the captured log. + + Args: + profile: Profile entry naming the preset and its backend settings. + args: Parsed command-line arguments. + + Returns: + Timing statistics from :func:`parse_log`, or ``False`` if the run failed. + """ + warp_cache_path = os.path.join(OUTPUT_PATH, "warp-cache") + + env = { + "NEWTON_USE_CUDA_GRAPH": "0", + "ISAACLAB_RENDER_PROFILE": "1", + "BENCHMARK_SAVE_IMAGE": "1" if args.save_image else "0", + "BENCHMARK_RENDER_RESOLUTION": f"{args.resolution}", + "WARP_CACHE_PATH": warp_cache_path, + } + + # Kernel compilation lands in the warm-up frames that FRAME_PADDING discards, so a warm cache + # does not reach the measured window -- it only saves a recompile per profile. + if not args.keep_warp_cache and os.path.exists(warp_cache_path): + shutil.rmtree(warp_cache_path) + + profile_name: str = profile["name"] + preset: str = profile["preset"] + renderer: str = preset.split(",")[0] + + if renderer == OVRTX_RENDERER: + env["LD_PRELOAD"] = os.path.join( + site.getsitepackages()[0], "ovrtx/bin/plugins/omni.client.lib/libomniclient.so" + ) + env["CUDA_VISIBLE_DEVICES"] = "0" + env["OMNI_KIT_ACCEPT_EULA"] = "YES" + env["OVRTX_rtx_post_tonemap_op"] = "0" + env["OVRTX_rtx_minimal_useMinimalPipeline"] = "1" if profile["settings"]["min-pipe"] else "0" + + if renderer == NEWTON_RENDERER: + env["NEWTON_BVH_SCENE"] = profile["settings"]["tlas"] + env["NEWTON_BVH_GEOMETRY"] = profile["settings"]["blas"] + + os.makedirs(OUTPUT_PATH, exist_ok=True) + log_filename = os.path.join(OUTPUT_PATH, profile_name + ".log") + + cmd = [ + sys.executable, + str(RUNTIME_SCRIPT), + "--task", + args.task, + "--num_envs", + f"{args.num_envs}", + "--warmup_steps", + "0", + "--num_steps", + f"{args.num_frames + FRAME_PADDING * 2}", + "--output_path", + OUTPUT_PATH, + f"presets={preset}", + ] + + with open(log_filename, "w") as file: + process = subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env=dict(os.environ) | env, text=True + ) + + for line in process.stdout: + file.write(line) + if args.verbose: + log_stream.write(f"\x1b[90m{line}\x1b[0m") + log_stream.flush() + + return_code = process.wait() + if return_code != 0: + log(f"Failed with exit code {process.returncode}, see {log_filename} for details.") + return False + + return parse_log(log_filename, args.num_frames) + + +def _build_arg_parser() -> argparse.ArgumentParser: + """Build the CLI parser, kept separate from module import so tests can import pure helpers.""" + parser = argparse.ArgumentParser("IsaacLab Benchmark: Sweep Franka Cabinet") + parser.add_argument("--num_frames", type=int, default="20", help="Number of frames to render") + parser.add_argument("--num_envs", type=int, default="1024", help="Number of environments to render") + parser.add_argument("--resolution", type=int, default="256", help="Render resolution") + parser.add_argument("--task", default=TASK_NAME, help="Gym task id to profile") + parser.add_argument( + "--keep_warp_cache", + action="store_true", + help="Reuse the warp kernel cache instead of recompiling per profile", + ) + parser.add_argument("-v", "--verbose", action="store_true", help="Verbose output") + parser.add_argument("-s", "--save_image", action="store_true", help="Save image for debugging purposes") + parser.add_argument("--json", action="store_true", help="Write results to stdout as JSON instead of a table") + parser.add_argument("profile", nargs="*", help="Profiles to run") + return parser + + +def main() -> None: + """Parse CLI arguments, run the selected profiles, and report their results.""" + global log_stream + + args = _build_arg_parser().parse_args() + + # Keep stdout free of anything but the JSON document. + if args.json: + log_stream = sys.stderr + + if not args.profile: + if args.json: + print(json.dumps({"available_profiles": [profile["name"] for profile in PROFILES]}, indent=2)) + else: + log("Available profiles:") + for profile in PROFILES: + log(" " + profile["name"]) + exit(0) + + matched_names = set() + for profile_name in args.profile: + matches = [profile["name"] for profile in PROFILES if fnmatch.fnmatch(profile["name"], profile_name)] + if not matches: + print(f"No profile found matching: {profile_name}", file=sys.stderr) + exit(1) + matched_names.update(matches) + + # Run in declaration order so a given selection always reports in the same order. + selected_profiles = [profile for profile in PROFILES if profile["name"] in matched_names] + + all_results = {} + for profile in selected_profiles: + log(f"profile: {profile['name']}") + log(f" preset: {profile['preset']}") + for key, value in profile["settings"].items(): + log(f" {key}: {value}") + + try: + all_results[profile["name"]] = run_profile(profile, args) + except KeyboardInterrupt: + break + + if results := all_results[profile["name"]]: + log(f" size: {results['size']}") + for key, value in results.items(): + if key != "size": + log(f" {key}: {value:.2f}ms") + log("") + + # A KeyboardInterrupt leaves the remaining profiles unrun; report only what completed. + records = [ + build_record(profile, all_results[profile["name"]], args.num_envs, args.resolution) + for profile in selected_profiles + if profile["name"] in all_results + ] + benchmark_failed = any(record["status"] == "failed" for record in records) + + if args.json: + print( + json.dumps( + { + "task": args.task, + "num_envs": args.num_envs, + "num_frames": args.num_frames, + "resolution": args.resolution, + "profiles": records, + }, + indent=2, + ) + ) + else: + separator = "|------------------------------------------|------|--------------|--------------|--------------|--------------|--------------|--------------|" # noqa: E501 + log("") + log( + "| PROFILE | SIZE | PIXEL / SEC | MEDIAN | MEAN | MIN | MAX | STDEV |" # noqa: E501 + ) + log(separator) + for record in records: + if record["status"] == "ok": + gpxs = record["pixels_per_second"] / 1e9 + log( + f"| {record['name']:<40} | {record['size']:>4} | {gpxs:>6.2f} Gpx/s | {record['median_ms']:>10.2f}ms | {record['mean_ms']:>10.2f}ms | {record['min_ms']:>10.2f}ms | {record['max_ms']:>10.2f}ms | {record['stdev_ms']:>10.2f}ms |" # noqa: E501 + ) + else: + log(f"| {record['name']:<40} | FAILED {record['log']:<87} |") + log(separator) + log("") + + if benchmark_failed: + exit(1) + exit(0) + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmarks/test/test_benchmark_renderer.py b/scripts/benchmarks/test/test_benchmark_renderer.py new file mode 100644 index 000000000000..05c57d63ff5b --- /dev/null +++ b/scripts/benchmarks/test/test_benchmark_renderer.py @@ -0,0 +1,136 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Unit tests for the pure helpers in ``scripts/benchmarks/benchmark_renderer.py``. + +The script drives a rendering backend end-to-end, so that path is not covered here. These tests +exercise the record-building, log-parsing, and CLI-parsing logic that does not require a GPU. +""" + +import importlib.util +import json +import subprocess +import sys +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parents[3] +SCRIPT_PATH = ROOT / "scripts" / "benchmarks" / "benchmark_renderer.py" + + +def _load_module(): + """Import ``benchmark_renderer.py`` as a module, bypassing its CLI entry point.""" + spec = importlib.util.spec_from_file_location("benchmark_renderer", SCRIPT_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def benchmark_renderer(): + return _load_module() + + +def test_render_scope_matches_render_context(benchmark_renderer): + """The script's timer name must match the one the renderer prints, or nothing is parsed.""" + from isaaclab.renderers.render_context import RENDER_PROFILE_SCOPE + + assert benchmark_renderer.RENDER_SCOPE == RENDER_PROFILE_SCOPE + + +def test_pixels_per_second(benchmark_renderer): + """Throughput is env count times tile area, scaled from milliseconds to seconds.""" + assert benchmark_renderer.pixels_per_second(median_ms=1000.0, num_envs=4, resolution=256) == pytest.approx( + 4 * 256 * 256 + ) + + +def test_build_record_ok(benchmark_renderer): + profile = {"name": "p", "preset": "newton_renderer,rgb", "settings": {"tlas": "sah"}} + results = {"size": 20, "median": 2.0, "mean": 2.1, "min": 1.5, "max": 3.0, "stdev": 0.2} + + record = benchmark_renderer.build_record(profile, results, num_envs=4, resolution=256) + + assert record["status"] == "ok" + assert record["size"] == 20 + assert record["median_ms"] == 2.0 + assert record["pixels_per_second"] == pytest.approx(benchmark_renderer.pixels_per_second(2.0, 4, 256)) + + +def test_build_record_failed(benchmark_renderer): + profile = {"name": "p", "preset": "newton_renderer,rgb", "settings": {"tlas": "sah"}} + + record = benchmark_renderer.build_record(profile, None, num_envs=4, resolution=256) + + assert record == { + "name": "p", + "preset": "newton_renderer,rgb", + "settings": {"tlas": "sah"}, + "status": "failed", + "log": str(Path(benchmark_renderer.OUTPUT_PATH) / "p.log"), + } + + +_RENDER_SCOPE = "IsaacLab::Renderer::render" + + +def _write_log(path: Path, timings_ms: list[float]) -> None: + """Write a synthetic run log with one ``wp.ScopedTimer`` print line per timing. + + Interleaves unrelated lines to mimic real subprocess output (warp init banner, other timers), + so the parser is exercised against noise rather than a file with only matching lines. + """ + lines = ["Warp 1.17.0 initialized:", "SomeOtherScope took 0.10 ms"] + for value in timings_ms: + lines.append(f"{_RENDER_SCOPE} took {value:.2f} ms") + path.write_text("\n".join(lines) + "\n") + + +def test_parse_log_skips_padding_and_keeps_num_frames(benchmark_renderer, tmp_path): + """Only the ``num_frames`` timings after :data:`FRAME_PADDING` warm-up frames are summarized.""" + log_path = tmp_path / "profile.log" + padding = benchmark_renderer.FRAME_PADDING + # padding warm-up frames of 1ms each, then 4 measured frames of 2ms each (only 3 are kept). + timings = [1.0] * padding + [2.0] * 4 + _write_log(log_path, timings) + + results = benchmark_renderer.parse_log(str(log_path), num_frames=3) + + assert results["size"] == 3 + assert results["median"] == pytest.approx(2.0) + assert results["min"] == pytest.approx(2.0) + assert results["max"] == pytest.approx(2.0) + + +def test_parse_log_returns_none_without_matching_lines(benchmark_renderer, tmp_path): + """A log with no ``RENDER_SCOPE`` timings means profiling was never enabled for that run.""" + log_path = tmp_path / "profile.log" + _write_log(log_path, []) + + assert benchmark_renderer.parse_log(str(log_path), num_frames=3) is None + + +def _run_cli(args: list[str]) -> subprocess.CompletedProcess: + return subprocess.run( + [sys.executable, str(SCRIPT_PATH), *args], capture_output=True, text=True, cwd=ROOT, timeout=30 + ) + + +def test_cli_lists_available_profiles_as_json(benchmark_renderer): + """With no profile selected, the CLI reports every declared profile name and exits cleanly.""" + result = _run_cli(["--json"]) + + assert result.returncode == 0 + payload = json.loads(result.stdout) + assert payload["available_profiles"] == [profile["name"] for profile in benchmark_renderer.PROFILES] + + +def test_cli_rejects_unmatched_profile_glob(): + """An unmatched profile pattern fails fast, before any profiling subprocess is launched.""" + result = _run_cli(["does-not-exist-*"]) + + assert result.returncode == 1 + assert "No profile found matching: does-not-exist-*" in result.stderr diff --git a/scripts/demos/bin_packing.py b/scripts/demos/bin_packing.py index ea14401cb6c7..3f8cf063e2a0 100644 --- a/scripts/demos/bin_packing.py +++ b/scripts/demos/bin_packing.py @@ -54,6 +54,7 @@ from random import Random import torch +import warp as wp import isaaclab.sim as sim_utils import isaaclab.utils.math as math_utils @@ -175,18 +176,18 @@ def spawn_grocery( return prim -def grocery_cfg(slot: int) -> RigidObjectCfg: +def grocery_cfg(slot: int) -> AssetBaseCfg: """Declare the grocery that occupies one slot of the spawn grid. Args: slot: Index into the spawn grid built by :func:`grocery_spawn_poses`. Returns: - The rigid-object configuration for that slot. + The spawn-only asset configuration for that slot. """ - return RigidObjectCfg( + return AssetBaseCfg( prim_path=f"{{ENV_REGEX_NS}}/Groceries/Grocery_{slot:02d}", - init_state=RigidObjectCfg.InitialStateCfg(pos=tuple(grocery_spawn_poses()[slot, :3].tolist())), + init_state=AssetBaseCfg.InitialStateCfg(pos=tuple(grocery_spawn_poses()[slot, :3].tolist())), spawn=sim_utils.UsdFileCfg(func=spawn_grocery, usd_path=GROCERY_USDS[slot % len(GROCERY_USDS)]), ) @@ -248,30 +249,30 @@ class BinPackingSceneCfg(InteractiveSceneCfg): ) # grocery slots, one per object that may end up in the bin - grocery_00: RigidObjectCfg = grocery_cfg(0) - grocery_01: RigidObjectCfg = grocery_cfg(1) - grocery_02: RigidObjectCfg = grocery_cfg(2) - grocery_03: RigidObjectCfg = grocery_cfg(3) - grocery_04: RigidObjectCfg = grocery_cfg(4) - grocery_05: RigidObjectCfg = grocery_cfg(5) - grocery_06: RigidObjectCfg = grocery_cfg(6) - grocery_07: RigidObjectCfg = grocery_cfg(7) - grocery_08: RigidObjectCfg = grocery_cfg(8) - grocery_09: RigidObjectCfg = grocery_cfg(9) - grocery_10: RigidObjectCfg = grocery_cfg(10) - grocery_11: RigidObjectCfg = grocery_cfg(11) - grocery_12: RigidObjectCfg = grocery_cfg(12) - grocery_13: RigidObjectCfg = grocery_cfg(13) - grocery_14: RigidObjectCfg = grocery_cfg(14) - grocery_15: RigidObjectCfg = grocery_cfg(15) - grocery_16: RigidObjectCfg = grocery_cfg(16) - grocery_17: RigidObjectCfg = grocery_cfg(17) - grocery_18: RigidObjectCfg = grocery_cfg(18) - grocery_19: RigidObjectCfg = grocery_cfg(19) - grocery_20: RigidObjectCfg = grocery_cfg(20) - grocery_21: RigidObjectCfg = grocery_cfg(21) - grocery_22: RigidObjectCfg = grocery_cfg(22) - grocery_23: RigidObjectCfg = grocery_cfg(23) + grocery_00: AssetBaseCfg = grocery_cfg(0) + grocery_01: AssetBaseCfg = grocery_cfg(1) + grocery_02: AssetBaseCfg = grocery_cfg(2) + grocery_03: AssetBaseCfg = grocery_cfg(3) + grocery_04: AssetBaseCfg = grocery_cfg(4) + grocery_05: AssetBaseCfg = grocery_cfg(5) + grocery_06: AssetBaseCfg = grocery_cfg(6) + grocery_07: AssetBaseCfg = grocery_cfg(7) + grocery_08: AssetBaseCfg = grocery_cfg(8) + grocery_09: AssetBaseCfg = grocery_cfg(9) + grocery_10: AssetBaseCfg = grocery_cfg(10) + grocery_11: AssetBaseCfg = grocery_cfg(11) + grocery_12: AssetBaseCfg = grocery_cfg(12) + grocery_13: AssetBaseCfg = grocery_cfg(13) + grocery_14: AssetBaseCfg = grocery_cfg(14) + grocery_15: AssetBaseCfg = grocery_cfg(15) + grocery_16: AssetBaseCfg = grocery_cfg(16) + grocery_17: AssetBaseCfg = grocery_cfg(17) + grocery_18: AssetBaseCfg = grocery_cfg(18) + grocery_19: AssetBaseCfg = grocery_cfg(19) + grocery_20: AssetBaseCfg = grocery_cfg(20) + grocery_21: AssetBaseCfg = grocery_cfg(21) + grocery_22: AssetBaseCfg = grocery_cfg(22) + grocery_23: AssetBaseCfg = grocery_cfg(23) # A slot claimed by some layout but active in none of the layouts the environments # actually receive is never spawned, which the scene rejects. Environments cycle @@ -297,30 +298,33 @@ class BinPackingSceneCfg(InteractiveSceneCfg): def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene) -> None: """Runs the simulation loop that coordinates spawn randomization and stepping. - A heterogeneous scene gives each grocery slot its own physics view whose - instance count is the number of environments containing that slot, not - ``num_envs``. + A flat physics view manages every grocery present in the heterogeneous scene. Returns: None: The simulator side-effects are applied through ``scene`` and ``sim``. """ device = scene.device - # Resolve each spawned slot once: (asset, env ids owning an instance, world spawn poses). - # The spawn pose is the asset's default (init-state) pose, which carries the jittered - # grid position sampled by :func:`grocery_spawn_poses` at scene-config build time. - groceries = [] - for name in GROCERY_NAMES: - asset = scene[name] - prim_paths = getattr(asset.root_view, "prim_paths", None) - if prim_paths is not None and len(prim_paths) == asset.num_instances: - env_ids = [int(path.partition("/env_")[2].partition("/")[0]) for path in prim_paths] - else: - # Homogeneous fallback: one instance per environment. - env_ids = range(asset.num_instances) - env_ids = torch.tensor(list(env_ids), dtype=torch.long, device=device) - spawn_pose_w = asset.data.default_root_pose.torch.clone() - spawn_pose_w[:, :3] += scene.env_origins[env_ids] - groceries.append((asset, env_ids, spawn_pose_w)) + physics_sim_view = sim.physics_manager.get_physics_sim_view() + root_view = physics_sim_view.create_rigid_body_view("/World/envs/env_*/Groceries/Grocery_*") + if root_view.count == 0: + raise RuntimeError("The flat grocery view matched no rigid bodies.") + + env_ids = [] + slot_ids = [] + for path in root_view.prim_paths: + env_ids.append(int(path.partition("/env_")[2].partition("/")[0])) + slot_ids.append(int(path.rpartition("/Grocery_")[2])) + env_ids_t = torch.tensor(env_ids, dtype=torch.long, device=device) + slot_ids_t = torch.tensor(slot_ids, dtype=torch.long, device=device) + env_origins = scene.env_origins[env_ids_t] + spawn_pose_w = grocery_spawn_poses(device)[slot_ids_t] + spawn_pose_w[:, :3] += env_origins + + all_indices_t = torch.arange(root_view.count, dtype=torch.int32, device=device) + cpu_indices_t = torch.arange(root_view.count, dtype=torch.int32, device="cpu") + all_indices = wp.from_torch(all_indices_t) + cpu_indices = wp.from_torch(cpu_indices_t) + print(f"[INFO] Grocery RigidBodyView: 1 flat view managing {root_view.count} rigid bodies") pose_ranges = torch.tensor([POSE_RANGE[axis] for axis in ("roll", "pitch", "yaw")], device=device) velocity_ranges = torch.tensor([VELOCITY_RANGE[axis] for axis in ("roll", "pitch", "yaw")], device=device) @@ -335,24 +339,22 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene) -> count = 0 # Retrieve positions with Timer("[INFO] Time to reset scene: "): - for asset, env_ids, spawn_pose_w in groceries: - num_instances = len(env_ids) - # Only randomize the [roll, pitch, yaw] rotation; keep the spawn position fixed. - noise = math_utils.sample_uniform( - pose_ranges[:, 0], pose_ranges[:, 1], (num_instances, 3), device=device - ) - pose = spawn_pose_w.clone() - noise_quat = math_utils.quat_from_euler_xyz(noise[:, 0], noise[:, 1], noise[:, 2]) - pose[:, 3:7] = math_utils.quat_mul(pose[:, 3:7], noise_quat) - asset.write_root_pose_to_sim_index(root_pose=pose) - # Only randomize the angular velocity; leave the linear velocity at zero. - root_velocity = torch.zeros((num_instances, 6), device=device) - root_velocity[:, 3:6] = math_utils.sample_uniform( - velocity_ranges[:, 0], velocity_ranges[:, 1], (num_instances, 3), device=device - ) - asset.write_root_velocity_to_sim_index(root_velocity=root_velocity) - # Vary the mass so the drop behavior differs between resets. - asset.set_masses_index(masses=torch.rand((num_instances, 1), device=device) * 0.2 + 0.2) + shape = (root_view.count,) + # Only randomize the [roll, pitch, yaw] rotation; keep the spawn position fixed. + noise = math_utils.sample_uniform(pose_ranges[:, 0], pose_ranges[:, 1], (*shape, 3), device=device) + pose = spawn_pose_w.clone() + noise_quat = math_utils.quat_from_euler_xyz(noise[..., 0], noise[..., 1], noise[..., 2]) + pose[..., 3:7] = math_utils.quat_mul(pose[..., 3:7], noise_quat) + root_view.set_transforms(wp.from_torch(pose.contiguous()), all_indices) + # Only randomize the angular velocity; leave the linear velocity at zero. + body_velocity = torch.zeros((*shape, 6), device=device) + body_velocity[..., 3:6] = math_utils.sample_uniform( + velocity_ranges[:, 0], velocity_ranges[:, 1], (*shape, 3), device=device + ) + root_view.set_velocities(wp.from_torch(body_velocity.contiguous()), all_indices) + # Vary the mass so the drop behavior differs between resets. + masses = torch.rand((root_view.count, 1), device="cpu") * 0.2 + 0.2 + root_view.set_masses(wp.from_torch(masses), cpu_indices) scene.reset() # Write data to sim @@ -360,14 +362,21 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene) -> # Perform step sim.step() # Bring out-of-bounds objects back to the bin. - for asset, env_ids, spawn_pose_w in groceries: - xy = asset.data.root_link_pose_w.torch[:, :2] - scene.env_origins[env_ids, :2] - stray = torch.nonzero(~((xy >= bounds_xy[0]) & (xy <= bounds_xy[1])).all(dim=-1)).flatten() - if stray.numel(): - asset.write_root_pose_to_sim_index(root_pose=spawn_pose_w[stray], env_ids=stray) - asset.write_root_velocity_to_sim_index( - root_velocity=torch.zeros((stray.numel(), 6), device=device), env_ids=stray - ) + transforms = root_view.get_transforms() + transforms_t = wp.to_torch(transforms) if isinstance(transforms, wp.array) else transforms + xy = transforms_t[:, :2] - env_origins[:, :2] + stray = torch.nonzero(~((xy >= bounds_xy[0]) & (xy <= bounds_xy[1])).all(dim=-1)).flatten() + if stray.numel(): + stray_indices = wp.from_torch(stray.to(dtype=torch.int32)) + corrected_transforms = transforms_t.clone() + corrected_transforms[stray] = spawn_pose_w[stray] + velocities = root_view.get_velocities() + corrected_velocities = ( + wp.to_torch(velocities).clone() if isinstance(velocities, wp.array) else velocities.clone() + ) + corrected_velocities[stray] = 0.0 + root_view.set_transforms(wp.from_torch(corrected_transforms.contiguous()), stray_indices) + root_view.set_velocities(wp.from_torch(corrected_velocities.contiguous()), stray_indices) # Increment counter count += 1 # Update buffers diff --git a/scripts/demos/mpm/newton_mpm_granular.py b/scripts/demos/mpm/newton_mpm_granular.py index 965cbe704e44..a7c009b5e516 100644 --- a/scripts/demos/mpm/newton_mpm_granular.py +++ b/scripts/demos/mpm/newton_mpm_granular.py @@ -146,7 +146,7 @@ class GranularSceneCfg(InteractiveSceneCfg): ground = AssetBaseCfg( prim_path="/World/Ground", - spawn=sim_utils.GroundPlaneCfg(size=(12.0, 12.0), color=(0.30, 0.30, 0.30)), + spawn=sim_utils.GroundPlaneCfg(), ) dome_light = AssetBaseCfg( diff --git a/scripts/demos/mpm/newton_mpm_twoway_coupling.py b/scripts/demos/mpm/newton_mpm_twoway_coupling.py index 04e44b753c99..2c60956c0399 100644 --- a/scripts/demos/mpm/newton_mpm_twoway_coupling.py +++ b/scripts/demos/mpm/newton_mpm_twoway_coupling.py @@ -230,7 +230,7 @@ class CoupledSceneCfg(InteractiveSceneCfg): ground = AssetBaseCfg( prim_path="/World/Ground", - spawn=sim_utils.GroundPlaneCfg(size=(12.0, 12.0), color=(0.30, 0.30, 0.30)), + spawn=sim_utils.GroundPlaneCfg(), init_state=AssetBaseCfg.InitialStateCfg(pos=(0.0, 0.0, -wall_t)), ) dome_light = AssetBaseCfg( diff --git a/scripts/demos/mpm/snowball_smash.py b/scripts/demos/mpm/snowball_smash.py index 137398441cd0..0be3833d4aed 100644 --- a/scripts/demos/mpm/snowball_smash.py +++ b/scripts/demos/mpm/snowball_smash.py @@ -279,7 +279,7 @@ class SnowballSmashSceneCfg(InteractiveSceneCfg): # the MPM model view. ground = AssetBaseCfg( prim_path="/World/Ground", - spawn=sim_utils.GroundPlaneCfg(size=(12.0, 12.0), color=(0.32, 0.34, 0.38)), + spawn=sim_utils.GroundPlaneCfg(), ) # The co-located hidden kinematic slab belongs only to the MPM entry. diff --git a/scripts/imitation_learning/locomanipulation_sdg/generate_data.py b/scripts/imitation_learning/locomanipulation_sdg/generate_data.py index f71916309190..399f89be47c3 100644 --- a/scripts/imitation_learning/locomanipulation_sdg/generate_data.py +++ b/scripts/imitation_learning/locomanipulation_sdg/generate_data.py @@ -95,13 +95,13 @@ "--background_usd_path", type=str, default=None, - help="Path to the USD file for the background asset", + help="Path to the NuRec USD file.", ) parser.add_argument( "--background_occupancy_yaml_file", type=str, default=None, - help="Path to the occupancy map YAML file for the background asset", + help="Path to the NuRec occupancy map YAML file.", ) parser.add_argument( "--high_res_video", @@ -126,7 +126,7 @@ # forward unrecognized args as Hydra-style task config overrides args_cli, hydra_overrides = parser.parse_known_args() -app_launcher = AppLauncher(args_cli) +app_launcher = AppLauncher(args_cli, enable_cameras=True) simulation_app = app_launcher.app import enum @@ -358,7 +358,9 @@ def project_robot_state_into_env(env: LocomanipulationSDGEnv, input_episode_data ) default_vel = env.scene["robot"].data.default_root_vel.torch.clone() default_vel[0] = torch.zeros(6, device=env.device) - env.scene["robot"].data.default_root_vel.warp.assign(wp.from_torch(default_vel.to(env.device).contiguous())) + env.scene["robot"].data.default_root_vel.warp.assign( + wp.from_torch(default_vel.to(env.device).contiguous()).view(wp.spatial_vectorf) + ) robot_state = recording_initial_state["articulation"]["robot"] joint_position = robot_state["joint_position"][0].to(env.device) @@ -409,7 +411,9 @@ def project_object_state_into_env(env: LocomanipulationSDGEnv, input_episode_dat ) default_vel = env.scene["object"].data.default_root_vel.torch.clone() default_vel[0] = torch.zeros(6, device=env.device) - env.scene["object"].data.default_root_vel.warp.assign(wp.from_torch(default_vel.to(env.device).contiguous())) + env.scene["object"].data.default_root_vel.warp.assign( + wp.from_torch(default_vel.to(env.device).contiguous()).view(wp.spatial_vectorf) + ) return new_object_pose @@ -438,8 +442,7 @@ def setup_navigation_scene( if background_fixture is not None: occupancy_map = background_fixture.get_occupancy_map() fixtures = [env.get_start_fixture(), env.get_end_fixture()] + env.get_obstacle_fixtures() - if not randomize_placement: - raise ValueError("randomize_placement needs to be True when background_usd_path is provided") + randomize_placement = True else: occupancy_map = merge_occupancy_maps( [ @@ -484,17 +487,20 @@ def setup_navigation_scene( env.obs_buf = env.observation_manager.compute(update_history=True) nav_map = occupancy_map.buffered_meters(0.15) - nav_fs = nav_map.freespace_mask() - start_pos = env.get_base().get_pose()[0, :2].detach().cpu().numpy() - start_px = nav_map.world_to_pixel_numpy(start_pos[None])[0].astype(int) - sx, sy = int(start_px[0]), int(start_px[1]) - - # Clear the buffer zone around the robot's start pixel if it falls inside it. - # The robot stands adjacent to the start table so its position sits within the 0.15m buffer. - if 0 <= sy < nav_fs.shape[0] and 0 <= sx < nav_fs.shape[1] and not nav_fs[sy, sx]: + + def clear_nav_map_buffer(world_xy: torch.Tensor | np.ndarray): + nav_fs = nav_map.freespace_mask() + pixel_xy = nav_map.world_to_pixel_numpy(np.asarray(world_xy)[None])[0].astype(int) + x_px, y_px = int(pixel_xy[0]), int(pixel_xy[1]) + if not (0 <= y_px < nav_fs.shape[0] and 0 <= x_px < nav_fs.shape[1]) or nav_fs[y_px, x_px]: + return clear_r = int(0.15 / nav_map.resolution) yy, xx = np.ogrid[: nav_map.data.shape[0], : nav_map.data.shape[1]] - nav_map.data[(xx - sx) ** 2 + (yy - sy) ** 2 <= clear_r**2] = OccupancyMapDataValue.FREESPACE + nav_map.data[(xx - x_px) ** 2 + (yy - y_px) ** 2 <= clear_r**2] = OccupancyMapDataValue.FREESPACE + + # The robot start and approach goal can sit adjacent to their tables, inside the 0.15 m buffer. + clear_nav_map_buffer(env.get_base().get_pose()[0, :2].detach().cpu().numpy()) + clear_nav_map_buffer(base_goal_approach.get_pose()[0, :2].detach().cpu().numpy()) base_path = plan_path(start=env.get_base(), end=base_goal_approach, occupancy_map=nav_map) @@ -899,6 +905,11 @@ def replay( print("Failed to setup navigation scene", flush=True) return False + # Background placement can move fixtures and re-project robot/object root poses + # after reset_to records initial_state. Re-record from the projected scene. + env.recorder_manager.reset(env_ids=[0]) + env.recorder_manager.record_post_reset(env_ids=[0]) + if sensor_camera_view: _set_sensor_camera_view() @@ -1001,7 +1012,17 @@ def replay( env_cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_SUCCEEDED_ONLY env_cfg.recorders.export_in_record_pre_reset = True - if args_cli.background_usd_path is not None and args_cli.background_occupancy_yaml_file is not None: + nurec_scene_used = ( + args_cli.background_usd_path is not None or args_cli.background_occupancy_yaml_file is not None + ) + if nurec_scene_used and ( + args_cli.background_usd_path is None or args_cli.background_occupancy_yaml_file is None + ): + raise ValueError( + "Both --background_usd_path and --background_occupancy_yaml_file must be provided for NuRec scenes." + ) + + if nurec_scene_used: env_cfg.background_usd_path = args_cli.background_usd_path env_cfg.background_occupancy_yaml_file = args_cli.background_occupancy_yaml_file @@ -1048,7 +1069,7 @@ def replay( following_offset=args_cli.following_offset, angle_threshold=args_cli.angle_threshold, approach_distance=args_cli.approach_distance, - randomize_placement=args_cli.randomize_placement, + randomize_placement=args_cli.randomize_placement or nurec_scene_used, sensor_camera_view=args_cli.sensor_camera_view, ) diff --git a/source/isaaclab/changelog.d/mimic-max-num-failures.rst b/source/isaaclab/changelog.d/mimic-max-num-failures.rst new file mode 100644 index 000000000000..9f811b3a0612 --- /dev/null +++ b/source/isaaclab/changelog.d/mimic-max-num-failures.rst @@ -0,0 +1,8 @@ +Fixed +^^^^^ + +* Fixed :attr:`~isaaclab.envs.mimic_env_cfg.DataGenConfig.max_num_failures` being ignored. The field + documented a cap on failed generation attempts but was never read, so a run with + :attr:`~isaaclab.envs.mimic_env_cfg.DataGenConfig.generation_guarantee` enabled retried without + bound on a task with a low success rate. Its default is now ``None`` (no limit) and setting it to + an integer stops generation once that many attempts have failed. diff --git a/source/isaaclab/changelog.d/render-nvtx-profile-scope.rst b/source/isaaclab/changelog.d/render-nvtx-profile-scope.rst new file mode 100644 index 000000000000..df22fa38fe68 --- /dev/null +++ b/source/isaaclab/changelog.d/render-nvtx-profile-scope.rst @@ -0,0 +1,8 @@ +Added +^^^^^ + +* Added a synchronized timer around the renderer call inside + :meth:`~isaaclab.renderers.render_context.RenderContext.render_into_camera`, gated by the ``ISAACLAB_RENDER_PROFILE`` + environment variable, so any rendering backend can be profiled through the same scope name + (:data:`~isaaclab.renderers.render_context.RENDER_PROFILE_SCOPE`). When enabled, each render prints its elapsed + time to the log. diff --git a/source/isaaclab/isaaclab/envs/mimic_env_cfg.py b/source/isaaclab/isaaclab/envs/mimic_env_cfg.py index a98c34bf5f90..990c792b895c 100644 --- a/source/isaaclab/isaaclab/envs/mimic_env_cfg.py +++ b/source/isaaclab/isaaclab/envs/mimic_env_cfg.py @@ -38,8 +38,21 @@ class DataGenConfig: Keeping failed demonstrations is useful for visualizing and debugging low success rates. """ - max_num_failures: int = 50 - """Maximum number of failures allowed before stopping generation.""" + max_num_failures: int | None = None + """Maximum number of failed generation attempts before stopping, or None for no limit. + + Only applies together with :attr:`generation_guarantee`. With the guarantee enabled, generation + keeps retrying until :attr:`generation_num_trials` demos succeed, so a task whose success rate is + low can run for an unbounded number of attempts; this caps that. Defaults to None so the + guarantee keeps its usual meaning unless a limit is asked for. + + With the guarantee disabled, generation already stops after :attr:`generation_num_trials` + attempts and this field is ignored, so setting it cannot cut a fixed-attempt run short. + + The bound is read once per simulation step. Attempts that end on the step that crosses it are + already complete, so a run over ``num_envs`` parallel environments can record up to + ``num_envs - 1`` failures beyond the bound; it is exact whenever attempts end on separate steps. + """ seed: int = 1 """Seed for randomization to ensure reproducibility.""" diff --git a/source/isaaclab/isaaclab/renderers/render_context.py b/source/isaaclab/isaaclab/renderers/render_context.py index aa132a98b2d8..c08fe2e7f143 100644 --- a/source/isaaclab/isaaclab/renderers/render_context.py +++ b/source/isaaclab/isaaclab/renderers/render_context.py @@ -8,6 +8,7 @@ from __future__ import annotations import logging +import os from typing import Any import torch @@ -21,6 +22,22 @@ logger = logging.getLogger(__name__) +RENDER_PROFILE_SCOPE = "IsaacLab::Renderer::render" +"""Name of the timed scope bracketing :meth:`BaseRenderer.render`, emitted when render profiling is on. + +Every backend renders through the same call, so a profile can compare them under one scope name +instead of one internal name per backend. ``wp.ScopedTimer`` prints one ``" took X.XX ms"`` +line per call, which ``scripts/benchmarks/benchmark_renderer.py`` parses back out of the run log. +""" + +_RENDER_PROFILE_ENABLED = os.environ.get("ISAACLAB_RENDER_PROFILE", "0") != "0" +"""Whether to time and print :data:`RENDER_PROFILE_SCOPE`, read once from ``ISAACLAB_RENDER_PROFILE``. + +Off by default because the timer synchronizes the device on entry and exit. That is what lets it +measure completed device work rather than submitted work, but it also removes CPU/GPU overlap, so +an enabled run is a profiling aid and not a throughput measurement. +""" + @wp.kernel(enable_backward=False) def _write_material( @@ -340,9 +357,20 @@ def render_into_camera( camera_data: CameraData, physics_step_count: int, ) -> None: - """Sync scene state, render, and read outputs into ``camera_data``.""" + """Sync scene state, render, and read outputs into ``camera_data``. + + Only the render itself is bracketed by :data:`RENDER_PROFILE_SCOPE`, so a profile + attributes neither the scene-state sync before it nor the output readback after it to + rendering. See :data:`_RENDER_PROFILE_ENABLED` for how to turn the timer on. + """ self.update_scene_state(physics_step_count) - renderer.render(render_data) + with wp.ScopedTimer( + RENDER_PROFILE_SCOPE, + active=_RENDER_PROFILE_ENABLED, + print=True, + synchronize=True, + ): + renderer.render(render_data) renderer.read_output(render_data, camera_data) def reset_stage_prepare_flag(self) -> None: diff --git a/source/isaaclab/test/renderers/test_simulation_render_context.py b/source/isaaclab/test/renderers/test_simulation_render_context.py index 7aff3c4c9099..eed7bd67dd1f 100644 --- a/source/isaaclab/test/renderers/test_simulation_render_context.py +++ b/source/isaaclab/test/renderers/test_simulation_render_context.py @@ -7,10 +7,12 @@ from __future__ import annotations +import re from typing import Any, cast import pytest +from isaaclab.renderers import render_context from isaaclab.renderers.base_renderer import BaseRenderer from isaaclab.renderers.output_contract import RenderBufferKind, RenderBufferSpec from isaaclab.renderers.render_context import RenderContext @@ -207,6 +209,38 @@ def test_render_into_camera_calls_update_render_read_order(): assert events == ["ut", "geo", "render", "read", "render", "read"] +def test_render_into_camera_call_order_unaffected_by_render_profile_flag(monkeypatch): + """Enabling ``ISAACLAB_RENDER_PROFILE`` only wraps the render call in a printed timer. + + The wrapping must not change the call order or drop the render call itself. + """ + monkeypatch.setattr(render_context, "_RENDER_PROFILE_ENABLED", True) + + ctx = RenderContext() + events: list[str] = [] + cfg = IsaacRtxRendererCfg() + fake = _FakeBackend(event_log=events) + _set_entries(ctx, (cfg, fake)) + + ctx.render_into_camera(cast(BaseRenderer, fake), object(), CameraData(), physics_step_count=1) + + assert events == ["ut", "geo", "render", "read"] + + +def test_render_into_camera_prints_timing_line_when_render_profile_enabled(monkeypatch, capsys): + """The printed line must match the format ``scripts/benchmarks/benchmark_renderer.py`` parses.""" + monkeypatch.setattr(render_context, "_RENDER_PROFILE_ENABLED", True) + + ctx = RenderContext() + cfg = IsaacRtxRendererCfg() + fake = _FakeBackend() + _set_entries(ctx, (cfg, fake)) + + ctx.render_into_camera(cast(BaseRenderer, fake), object(), CameraData(), physics_step_count=1) + + assert re.search(rf"{re.escape(render_context.RENDER_PROFILE_SCOPE)} took [\d.]+ ms", capsys.readouterr().out) + + def test_reset_stage_prepare_flag_allows_second_prepare_stage(): """After reset_stage_prepare_flag, ensure_prepare_stage invokes prepare_stage again.""" ctx = RenderContext() diff --git a/source/isaaclab_assets/changelog.d/fix-newton-ball-joint-dof-space.rst b/source/isaaclab_assets/changelog.d/fix-newton-ball-joint-dof-space.rst new file mode 100644 index 000000000000..fb2e3c7b3922 --- /dev/null +++ b/source/isaaclab_assets/changelog.d/fix-newton-ball-joint-dof-space.rst @@ -0,0 +1,6 @@ +Changed +^^^^^^^ + +* Changed :obj:`~isaaclab_assets.robots.agility.DIGIT_V4_CFG` to disable colliders on the + RealSense camera mount's decoration meshes and to split actuators into a stable-armature and a + low-armature group, matching what the Digit velocity tasks previously configured per-task. diff --git a/source/isaaclab_assets/isaaclab_assets/robots/agility.py b/source/isaaclab_assets/isaaclab_assets/robots/agility.py index 2c85a42ec681..ed67e16dfba0 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/agility.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/agility.py @@ -6,6 +6,7 @@ import isaaclab.sim as sim_utils from isaaclab.actuators import ImplicitActuatorCfg from isaaclab.assets.articulation import ArticulationCfg +from isaaclab.sim.schemas import UsdPhysicsCollisionCfg from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR LEG_JOINT_NAMES = [ @@ -19,21 +20,50 @@ ARM_JOINT_NAMES = [".*_arm_.*"] +# On MJWarp these ten joints diverge at the armature the USD authors. With the damping the asset +# carries (c = 57.3) and the velocity tasks' substep (h = sim.dt / num_substeps = 0.0025 s), the +# observed threshold is c*h/I > 2: wrist_yaw at I = 0.01822 [kg m^2] gives 7.86 and grew 6.9x per +# substep, the other eight at 0.05228 give 2.74, and every joint at or above I = 0.0716 was stable. +# 0.10 leaves margin. Re-measure if sim.dt, num_substeps or the asset's damping change. +# Together these cover exactly ``LEG_JOINT_NAMES + ARM_JOINT_NAMES``. +_LOW_ARMATURE_JOINT_NAMES = [".*_arm_wrist_.*", ".*_toe_a", ".*_toe_b"] +_STABLE_ARMATURE_JOINT_NAMES = [ + ".*_hip_roll", + ".*_hip_yaw", + ".*_hip_pitch", + ".*_knee", + ".*_arm_shoulder_.*", + ".*_arm_elbow", +] +_MIN_STABLE_ARMATURE = 0.10 + DIGIT_V4_CFG = ArticulationCfg( spawn=sim_utils.UsdFileCfg( usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/Agility/Digit/digit_v4.usd", activate_contact_sensors=True, + # digit_v4.usd applies CollisionAPI to 32 prims, every one a decoration mesh on a + # RealSense camera mount -- glass, USB-C, case halves. They are 32 of the robot's 55 + # collision shapes while the arms, hips and rods carry none, and produced 3e7 N contact + # forces on bodies 1.4 m above the ground. Colliders on a camera's glass are an authoring + # error on either backend. + collision_props={"/.*camera_mount/.*/Visual/.*": [UsdPhysicsCollisionCfg(collision_enabled=False)]}, ), init_state=ArticulationCfg.InitialStateCfg( pos=(0.0, 0.0, 1.05), ), soft_joint_pos_limit_factor=0.9, actuators={ - "all": ImplicitActuatorCfg( - joint_names_expr=".*", + "legs_arms": ImplicitActuatorCfg( + joint_names_expr=_STABLE_ARMATURE_JOINT_NAMES, + stiffness=None, + damping=None, + ), + "low_armature": ImplicitActuatorCfg( + joint_names_expr=_LOW_ARMATURE_JOINT_NAMES, stiffness=None, damping=None, + armature=_MIN_STABLE_ARMATURE, ), }, ) diff --git a/source/isaaclab_mimic/changelog.d/locomanip-nurec-sdg.rst b/source/isaaclab_mimic/changelog.d/locomanip-nurec-sdg.rst new file mode 100644 index 000000000000..5c13727e221a --- /dev/null +++ b/source/isaaclab_mimic/changelog.d/locomanip-nurec-sdg.rst @@ -0,0 +1,7 @@ +Fixed +^^^^^ + +* Fixed locomanipulation SDG generation with NuRec backgrounds by enabling camera capture, applying Isaac RTX + Gaussian renderer settings, syncing randomized fixture poses, and recording the projected scene state after + placement. ``--high_res_video`` now records RGB observations at 512x320 instead of 960x540; update MP4 + conversion dimensions and model input shapes accordingly. diff --git a/source/isaaclab_mimic/changelog.d/mimic-max-num-failures.rst b/source/isaaclab_mimic/changelog.d/mimic-max-num-failures.rst new file mode 100644 index 000000000000..bffb2ae2ed24 --- /dev/null +++ b/source/isaaclab_mimic/changelog.d/mimic-max-num-failures.rst @@ -0,0 +1,14 @@ +Fixed +^^^^^ + +* Fixed the Mimic generation loop never bounding a run by failure count. ``env_loop`` now stops when + ``datagen_config.max_num_failures`` failed attempts have accumulated, alongside the existing stop + on enough successes or attempts. + +Removed +^^^^^^^ + +* Removed the ``datagen_config.max_num_failures = 25`` assignment from the shipped Mimic environment + configs. The field was never read when those lines were written, so honouring it now would newly + cap every shipped task at 25 failed attempts and cut short any run asking for a large number of + demos. Set the field explicitly to opt into a cap. diff --git a/source/isaaclab_mimic/isaaclab_mimic/datagen/generation.py b/source/isaaclab_mimic/isaaclab_mimic/datagen/generation.py index 5001726e50c3..687a74f0cc4f 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/datagen/generation.py +++ b/source/isaaclab_mimic/isaaclab_mimic/datagen/generation.py @@ -140,6 +140,18 @@ def env_loop( print(f"Reached {generation_num_trials} successes/attempts. Exiting.") break + # with the success guarantee on, nothing else bounds the run: a task that rarely + # succeeds retries forever. max_num_failures is the opt-in bound on that. Without the + # guarantee the check above already stops on generation_num_trials attempts, so the + # cap has nothing to add and must not cut a fixed-attempt run short. + max_num_failures = env.cfg.datagen_config.max_num_failures + if generation_guarantee and max_num_failures is not None and num_failures >= max_num_failures: + print( + f"Reached {num_failures} failures (max_num_failures={max_num_failures}) after" + f" {num_success}/{generation_num_trials} successes. Exiting." + ) + break + # check that simulation is stopped or not if env.sim.is_stopped(): break diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/agibot_place_toy2box_mimic_env_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/envs/agibot_place_toy2box_mimic_env_cfg.py index 43085361f950..0df362689ec5 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/agibot_place_toy2box_mimic_env_cfg.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/agibot_place_toy2box_mimic_env_cfg.py @@ -32,7 +32,6 @@ def __post_init__(self): self.datagen_config.generation_select_src_per_subtask = True self.datagen_config.generation_transform_first_robot_pose = False self.datagen_config.generation_interpolate_from_last_target_pose = True - self.datagen_config.max_num_failures = 25 self.datagen_config.seed = 1 # The following are the subtask configurations for the stack task. diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/agibot_place_upright_mug_mimic_env_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/envs/agibot_place_upright_mug_mimic_env_cfg.py index 31bf8fe81f3a..b645c4fcf2ab 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/agibot_place_upright_mug_mimic_env_cfg.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/agibot_place_upright_mug_mimic_env_cfg.py @@ -29,7 +29,6 @@ def __post_init__(self): self.datagen_config.generation_select_src_per_subtask = True self.datagen_config.generation_transform_first_robot_pose = False self.datagen_config.generation_interpolate_from_last_target_pose = True - self.datagen_config.max_num_failures = 25 self.datagen_config.seed = 1 # The following are the subtask configurations for the stack task. diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/exhaustpipe_gr1t2_mimic_env_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/envs/exhaustpipe_gr1t2_mimic_env_cfg.py index 78df5903d761..4c7192a6307f 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/exhaustpipe_gr1t2_mimic_env_cfg.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/exhaustpipe_gr1t2_mimic_env_cfg.py @@ -30,7 +30,6 @@ def __post_init__(self): self.datagen_config.generation_joint_pos = False self.datagen_config.generation_transform_first_robot_pose = False self.datagen_config.generation_interpolate_from_last_target_pose = True - self.datagen_config.max_num_failures = 25 self.datagen_config.num_demo_to_render = 10 self.datagen_config.num_fail_demo_to_render = 25 self.datagen_config.seed = 10 diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/franka_bin_stack_ik_rel_mimic_env_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/envs/franka_bin_stack_ik_rel_mimic_env_cfg.py index ef79f102f330..3e35749f9c33 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/franka_bin_stack_ik_rel_mimic_env_cfg.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/franka_bin_stack_ik_rel_mimic_env_cfg.py @@ -28,7 +28,6 @@ def __post_init__(self): self.datagen_config.generation_transform_first_robot_pose = False self.datagen_config.generation_interpolate_from_last_target_pose = True self.datagen_config.generation_relative = True - self.datagen_config.max_num_failures = 25 self.datagen_config.seed = 1 # The following are the subtask configurations for the stack task. diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_abs_mimic_env_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_abs_mimic_env_cfg.py index 974e0895b984..537541ffae01 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_abs_mimic_env_cfg.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_abs_mimic_env_cfg.py @@ -27,7 +27,6 @@ def __post_init__(self): self.datagen_config.generation_select_src_per_subtask = True self.datagen_config.generation_transform_first_robot_pose = False self.datagen_config.generation_interpolate_from_last_target_pose = True - self.datagen_config.max_num_failures = 25 self.datagen_config.seed = 1 # The following are the subtask configurations for the stack task. diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_blueprint_mimic_env_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_blueprint_mimic_env_cfg.py index cda843307acf..2773e8139013 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_blueprint_mimic_env_cfg.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_blueprint_mimic_env_cfg.py @@ -29,7 +29,6 @@ def __post_init__(self): self.datagen_config.generation_select_src_per_subtask = True self.datagen_config.generation_transform_first_robot_pose = False self.datagen_config.generation_interpolate_from_last_target_pose = True - self.datagen_config.max_num_failures = 25 self.datagen_config.seed = 1 # The following are the subtask configurations for the stack task. diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_mimic_env_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_mimic_env_cfg.py index a3ce6fc310a2..6f84e90bf97a 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_mimic_env_cfg.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_mimic_env_cfg.py @@ -31,7 +31,6 @@ def __post_init__(self): self.datagen_config.generation_transform_first_robot_pose = False self.datagen_config.generation_interpolate_from_last_target_pose = True self.datagen_config.generation_relative = True - self.datagen_config.max_num_failures = 25 self.datagen_config.seed = 1 # The following are the subtask configurations for the stack task. diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_skillgen_env_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_skillgen_env_cfg.py index cd7e3e472c6e..9b1cb235b6c6 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_skillgen_env_cfg.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_skillgen_env_cfg.py @@ -33,7 +33,6 @@ def __post_init__(self): self.datagen_config.generation_transform_first_robot_pose = False self.datagen_config.generation_interpolate_from_last_target_pose = True self.datagen_config.generation_relative = True - self.datagen_config.max_num_failures = 25 self.datagen_config.seed = 1 # The following are the subtask configurations for the stack task. diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_visuomotor_cosmos_mimic_env_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_visuomotor_cosmos_mimic_env_cfg.py index 83a727ee6022..f188a9ecffe9 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_visuomotor_cosmos_mimic_env_cfg.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_visuomotor_cosmos_mimic_env_cfg.py @@ -30,7 +30,6 @@ def __post_init__(self): self.datagen_config.generation_transform_first_robot_pose = False self.datagen_config.generation_interpolate_from_last_target_pose = True self.datagen_config.generation_relative = True - self.datagen_config.max_num_failures = 25 self.datagen_config.seed = 1 # The following are the subtask configurations for the stack task. diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_visuomotor_mimic_env_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_visuomotor_mimic_env_cfg.py index e6ca8a12cd02..050128cf9a93 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_visuomotor_mimic_env_cfg.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_visuomotor_mimic_env_cfg.py @@ -30,7 +30,6 @@ def __post_init__(self): self.datagen_config.generation_transform_first_robot_pose = False self.datagen_config.generation_interpolate_from_last_target_pose = True self.datagen_config.generation_relative = True - self.datagen_config.max_num_failures = 25 self.datagen_config.seed = 1 # The following are the subtask configurations for the stack task. diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/galbot_stack_rmp_abs_mimic_env_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/envs/galbot_stack_rmp_abs_mimic_env_cfg.py index 1b3d544444cb..ef94781bc961 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/galbot_stack_rmp_abs_mimic_env_cfg.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/galbot_stack_rmp_abs_mimic_env_cfg.py @@ -31,7 +31,6 @@ def __post_init__(self): self.datagen_config.generation_select_src_per_subtask = True self.datagen_config.generation_transform_first_robot_pose = False self.datagen_config.generation_interpolate_from_last_target_pose = True - self.datagen_config.max_num_failures = 25 self.datagen_config.seed = 1 # The following are the subtask configurations for the stack task. @@ -157,7 +156,6 @@ def __post_init__(self): self.datagen_config.generation_select_src_per_subtask = True self.datagen_config.generation_transform_first_robot_pose = False self.datagen_config.generation_interpolate_from_last_target_pose = True - self.datagen_config.max_num_failures = 25 self.datagen_config.seed = 1 # The following are the subtask configurations for the stack task. diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/galbot_stack_rmp_rel_mimic_env_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/envs/galbot_stack_rmp_rel_mimic_env_cfg.py index be87d9080743..8def41266475 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/galbot_stack_rmp_rel_mimic_env_cfg.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/galbot_stack_rmp_rel_mimic_env_cfg.py @@ -31,7 +31,6 @@ def __post_init__(self): self.datagen_config.generation_select_src_per_subtask = True self.datagen_config.generation_transform_first_robot_pose = False self.datagen_config.generation_interpolate_from_last_target_pose = True - self.datagen_config.max_num_failures = 25 self.datagen_config.seed = 1 # The following are the subtask configurations for the stack task. @@ -155,7 +154,6 @@ def __post_init__(self): self.datagen_config.generation_select_src_per_subtask = True self.datagen_config.generation_transform_first_robot_pose = False self.datagen_config.generation_interpolate_from_last_target_pose = True - self.datagen_config.max_num_failures = 25 self.datagen_config.seed = 1 # The following are the subtask configurations for the stack task. diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/locomanipulation_g1_mimic_env_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/envs/locomanipulation_g1_mimic_env_cfg.py index 5ddcf3d26c3f..71ad54c594ab 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/locomanipulation_g1_mimic_env_cfg.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/locomanipulation_g1_mimic_env_cfg.py @@ -30,7 +30,6 @@ def __post_init__(self): self.datagen_config.generation_joint_pos = False self.datagen_config.generation_transform_first_robot_pose = False self.datagen_config.generation_interpolate_from_last_target_pose = True - self.datagen_config.max_num_failures = 25 self.datagen_config.num_demo_to_render = 10 self.datagen_config.num_fail_demo_to_render = 25 self.datagen_config.seed = 1 diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/nutpour_gr1t2_mimic_env_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/envs/nutpour_gr1t2_mimic_env_cfg.py index e145257e2fbe..dc12d03f49e5 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/nutpour_gr1t2_mimic_env_cfg.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/nutpour_gr1t2_mimic_env_cfg.py @@ -28,7 +28,6 @@ def __post_init__(self): self.datagen_config.generation_joint_pos = False self.datagen_config.generation_transform_first_robot_pose = False self.datagen_config.generation_interpolate_from_last_target_pose = True - self.datagen_config.max_num_failures = 25 self.datagen_config.num_demo_to_render = 10 self.datagen_config.num_fail_demo_to_render = 25 self.datagen_config.seed = 10 diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/pickplace_gr1t2_mimic_env_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/envs/pickplace_gr1t2_mimic_env_cfg.py index 9a987a3b0f21..b220625bdb46 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/pickplace_gr1t2_mimic_env_cfg.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/pickplace_gr1t2_mimic_env_cfg.py @@ -28,7 +28,6 @@ def __post_init__(self): self.datagen_config.generation_joint_pos = False self.datagen_config.generation_transform_first_robot_pose = False self.datagen_config.generation_interpolate_from_last_target_pose = True - self.datagen_config.max_num_failures = 25 self.datagen_config.num_demo_to_render = 10 self.datagen_config.num_fail_demo_to_render = 25 self.datagen_config.seed = 1 diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/pickplace_gr1t2_waist_enabled_mimic_env_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/envs/pickplace_gr1t2_waist_enabled_mimic_env_cfg.py index 6ed48a7590ca..bfae20234e67 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/pickplace_gr1t2_waist_enabled_mimic_env_cfg.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/pickplace_gr1t2_waist_enabled_mimic_env_cfg.py @@ -30,7 +30,6 @@ def __post_init__(self): self.datagen_config.generation_joint_pos = False self.datagen_config.generation_transform_first_robot_pose = False self.datagen_config.generation_interpolate_from_last_target_pose = True - self.datagen_config.max_num_failures = 25 self.datagen_config.num_demo_to_render = 10 self.datagen_config.num_fail_demo_to_render = 25 self.datagen_config.seed = 1 diff --git a/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/envs/g1_locomanipulation_sdg_env.py b/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/envs/g1_locomanipulation_sdg_env.py index 72be2a63c075..a981caefb9a7 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/envs/g1_locomanipulation_sdg_env.py +++ b/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/envs/g1_locomanipulation_sdg_env.py @@ -5,6 +5,7 @@ import numpy as np import torch +from isaaclab_physx.renderers import IsaacRtxRendererCfg, IsaacRtxRendererGlobalSettingsCfg import isaaclab.sim as sim_utils from isaaclab.assets import AssetBaseCfg @@ -35,6 +36,11 @@ NUM_FORKLIFTS = 0 NUM_BOXES = 0 +ISAAC_RTX_GAUSSIAN_CAMERA_RENDERER_CARB_SETTINGS = { + "/rtx/rtpt/gaussian/accumulatedDepth/allHits/enabled": True, + "/rtx/rtpt/gaussian/accumulatedAlbedo/enabled": True, + "/rtx/rtpt/gaussian/maxGaussiansToAccumulate": 360, +} @configclass @@ -52,7 +58,17 @@ class G1LocomanipulationSDGSceneCfg(LocomanipulationG1SceneCfg): ), ) - def add_robot_pov_cam(self, height, width): + def add_robot_pov_cam(self, height, width, use_nurec_renderer_settings: bool = False): + camera_kwargs = {} + if use_nurec_renderer_settings: + camera_kwargs["renderer_cfg"] = IsaacRtxRendererCfg( + global_settings=IsaacRtxRendererGlobalSettingsCfg( + enable_dl_denoiser=True, + antialiasing_mode="DLSS", + carb_settings=ISAAC_RTX_GAUSSIAN_CAMERA_RENDERER_CARB_SETTINGS, + ) + ) + robot_pov_cam = CameraCfg( prim_path="{ENV_REGEX_NS}/Robot/torso_link/d435_link/camera", update_period=0.0, @@ -61,6 +77,7 @@ def add_robot_pov_cam(self, height, width): data_types=["rgb"], spawn=sim_utils.PinholeCameraCfg(focal_length=8.0, clipping_range=(0.1, 20.0)), offset=CameraCfg.OffsetCfg(pos=(0.0, 0.0, 0.0), rot=(0.0, -0.1736482, 0.0, 0.9848078), convention="world"), + **camera_kwargs, ) setattr(self, "robot_pov_cam", robot_pov_cam) @@ -176,10 +193,10 @@ def __init__(self, cfg: G1LocomanipulationSDGEnvCfg, **kwargs): self._num_boxes = NUM_BOXES set_ground_invisible = False - if cfg.high_res_video: - cfg.scene.add_robot_pov_cam(540, 960) - else: - cfg.scene.add_robot_pov_cam(160, 256) + camera_height, camera_width = (320, 512) if cfg.high_res_video else (160, 256) + cfg.scene.add_robot_pov_cam( + camera_height, camera_width, use_nurec_renderer_settings=cfg.background_usd_path is not None + ) cfg.scene.add_forklifts(self._num_forklifts) cfg.scene.add_boxes(self._num_boxes) diff --git a/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py b/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py index 0b2506f48b3b..d6a4780f2513 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py +++ b/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py @@ -11,7 +11,7 @@ import isaaclab.utils.math as math_utils from isaaclab.assets import AssetBaseCfg -from isaaclab.sim.views import FrameView +from isaaclab.sim.views import UsdFrameView from .occupancy_map_utils import OccupancyMap, intersect_occupancy_maps from .transform_utils import transform_mul @@ -101,10 +101,11 @@ class SceneAsset(HasPose): def __init__(self, scene, entity_name: str): self.scene = scene self.entity_name = entity_name - self._xform_view: FrameView | None = None + self._xform_view: UsdFrameView | None = None + self._cached_pose: torch.Tensor | None = None - def _get_xform_view(self) -> FrameView: - """Return the FrameView for this asset, building it on demand. + def _get_xform_view(self) -> UsdFrameView: + """Return the USD transform view for this asset, building it on demand. Static scene assets carry no runtime view (``scene[name]`` returns the spawned configuration), so the pose view is created lazily here and cached on this @@ -114,23 +115,27 @@ def _get_xform_view(self) -> FrameView: if self._xform_view is None or self._xform_view.count == 0: if self._xform_view is not None: self._xform_view.close() + self._cached_pose = None entity = self.scene[self.entity_name] prim_path = ( entity.prim_path if isinstance(entity, AssetBaseCfg) else getattr(entity, "_usd_view", entity)._prim_path ) - self._xform_view = FrameView(prim_path, device=self.scene.device) + self._xform_view = UsdFrameView(prim_path, device=self.scene.device) return self._xform_view def get_pose(self): """Get the 3D pose of the entity.""" + if self._cached_pose is not None: + return self._cached_pose.clone() xform_prim = self._get_xform_view() pos_w, quat_w = xform_prim.get_world_poses() position = pos_w.torch orientation = quat_w.torch pose = torch.cat([position, orientation], dim=-1) - return pose + self._cached_pose = pose.detach().clone() + return self._cached_pose.clone() def set_pose(self, pose: torch.Tensor): """Set the 3D pose of the entity.""" @@ -139,6 +144,16 @@ def set_pose(self, pose: torch.Tensor): orientation = pose[..., 3:] with xform_prim.xform_world_space_writer() as writer: writer.set_poses(wp.from_torch(position.contiguous()), wp.from_torch(orientation.contiguous()), None) + self._cached_pose = pose.detach().clone() + entity = self.scene[self.entity_name] + if isinstance(entity, AssetBaseCfg): + # Static assets have no runtime data object, so keep their config pose aligned. + pose_cfg = pose[0] if pose.ndim > 1 else pose + env_origin = self.scene.env_origins[0].to(pose_cfg.device) + pos = (pose_cfg[:3] - env_origin).detach().cpu().tolist() + rot = pose_cfg[3:].detach().cpu().tolist() + entity.init_state.pos = tuple(float(value) for value in pos) + entity.init_state.rot = tuple(float(value) for value in rot) class RelativePose(HasPose): diff --git a/source/isaaclab_mimic/test/test_generation_failure_cap.py b/source/isaaclab_mimic/test/test_generation_failure_cap.py new file mode 100644 index 000000000000..c029c1addf8b --- /dev/null +++ b/source/isaaclab_mimic/test/test_generation_failure_cap.py @@ -0,0 +1,150 @@ +# Copyright (c) 2024-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the ``max_num_failures`` bound in the Mimic generation loop. + +``env_loop`` is driven with a fake environment whose ``step`` consumes one scripted attempt outcome +per call and feeds the action queue back, so the loop's own termination logic runs unmodified +without a simulator, data generator, or dataset. A step-count fuse turns an unbounded loop into a +test failure instead of a hang. +""" + +from isaaclab.app import AppLauncher + +# launch omniverse app +simulation_app = AppLauncher(headless=True).app + +import asyncio +from types import SimpleNamespace + +import pytest +import torch + +from isaaclab_mimic.datagen import generation + +FUSE_STEPS = 40 + + +class _Fuse(Exception): + """Raised by the fake environment once the loop has run for longer than any test expects.""" + + +def _run( + outcomes, + *, + max_num_failures, + generation_num_trials=3, + generation_guarantee=True, + num_envs=1, + attempts_per_step=1, +): + """Run ``env_loop`` over scripted attempt outcomes; return how it ended and the final counters. + + ``attempts_per_step`` is how many of the ``num_envs`` generators finish an attempt on the same + step, which is what decides whether the bound is read before or after the extra attempts land. + + Returns: + A tuple ``(how, num_success, num_failures, num_attempts)`` where ``how`` is ``"exited"`` if + the loop returned on its own and ``"fuse"`` if it was still running after ``FUSE_STEPS``. + """ + generation.num_success = generation.num_failures = generation.num_attempts = 0 + loop = asyncio.get_event_loop() + action_queue: asyncio.Queue = asyncio.Queue() + reset_queue: asyncio.Queue = asyncio.Queue() + for env_id in range(num_envs): + action_queue.put_nowait((env_id, torch.zeros(7))) + scripted = iter(outcomes) + steps = {"n": 0} + + def step(actions): + steps["n"] += 1 + if steps["n"] > FUSE_STEPS: + raise _Fuse(f"env_loop still running after {FUSE_STEPS} steps") + for _ in range(attempts_per_step): + if next(scripted): + generation.num_success += 1 + else: + generation.num_failures += 1 + generation.num_attempts += 1 + for env_id in range(num_envs): + action_queue.put_nowait((env_id, torch.zeros(7))) + + env = SimpleNamespace( + num_envs=num_envs, + device="cpu", + action_space=SimpleNamespace(shape=(num_envs, 7)), + step=step, + reset=lambda env_ids=None: None, + close=lambda: None, + sim=SimpleNamespace(is_stopped=lambda: False), + cfg=SimpleNamespace( + datagen_config=SimpleNamespace( + generation_guarantee=generation_guarantee, + generation_num_trials=generation_num_trials, + max_num_failures=max_num_failures, + ) + ), + ) + try: + generation.env_loop(env, reset_queue, action_queue, None, loop) + how = "exited" + except _Fuse: + how = "fuse" + return how, generation.num_success, generation.num_failures, generation.num_attempts + + +def test_failure_cap_stops_a_run_that_never_succeeds(): + how, succ, fail, attempts = _run([False] * 100, max_num_failures=5) + assert how == "exited" + assert (succ, fail, attempts) == (0, 5, 5) + + +def test_no_cap_by_default_keeps_the_success_guarantee(): + """``None`` means what the guarantee always meant: keep going until enough demos succeed.""" + how, _, fail, _ = _run([False] * 100, max_num_failures=None) + assert how == "fuse" + assert fail == FUSE_STEPS + + +def test_enough_successes_still_end_the_run_first(): + outcomes = [False, True, False, True, True] + [False] * 50 + how, succ, fail, _ = _run(outcomes, max_num_failures=100, generation_num_trials=3) + assert how == "exited" + assert (succ, fail) == (3, 2) + + +def test_cap_reached_before_enough_successes_ends_the_run(): + outcomes = [False, True, False, True, False, False] + [True] * 50 + how, succ, fail, _ = _run(outcomes, max_num_failures=4, generation_num_trials=3) + assert how == "exited" + assert (succ, fail) == (2, 4) + + +@pytest.mark.parametrize("max_num_failures", [None, 3, 100]) +def test_attempt_based_termination_is_unchanged(max_num_failures): + """With the guarantee off the run stops on ``generation_num_trials`` attempts, cap or no cap. + + A cap of 3 against 10 requested attempts is the case that matters: the fixed-attempt contract + says the run delivers the attempts it was asked for, so the cap must not end it at 3. + """ + how, _, _, attempts = _run( + [False] * 100, max_num_failures=max_num_failures, generation_num_trials=10, generation_guarantee=False + ) + assert how == "exited" + assert attempts == 10 + + +def test_bound_is_exact_when_attempts_end_on_separate_steps(): + """Four environments, one attempt landing per step: the bound is read between every attempt.""" + how, _, fail, _ = _run([False] * 100, max_num_failures=5, num_envs=4, attempts_per_step=1) + assert how == "exited" + assert fail == 5 + + +def test_attempts_ending_together_overshoot_the_bound_by_at_most_num_envs_minus_one(): + """Attempts that end on the step that crosses the bound are already complete and still count.""" + how, _, fail, _ = _run([False] * 100, max_num_failures=5, num_envs=4, attempts_per_step=4) + assert how == "exited" + assert 5 < fail <= 5 + 4 - 1 diff --git a/source/isaaclab_newton/changelog.d/hehu-newton-ball-joint-dof-space.rst b/source/isaaclab_newton/changelog.d/hehu-newton-ball-joint-dof-space.rst new file mode 100644 index 000000000000..c9190865df83 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/hehu-newton-ball-joint-dof-space.rst @@ -0,0 +1,18 @@ +Fixed +^^^^^ + +* Fixed :attr:`~isaaclab_newton.assets.articulation.articulation_data.ArticulationData.joint_pos` + exposing Newton's ``joint_q`` -- joint *coordinate* space -- rather than DOF positions. A ball + joint occupies 4 quaternion components against 3 DOFs, so on an articulation containing one the + array was wider than ``num_joints`` while ``joint_names``, ``joint_vel``, ``default_joint_pos`` + and the joint gains stayed in DOF space, and every consumer indexing them with the same joint ids + read a different joint past the first ball joint. Added + :class:`~isaaclab_newton.assets.articulation.joint_coordinates.JointCoordinateMap`, which converts + between the two spaces on read and on write. Articulations whose joints all have one coordinate + per DOF keep the existing zero-copy view and are unaffected. + +* Fixed joint position targets being written in DOF space into Newton's coordinate-layout + ``joint_target_q``. That array follows ``newton.use_coord_layout_targets``, which defaults to + ``True`` from Newton 1.6, so on a ball-jointed articulation the actuators were writing 50 DOF + targets into a 56-wide coordinate array. Targets now go through a DOF-shaped staging buffer and + the same coordinate map. diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/actuator_control.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/actuator_control.py index 9921ca150915..f69b6ad35068 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/actuator_control.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/actuator_control.py @@ -22,6 +22,7 @@ from isaaclab.assets.articulation import ordering_kernels from isaaclab.sim.schemas.schemas_actuators import _validate_newton_native_actuator_cfgs +from isaaclab_newton.assets.articulation.joint_coordinates import scatter_joint_coordinates from isaaclab_newton.physics import NewtonManager as SimulationManager if TYPE_CHECKING: @@ -93,6 +94,9 @@ def finalize_native_actuators(self, collection: ActuatorCollection) -> NewtonAct ) def _post_actuator() -> None: + # Telemetry reads _sim_bind_joint_pos inside the decimation loop, ahead of the + # post-step gather, so the DOF-space view is re-derived here. + articulation._data._gather_joint_coordinates() wp.launch( actuator_kernels.sync_torque_telemetry, dim=(self.num_instances, self.num_joints), @@ -134,7 +138,9 @@ def compute_native_actuators(self, collection: ActuatorCollection, dt: float) -> return self._native_actuator_path_active def submit_commands(self, collection: ActuatorCollection) -> None: + """Publish the collection's targets to the backend arrays.""" articulation = self._articulation + data = articulation.data if self._native_actuator_path_active: # Newton consumes raw explicit-actuator targets through joint_act. user_effort = collection._joint_effort_target @@ -143,12 +149,6 @@ def submit_commands(self, collection: ActuatorCollection) -> None: write_pos_target = True write_vel_target = True write_joint_act = True - if not articulation.data.has_joint_ordering: - articulation.data._sim_bind_joint_position_target.assign(collection._joint_pos_target) - articulation.data._sim_bind_joint_velocity_target.assign(collection._joint_vel_target) - articulation.data._sim_bind_joint_act.assign(collection._joint_effort_target) - articulation.data._sim_bind_joint_effort.assign(collection._joint_effort_target) - return else: # Lab executors publish processed targets; only implicit joints use # the backend position and velocity drives. @@ -158,28 +158,43 @@ def submit_commands(self, collection: ActuatorCollection) -> None: write_pos_target = collection.has_implicit_actuators write_vel_target = collection.has_implicit_actuators write_joint_act = False - if not articulation.data.has_joint_ordering: - articulation.data._sim_bind_joint_effort.assign(collection._joint_effort_target_sim) - if collection.has_implicit_actuators: - articulation.data._sim_bind_joint_position_target.assign(collection._joint_pos_target_sim) - articulation.data._sim_bind_joint_velocity_target.assign(collection._joint_vel_target_sim) - return - - ordering_kernels.launch_reorder_joint_targets_user_to_backend( - user_effort=user_effort, - user_pos_target=user_pos_target, - user_vel_target=user_vel_target, - backend_to_user=articulation._joint_backend_to_user_map(), - write_effort=True, - write_pos_target=write_pos_target, - write_vel_target=write_vel_target, - write_joint_act=write_joint_act, - backend_effort=articulation.data._sim_bind_joint_effort, - backend_pos_target=articulation.data._sim_bind_joint_position_target, - backend_vel_target=articulation.data._sim_bind_joint_velocity_target, - backend_joint_act=articulation.data._sim_bind_joint_act, - device=self.device, - ) + + if data.has_joint_ordering: + ordering_kernels.launch_reorder_joint_targets_user_to_backend( + user_effort=user_effort, + user_pos_target=user_pos_target, + user_vel_target=user_vel_target, + backend_to_user=articulation._joint_backend_to_user_map(), + write_effort=True, + write_pos_target=write_pos_target, + write_vel_target=write_vel_target, + write_joint_act=write_joint_act, + backend_effort=data._sim_bind_joint_effort, + backend_pos_target=data._sim_bind_joint_position_target, + backend_vel_target=data._sim_bind_joint_velocity_target, + backend_joint_act=data._sim_bind_joint_act, + device=self.device, + ) + else: + data._sim_bind_joint_effort.assign(user_effort) + if write_pos_target: + data._sim_bind_joint_position_target.assign(user_pos_target) + if write_vel_target: + data._sim_bind_joint_velocity_target.assign(user_vel_target) + if write_joint_act: + data._sim_bind_joint_act.assign(user_effort) + + # Newton takes coordinate-layout position targets from 1.6 on; every branch above writes + # DOF-indexed targets, so the staging buffer has to be scattered across unconditionally. + # Sequenced after the writes above (not a try/finally) so an exception mid-write leaves the + # staging buffer unflushed rather than scattering a half-written buffer into joint_target_q. + if articulation.data._joint_targets_need_conversion: + scatter_joint_coordinates( + articulation.data._joint_coord_map, + articulation.data._sim_bind_joint_position_target, + articulation.data._sim_bind_joint_target_coords, + articulation._ALL_ENV_MASK, + ) def reset_native_actuators(self, env_ids: Sequence[int] | slice) -> None: if self._native_actuator_path_active and SimulationManager._adapter is not None: diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index 53b23d09ee06..84d107b7e347 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -12,7 +12,7 @@ import re import warnings from collections.abc import Sequence -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import numpy as np import torch @@ -36,6 +36,7 @@ from isaaclab_newton.assets import kernels as shared_kernels from isaaclab_newton.assets.articulation import kernels as articulation_kernels +from isaaclab_newton.assets.articulation.joint_coordinates import scatter_joint_coordinates from isaaclab_newton.physics import NewtonManager as SimulationManager from .actuator_control import NewtonActuatorControl @@ -1200,6 +1201,14 @@ def write_joint_state_to_sim_index( ], device=self.device, ) + # The write landed in DOF space; push it back into Newton's joint coordinates. + if self.data._joint_coord_map.required: + scatter_joint_coordinates( + self.data._joint_coord_map, + self.data._sim_bind_joint_pos, + self.data._sim_bind_joint_coords, + self._env_ids_to_mask(env_ids), + ) # Let the data class handle the invalidation of the pose and velocity related properties. if not skip_forward: self.data._reset_pose(env_ids=env_ids) @@ -1266,6 +1275,14 @@ def write_joint_state_to_sim_mask( ], device=self.device, ) + # The write landed in DOF space; push it back into Newton's joint coordinates. + if self.data._joint_coord_map.required: + scatter_joint_coordinates( + self.data._joint_coord_map, + self.data._sim_bind_joint_pos, + self.data._sim_bind_joint_coords, + env_mask, + ) # Let the data class handle the invalidation of the pose and velocity related properties. if not skip_forward: self.data._reset_pose(env_mask=env_mask) @@ -1319,6 +1336,14 @@ def write_joint_position_to_sim_index( self.data._sim_bind_joint_pos, device=self.device, ) + # The write landed in DOF space; push it back into Newton's joint coordinates. + if self.data._joint_coord_map.required: + scatter_joint_coordinates( + self.data._joint_coord_map, + self.data._sim_bind_joint_pos, + self.data._sim_bind_joint_coords, + self._env_ids_to_mask(env_ids), + ) # Let the data class handle the invalidation of pose- and velocity-dependent properties. if not skip_forward: self.data._reset_pose(env_ids=env_ids) @@ -1369,6 +1394,14 @@ def write_joint_position_to_sim_mask( self.data._sim_bind_joint_pos, device=self.device, ) + # The write landed in DOF space; push it back into Newton's joint coordinates. + if self.data._joint_coord_map.required: + scatter_joint_coordinates( + self.data._joint_coord_map, + self.data._sim_bind_joint_pos, + self.data._sim_bind_joint_coords, + env_mask, + ) # Let the data class handle the invalidation of pose- and velocity-dependent properties. if not skip_forward: self.data._reset_pose(env_mask=env_mask) @@ -2249,14 +2282,24 @@ def write_joint_viscous_friction_coefficient_to_sim_mask( @staticmethod @wp.kernel(enable_backward=False) - def _build_env_mask_kernel(mask: wp.array(dtype=wp.bool), indices: wp.array(dtype=wp.int32)): + def _build_env_mask_kernel(mask: wp.array(dtype=wp.bool), indices: wp.array(dtype=Any)): i = wp.tid() - mask[indices[i]] = True + mask[wp.int32(indices[i])] = True - def _env_ids_to_mask(self, env_ids: wp.array) -> wp.array: - """Convert warp env_ids to a boolean Warp mask.""" + def _env_ids_to_mask(self, env_ids: wp.array | torch.Tensor) -> wp.array: + """Convert env_ids to a boolean Warp mask. + + Args: + env_ids: Environment indices as returned by :meth:`_resolve_env_ids`, which may be a + warp array or a torch tensor of any integer width. + + Returns: + A per-environment boolean mask. + """ if env_ids is self._ALL_INDICES: return self._ALL_ENV_MASK + if isinstance(env_ids, torch.Tensor): + env_ids = wp.from_torch(env_ids) mask = wp.zeros(self.num_instances, dtype=wp.bool, device=self.device) wp.launch(self._build_env_mask_kernel, dim=env_ids.shape[0], inputs=[mask, env_ids], device=self.device) return mask @@ -3444,13 +3487,16 @@ def _create_buffers(self): ) # Republish the Tier-1 backend->user state shadows inside the stepped # (and captured) region after the last solver substep. Registering only - # when ordering is non-identity keeps identity-ordering scenes at zero + # when ordering is non-identity or a ball joint needs the coordinate + # gather keeps a plain identity-ordering, non-ball-joint scene at zero # overhead (empty callback list). The reorders are then recorded into # every captured graph, so passthrough state getters never replay stale. # The stored handle is the exact bound method ``_clear_callbacks`` later # deregisters. + # A ball-jointed articulation also needs the slot: its DOF-space joint_pos is derived, not + # sim-bound, so it has to be republished after every step just like the ordering shadows. self._post_step_callback = None - if self.data.has_joint_ordering or self.data.has_body_ordering: + if self.data.has_joint_ordering or self.data.has_body_ordering or self.data._joint_coord_map.required: self._post_step_callback = self._data._refresh_user_order_state SimulationManager.register_post_step_callback(self._post_step_callback) # tendon names are set in _process_tendons function diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py index 04a29718fc3e..b17e37f3c782 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py @@ -22,6 +22,10 @@ from isaaclab_newton.assets import kernels as shared_kernels from isaaclab_newton.assets.articulation import kernels as articulation_kernels +from isaaclab_newton.assets.articulation.joint_coordinates import ( + build_ball_joint_coordinate_map, + gather_joint_coordinates, +) from isaaclab_newton.physics import NewtonManager as SimulationManager if TYPE_CHECKING: @@ -1591,7 +1595,23 @@ def _create_simulation_bindings(self) -> None: "joint_effort_limit", SimulationManager.get_model() )[:, 0] # -- joint states - self._sim_bind_joint_pos = self._root_view.get_dof_positions(SimulationManager.get_state_0())[:, 0] + # ``get_dof_positions`` returns Newton's ``joint_q``, which is *coordinate* space: a ball + # joint occupies 4 quaternion components against 3 DOFs, so the array is wider than + # ``num_joints`` whenever the articulation has one. IsaacLab addresses joints by DOF + # index everywhere, so keep the coordinate array separate and publish a DOF-space view. + self._sim_bind_joint_coords = self._root_view.get_dof_positions(SimulationManager.get_state_0())[:, 0] + # The view's per-joint counts are already in the column order of the array above and + # already exclude the free root, fixed joints and loop-closing joints. + self._joint_coord_map = build_ball_joint_coordinate_map( + self._root_view.joint_coord_counts, self._root_view.joint_dof_counts, self.device + ) + if self._joint_coord_map.required: + self._sim_bind_joint_pos = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + gather_joint_coordinates(self._joint_coord_map, self._sim_bind_joint_coords, self._sim_bind_joint_pos) + else: + self._sim_bind_joint_pos = self._sim_bind_joint_coords self._sim_bind_joint_vel = self._root_view.get_dof_velocities(SimulationManager.get_state_0())[:, 0] # -- joint commands (sent to the simulation) self._sim_bind_joint_effort = self._root_view.get_attribute("joint_f", SimulationManager.get_control())[ @@ -1601,6 +1621,19 @@ def _create_simulation_bindings(self) -> None: self._sim_bind_joint_position_target = self._root_view.get_attribute( "joint_target_q", SimulationManager.get_control() )[:, 0] + # ``joint_target_q`` follows ``newton.use_coord_layout_targets``, which defaults to + # True from Newton 1.6. Under that layout the array is coordinate-shaped, exactly like + # ``joint_q``, so actuators keep writing DOF-indexed targets into a staging buffer and + # the same map scatters them across. + self._sim_bind_joint_target_coords = self._sim_bind_joint_position_target + self._joint_targets_need_conversion = self._sim_bind_joint_target_coords.shape[1] != self._num_joints + if self._joint_targets_need_conversion: + self._sim_bind_joint_position_target = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + gather_joint_coordinates( + self._joint_coord_map, self._sim_bind_joint_target_coords, self._sim_bind_joint_position_target + ) self._sim_bind_joint_velocity_target = self._root_view.get_attribute( "joint_target_qd", SimulationManager.get_control() )[:, 0] @@ -1629,13 +1662,20 @@ def _create_simulation_bindings(self) -> None: self._sim_bind_joint_effort_limits_sim = wp.zeros( (self._num_instances, 0), dtype=wp.float32, device=self.device ) - self._sim_bind_joint_pos = wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) + # No joints: the view reports no counts, so the map has nothing to convert. + self._sim_bind_joint_coords = self._sim_bind_joint_pos = wp.zeros( + (self._num_instances, 0), dtype=wp.float32, device=self.device + ) + self._joint_coord_map = build_ball_joint_coordinate_map( + self._root_view.joint_coord_counts, self._root_view.joint_dof_counts, self.device + ) self._sim_bind_joint_vel = wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) self._sim_bind_joint_effort = wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) self._sim_bind_joint_act = wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) self._sim_bind_joint_position_target = wp.zeros( (self._num_instances, 0), dtype=wp.float32, device=self.device ) + self._joint_targets_need_conversion = False self._sim_bind_joint_velocity_target = wp.zeros( (self._num_instances, 0), dtype=wp.float32, device=self.device ) @@ -2157,8 +2197,18 @@ def _refresh_user_order_body_state(self) -> None: outputs=[self._body_link_pose_w_user, self._body_com_vel_w_user], ) + def _gather_joint_coordinates(self) -> None: + """Re-derive the DOF-space joint positions from Newton's coordinate array. + + A no-op unless the articulation has a ball joint; without one ``_sim_bind_joint_pos`` is a + zero-copy view onto ``joint_q`` and is always current. + """ + if self._joint_coord_map.required: + gather_joint_coordinates(self._joint_coord_map, self._sim_bind_joint_coords, self._sim_bind_joint_pos) + def _refresh_user_order_state(self) -> None: - """Republish all Tier-1 user-order state shadows from live backend state. + """Republish all Tier-1 user-order state shadows from live backend state, and gather any + ball-joint DOF positions ahead of them. Registered as a post-step callback (see :meth:`isaaclab_newton.physics.NewtonManager.register_post_step_callback`) @@ -2166,8 +2216,12 @@ def _refresh_user_order_state(self) -> None: the last solver substep. With no Python freshness guard the launches are recorded into every captured graph and replayed on each tick, so the passthrough ``joint_pos`` / ``joint_vel`` / ``body_link_pose_w`` / - ``body_com_vel_w`` shadows behave exactly like sim-bound memory. + ``body_com_vel_w`` shadows behave exactly like sim-bound memory. For an + identity-ordered ball-joint articulation the coordinate gather is the only + work this method does -- the two reorder calls below are no-ops. """ + # Ahead of the reorder: the user-order shadows have to gather from a current DOF buffer. + self._gather_joint_coordinates() self._refresh_user_order_joint_state() self._refresh_user_order_body_state() diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/joint_coordinates.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/joint_coordinates.py new file mode 100644 index 000000000000..cb7f041217bb --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/joint_coordinates.py @@ -0,0 +1,232 @@ +# 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 + +"""Conversion between Newton's joint coordinate space and Isaac Lab's DOF space. + +Newton stores a ball joint as a 4-component unit quaternion against 3 DOFs, so an articulation +containing one has more joint coordinates than DOFs. Free and distance joints have the same kind +of mismatch, but Isaac Lab's view excludes free joints and this module rejects distance joints, so +ball is the only mismatched layout this map converts. Every other joint type has one coordinate +per DOF, so the two spaces coincide for most assets and :func:`build_ball_joint_coordinate_map` +reports ``required = False`` for them. + +Isaac Lab addresses joints by DOF index throughout -- ``joint_names``, ``find_joints``, +``SceneEntityCfg.joint_ids`` -- so joint positions have to be exposed in DOF space to stay +consistent with ``joint_vel``, ``default_joint_pos`` and the joint gains, all of which already are. + +The rotation vector is the representation consistent with ``joint_qd``, which holds angular velocity +in the joint frame. PhysX also reports spherical-joint positions as an axis-angle vector projected +onto the DOF axes; sign and frame parity with PhysX has not been verified numerically. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import warp as wp + +_BALL_LAYOUT = (4, 3) +"""``(coordinates, DOFs)`` of a ball joint -- the only mismatched layout this map converts; free +joints are excluded from Lab's view, and distance joints (7 against 6) are rejected.""" + + +@wp.kernel(enable_backward=False) +def gather_single_coord_dofs( + coords: wp.array2d(dtype=wp.float32), + dof_index: wp.array(dtype=wp.int32), + coord_index: wp.array(dtype=wp.int32), + dofs: wp.array2d(dtype=wp.float32), +): + """Copy joints whose coordinate count equals their DOF count.""" + env, i = wp.tid() + dofs[env, dof_index[i]] = coords[env, coord_index[i]] + + +@wp.kernel(enable_backward=False) +def gather_ball_dofs( + coords: wp.array2d(dtype=wp.float32), + dof_index: wp.array(dtype=wp.int32), + coord_index: wp.array(dtype=wp.int32), + dofs: wp.array2d(dtype=wp.float32), +): + """Quaternion -> rotation vector for each ball joint. + + ``wp.quat_to_axis_angle`` resolves the double cover itself -- it flips the axis with the sign of + ``w`` and returns an angle in ``[0, pi]`` -- so a fixed pose always decodes to the same vector. + """ + env, i = wp.tid() + c = coord_index[i] + d = dof_index[i] + axis, angle = wp.quat_to_axis_angle( + wp.quat(coords[env, c + 0], coords[env, c + 1], coords[env, c + 2], coords[env, c + 3]) + ) + rotvec = axis * angle + dofs[env, d + 0] = rotvec[0] + dofs[env, d + 1] = rotvec[1] + dofs[env, d + 2] = rotvec[2] + + +@wp.kernel(enable_backward=False) +def scatter_single_coord_dofs( + dofs: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), + dof_index: wp.array(dtype=wp.int32), + coord_index: wp.array(dtype=wp.int32), + coords: wp.array2d(dtype=wp.float32), +): + """Inverse of :func:`gather_single_coord_dofs`, for the selected environments.""" + env, i = wp.tid() + if env_mask[env]: + coords[env, coord_index[i]] = dofs[env, dof_index[i]] + + +@wp.kernel(enable_backward=False) +def scatter_ball_dofs( + dofs: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), + dof_index: wp.array(dtype=wp.int32), + coord_index: wp.array(dtype=wp.int32), + coords: wp.array2d(dtype=wp.float32), +): + """Rotation vector -> quaternion for each ball joint, for the selected environments.""" + env, i = wp.tid() + if not env_mask[env]: + return + c = coord_index[i] + d = dof_index[i] + rotvec = wp.vec3(dofs[env, d + 0], dofs[env, d + 1], dofs[env, d + 2]) + angle = wp.length(rotvec) + q = wp.quat_identity() + if angle > 1.0e-9: + q = wp.quat_from_axis_angle(rotvec / angle, angle) + coords[env, c + 0] = q[0] + coords[env, c + 1] = q[1] + coords[env, c + 2] = q[2] + coords[env, c + 3] = q[3] + + +@dataclass(frozen=True) +class BallJointCoordinateMap: + """Index tables mapping an articulation view's ball-joint coordinates to their DOFs and back. + + A plain data record, not an object with behavior -- :func:`build_ball_joint_coordinate_map` + produces one, and :func:`gather_joint_coordinates` / :func:`scatter_joint_coordinates` consume + one. Fields are only meaningful when ``required`` is True; the ``as_wp`` conversion in + :func:`build_ball_joint_coordinate_map` is skipped otherwise, so on a required=False map the + remaining fields are placeholder empty arrays. The frozen dataclass prevents field + reassignment, but the Warp arrays remain mutable objects. + """ + + required: bool + """Whether the articulation has a ball joint and the remaining fields are populated.""" + + single_dof: wp.array + """DOF indices of joints whose coordinate count equals their DOF count, dtype ``wp.int32``.""" + + single_coord: wp.array + """Matching coordinate indices for :attr:`single_dof`, dtype ``wp.int32``.""" + + ball_dof: wp.array + """First DOF index of each ball joint, dtype ``wp.int32``.""" + + ball_coord: wp.array + """First coordinate index of each ball joint, dtype ``wp.int32``.""" + + +def build_ball_joint_coordinate_map(coord_counts: list[int], dof_counts: list[int], device) -> BallJointCoordinateMap: + """Build the index tables mapping an articulation view's joint coordinates to its DOFs and back. + + Built from the view's own per-joint counts, which are in the column order of + :meth:`~newton.ArticulationView.get_dof_positions` and already exclude the free root joint, + fixed joints and loop-closing joints. + + Args: + coord_counts: Coordinates per selected joint (``ArticulationView.joint_coord_counts``). + dof_counts: DOFs per selected joint (``ArticulationView.joint_dof_counts``). + device: Device to allocate the index tables on. + + Returns: + The index tables, with ``required`` set when the articulation has a ball joint. + + Raises: + NotImplementedError: If a joint's coordinate and DOF counts differ in any way other than the + ball-joint layout, which would otherwise be silently decoded as a quaternion. + """ + single_dof: list[int] = [] + single_coord: list[int] = [] + ball_dof: list[int] = [] + ball_coord: list[int] = [] + + coord, dof = 0, 0 + for n_coords, n_dofs in zip(coord_counts, dof_counts, strict=True): + if n_coords == n_dofs: + single_dof.extend(range(dof, dof + n_dofs)) + single_coord.extend(range(coord, coord + n_coords)) + elif (n_coords, n_dofs) == _BALL_LAYOUT: + ball_dof.append(dof) + ball_coord.append(coord) + else: + raise NotImplementedError( + f"Joint with {n_coords} coordinates against {n_dofs} DOFs has no coordinate" + " conversion; only ball joints (4 against 3) are supported." + ) + coord += n_coords + dof += n_dofs + + required = bool(ball_dof) + if not required: + empty = wp.array([], dtype=wp.int32, device=device) + return BallJointCoordinateMap(False, empty, empty, empty, empty) + as_wp = lambda values: wp.array(values, dtype=wp.int32, device=device) # noqa: E731 + return BallJointCoordinateMap(True, as_wp(single_dof), as_wp(single_coord), as_wp(ball_dof), as_wp(ball_coord)) + + +def gather_joint_coordinates(coord_map: BallJointCoordinateMap, coords: wp.array, dofs: wp.array) -> None: + """Write the DOF-space view of ``coords`` into ``dofs``. + + Args: + coord_map: Index tables from :func:`build_ball_joint_coordinate_map`. + coords: Newton's joint coordinate array for this view. + dofs: DOF-space destination [rad or m, depending on joint type]. + """ + num_envs = coords.shape[0] + for kernel, dof_index, coord_index in ( + (gather_single_coord_dofs, coord_map.single_dof, coord_map.single_coord), + (gather_ball_dofs, coord_map.ball_dof, coord_map.ball_coord), + ): + wp.launch( + kernel, + dim=(num_envs, dof_index.shape[0]), + inputs=[coords, dof_index, coord_index, dofs], + device=coords.device, + ) + + +def scatter_joint_coordinates( + coord_map: BallJointCoordinateMap, dofs: wp.array, coords: wp.array, env_mask: wp.array +) -> None: + """Write ``dofs`` back into ``coords`` for the selected environments. + + Scoping this to the written environments matters: resets are staggered, and ``dofs`` for + the other environments only holds the last post-step gather. An all-environment scatter + would overwrite their live ``joint_q`` with an exp(log(q)) round trip of itself, perturbing + state that was never written. + + Args: + coord_map: Index tables from :func:`build_ball_joint_coordinate_map`. + dofs: DOF-space joint positions [rad or m, depending on joint type]. + coords: Newton's joint coordinate array to write into. + env_mask: Per-environment boolean selection of the environments that were written. + """ + for kernel, dof_index, coord_index in ( + (scatter_single_coord_dofs, coord_map.single_dof, coord_map.single_coord), + (scatter_ball_dofs, coord_map.ball_dof, coord_map.ball_coord), + ): + wp.launch( + kernel, + dim=(env_mask.shape[0], dof_index.shape[0]), + inputs=[dofs, env_mask, dof_index, coord_index, coords], + device=coords.device, + ) diff --git a/source/isaaclab_newton/isaaclab_newton/test/fixtures/views/mock_articulation_view.py b/source/isaaclab_newton/isaaclab_newton/test/fixtures/views/mock_articulation_view.py index 3820642a0422..7300f3f2175b 100644 --- a/source/isaaclab_newton/isaaclab_newton/test/fixtures/views/mock_articulation_view.py +++ b/source/isaaclab_newton/isaaclab_newton/test/fixtures/views/mock_articulation_view.py @@ -240,6 +240,7 @@ def __init__( joint_names: list[str] | None = None, body_names: list[str] | None = None, tendon_names: list[str] | None = None, + joint_coord_counts: list[int] | None = None, ): """Initialize the mock Newton articulation view. @@ -253,6 +254,10 @@ def __init__( joint_names: Names of joints. Defaults to ["joint_0", ...]. body_names: Names of bodies. Defaults to ["body_0", ...]. tendon_names: Names of fixed tendons. Defaults to ["tendon_0", ...]. + joint_coord_counts: Coordinate count per selected joint, e.g. ``[1, 4, 1]`` for a + three-joint articulation with a ball joint in the middle. Must sum with its + matching per-joint DOF counts (inferred as 3 in place of any ``4``, 1 otherwise) + to ``num_joints`` DOFs. Defaults to one coordinate per DOF -- no ball joints. """ self._count = num_instances self._link_count = num_bodies @@ -261,6 +266,10 @@ def __init__( self._device = device self._is_fixed_base = is_fixed_base self._noop_setters = False + self._joint_coord_counts_override = joint_coord_counts + # Real Newton's ``get_dof_positions`` returns ``joint_q`` (coordinate space), wider than the + # DOF count whenever a ball joint is present; fake that width here too. + self._joint_coord_count_total = sum(joint_coord_counts) if joint_coord_counts is not None else num_joints # Set joint and body names self._joint_dof_names = joint_names if joint_names is not None else [f"joint_{i}" for i in range(num_joints)] @@ -320,6 +329,20 @@ def joint_dof_count(self) -> int: """Number of DOFs (joints) per instance.""" return self._joint_dof_count + @property + def joint_dof_counts(self) -> list[int]: + """DOF count per selected joint. One DOF per joint unless a ball layout was requested.""" + if self._joint_coord_counts_override is None: + return [1] * self._joint_dof_count + return [3 if n_coords == 4 else 1 for n_coords in self._joint_coord_counts_override] + + @property + def joint_coord_counts(self) -> list[int]: + """Coordinate count per selected joint. Equal to :attr:`joint_dof_counts` unless a ball layout was requested.""" + if self._joint_coord_counts_override is None: + return [1] * self._joint_dof_count + return list(self._joint_coord_counts_override) + @property def is_fixed_base(self) -> bool: """Whether the articulation has a fixed base.""" @@ -396,10 +419,10 @@ def _ensure_link_velocities(self) -> wp.array | None: return self._link_velocities def _ensure_dof_positions(self) -> wp.array: - """Lazily create DOF positions.""" + """Lazily create DOF positions, in coordinate space (see ``_joint_coord_count_total``).""" if self._dof_positions is None: self._dof_positions = wp.zeros( - (self._count, 1, self._joint_dof_count), dtype=wp.float32, device=self._device + (self._count, 1, self._joint_coord_count_total), dtype=wp.float32, device=self._device ) return self._dof_positions @@ -719,8 +742,10 @@ def set_random_mock_data(self) -> None: link_vel_np = np.random.randn(N, 1, L, 6).astype(np.float32) self._link_velocities = wp.array(link_vel_np, dtype=wp.spatial_vectorf, device=dev) - # DOF state - self._dof_positions = wp.array(np.random.randn(N, 1, J).astype(np.float32), dtype=wp.float32, device=dev) + # DOF state -- positions are coordinate-space width (see ``_joint_coord_count_total``), + # velocities are always DOF-space width. + C = self._joint_coord_count_total + self._dof_positions = wp.array(np.random.randn(N, 1, C).astype(np.float32), dtype=wp.float32, device=dev) self._dof_velocities = wp.array(np.random.randn(N, 1, J).astype(np.float32), dtype=wp.float32, device=dev) # Body properties diff --git a/source/isaaclab_newton/test/assets/test_joint_coordinates.py b/source/isaaclab_newton/test/assets/test_joint_coordinates.py new file mode 100644 index 000000000000..216ca4b7c5ee --- /dev/null +++ b/source/isaaclab_newton/test/assets/test_joint_coordinates.py @@ -0,0 +1,210 @@ +# 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 + +"""Kitless tests for the Newton joint coordinate/DOF conversion.""" + +from unittest.mock import MagicMock + +import numpy as np +import pytest +import warp as wp +from isaaclab_newton.assets.articulation.joint_coordinates import ( + BallJointCoordinateMap, + build_ball_joint_coordinate_map, + gather_joint_coordinates, + scatter_joint_coordinates, +) +from scipy.spatial.transform import Rotation + +# One revolute, one ball, one revolute -- the layout that hides an off-by-one when the tables are +# built by walking the model instead of the view's own per-joint counts. +COORD_COUNTS = [1, 4, 1] +DOF_COUNTS = [1, 3, 1] + +# Two balls, straddling a revolute on each side -- the second ball's offsets (ball_coord[1] = 5, +# ball_dof[1] = 4) are not exercised by the single-ball layout above. +TWO_BALL_COORD_COUNTS = [1, 4, 4, 1] +TWO_BALL_DOF_COUNTS = [1, 3, 3, 1] + + +def _map() -> BallJointCoordinateMap: + return build_ball_joint_coordinate_map(COORD_COUNTS, DOF_COUNTS, "cpu") + + +def test_tables_cover_every_dof() -> None: + """Every DOF of every joint must be tabulated exactly once.""" + m = _map() + covered = list(m.single_dof.numpy()) + [b + k for b in m.ball_dof.numpy() for k in range(3)] + assert sorted(covered) == list(range(sum(DOF_COUNTS))) + assert sorted(list(m.single_coord.numpy()) + [b + k for b in m.ball_coord.numpy() for k in range(4)]) == list( + range(sum(COORD_COUNTS)) + ) + + +def test_two_ball_tables_cover_every_dof() -> None: + """A second ball joint's offsets are not a simple repeat of the first's.""" + m = build_ball_joint_coordinate_map(TWO_BALL_COORD_COUNTS, TWO_BALL_DOF_COUNTS, "cpu") + assert list(m.ball_dof.numpy()) == [1, 4] + assert list(m.ball_coord.numpy()) == [1, 5] + covered = list(m.single_dof.numpy()) + [b + k for b in m.ball_dof.numpy() for k in range(3)] + assert sorted(covered) == list(range(sum(TWO_BALL_DOF_COUNTS))) + assert sorted(list(m.single_coord.numpy()) + [b + k for b in m.ball_coord.numpy() for k in range(4)]) == list( + range(sum(TWO_BALL_COORD_COUNTS)) + ) + + +def test_zero_rotation_vector_scatters_to_identity_quaternion() -> None: + """The ``angle > 1e-9`` guard's else-branch -- taken on every reset of a passive ball joint at + ``default_joint_pos = 0`` -- must produce the identity quaternion, not ``0/0``.""" + m = _map() + n_dofs, n_coords = sum(DOF_COUNTS), sum(COORD_COUNTS) + dofs = wp.zeros((1, n_dofs), dtype=wp.float32, device="cpu") + coords = wp.zeros((1, n_coords), dtype=wp.float32, device="cpu") + mask = wp.array(np.array([True]), dtype=wp.bool, device="cpu") + + scatter_joint_coordinates(m, dofs, coords, mask) + + ball_coord = int(m.ball_coord.numpy()[0]) + np.testing.assert_array_equal(coords.numpy()[0, ball_coord : ball_coord + 4], [0.0, 0.0, 0.0, 1.0]) + + +@pytest.mark.parametrize("sign", [1.0, -1.0]) +def test_identity_quaternion_gathers_to_zero_rotation_vector(sign: float) -> None: + """The inverse of the guard above: the identity quaternion (either hemisphere) decodes to the + zero rotation vector, not a NaN from a degenerate axis normalization.""" + m = _map() + n_dofs, n_coords = sum(DOF_COUNTS), sum(COORD_COUNTS) + coords_np = np.zeros((1, n_coords), dtype=np.float32) + ball_coord = int(m.ball_coord.numpy()[0]) + coords_np[0, ball_coord : ball_coord + 4] = np.array([0.0, 0.0, 0.0, 1.0]) * sign + out = wp.zeros((1, n_dofs), dtype=wp.float32, device="cpu") + + gather_joint_coordinates(m, wp.array(coords_np, dtype=wp.float32, device="cpu"), out) + + ball_dof = int(m.ball_dof.numpy()[0]) + np.testing.assert_array_equal(out.numpy()[0, ball_dof : ball_dof + 3], [0.0, 0.0, 0.0]) + + +def test_joint_pos_is_dof_shaped_on_a_ball_jointed_mock() -> None: + """Regression for the original bug: :attr:`ArticulationData.joint_pos` must be DOF-shaped + even when the underlying view's ``joint_q`` is wider (a ball joint). This exercises + ``_create_simulation_bindings`` end to end, unlike the coordinate-table-only tests + above -- it fails on the pre-fix binding, which read ``get_dof_positions()`` (coordinate + space) straight into ``joint_pos`` with no conversion, so ``joint_pos.shape[1]`` would be + ``sum(TWO_BALL_COORD_COUNTS) == 10`` instead of ``sum(TWO_BALL_DOF_COUNTS) == 8``. + """ + import isaaclab_newton.assets.articulation.articulation_data as newton_data_module + from isaaclab_newton.assets.articulation.articulation_data import ArticulationData + from isaaclab_newton.test.fixtures.views import MockNewtonArticulationView + + num_instances = 2 + num_dofs = sum(TWO_BALL_DOF_COUNTS) + num_bodies = len(TWO_BALL_DOF_COUNTS) + 1 + + mock_view = MockNewtonArticulationView( + num_instances=num_instances, + num_bodies=num_bodies, + num_joints=num_dofs, + device="cpu", + is_fixed_base=True, + joint_coord_counts=TWO_BALL_COORD_COUNTS, + ) + mock_view.set_random_mock_data() + mock_view._noop_setters = True + + mock_model = MagicMock() + mock_model.world_count = num_instances + mock_model.gravity = wp.array( + np.tile(np.array([[0.0, 0.0, -9.81]], dtype=np.float32), (num_instances + 1, 1)), + dtype=wp.vec3f, + device="cpu", + ) + mock_model.articulation_count = num_instances + mock_model.max_joints_per_articulation = num_bodies + mock_model.max_dofs_per_articulation = num_dofs + mock_model.joint_dof_count = num_instances * num_dofs + mock_model.body_count = num_instances * num_bodies + + mock_manager = MagicMock() + mock_manager.get_model.return_value = mock_model + mock_manager.get_state_0.return_value = MagicMock() + mock_manager.get_state_1.return_value = MagicMock() + mock_manager.get_control.return_value = MagicMock() + + original_sim_manager = newton_data_module.SimulationManager + newton_data_module.SimulationManager = mock_manager + try: + data = ArticulationData(mock_view, "cpu") + finally: + newton_data_module.SimulationManager = original_sim_manager + + assert data.joint_pos.torch.shape[1] == num_dofs + + +def test_map_is_inert_without_ball_joints() -> None: + """An articulation whose joints all have one coordinate per DOF needs no conversion.""" + assert not build_ball_joint_coordinate_map([1, 1, 1], [1, 1, 1], "cpu").required + assert not build_ball_joint_coordinate_map([], [], "cpu").required + + +def test_unsupported_layout_is_rejected() -> None: + """A distance joint (7 coordinates, 6 DOFs) must not be decoded as a quaternion.""" + with pytest.raises(NotImplementedError, match="7 coordinates against 6 DOFs"): + build_ball_joint_coordinate_map([7], [6], "cpu") + + +@pytest.mark.parametrize("num_envs", [1, 2]) +def test_scatter_then_gather_round_trips(num_envs: int) -> None: + """DOF values survive a trip through coordinate space, at one environment and at two.""" + m = _map() + n_dofs, n_coords = sum(DOF_COUNTS), sum(COORD_COUNTS) + rng = np.random.default_rng(0) + dofs_np = rng.uniform(-0.7, 0.7, size=(num_envs, n_dofs)).astype(np.float32) + dofs = wp.array(dofs_np, dtype=wp.float32, device="cpu") + coords = wp.zeros((num_envs, n_coords), dtype=wp.float32, device="cpu") + mask = wp.array(np.ones(num_envs, dtype=bool), dtype=wp.bool, device="cpu") + + scatter_joint_coordinates(m, dofs, coords, mask) + # The quaternion the scatter wrote must match an independent exp map. + ball_coord = int(m.ball_coord.numpy()[0]) + ball_dof = int(m.ball_dof.numpy()[0]) + for env in range(num_envs): + expected = Rotation.from_rotvec(dofs_np[env, ball_dof : ball_dof + 3]).as_quat() + np.testing.assert_allclose(coords.numpy()[env, ball_coord : ball_coord + 4], expected, atol=1e-6) + + out = wp.zeros((num_envs, n_dofs), dtype=wp.float32, device="cpu") + gather_joint_coordinates(m, coords, out) + np.testing.assert_allclose(out.numpy(), dofs_np, atol=1e-5) + + +def test_gather_is_invariant_to_quaternion_sign() -> None: + """``q`` and ``-q`` are the same rotation and must decode to the same rotation vector.""" + m = _map() + n_dofs, n_coords = sum(DOF_COUNTS), sum(COORD_COUNTS) + c = int(m.ball_coord.numpy()[0]) + base = np.zeros((1, n_coords), dtype=np.float32) + base[0, c : c + 4] = Rotation.from_rotvec([0.2, -0.5, 0.1]).as_quat() + + decoded = [] + for sign in (1.0, -1.0): + coords = base.copy() + coords[0, c : c + 4] *= sign + out = wp.zeros((1, n_dofs), dtype=wp.float32, device="cpu") + gather_joint_coordinates(m, wp.array(coords, dtype=wp.float32, device="cpu"), out) + decoded.append(out.numpy().copy()) + np.testing.assert_allclose(decoded[0], decoded[1], atol=1e-6) + + +def test_scatter_only_touches_masked_environments() -> None: + """Resets are staggered, so an unmasked environment's coordinates must not move.""" + m = _map() + n_dofs, n_coords = sum(DOF_COUNTS), sum(COORD_COUNTS) + coords_np = np.full((2, n_coords), 0.25, dtype=np.float32) + coords = wp.array(coords_np, dtype=wp.float32, device="cpu") + dofs = wp.array(np.full((2, n_dofs), 0.3, dtype=np.float32), dtype=wp.float32, device="cpu") + + scatter_joint_coordinates(m, dofs, coords, wp.array(np.array([True, False]), dtype=wp.bool, device="cpu")) + assert not np.allclose(coords.numpy()[0], coords_np[0]) + np.testing.assert_array_equal(coords.numpy()[1], coords_np[1]) diff --git a/source/isaaclab_tasks/changelog.d/benchmark-task-package.minor.rst b/source/isaaclab_tasks/changelog.d/benchmark-task-package.minor.rst new file mode 100644 index 000000000000..43b0be9ed46c --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/benchmark-task-package.minor.rst @@ -0,0 +1,10 @@ +Changed +^^^^^^^ + +* Moved benchmark-only tasks into a new ``isaaclab_tasks.benchmark`` package so they are no longer mixed in with + the trainable ``core`` and ``contrib`` task families. The ``IsaacContrib-Reorient-Cube-Shadow-Camera-Benchmark-Direct`` + task id is unchanged, but its configuration moved from + ``isaaclab_tasks.contrib.reorient.config.shadow_hand.shadow_hand_camera_benchmark_env_cfg`` to + ``isaaclab_tasks.benchmark.shadow_hand_camera.shadow_hand_camera_benchmark_env_cfg``. Code that imports the + configuration class directly must update the import path; code that resolves the task through ``gym.make`` or + ``load_cfg_from_registry`` needs no change. diff --git a/source/isaaclab_tasks/changelog.d/fix-franka-lift-collider-variant.rst b/source/isaaclab_tasks/changelog.d/fix-franka-lift-collider-variant.rst new file mode 100644 index 000000000000..c52a004b6e7e --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/fix-franka-lift-collider-variant.rst @@ -0,0 +1,5 @@ +Fixed +^^^^^ + +* Fixed Franka Lift and Reorient reset sampling with updated Franka assets by explicitly selecting the convex-hull + arm colliders used by the tasks' clearance criteria. diff --git a/source/isaaclab_tasks/changelog.d/fix-franka-pour-convex-hulls.rst b/source/isaaclab_tasks/changelog.d/fix-franka-pour-convex-hulls.rst new file mode 100644 index 000000000000..3826c316199c --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/fix-franka-pour-convex-hulls.rst @@ -0,0 +1,5 @@ +Fixed +^^^^^ + +* Fixed the Franka Pour task selecting the primitive robot collider variant, + which omitted the arm collision meshes required by the task. diff --git a/source/isaaclab_tasks/changelog.d/fix-skip-drlegs-walk-contrib-test.skip b/source/isaaclab_tasks/changelog.d/fix-skip-drlegs-walk-contrib-test.skip new file mode 100644 index 000000000000..e29cd024514a --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/fix-skip-drlegs-walk-contrib-test.skip @@ -0,0 +1,3 @@ +Test-only change: skip ``IsaacContrib-DrLegs-Walk`` (and its deprecated alias) in the contributed +environment smoke test because the Kamino solver intermittently diverges to NaN under random actions. +No user-visible behavior change. diff --git a/source/isaaclab_tasks/changelog.d/hehu-digit-newton-velocity-tasks.rst b/source/isaaclab_tasks/changelog.d/hehu-digit-newton-velocity-tasks.rst new file mode 100644 index 000000000000..db6ccfdd69ae --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/hehu-digit-newton-velocity-tasks.rst @@ -0,0 +1,17 @@ +Added +^^^^^ + +* Added a ``newton_mjwarp`` branch to ``DigitPhysicsCfg``, so the Digit velocity tasks run on the + Newton backend with ``presets=newton_mjwarp``. Gated on that preset: ``self_collision_enabled`` + (the asset authors ``enabledSelfCollisions=False`` and ``DIGIT_V4_CFG`` sets no + ``articulation_props``, so Newton filters all 253 intra-articulation shape pairs). An armature + floor for the ten joints below MJWarp's observed stability threshold (expressed as a second + actuator group) and ``entropy_coef = 0.005`` apply on **both** backends, consistent with #7607's + direction of not special-casing PhysX for values that do not hurt it there. + +Fixed +^^^^^ + +* Fixed 32 ``CollisionAPI`` prims on Digit's RealSense camera decoration meshes -- glass, USB-C and + case halves -- becoming collision shapes. This applies on **both** backends, so the PhysX contact + behaviour of the Digit tasks changes: those 32 shapes no longer collide. diff --git a/source/isaaclab_tasks/changelog.d/multitask-manipulation.minor.rst b/source/isaaclab_tasks/changelog.d/multitask-manipulation.minor.rst new file mode 100644 index 000000000000..eea0c7cd6865 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/multitask-manipulation.minor.rst @@ -0,0 +1,8 @@ +Added +^^^^^ + +* Added the contributed ``IsaacContrib-Multitask-Manipulation`` manager-based task for training one policy across heterogeneous + OpenArm lift, Franka cabinet, and UR10 reach environments using selection-aware physics-view indexing, including + task-specific Gaussian policy and value heads, task-wise PPO advantage normalization, bounded cabinet joint targets, + cabinet state-safety resets, task-balanced OpenArm reward scaling, and playback visualization with shortened + evaluation horizons and task-appropriate command markers. diff --git a/source/isaaclab_tasks/changelog.d/render-benchmark-task.rst b/source/isaaclab_tasks/changelog.d/render-benchmark-task.rst new file mode 100644 index 000000000000..ccbd42cd669f --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/render-benchmark-task.rst @@ -0,0 +1,5 @@ +Added +^^^^^ + +* Added the ``Isaac-RenderBenchmark-Franka-Cabinet`` direct task, an animated Franka-and-cabinet scene with no + policy or reward used to compare rendering backends under ``scripts/benchmarks/benchmark_renderer.py``. diff --git a/source/isaaclab_tasks/changelog.d/unify-velocity-solver-inputs.rst b/source/isaaclab_tasks/changelog.d/unify-velocity-solver-inputs.rst new file mode 100644 index 000000000000..c29d3af3d725 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/unify-velocity-solver-inputs.rst @@ -0,0 +1,7 @@ +Changed +^^^^^^^ + +* Unified rough-velocity task inputs across physics backends by removing MJWarp-only actuator armatures, + using 5,000 G1 training iterations for every backend, and representing shared base-COM randomization as a + plain event. Downstream configurations that require the former backend-specific behavior should set it + explicitly. diff --git a/source/isaaclab_tasks/isaaclab_tasks/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/__init__.py index 02cb223e544c..19ec4b1f023f 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/__init__.py @@ -10,6 +10,9 @@ - ``core``: Core task families maintained as part of Isaac Lab. - ``contrib``: Contributed task families. These may depend on ``core`` tasks, but ``core`` tasks never depend on ``contrib`` tasks. +- ``benchmark``: Benchmark-only task families used to measure simulation and rendering + throughput. These may depend on ``core`` and ``contrib`` tasks, but neither depends + on ``benchmark`` tasks. - ``utils``: These include utility functions for the tasks. """ diff --git a/source/isaaclab_tasks/isaaclab_tasks/benchmark/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/benchmark/__init__.py new file mode 100644 index 000000000000..4bf835f15131 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/benchmark/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Benchmark-only task families. + +These tasks exist to measure simulation and rendering throughput rather than to train +policies. They may depend on ``core`` and ``contrib`` tasks, but neither depends on them. +""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/benchmark/render_benchmark/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/benchmark/render_benchmark/__init__.py new file mode 100644 index 000000000000..7ab3de636b7a --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/benchmark/render_benchmark/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Render-only benchmark scene for OVRTX / Newton-Warp comparison.""" + +import gymnasium as gym + +gym.register( + id="Isaac-RenderBenchmark-Franka-Cabinet", + entry_point=f"{__name__}.render_benchmark_env:RenderBenchmarkEnv", + disable_env_checker=True, + kwargs={"env_cfg_entry_point": f"{__name__}.render_benchmark_env_cfg:RenderBenchmarkFrankaCabinetEnvCfg"}, +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/benchmark/render_benchmark/render_benchmark_env.py b/source/isaaclab_tasks/isaaclab_tasks/benchmark/render_benchmark/render_benchmark_env.py new file mode 100644 index 000000000000..4c71338e4f81 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/benchmark/render_benchmark/render_benchmark_env.py @@ -0,0 +1,190 @@ +# 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 + +"""Direct environment for render-only benchmarking.""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +import torch + +import isaaclab.sim as sim_utils +from isaaclab import cloner +from isaaclab.assets import Articulation, RigidObject, RigidObjectCfg +from isaaclab.envs import DirectRLEnv +from isaaclab.sensors import Camera, save_images_to_file + +if TYPE_CHECKING: + from .render_benchmark_env_cfg import RenderBenchmarkFrankaCabinetEnvCfg + + +class RenderBenchmarkEnv(DirectRLEnv): + """Environment that animates its articulations so a renderer can be profiled on them. + + Every articulation in :attr:`RenderBenchmarkFrankaCabinetEnvCfg.articulations` is driven by a + sinusoid around its default joint positions. Actions are ignored and rewards are zero: the + only output that matters is the camera image, and the only cost that matters is the time + spent producing it. + """ + + cfg: RenderBenchmarkFrankaCabinetEnvCfg + + # --- scene construction -------------------------------------------------- + + def _setup_scene(self): + # Per-(env, joint) sinusoid phases [rad] and the elapsed animation time [s]. The phases + # are filled on the first step; see :meth:`_sample_animation_phases`. + self._anim_phases: dict[str, torch.Tensor] | None = None + self._anim_time: float = 0.0 + + self._articulations: dict[str, Articulation] = { + name: Articulation(articulation_cfg) for name, articulation_cfg in self.cfg.articulations.items() + } + self._tiled_camera = Camera(self.cfg.tiled_camera) + self._ground = RigidObject(self._ground_cfg()) + self._spawn_lights() + + # Assets must exist before planning: the plan keys its rows off the cfgs they queued. + src, dest = "/World/envs/env_0", "/World/envs/env_{}" + pos = cloner.grid_transforms(self.scene.num_envs, self.scene.cfg.env_spacing)[0] + plan = cloner.clone_plan_from_env_0(src, dest, self.scene.num_envs, pos) + cloner.replicate(plan) + + # PhysX replication requires explicit collision filtering between environments. + if "physx" in self.scene.physics_backend: + self.scene.filter_collisions(global_prim_paths=[]) + + for name, articulation in self._articulations.items(): + self.scene.articulations[name] = articulation + self.scene.rigid_objects["ground"] = self._ground + self.scene.sensors["tiled_camera"] = self._tiled_camera + + def _ground_cfg(self) -> RigidObjectCfg: + """Build the per-environment ground cuboid. + + A flat cuboid routed through :class:`~isaaclab.assets.RigidObject` keeps the comparison + apples-to-apples: a USD ground plane is invisible to Warp's renderer, which only knows + about simulation meshes, so the two renderers would otherwise disagree on primary-ray + misses and shadow-ray counts. + + The tile is clamped to the env spacing so neighboring clones tile edge-to-edge instead of + overlapping. An oversized tile would leave every environment's camera looking at several + redundant, coplanar ground meshes -- extra geometry that neither backend needs to resolve + in the real scene, and whose BVH/intersection cost grows with the environment count. + """ + spacing = self.scene.cfg.env_spacing + size_x = min(self.cfg.ground_size[0], spacing) + size_y = min(self.cfg.ground_size[1], spacing) + thickness = self.cfg.ground_thickness + return RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Ground", + spawn=sim_utils.CuboidCfg( + size=(size_x, size_y, thickness), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=self.cfg.ground_color, metallic=0.0), + rigid_props=sim_utils.RigidBodyPropertiesCfg(disable_gravity=True, kinematic_enabled=True), + mass_props=sim_utils.MassPropertiesCfg(mass=1.0), + collision_props=sim_utils.CollisionPropertiesCfg(), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, self.cfg.ground_top_z - thickness / 2.0)), + ) + + def _spawn_lights(self) -> None: + """Spawn the ambient dome light and, when configured, the directional light above it.""" + dome_cfg = sim_utils.DomeLightCfg(intensity=self.cfg.dome_light_intensity, color=(0.75, 0.75, 0.75)) + dome_cfg.func("/World/Light", dome_cfg) + if self.cfg.light_cfg is not None: + self.cfg.light_cfg.func( + "/World/LightDirectional", + self.cfg.light_cfg, + orientation=self.cfg.light_orientation, + ) + + # --- joint animation ----------------------------------------------------- + + def _pre_physics_step(self, actions: torch.Tensor) -> None: + if self.cfg.joint_animation_amplitude <= 0.0: + return + if self._anim_phases is None: + self._anim_phases = self._sample_animation_phases() + self._animate_joints() + + def _apply_action(self) -> None: + pass + + def _sample_animation_phases(self) -> dict[str, torch.Tensor]: + """Derive a phase offset [rad] for every (env, joint) pair, keyed by articulation name. + + The phases come from a deterministic hash of the indices rather than the torch RNG so a + given frame holds the same pose in every run. The default RNG state varies between + processes, which would leave an OVRTX run and a Warp run rendering different poses at the + same frame even though the solver itself is deterministic, making their timings + incomparable. Multiplying by two coprime primes spreads the phases pseudo-uniformly. + + Deferred to the first step because it reads articulation data, which the simulation only + populates after :meth:`_setup_scene` returns. + """ + phases = {} + for name, articulation in self._articulations.items(): + default_pos = articulation.data.default_joint_pos.torch + num_envs, num_joints = default_pos.shape + env_idx = torch.arange(num_envs, device=self.device, dtype=default_pos.dtype).unsqueeze(1) + joint_idx = torch.arange(num_joints, device=self.device, dtype=default_pos.dtype).unsqueeze(0) + phases[name] = ((env_idx * 7919.0 + joint_idx * 6553.0) % 10007.0) * (2.0 * math.pi / 10007.0) + return phases + + def _animate_joints(self) -> None: + """Advance the animation clock and drive every joint to its sinusoidal target.""" + self._anim_time += self.cfg.sim.dt * self.cfg.decimation + omega = 2.0 * math.pi * self.cfg.joint_animation_freq_hz + for name, articulation in self._articulations.items(): + default_pos = articulation.data.default_joint_pos.torch + offset = self.cfg.joint_animation_amplitude * torch.sin(omega * self._anim_time + self._anim_phases[name]) + soft_limits = articulation.data.soft_joint_pos_limits.torch + target = torch.clamp(default_pos + offset, soft_limits[..., 0], soft_limits[..., 1]) + articulation.actuators.target_command.set_position_index(value=target) + + # --- DirectRLEnv plumbing ------------------------------------------------ + + def _get_observations(self) -> dict: + # Sensor buffers update lazily, so reading the camera's data is what drives the render. + # This access is the work the benchmark measures: keep it unconditional even when no + # image is written, or the profile records a scene that was never rendered. + output = self._tiled_camera.data.output + if self.cfg.write_image_to_file: + self._write_camera_image(output) + return {"policy": torch.zeros((self.num_envs, 1), device=self.device)} + + def _write_camera_image(self, output: dict) -> None: + """Dump the current camera frame to a PNG. + + Purely a debugging aid, so any failure is reported and swallowed rather than allowed to + abort a benchmark run. + + Args: + output: Rendered outputs from the tiled camera, keyed by data type. + """ + data_type = self.cfg.tiled_camera.data_types[0] + if not isinstance(output, dict) or data_type not in output: + return + try: + image = output[data_type] + if not torch.is_tensor(image): + image = torch.from_dlpack(image) + image = image.float() + # Renderers hand back either 0-255 or already-normalized data. + if image.max() > 1.5: + image = image / 255.0 + save_images_to_file(image[:, ..., :3], f"render_benchmark_{data_type}.{self.common_step_counter:06d}.png") + except Exception as error: + print(f"[render_benchmark] write_image_to_file failed: {error}") + + def _get_rewards(self) -> torch.Tensor: + return torch.zeros(self.num_envs, device=self.device) + + def _get_dones(self): + time_out = self.episode_length_buf >= self.max_episode_length - 1 + return torch.zeros_like(time_out), time_out diff --git a/source/isaaclab_tasks/isaaclab_tasks/benchmark/render_benchmark/render_benchmark_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/benchmark/render_benchmark/render_benchmark_env_cfg.py new file mode 100644 index 000000000000..d5459c316dd8 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/benchmark/render_benchmark/render_benchmark_env_cfg.py @@ -0,0 +1,213 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os + +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +from isaaclab_newton.renderers import NewtonWarpRendererCfg +from isaaclab_physx.physics import PhysxCfg + +import isaaclab.sim as sim_utils +from isaaclab.actuators import ImplicitActuatorCfg +from isaaclab.assets import ArticulationCfg +from isaaclab.envs import DirectRLEnvCfg +from isaaclab.physics import PhysxAutoCfg +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sensors import CameraCfg +from isaaclab.sim import SimulationCfg +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.configclass import configclass + +from isaaclab_tasks.utils import PresetCfg +from isaaclab_tasks.utils.presets import MultiBackendRendererCfg + +from isaaclab_assets.robots.franka import FRANKA_PANDA_HIGH_PD_CFG + + +@configclass +class RenderBenchmarkPhysicsCfg(PresetCfg): + """Physics backend presets. + + Pick via ``presets=newton_mjwarp`` (default) or ``presets=physx``, which resolves to the + concrete PhysX backend at launch. Use ``presets=isaacsim_physx`` to pin Isaac Sim PhysX. + The BVH constructors are read from the environment so ``benchmark_renderer.py`` can sweep + them without a separate preset per combination. + """ + + newton_mjwarp: NewtonCfg = NewtonCfg( + solver_cfg=MJWarpSolverCfg(solver="newton", integrator="implicitfast", njmax=200, nconmax=70), + num_substeps=2, + bvh_constructor_geometry=os.getenv("NEWTON_BVH_GEOMETRY", "cubql"), + bvh_constructor_gaussian=os.getenv("NEWTON_BVH_GAUSSIAN", "cubql"), + bvh_constructor_scene=os.getenv("NEWTON_BVH_SCENE", "sah"), + use_cuda_graph=os.getenv("NEWTON_USE_CUDA_GRAPH", "0") == "1", + ) + isaacsim_physx: PhysxCfg = PhysxCfg() + physx: PhysxAutoCfg = PhysxAutoCfg(isaacsim_physx=isaacsim_physx) + default = newton_mjwarp + + +@configclass +class RenderBenchmarkTiledCameraCfg(PresetCfg): + """Render-target presets — pick via ``presets=rgb`` (default), ``presets=depth``, and so on.""" + + @configclass + class BaseRenderBenchmarkCameraCfg(CameraCfg): + """Front view of the workspace, pitch-only (no yaw, no roll). + + The camera sits on the drawer-facing ``+X`` side of the cabinet, elevated and pitched + ~34.5 deg down so the workspace reads as a slight bird's-eye. ``rot`` is ``(qx, qy, qz, qw)`` + in ``convention="world"`` (``+X`` camera-forward): with ``qw=0`` and only ``qx``/``qz`` set it + is a 180 deg rotation about the ``(-X, 0, +Z)`` axis, aiming the camera at ``(0.4, 0, 0.4)``. + """ + + prim_path: str = "{ENV_REGEX_NS}/Camera" + offset: CameraCfg.OffsetCfg = CameraCfg.OffsetCfg( + pos=(2.0, 0.0, 1.5), rot=(-0.296, 0.0, 0.955, 0.0), convention="world" + ) + data_types: list[str] = [] + spawn: sim_utils.PinholeCameraCfg = sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.05, 50.0) + ) + width: int = int(os.getenv("BENCHMARK_RENDER_RESOLUTION", "256")) + height: int = int(os.getenv("BENCHMARK_RENDER_RESOLUTION", "256")) + renderer_cfg: MultiBackendRendererCfg = MultiBackendRendererCfg( + newton_renderer=NewtonWarpRendererCfg(enable_shadows=True), + ) + + default = BaseRenderBenchmarkCameraCfg(data_types=["rgb"]) + rgb = default + albedo = BaseRenderBenchmarkCameraCfg(data_types=["albedo"]) + depth = BaseRenderBenchmarkCameraCfg(data_types=["depth"]) + simple_shading_constant_diffuse = BaseRenderBenchmarkCameraCfg(data_types=["simple_shading_constant_diffuse"]) + simple_shading_diffuse_mdl = BaseRenderBenchmarkCameraCfg(data_types=["simple_shading_diffuse_mdl"]) + simple_shading_full_mdl = BaseRenderBenchmarkCameraCfg(data_types=["simple_shading_full_mdl"]) + + +@configclass +class RenderBenchmarkFrankaCabinetEnvCfg(DirectRLEnvCfg): + """Franka Panda and Sektion cabinet, animated for renderer benchmarking. + + The cabinet contributes four articulated joints (two drawers, two doors) on top of the + Franka's seven, and a sinusoidal animation drives all of them so every rendered frame has + moving articulated geometry rather than a static scene. There is no policy: actions are + ignored, rewards are zero, and the episode only ends on time-out. + + A mirrored layout of the canonical ``Isaac-Franka-Cabinet-Direct-v0`` task, which places the + Franka at the origin facing its default ``+X``: here the Franka sits at ``(1.0, 0, 0)`` rotated + 180 deg about Z (facing ``-X``, toward the cabinet), and the cabinet sits at the origin in its + default USD orientation. + """ + + decimation: int = 2 + episode_length_s: float = 60.0 + + action_space: int = 1 + observation_space: int = 1 + state_space: int = 0 + + sim: SimulationCfg = SimulationCfg(dt=1.0 / 120.0, render_interval=2, physics=RenderBenchmarkPhysicsCfg()) + + scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=4, env_spacing=3.0, replicate_physics=True) + + tiled_camera: RenderBenchmarkTiledCameraCfg = RenderBenchmarkTiledCameraCfg() + + articulations: dict[str, ArticulationCfg] = { + # High-PD variant so the joints track the sinusoidal targets smoothly. + "robot": FRANKA_PANDA_HIGH_PD_CFG.replace( + prim_path="{ENV_REGEX_NS}/Robot", + init_state=FRANKA_PANDA_HIGH_PD_CFG.init_state.replace( + pos=(1.0, 0.0, 0.0), + rot=(0.0, 0.0, 1.0, 0.0), # 180 deg about Z, xyzw + ), + ), + # Loaded as an ArticulationCfg rather than a static USD reference so its four joints + # animate and it clones to every env through the standard Isaac Lab path. + "cabinet": ArticulationCfg( + prim_path="{ENV_REGEX_NS}/Cabinet", + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Sektion_Cabinet/sektion_cabinet_instanceable.usd", + activate_contact_sensors=False, + ), + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.0, 0.0, 0.4), + rot=(0.0, 0.0, 0.0, 1.0), # identity: cabinet's default USD orientation + joint_pos={ + "door_left_joint": 0.0, + "door_right_joint": 0.0, + "drawer_bottom_joint": 0.0, + "drawer_top_joint": 0.0, + }, + ), + actuators={ + "drawers": ImplicitActuatorCfg( + joint_names_expr=["drawer_top_joint", "drawer_bottom_joint"], + joint_effort_limit=87.0, + stiffness=10.0, + damping=1.0, + ), + "doors": ImplicitActuatorCfg( + joint_names_expr=["door_left_joint", "door_right_joint"], + joint_effort_limit=87.0, + stiffness=10.0, + damping=2.5, + ), + }, + ), + } + """Articulations spawned into every environment, keyed by scene name.""" + + ground_top_z: float = 0.0 + """Height of the ground's top surface [m].""" + + ground_size: tuple[float, float] = (50.0, 50.0) + """Requested extent of the per-environment ground cuboid in XY [m]. + + Clamped to :attr:`scene.env_spacing` when the cuboid is built, so neighboring environments' + ground tiles meet at the boundary instead of overlapping. + """ + + ground_thickness: float = 0.1 + """Thickness of the ground cuboid along Z [m].""" + + ground_color: tuple[float, float, float] = (0.5, 0.5, 0.5) + """Diffuse color of the ground, as linear RGB in ``[0, 1]``.""" + + dome_light_intensity: float = 2000.0 + """Intensity of the ambient dome light. + + Newton's renderer has no tone mapping or ambient defaults of its own, so without this + unlit surfaces render pure black. + """ + + light_cfg: sim_utils.LightCfg | None = sim_utils.DistantLightCfg( + intensity=200.0, + exposure=0.0, + angle=0.0, + color=(1.0, 1.0, 1.0), + normalize=True, + ) + """Directional light spawned on top of the ambient dome light, or ``None`` for dome only.""" + + light_orientation: tuple[float, float, float, float] = (0.3251, 0.3251, 0.0, 0.8881) + """Orientation of :attr:`light_cfg` as ``(qx, qy, qz, qw)``. + + Rotates a USD ``DistantLight``'s default ``-Z`` onto ``(-0.57735, 0.57735, -0.57735)``, the + direction Warp's renderer hard-codes, so both renderers light the scene identically. + """ + + joint_animation_amplitude: float = 0.4 + """Peak sinusoidal offset from each joint's default position [m or rad, depending on joint type]. + + Values at or below zero disable the animation, leaving a static scene. + """ + + joint_animation_freq_hz: float = 0.35 + """Frequency of the joint animation [Hz].""" + + write_image_to_file: bool = os.getenv("BENCHMARK_SAVE_IMAGE", "0") == "1" + """Whether to dump each rendered frame to a PNG, for eyeballing renderer output.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/benchmark/shadow_hand_camera/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/benchmark/shadow_hand_camera/__init__.py new file mode 100644 index 000000000000..0c377d8ded6d --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/benchmark/shadow_hand_camera/__init__.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 + +"""Rendering-throughput benchmark variant of the Shadow Hand camera task.""" + +import gymnasium as gym + +from isaaclab_tasks.core.reorient.config.shadow_hand import agents as core_agents + +gym.register( + id="IsaacContrib-Reorient-Cube-Shadow-Camera-Benchmark-Direct", + entry_point="isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_direct_camera_env:ShadowHandCameraEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.shadow_hand_camera_benchmark_env_cfg:ShadowHandCameraBenchmarkEnvCfg", + "rsl_rl_cfg_entry_point": f"{core_agents.__name__}.rsl_rl_ppo_cfg:ShadowHandCameraFFPPORunnerCfg", + "rl_games_cfg_entry_point": f"{core_agents.__name__}:rl_games_ppo_camera_cfg.yaml", + }, +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/reorient/config/shadow_hand/shadow_hand_camera_benchmark_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/benchmark/shadow_hand_camera/shadow_hand_camera_benchmark_env_cfg.py similarity index 100% rename from source/isaaclab_tasks/isaaclab_tasks/contrib/reorient/config/shadow_hand/shadow_hand_camera_benchmark_env_cfg.py rename to source/isaaclab_tasks/isaaclab_tasks/benchmark/shadow_hand_camera/shadow_hand_camera_benchmark_env_cfg.py diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env_cfg.py index 2662affebd7a..775aeb222ec8 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env_cfg.py @@ -339,6 +339,7 @@ class PourSceneCfg(InteractiveSceneCfg): ) robot = FRANKA_PANDA_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") robot.spawn.usd_path = FRANKA_POUR_ROBOT_USD_PATH + robot.spawn.variants = {"Colliders": "convex_hulls"} robot.spawn.func = spawn_franka_with_arm_collisions robot.spawn.articulation_props.enabled_self_collisions = True robot.actuators = { diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/__init__.py new file mode 100644 index 000000000000..173a089f23c7 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/__init__.py @@ -0,0 +1,22 @@ +# 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 + +"""Contributed heterogeneous fixed-arm manipulation environments.""" + +import gymnasium as gym + +from . import agents + + +gym.register( + id="IsaacContrib-Multitask-Manipulation", + entry_point=f"{__name__}.multitask_env:MultitaskManipulationEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.multitask_env_cfg:MultitaskManipulationEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:MultitaskManipulationPPORunnerCfg", + "default_agent": "rsl_rl", + }, +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/agents/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/agents/__init__.py new file mode 100644 index 000000000000..f92c3e46e50d --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/agents/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Agent configurations for contributed heterogeneous manipulation training.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/agents/models.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/agents/models.py new file mode 100644 index 000000000000..fd3e3d4725e2 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/agents/models.py @@ -0,0 +1,459 @@ +# 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 + +"""Task-headed RSL-RL models for contributed heterogeneous manipulation.""" + +from __future__ import annotations + +import copy +import math +from collections.abc import Sequence + +import torch +import torch.nn as nn +import torch.nn.functional as F +from rsl_rl.models import MLPModel +from rsl_rl.modules import HiddenState +from rsl_rl.modules.distribution import Distribution +from rsl_rl.utils import resolve_callable, resolve_nn_activation, unpad_trajectories +from tensordict import TensorDict +from torch.distributions import Normal + + +def _positive_dimensions(dimensions: Sequence[int], label: str) -> tuple[int, ...]: + """Return dimensions as a tuple after validating that every entry is positive.""" + dimensions = tuple(dimensions) + if not dimensions or any(dimension <= 0 for dimension in dimensions): + raise ValueError(f"Expected positive {label} dimensions, got {dimensions}.") + return dimensions + + +def _task_action_dimensions(dimensions: Sequence[int], output_dim: int) -> tuple[int, ...]: + """Validate task action dimensions against the fixed global action width.""" + dimensions = _positive_dimensions(dimensions, "task action") + if sum(dimensions) != output_dim: + raise ValueError(f"Task action dimensions {dimensions} sum to {sum(dimensions)}, expected {output_dim}.") + return dimensions + + +def _task_encoding_slice(encoding_slice: Sequence[int], input_dim: int, task_count: int) -> tuple[int, int]: + """Validate and return the task one-hot slice within a model input.""" + if len(encoding_slice) != 2: + raise ValueError(f"Expected a two-element task encoding slice, got {encoding_slice}.") + resolved_slice = (int(encoding_slice[0]), int(encoding_slice[1])) + start, end = resolved_slice + if start < 0 or end > input_dim or end <= start: + raise ValueError(f"Task encoding slice {resolved_slice} is outside the {input_dim}D model input.") + if end - start != task_count: + raise ValueError(f"Task encoding slice {resolved_slice} has dimension {end - start}, expected {task_count}.") + return resolved_slice + + +def _task_encoding(obs: TensorDict, obs_groups: Sequence[str], encoding_slice: tuple[int, int]) -> torch.Tensor: + """Extract the raw task one-hot before observation normalization.""" + model_input = torch.cat([obs[obs_group] for obs_group in obs_groups], dim=-1) + return model_input[..., encoding_slice[0] : encoding_slice[1]] + + +def _validate_task_encoding(task_encoding: torch.Tensor) -> None: + """Validate that each sample contains exactly one one-hot task identity.""" + is_binary = torch.logical_or(task_encoding == 0.0, task_encoding == 1.0).all() + has_one_task = (task_encoding.sum(dim=-1) == 1.0).all() + if not bool(is_binary and has_one_task): + raise ValueError("Task encoding observations must be one-hot with exactly one active task per sample.") + + +class TaskHeadedGaussianDistribution(Distribution): + """Diagonal Gaussian distribution routed over contiguous task action heads.""" + + def __init__( + self, + output_dim: int, + task_action_dims: Sequence[int], + init_std: float = 1.0, + std_range: tuple[float, float] = (1e-6, 1e6), + std_type: str = "scalar", + learn_std: bool = True, + ) -> None: + """Initialize task-specific Gaussian standard deviations. + + Args: + output_dim: Fixed global action dimension. + task_action_dims: Action dimensions of the contiguous task heads. + init_std: Initial standard deviation for every active action dimension. + std_range: Inclusive standard-deviation clamp range. + std_type: Standard-deviation parameterization, either ``"scalar"`` or ``"log"``. + learn_std: Whether the task-specific standard deviations are learnable. + + Raises: + ValueError: If the task dimensions or standard-deviation configuration are invalid. + """ + super().__init__(output_dim) + self.task_action_dims = _task_action_dimensions(task_action_dims, output_dim) + if init_std <= 0.0: + raise ValueError(f"Expected a positive initial standard deviation, got {init_std}.") + if std_range[0] <= 0.0 or std_range[0] > std_range[1]: + raise ValueError(f"Expected a positive ordered standard-deviation range, got {std_range}.") + if std_type not in ("scalar", "log"): + raise ValueError(f"Unknown standard deviation type: {std_type}. Should be 'scalar' or 'log'.") + + self.std_type = std_type + self.std_range = std_range + self.log_std_range = (math.log(std_range[0]), math.log(std_range[1])) + initial_value = init_std if std_type == "scalar" else math.log(init_std) + self._std_params = nn.ParameterList( + [nn.Parameter(torch.full((dim,), initial_value), requires_grad=learn_std) for dim in self.task_action_dims] + ) + + action_masks = torch.zeros(len(self.task_action_dims), output_dim) + start = 0 + for task_id, action_dim in enumerate(self.task_action_dims): + action_masks[task_id, start : start + action_dim] = 1.0 + start += action_dim + self.register_buffer("_task_action_masks", action_masks) + + self._distribution: Normal | None = None + self._action_mask: torch.Tensor | None = None + Normal.set_default_validate_args(False) + + def update(self, mlp_output: torch.Tensor, task_encoding: torch.Tensor | None = None) -> None: + """Update the routed distribution from global means and one-hot task identities.""" + if task_encoding is None: + raise ValueError("TaskHeadedGaussianDistribution requires a task encoding when it is updated.") + if task_encoding.shape[-1] != len(self.task_action_dims): + raise ValueError( + f"Expected {len(self.task_action_dims)} task encoding dimensions, got {task_encoding.shape[-1]}." + ) + + action_mask = task_encoding.to(dtype=mlp_output.dtype) @ self._task_action_masks.to(dtype=mlp_output.dtype) + std = self._global_std().expand_as(mlp_output) + # Unit standard deviation keeps the inactive Normal parameters numerically valid; all probability terms mask it. + distribution_std = std * action_mask + (1.0 - action_mask) + self._distribution = Normal(mlp_output, distribution_std) + self._action_mask = action_mask + + def sample(self) -> torch.Tensor: + """Sample the active task head and return zero in every inactive global dimension.""" + distribution, action_mask = self._require_state() + return distribution.sample() * action_mask + + def deterministic_output(self, mlp_output: torch.Tensor) -> torch.Tensor: + """Return the routed means produced by the task-headed model.""" + return mlp_output + + def as_deterministic_output_module(self) -> nn.Module: + """Return an identity module because the model already routes its deterministic output.""" + return nn.Identity() + + @property + def input_dim(self) -> int: + """Return the fixed global mean dimension.""" + return self.output_dim + + @property + def mean(self) -> torch.Tensor: + """Return the current routed global action mean.""" + return self._require_state()[0].mean + + @property + def std(self) -> torch.Tensor: + """Return the concatenated standard deviations of every task head.""" + return self._global_std() + + @property + def entropy(self) -> torch.Tensor: + """Return entropy summed only over the selected task head.""" + distribution, action_mask = self._require_state() + return (distribution.entropy() * action_mask).sum(dim=-1) + + @property + def params(self) -> tuple[torch.Tensor, ...]: + """Return mean, numerically valid standard deviation, and the active action mask.""" + distribution, action_mask = self._require_state() + return distribution.mean, distribution.stddev, action_mask + + def log_prob(self, outputs: torch.Tensor) -> torch.Tensor: + """Return log-probability summed only over the selected task head.""" + distribution, action_mask = self._require_state() + return (distribution.log_prob(outputs) * action_mask).sum(dim=-1) + + def kl_divergence( + self, + old_params: tuple[torch.Tensor, ...], + new_params: tuple[torch.Tensor, ...], + ) -> torch.Tensor: + """Return KL divergence summed only over the task head stored with the old policy.""" + old_mean, old_std, old_action_mask = old_params + new_mean, new_std, _ = new_params + divergence = torch.distributions.kl_divergence(Normal(old_mean, old_std), Normal(new_mean, new_std)) + return (divergence * old_action_mask).sum(dim=-1) + + def _global_std(self) -> torch.Tensor: + """Concatenate the independently parameterized task-head standard deviations.""" + if self.std_type == "scalar": + stds = [param.clamp(self.std_range[0], self.std_range[1]) for param in self._std_params] + else: + stds = [torch.exp(param.clamp(self.log_std_range[0], self.log_std_range[1])) for param in self._std_params] + return torch.cat(stds) + + def _require_state(self) -> tuple[Normal, torch.Tensor]: + """Return the current distribution and mask or reject access before an update.""" + if self._distribution is None or self._action_mask is None: + raise RuntimeError("The task-headed Gaussian distribution has not been updated.") + return self._distribution, self._action_mask + + +class TaskHeadedMLPModel(MLPModel): + """Shared observation backbone with task-specific Gaussian action heads.""" + + def __init__( + self, + obs: TensorDict, + obs_groups: dict[str, list[str]], + obs_set: str, + output_dim: int, + hidden_dims: tuple[int, ...] | list[int] = (256, 256, 256), + activation: str = "elu", + obs_normalization: bool = False, + distribution_cfg: dict | None = None, + task_action_dims: tuple[int, ...] | list[int] = (8, 8, 6), + task_encoding_slice: tuple[int, int] | list[int] = (0, 3), + ) -> None: + """Initialize the shared backbone and contiguous task action heads. + + Args: + obs: Observation dictionary used to infer model input dimensions. + obs_groups: Observation groups assigned to each model set. + obs_set: Observation set consumed by this model. + output_dim: Fixed global action dimension. + hidden_dims: Shared-backbone layer dimensions. + activation: Shared-backbone activation function. + obs_normalization: Whether to normalize the shared observation input. + distribution_cfg: Task-headed Gaussian distribution configuration. + task_action_dims: Action dimensions ordered by task identity. + task_encoding_slice: Half-open slice of the raw model input containing the task one-hot. + + Raises: + TypeError: If the configured distribution is not task-headed. + ValueError: If dimensions or the initial task encoding are invalid. + """ + hidden_dims = _positive_dimensions(hidden_dims, "shared-backbone") + + super().__init__( + obs=obs, + obs_groups=obs_groups, + obs_set=obs_set, + output_dim=output_dim, + hidden_dims=hidden_dims, + activation=activation, + obs_normalization=obs_normalization, + distribution_cfg=None, + ) + + self.task_action_dims = _task_action_dimensions(task_action_dims, output_dim) + self.task_encoding_slice = _task_encoding_slice(task_encoding_slice, self.obs_dim, len(self.task_action_dims)) + _validate_task_encoding(self._task_encoding(obs)) + + # Replace the standard monolithic MLP with an activated shared backbone and task-local linear heads. + del self.mlp + self.backbone = _make_backbone(self.obs_dim, hidden_dims, activation) + self.action_heads = nn.ModuleList([nn.Linear(hidden_dims[-1], dim) for dim in self.task_action_dims]) + + if distribution_cfg is None: + raise ValueError("TaskHeadedMLPModel requires a task-headed distribution configuration.") + distribution_cfg = dict(distribution_cfg) + distribution_class: type[Distribution] = resolve_callable(distribution_cfg.pop("class_name")) # type: ignore + if not issubclass(distribution_class, TaskHeadedGaussianDistribution): + raise TypeError( + f"TaskHeadedMLPModel requires a TaskHeadedGaussianDistribution, got {distribution_class.__name__}." + ) + self.distribution = distribution_class( + output_dim, + task_action_dims=self.task_action_dims, + **distribution_cfg, + ) + + def forward( + self, + obs: TensorDict, + masks: torch.Tensor | None = None, + hidden_state: HiddenState = None, + stochastic_output: bool = False, + ) -> torch.Tensor: + """Route shared features through the task head selected by the raw task one-hot.""" + obs = unpad_trajectories(obs, masks) if masks is not None and not self.is_recurrent else obs + latent = self.backbone(self.get_latent(obs, masks, hidden_state)) + task_ids = self._task_encoding(obs).argmax(dim=-1) + task_encoding = F.one_hot(task_ids, num_classes=len(self.task_action_dims)).to(dtype=latent.dtype) + output = torch.cat( + [ + head(latent) * task_encoding[..., task_id : task_id + 1] + for task_id, head in enumerate(self.action_heads) + ], + dim=-1, + ) + if stochastic_output: + self.distribution.update(output, task_encoding) + return self.distribution.sample() + return output + + def as_jit(self) -> nn.Module: + """Return a TorchScript-compatible deterministic task-headed policy.""" + return _TorchTaskHeadedMLPModel(self) + + def as_onnx(self, verbose: bool = False) -> nn.Module: + """Return an ONNX-compatible deterministic task-headed policy.""" + return _OnnxTaskHeadedMLPModel(self, verbose) + + def _task_encoding(self, obs: TensorDict) -> torch.Tensor: + """Extract the raw task one-hot before observation normalization.""" + return _task_encoding(obs, self.obs_groups, self.task_encoding_slice) + + +class TaskHeadedValueModel(MLPModel): + """Shared observation backbone with task-specific scalar value heads.""" + + def __init__( + self, + obs: TensorDict, + obs_groups: dict[str, list[str]], + obs_set: str, + output_dim: int, + hidden_dims: tuple[int, ...] | list[int] = (256, 256, 256), + activation: str = "elu", + obs_normalization: bool = False, + distribution_cfg: dict | None = None, + task_head_count: int = 3, + task_encoding_slice: tuple[int, int] | list[int] = (0, 3), + ) -> None: + """Initialize a shared value backbone and one scalar head per task. + + Args: + obs: Observation dictionary used to infer model input dimensions. + obs_groups: Observation groups assigned to each model set. + obs_set: Observation set consumed by this model. + output_dim: Critic output dimension, which must be one. + hidden_dims: Shared-backbone layer dimensions. + activation: Shared-backbone activation function. + obs_normalization: Whether to normalize the shared observation input. + distribution_cfg: Output distribution configuration, which must be ``None`` for a critic. + task_head_count: Number of task-specific scalar value heads. + task_encoding_slice: Half-open slice of the raw model input containing the task one-hot. + + Raises: + ValueError: If the value, backbone, task-head, or task-encoding configuration is invalid. + """ + if output_dim != 1: + raise ValueError(f"TaskHeadedValueModel requires a scalar output, got {output_dim}.") + hidden_dims = _positive_dimensions(hidden_dims, "shared-backbone") + if task_head_count <= 0: + raise ValueError(f"Expected a positive task head count, got {task_head_count}.") + if distribution_cfg is not None: + raise ValueError("TaskHeadedValueModel does not support an output distribution.") + + super().__init__( + obs=obs, + obs_groups=obs_groups, + obs_set=obs_set, + output_dim=output_dim, + hidden_dims=hidden_dims, + activation=activation, + obs_normalization=obs_normalization, + distribution_cfg=None, + ) + + self.task_head_count = task_head_count + self.task_encoding_slice = _task_encoding_slice(task_encoding_slice, self.obs_dim, task_head_count) + _validate_task_encoding(self._task_encoding(obs)) + + del self.mlp + self.backbone = _make_backbone(self.obs_dim, hidden_dims, activation) + self.value_heads = nn.ModuleList([nn.Linear(hidden_dims[-1], 1) for _ in range(task_head_count)]) + + def forward( + self, + obs: TensorDict, + masks: torch.Tensor | None = None, + hidden_state: HiddenState = None, + stochastic_output: bool = False, + ) -> torch.Tensor: + """Return the scalar value from the head selected by the raw task one-hot.""" + del stochastic_output + obs = unpad_trajectories(obs, masks) if masks is not None and not self.is_recurrent else obs + latent = self.backbone(self.get_latent(obs, masks, hidden_state)) + task_ids = self._task_encoding(obs).argmax(dim=-1, keepdim=True) + task_values = torch.cat([head(latent) for head in self.value_heads], dim=-1) + return torch.gather(task_values, dim=-1, index=task_ids) + + def _task_encoding(self, obs: TensorDict) -> torch.Tensor: + """Extract the raw task one-hot before observation normalization.""" + return _task_encoding(obs, self.obs_groups, self.task_encoding_slice) + + +class _TorchTaskHeadedMLPModel(nn.Module): + """TorchScript-compatible deterministic task-headed policy.""" + + def __init__(self, model: TaskHeadedMLPModel) -> None: + """Copy inference modules and task routing metadata from a trained model.""" + super().__init__() + self.obs_normalizer = copy.deepcopy(model.obs_normalizer) + self.backbone = copy.deepcopy(model.backbone) + self.action_heads = copy.deepcopy(model.action_heads) + self.task_encoding_start = model.task_encoding_slice[0] + self.task_encoding_end = model.task_encoding_slice[1] + + def forward(self, obs: torch.Tensor) -> torch.Tensor: + """Return a fixed-width action containing only the selected task head.""" + task_ids = obs[..., self.task_encoding_start : self.task_encoding_end].argmax(dim=-1) + latent = self.backbone(self.obs_normalizer(obs)) + outputs: list[torch.Tensor] = [] + for task_id, head in enumerate(self.action_heads): + selected = (task_ids == task_id).unsqueeze(-1).to(dtype=latent.dtype) + outputs.append(head(latent) * selected) + return torch.cat(outputs, dim=-1) + + @torch.jit.export + def reset(self) -> None: + """Reset recurrent export state (no-op for the feed-forward model).""" + pass + + +class _OnnxTaskHeadedMLPModel(_TorchTaskHeadedMLPModel): + """ONNX-compatible deterministic task-headed policy.""" + + is_recurrent: bool = False + + def __init__(self, model: TaskHeadedMLPModel, verbose: bool) -> None: + """Copy a task-headed model and record ONNX export metadata.""" + super().__init__(model) + self.verbose = verbose + self.input_size = model.obs_dim + + def get_dummy_inputs(self) -> tuple[torch.Tensor]: + """Return a representative task-valid observation for ONNX tracing.""" + dummy_obs = torch.zeros(1, self.input_size) + dummy_obs[0, self.task_encoding_start] = 1.0 + return (dummy_obs,) + + @property + def input_names(self) -> list[str]: + """Return ONNX input tensor names.""" + return ["obs"] + + @property + def output_names(self) -> list[str]: + """Return ONNX output tensor names.""" + return ["actions"] + + +def _make_backbone(input_dim: int, hidden_dims: Sequence[int], activation: str) -> nn.Sequential: + """Create an activated MLP whose final hidden layer is shared by every task head.""" + layers: list[nn.Module] = [] + layer_input_dim = input_dim + for hidden_dim in hidden_dims: + layers.append(nn.Linear(layer_input_dim, hidden_dim)) + layers.append(resolve_nn_activation(activation)) + layer_input_dim = hidden_dim + return nn.Sequential(*layers) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/agents/ppo.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/agents/ppo.py new file mode 100644 index 000000000000..370efabeba3f --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/agents/ppo.py @@ -0,0 +1,108 @@ +# 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 + +"""Task-balanced PPO for contributed heterogeneous manipulation.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import torch +from rsl_rl.algorithms import PPO +from rsl_rl.models import MLPModel +from rsl_rl.storage import RolloutStorage +from tensordict import TensorDict + + +class TaskBalancedPPO(PPO): + """PPO that normalizes rollout advantages independently for each task.""" + + def __init__( + self, + actor: MLPModel, + critic: MLPModel, + storage: RolloutStorage, + task_names: Sequence[str], + task_encoding_obs_group: str, + task_encoding_slice: Sequence[int], + **kwargs: Any, + ) -> None: + """Initialize task routing metadata and the base PPO algorithm. + + Args: + actor: Task-headed policy model. + critic: Task-headed value model. + storage: On-policy rollout storage. + task_names: Task names ordered by one-hot identity. + task_encoding_obs_group: Observation group containing the task one-hot. + task_encoding_slice: Half-open slice containing the task one-hot. + **kwargs: Base PPO keyword arguments. + + Raises: + ValueError: If task metadata is invalid or mixed mini-batch normalization is enabled. + """ + self.task_names = tuple(task_names) + if not self.task_names: + raise ValueError("Expected at least one task name.") + if len(task_encoding_slice) != 2: + raise ValueError(f"Expected a two-element task encoding slice, got {task_encoding_slice}.") + self.task_encoding_slice = (int(task_encoding_slice[0]), int(task_encoding_slice[1])) + encoding_start, encoding_end = self.task_encoding_slice + if encoding_start < 0 or encoding_end <= encoding_start: + raise ValueError(f"Expected an ordered non-negative task encoding slice, got {self.task_encoding_slice}.") + if encoding_end - encoding_start != len(self.task_names): + raise ValueError( + f"Task encoding slice {self.task_encoding_slice} has dimension {encoding_end - encoding_start}," + f" expected {len(self.task_names)}." + ) + if kwargs.get("normalize_advantage_per_mini_batch", False): + raise ValueError("TaskBalancedPPO is incompatible with mixed-task mini-batch advantage normalization.") + + self.task_encoding_obs_group = task_encoding_obs_group + self.task_statistics: dict[str, float] = {} + super().__init__(actor, critic, storage, **kwargs) + + def compute_returns(self, obs: TensorDict) -> None: + """Compute GAE targets and normalize advantages within each task rollout.""" + super().compute_returns(obs) + storage = self.storage + raw_advantages = storage.returns - storage.values + task_encoding = storage.observations[self.task_encoding_obs_group][ + ..., self.task_encoding_slice[0] : self.task_encoding_slice[1] + ] + self._validate_task_encoding(task_encoding) + task_ids = task_encoding.argmax(dim=-1) + + task_statistics = {} + for task_id, task_name in enumerate(self.task_names): + selected = task_ids == task_id + if not torch.any(selected): + raise ValueError(f"Rollout contains no samples for task '{task_name}'.") + task_advantages = raw_advantages[selected] + task_returns = storage.returns[selected] + advantage_mean = task_advantages.mean() + advantage_std = task_advantages.std(unbiased=False) + storage.advantages[selected] = (task_advantages - advantage_mean) / (advantage_std + 1.0e-8) + task_statistics[f"{task_name}_advantage_std"] = advantage_std.item() + task_statistics[f"{task_name}_return_std"] = task_returns.std(unbiased=False).item() + self.task_statistics = task_statistics + + def update(self) -> dict[str, float]: + """Run PPO updates and append per-task rollout scale diagnostics.""" + losses = super().update() + losses.update(self.task_statistics) + return losses + + def _validate_task_encoding(self, task_encoding: torch.Tensor) -> None: + """Validate one-hot task identities stored in the rollout.""" + if task_encoding.shape[-1] != len(self.task_names): + raise ValueError( + f"Expected {len(self.task_names)} task encoding dimensions, got {task_encoding.shape[-1]}." + ) + is_binary = torch.logical_or(task_encoding == 0.0, task_encoding == 1.0).all() + has_one_task = (task_encoding.sum(dim=-1) == 1.0).all() + if not bool(is_binary and has_one_task): + raise ValueError("Task encoding observations must be one-hot with exactly one active task per sample.") diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/__init__.py new file mode 100644 index 000000000000..b3618332aa8a --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""MDP terms for contributed heterogeneous manipulation training.""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/__init__.pyi new file mode 100644 index 000000000000..cd12356000bf --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/__init__.pyi @@ -0,0 +1,79 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +__all__ = [ + "CabinetOpenDrawerBonus", + "LiftGoalTracking", + "LiftSuccess", + "SelectedBinaryJointPositionActionCfg", + "SelectedJointPositionActionCfg", + "SelectedUniformPoseCommandCfg", + "articulation_state_invalid", + "cabinet_align_ee_handle", + "cabinet_align_grasp", + "cabinet_approach_ee_handle", + "cabinet_approach_gripper", + "cabinet_drawer_state", + "cabinet_ee_to_handle", + "cabinet_grasp_handle", + "cabinet_multi_stage_open", + "lift_ee_object_distance", + "lift_goal_tracking", + "lift_object_dropped", + "lift_object_height", + "lift_object_out_of_bounds", + "lift_object_position_b", + "lift_object_pose_b", + "reach_orientation_error", + "reach_position_error", + "reach_success", + "reset_multitask_scene", + "selected_action_l2", + "selected_action_rate_l2", + "selected_joint_pos_rel", + "selected_joint_vel_l2", + "selected_joint_vel_rel", + "task_encoding", + "task_time_out", +] + +from .actions_cfg import SelectedBinaryJointPositionActionCfg, SelectedJointPositionActionCfg +from .commands_cfg import SelectedUniformPoseCommandCfg +from .events import reset_multitask_scene +from .observations import ( + cabinet_drawer_state, + cabinet_ee_to_handle, + lift_object_position_b, + lift_object_pose_b, + selected_joint_pos_rel, + selected_joint_vel_rel, + task_encoding, +) +from .rewards import ( + CabinetOpenDrawerBonus, + LiftGoalTracking, + LiftSuccess, + cabinet_align_ee_handle, + cabinet_align_grasp, + cabinet_approach_ee_handle, + cabinet_approach_gripper, + cabinet_grasp_handle, + cabinet_multi_stage_open, + lift_ee_object_distance, + lift_goal_tracking, + lift_object_height, + reach_orientation_error, + reach_position_error, + selected_action_l2, + selected_action_rate_l2, + selected_joint_vel_l2, +) +from .terminations import ( + articulation_state_invalid, + lift_object_dropped, + lift_object_out_of_bounds, + reach_success, + task_time_out, +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/actions.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/actions.py new file mode 100644 index 000000000000..d2cb2af78a22 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/actions.py @@ -0,0 +1,130 @@ +# 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 + +"""Selection-aware joint action terms.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch + +from isaaclab.managers import ActionTerm + +from ..selection_utils import SceneEntitySelectionCfg + +if TYPE_CHECKING: + from isaaclab.assets import Articulation + from isaaclab.envs import ManagerBasedEnv + + from .actions_cfg import SelectedBinaryJointPositionActionCfg, SelectedJointPositionActionCfg + + +class _SelectedJointAction(ActionTerm): + """Resolve joints and maintain global inputs for a partial articulation view.""" + + cfg: SelectedBinaryJointPositionActionCfg | SelectedJointPositionActionCfg + _asset: Articulation + + def __init__( + self, cfg: SelectedBinaryJointPositionActionCfg | SelectedJointPositionActionCfg, env: ManagerBasedEnv + ) -> None: + super().__init__(cfg, env) + joint_ids, self._joint_names = self._asset.find_joints(cfg.joint_names, preserve_order=True, as_proxy=True) + self._joint_ids = joint_ids.torch + self._num_joints = len(joint_ids) + self._asset_cfg = SceneEntitySelectionCfg(cfg.asset_name) + self._asset_cfg.resolve(env.scene) + self._raw_actions = torch.zeros(env.num_envs, self.action_dim, device=self.device) + + @property + def raw_actions(self) -> torch.Tensor: + """Unprocessed global action batch.""" + return self._raw_actions + + @property + def processed_actions(self) -> torch.Tensor: + """Processed actions in asset physics-view order.""" + return self._processed_actions + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Clear raw actions for the reset global environments.""" + self._raw_actions[env_ids] = 0.0 + + +class SelectedJointPositionAction(_SelectedJointAction): + """Apply a global action batch to one articulation's selected physics-view rows.""" + + cfg: SelectedJointPositionActionCfg + _asset: Articulation + + def __init__(self, cfg: SelectedJointPositionActionCfg, env: ManagerBasedEnv) -> None: + super().__init__(cfg, env) + self._joint_limit_margin = cfg.joint_limit_margin + if self._joint_limit_margin is not None: + if not math.isfinite(self._joint_limit_margin) or self._joint_limit_margin < 0.0: + raise ValueError("joint_limit_margin must be finite and non-negative.") + limits = self._asset.data.soft_joint_pos_limits.torch[:, self._joint_ids] + if torch.any(limits[..., 0] + self._joint_limit_margin > limits[..., 1] - self._joint_limit_margin): + raise ValueError("joint_limit_margin leaves an empty target range for at least one joint.") + self._processed_actions = torch.zeros(len(self._asset_cfg.env_ids), self.action_dim, device=self.device) + + @property + def action_dim(self) -> int: + """Number of controlled joints.""" + return self._num_joints + + def process_actions(self, actions: torch.Tensor) -> None: + """Scale global actions and align them with the articulation physics view.""" + self._raw_actions[:] = actions + self._processed_actions = actions[self._asset_cfg.env_ids] * self.cfg.scale + if not self.cfg.relative: + self._processed_actions += self._asset.data.default_joint_pos.torch[:, self._joint_ids] + + def apply_actions(self) -> None: + """Write joint-position targets to the selected articulation instances.""" + target = self._processed_actions + if self.cfg.relative: + target = target + self._asset.data.joint_pos.torch[:, self._joint_ids] + if self._joint_limit_margin is not None: + limits = self._asset.data.soft_joint_pos_limits.torch[:, self._joint_ids] + target = torch.maximum( + torch.minimum(target, limits[..., 1] - self._joint_limit_margin), + limits[..., 0] + self._joint_limit_margin, + ) + self._asset.actuators.target_command.set_position_index(value=target, joint_ids=self._joint_ids) + + +class SelectedBinaryJointPositionAction(_SelectedJointAction): + """Map one binary global action to joint targets in asset physics-view order.""" + + cfg: SelectedBinaryJointPositionActionCfg + _asset: Articulation + + def __init__(self, cfg: SelectedBinaryJointPositionActionCfg, env: ManagerBasedEnv) -> None: + super().__init__(cfg, env) + self._processed_actions = torch.zeros(len(self._asset_cfg.env_ids), self._num_joints, device=self.device) + self._open_command = torch.full((self._num_joints,), cfg.open_command, device=self.device) + self._close_command = torch.full((self._num_joints,), cfg.close_command, device=self.device) + + @property + def action_dim(self) -> int: + """Binary action dimension.""" + return 1 + + def process_actions(self, actions: torch.Tensor) -> None: + """Choose open or closed targets from the sign of selected actions.""" + self._raw_actions[:] = actions + self._processed_actions = torch.where( + actions[self._asset_cfg.env_ids] < 0.0, self._close_command, self._open_command + ) + + def apply_actions(self) -> None: + """Write selected binary joint targets.""" + self._asset.actuators.target_command.set_position_index( + value=self._processed_actions, joint_ids=self._joint_ids + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/commands.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/commands.py new file mode 100644 index 000000000000..2f85cceb4472 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/commands.py @@ -0,0 +1,103 @@ +# 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 + +"""Selection-aware pose commands for contributed heterogeneous task batches.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch + +from isaaclab.managers import CommandTerm +from isaaclab.markers import VisualizationMarkers +from isaaclab.utils.math import combine_frame_transforms, quat_from_euler_xyz, quat_unique + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + from .commands_cfg import SelectedUniformPoseCommandCfg + + +class SelectedUniformPoseCommand(CommandTerm): + """Sample root-frame pose commands only where the configured scene entity exists.""" + + cfg: SelectedUniformPoseCommandCfg + + def __init__(self, cfg: SelectedUniformPoseCommandCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + self.pose_command = torch.zeros(self.num_envs, 7, device=self.device) + self._all_env_ids = torch.arange(self.num_envs, device=self.device) + self._reference_cfg = cfg.reference_cfg + self._tracked_cfg = cfg.tracked_cfg + self._reference_cfg.resolve(env.scene) + self._tracked_cfg.resolve(env.scene) + self._reference = env.scene[self._reference_cfg.name] + self._tracked = env.scene[self._tracked_cfg.name] + + @property + def command(self) -> torch.Tensor: + """Pose command, shape ``(num_envs, 7)``.""" + return self.pose_command + + def _resample_command(self, env_ids: Sequence[int]) -> None: + global_env_ids = self._all_env_ids[env_ids] + _, selected_env_ids = self._reference_cfg.select(global_env_ids) + ranges = self.cfg.ranges + sample = torch.empty(len(selected_env_ids), device=self.device) + self.pose_command[selected_env_ids, 0] = sample.uniform_(*ranges.pos_x) + self.pose_command[selected_env_ids, 1] = sample.uniform_(*ranges.pos_y) + self.pose_command[selected_env_ids, 2] = sample.uniform_(*ranges.pos_z) + euler = torch.empty(len(selected_env_ids), 3, device=self.device) + euler[:, 0].uniform_(*ranges.roll) + euler[:, 1].uniform_(*ranges.pitch) + euler[:, 2].uniform_(*ranges.yaw) + self.pose_command[selected_env_ids, 3:] = quat_unique(quat_from_euler_xyz(*euler.unbind(-1))) + + def _update_metrics(self) -> None: + pass + + def _update_command(self) -> None: + pass + + def _set_debug_vis_impl(self, debug_vis: bool) -> None: + if debug_vis and not hasattr(self, "goal_pose_visualizer"): + self.goal_pose_visualizer = VisualizationMarkers(self.cfg.goal_pose_visualizer_cfg) + self.current_pose_visualizer = VisualizationMarkers(self.cfg.current_pose_visualizer_cfg) + if not hasattr(self, "goal_pose_visualizer"): + return + self.goal_pose_visualizer.set_visibility(debug_vis) + self.current_pose_visualizer.set_visibility(debug_vis) + + def _debug_vis_callback(self, event) -> None: + if not self._reference.is_initialized or not self._tracked.is_initialized: + return + + reference_pose_w = self._reference.data.root_link_pose_w.torch + command = self.pose_command[self._reference_cfg.env_ids] + goal_pos_w, goal_quat_w = combine_frame_transforms( + reference_pose_w[:, :3], reference_pose_w[:, 3:7], command[:, :3], command[:, 3:] + ) + self.goal_pose_visualizer.visualize( + goal_pos_w, + goal_quat_w, + environment_ids=self._reference_cfg.env_ids, + ) + + if self._tracked_cfg.body_names is None: + tracked_pose_w = self._tracked.data.root_link_pose_w.torch + else: + tracked_pose_w = self._tracked.data.body_link_pose_w.torch[:, self._tracked_cfg.body_ids] + if tracked_pose_w.shape[1] != 1: + raise ValueError( + f"Expected '{self._tracked_cfg.name}' to select one tracked body, got {tracked_pose_w.shape[1]}." + ) + tracked_pose_w = tracked_pose_w[:, 0] + self.current_pose_visualizer.visualize( + tracked_pose_w[:, :3], + tracked_pose_w[:, 3:7], + environment_ids=self._tracked_cfg.env_ids, + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/events.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/events.py new file mode 100644 index 000000000000..ffc4ed42fd36 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/events.py @@ -0,0 +1,119 @@ +# 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 + +"""Selection-aware reset events for heterogeneous manipulation scenes.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from ..selection_utils import SceneEntitySelectionCfg + +if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject + from isaaclab.envs import ManagerBasedEnv + + +def _reset_root(env: ManagerBasedEnv, asset_cfg: SceneEntitySelectionCfg, env_ids: torch.Tensor) -> None: + """Write an asset's default root state for selected global environments.""" + asset: Articulation | RigidObject = env.scene[asset_cfg.name] + rows, selected_env_ids = asset_cfg.select(env_ids) + if rows.numel() == 0: + return + root_pose = asset.data.default_root_pose.torch[rows].clone() + root_pose[:, :3] += env.scene.env_origins[selected_env_ids] + asset.write_root_pose_to_sim_index(root_pose=root_pose, env_ids=rows) + asset.write_root_velocity_to_sim_index(root_velocity=asset.data.default_root_vel.torch[rows], env_ids=rows) + + +def _reset_joints_default(env: ManagerBasedEnv, asset_cfg: SceneEntitySelectionCfg, env_ids: torch.Tensor) -> None: + """Write default articulation joint states for selected global environments.""" + asset: Articulation = env.scene[asset_cfg.name] + rows, _ = asset_cfg.select(env_ids) + if rows.numel() == 0: + return + joint_pos = asset.data.default_joint_pos.torch[rows].clone() + joint_vel = asset.data.default_joint_vel.torch[rows].clone() + asset.write_joint_position_to_sim_index(position=joint_pos, env_ids=rows) + asset.write_joint_velocity_to_sim_index(velocity=joint_vel, env_ids=rows) + asset.actuators.target_command.set_position_index(value=joint_pos, env_ids=rows) + asset.actuators.target_command.set_velocity_index(value=joint_vel, env_ids=rows) + + +def _randomize_joint_offset( + env: ManagerBasedEnv, + asset_cfg: SceneEntitySelectionCfg, + env_ids: torch.Tensor, + position_range: tuple[float, float], +) -> None: + """Reset articulation joints with uniform offsets from their defaults.""" + asset: Articulation = env.scene[asset_cfg.name] + rows, _ = asset_cfg.select(env_ids) + if rows.numel() == 0: + return + joint_pos = asset.data.default_joint_pos.torch[rows].clone() + joint_pos += torch.empty_like(joint_pos).uniform_(*position_range) + limits = asset.data.soft_joint_pos_limits.torch[rows] + joint_pos.clamp_(limits[..., 0], limits[..., 1]) + joint_vel = torch.zeros_like(joint_pos) + asset.write_joint_position_to_sim_index(position=joint_pos, env_ids=rows) + asset.write_joint_velocity_to_sim_index(velocity=joint_vel, env_ids=rows) + asset.actuators.target_command.set_position_index(value=joint_pos, env_ids=rows) + + +def _randomize_joint_scale( + env: ManagerBasedEnv, + asset_cfg: SceneEntitySelectionCfg, + env_ids: torch.Tensor, + position_range: tuple[float, float], +) -> None: + """Reset articulation joints by uniformly scaling their default positions.""" + asset: Articulation = env.scene[asset_cfg.name] + rows, _ = asset_cfg.select(env_ids) + if rows.numel() == 0: + return + joint_pos = asset.data.default_joint_pos.torch[rows].clone() + joint_pos *= torch.empty_like(joint_pos).uniform_(*position_range) + limits = asset.data.soft_joint_pos_limits.torch[rows] + joint_pos.clamp_(limits[..., 0], limits[..., 1]) + joint_vel = torch.zeros_like(joint_pos) + asset.write_joint_position_to_sim_index(position=joint_pos, env_ids=rows) + asset.write_joint_velocity_to_sim_index(velocity=joint_vel, env_ids=rows) + asset.actuators.target_command.set_position_index(value=joint_pos, env_ids=rows) + + +def reset_multitask_scene( + env: ManagerBasedEnv, + env_ids: torch.Tensor, + root_asset_cfgs: tuple[SceneEntitySelectionCfg, ...], + lift_robot_cfg: SceneEntitySelectionCfg, + lift_object_cfg: SceneEntitySelectionCfg, + cabinet_robot_cfg: SceneEntitySelectionCfg, + cabinet_cfg: SceneEntitySelectionCfg, + reach_robot_cfg: SceneEntitySelectionCfg, +) -> None: + """Reset active assets and apply task-specific initial-state randomization.""" + for asset_cfg in root_asset_cfgs: + _reset_root(env, asset_cfg, env_ids) + + _reset_joints_default(env, lift_robot_cfg, env_ids) + _reset_joints_default(env, cabinet_cfg, env_ids) + _randomize_joint_offset(env, cabinet_robot_cfg, env_ids, (-0.1, 0.1)) + _randomize_joint_scale(env, reach_robot_cfg, env_ids, (0.75, 1.25)) + + lift_object: RigidObject = env.scene[lift_object_cfg.name] + object_rows, object_env_ids = lift_object_cfg.select(env_ids) + if object_rows.numel() == 0: + return + object_pose = lift_object.data.default_root_pose.torch[object_rows].clone() + object_pose[:, :3] += env.scene.env_origins[object_env_ids] + object_pose[:, 0] += torch.empty(len(object_rows), device=env.device).uniform_(-0.1, 0.1) + object_pose[:, 1] += torch.empty(len(object_rows), device=env.device).uniform_(-0.25, 0.25) + lift_object.write_root_pose_to_sim_index(root_pose=object_pose, env_ids=object_rows) + lift_object.write_root_velocity_to_sim_index( + root_velocity=torch.zeros(len(object_rows), 6, device=env.device), env_ids=object_rows + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/observations.py new file mode 100644 index 000000000000..0f7316e2533f --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/observations.py @@ -0,0 +1,112 @@ +# 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 + +"""Selection-aware observations for heterogeneous manipulation tasks.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from isaaclab.utils import math as math_utils + +from ..selection_utils import SceneEntitySelectionCfg +from .utils import _offset_body_pose + +if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject + from isaaclab.envs import ManagerBasedRLEnv + + +def task_encoding(env: ManagerBasedRLEnv, task_asset_cfgs: tuple[SceneEntitySelectionCfg, ...]) -> torch.Tensor: + """Return a one-hot identity for the selected manipulation task.""" + encoding = torch.zeros(env.num_envs, len(task_asset_cfgs), device=env.device) + for task_id, asset_cfg in enumerate(task_asset_cfgs): + encoding[asset_cfg.instance_ids >= 0, task_id] = 1.0 + return encoding + + +def selected_joint_pos_rel(env: ManagerBasedRLEnv, asset_cfg: SceneEntitySelectionCfg) -> torch.Tensor: + """Return relative joint positions in global environment order.""" + asset: Articulation = env.scene[asset_cfg.name] + values = ( + asset.data.joint_pos.torch[:, asset_cfg.joint_ids] - asset.data.default_joint_pos.torch[:, asset_cfg.joint_ids] + ) + return asset_cfg.scatter_to_envs(values) + + +def selected_joint_vel_rel(env: ManagerBasedRLEnv, asset_cfg: SceneEntitySelectionCfg) -> torch.Tensor: + """Return relative joint velocities in global environment order.""" + asset: Articulation = env.scene[asset_cfg.name] + values = ( + asset.data.joint_vel.torch[:, asset_cfg.joint_ids] - asset.data.default_joint_vel.torch[:, asset_cfg.joint_ids] + ) + return asset_cfg.scatter_to_envs(values) + + +def lift_object_pose_b( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntitySelectionCfg, + object_cfg: SceneEntitySelectionCfg, +) -> torch.Tensor: + """Return the lift object's pose in the Franka root frame.""" + robot: Articulation = env.scene[robot_cfg.name] + object_asset: RigidObject = env.scene[object_cfg.name] + robot_rows = robot_cfg.instance_ids[object_cfg.env_ids] + pos_b, quat_b = math_utils.subtract_frame_transforms( + robot.data.root_pos_w.torch[robot_rows], + robot.data.root_quat_w.torch[robot_rows], + object_asset.data.root_pos_w.torch, + object_asset.data.root_quat_w.torch, + ) + return object_cfg.scatter_to_envs(torch.cat((pos_b, quat_b), dim=-1)) + + +def lift_object_position_b( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntitySelectionCfg, + object_cfg: SceneEntitySelectionCfg, +) -> torch.Tensor: + """Return the lift object's position [m] in the OpenArm root frame.""" + robot: Articulation = env.scene[robot_cfg.name] + object_asset: RigidObject = env.scene[object_cfg.name] + robot_rows = robot_cfg.instance_ids[object_cfg.env_ids] + pos_b, _ = math_utils.subtract_frame_transforms( + robot.data.root_pos_w.torch[robot_rows], + robot.data.root_quat_w.torch[robot_rows], + object_asset.data.root_pos_w.torch, + ) + return object_cfg.scatter_to_envs(pos_b) + + +def cabinet_drawer_state(env: ManagerBasedRLEnv, cabinet_cfg: SceneEntitySelectionCfg) -> torch.Tensor: + """Return top-drawer position and velocity in global environment order.""" + cabinet: Articulation = env.scene[cabinet_cfg.name] + values = torch.cat( + ( + cabinet.data.joint_pos.torch[:, cabinet_cfg.joint_ids], + cabinet.data.joint_vel.torch[:, cabinet_cfg.joint_ids], + ), + dim=-1, + ) + return cabinet_cfg.scatter_to_envs(values) + + +def cabinet_ee_to_handle( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntitySelectionCfg, + cabinet_cfg: SceneEntitySelectionCfg, +) -> torch.Tensor: + """Return the vector [m] from the cabinet Franka TCP to the drawer handle.""" + env_ids, ee_pos_w, _ = _offset_body_pose(env, robot_cfg, (0.0, 0.0, 0.1034)) + _, handle_pos_w, _ = _offset_body_pose( + env, + cabinet_cfg, + (0.305, 0.0, 0.01), + (0.5, -0.5, -0.5, 0.5), + ) + values = handle_pos_w[cabinet_cfg.instance_ids[env_ids]] - ee_pos_w + return robot_cfg.scatter_to_envs(values) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/rewards.py new file mode 100644 index 000000000000..7b57ea79eb85 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/rewards.py @@ -0,0 +1,364 @@ +# 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 + +"""Selection-aware rewards for heterogeneous manipulation tasks.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from isaaclab.managers import ManagerTermBase, RewardTermCfg +from isaaclab.utils import math as math_utils + +from ..selection_utils import SceneEntitySelectionCfg +from .utils import _offset_body_pose + +if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject + from isaaclab.envs import ManagerBasedRLEnv + + +def _action_slice(env: ManagerBasedRLEnv, term_names: tuple[str, ...]) -> slice: + """Return the cached contiguous global action slice occupied by action terms.""" + cache = env.__dict__.setdefault("_multitask_action_slices", {}) + if term_names not in cache: + starts = {} + start = 0 + for name, dim in zip(env.action_manager.active_terms, env.action_manager.action_term_dim): + starts[name] = (start, start + dim) + start += dim + cache[term_names] = slice(starts[term_names[0]][0], starts[term_names[-1]][1]) + return cache[term_names] + + +def _lift_goal_error( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntitySelectionCfg, + object_cfg: SceneEntitySelectionCfg, + command_name: str, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return lift global IDs and object-to-goal position error [m].""" + robot: Articulation = env.scene[robot_cfg.name] + object_asset: RigidObject = env.scene[object_cfg.name] + robot_rows = robot_cfg.instance_ids[object_cfg.env_ids] + command = env.command_manager.get_command(command_name)[object_cfg.env_ids] + goal_pos_w, _ = math_utils.combine_frame_transforms( + robot.data.root_pos_w.torch[robot_rows], robot.data.root_quat_w.torch[robot_rows], command[:, :3] + ) + error = torch.linalg.norm(object_asset.data.root_pos_w.torch - goal_pos_w, dim=-1) + return object_cfg.env_ids, error + + +def lift_ee_object_distance( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntitySelectionCfg, + object_cfg: SceneEntitySelectionCfg, + std: float, +) -> torch.Tensor: + """Reward the lift OpenArm for bringing its TCP close to the object.""" + env_ids, tcp_pos_w, _ = _offset_body_pose(env, robot_cfg, (0.0, 0.0, 0.0)) + object_asset: RigidObject = env.scene[object_cfg.name] + distance = torch.linalg.norm( + object_asset.data.root_pos_w.torch[object_cfg.instance_ids[env_ids]] - tcp_pos_w, dim=-1 + ) + return robot_cfg.scatter_to_envs(1.0 - torch.tanh(distance / std)) + + +def lift_object_height( + env: ManagerBasedRLEnv, object_cfg: SceneEntitySelectionCfg, minimum_height: float +) -> torch.Tensor: + """Reward lift environments whose object rises above an environment-frame height [m].""" + object_asset: RigidObject = env.scene[object_cfg.name] + height = object_asset.data.root_pos_w.torch[:, 2] - env.scene.env_origins[object_cfg.env_ids, 2] + return object_cfg.scatter_to_envs((height > minimum_height).float()) + + +def lift_goal_tracking( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntitySelectionCfg, + object_cfg: SceneEntitySelectionCfg, + command_name: str, + std: float, + minimum_height: float, +) -> torch.Tensor: + """Reward lift-object target tracking after the object clears the table.""" + env_ids, error = _lift_goal_error(env, robot_cfg, object_cfg, command_name) + object_asset: RigidObject = env.scene[object_cfg.name] + height = object_asset.data.root_pos_w.torch[:, 2] - env.scene.env_origins[env_ids, 2] + reward = (height > minimum_height) * (1.0 - torch.tanh(error / std)) + return object_cfg.scatter_to_envs(reward) + + +class _LiftSuccessTerm(ManagerTermBase): + """Maintain and log sticky success state for lift reward terms.""" + + def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv) -> None: + super().__init__(cfg, env) + self.succeeded = torch.zeros(env.num_envs, dtype=torch.bool, device=env.device) + + def reset(self, env_ids: torch.Tensor) -> None: + """Log and clear success for reset lift environments.""" + object_cfg: SceneEntitySelectionCfg = self.cfg.params["object_cfg"] + selected = env_ids[object_cfg.instance_ids[env_ids] >= 0] + if selected.numel() == 0: + return + log = self._env.extras.setdefault("log", {}) + log["Metrics/lift_success_rate"] = self.succeeded[selected].float().mean().item() + self.succeeded[selected] = False + + +class LiftGoalTracking(_LiftSuccessTerm): + """Reward OpenArm goal tracking and log sticky per-episode success.""" + + def __call__( + self, + env: ManagerBasedRLEnv, + robot_cfg: SceneEntitySelectionCfg, + object_cfg: SceneEntitySelectionCfg, + command_name: str, + std: float, + minimum_height: float, + success_threshold: float, + ) -> torch.Tensor: + """Return coarse goal tracking and update sticky success state.""" + env_ids, error = _lift_goal_error(env, robot_cfg, object_cfg, command_name) + object_asset: RigidObject = env.scene[object_cfg.name] + height = object_asset.data.root_pos_w.torch[:, 2] - env.scene.env_origins[env_ids, 2] + is_lifted = height > minimum_height + self.succeeded[env_ids] |= is_lifted & (error < success_threshold) + reward = is_lifted * (1.0 - torch.tanh(error / std)) + return object_cfg.scatter_to_envs(reward) + + +class LiftSuccess(_LiftSuccessTerm): + """Reward lift-goal completion and log sticky episode success.""" + + def __call__( + self, + env: ManagerBasedRLEnv, + robot_cfg: SceneEntitySelectionCfg, + object_cfg: SceneEntitySelectionCfg, + command_name: str, + threshold: float, + ) -> torch.Tensor: + """Reward lift environments whose object reaches the sampled goal.""" + env_ids, error = _lift_goal_error(env, robot_cfg, object_cfg, command_name) + succeeded = error < threshold + self.succeeded[env_ids] |= succeeded + return object_cfg.scatter_to_envs(succeeded.float()) + + +def _cabinet_frames( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntitySelectionCfg, + cabinet_cfg: SceneEntitySelectionCfg, +) -> tuple[torch.Tensor, ...]: + """Return aligned cabinet-task TCP, handle, and fingertip poses.""" + env_ids, ee_pos, ee_quat = _offset_body_pose(env, robot_cfg, (0.0, 0.0, 0.1034)) + _, handle_pos, handle_quat = _offset_body_pose(env, cabinet_cfg, (0.305, 0.0, 0.01), (0.5, -0.5, -0.5, 0.5)) + _, left_pos, _ = _offset_body_pose(env, robot_cfg, (0.0, 0.0, 0.046), body_index=1) + _, right_pos, _ = _offset_body_pose(env, robot_cfg, (0.0, 0.0, 0.046), body_index=2) + cabinet_rows = cabinet_cfg.instance_ids[env_ids] + return env_ids, ee_pos, ee_quat, handle_pos[cabinet_rows], handle_quat[cabinet_rows], left_pos, right_pos + + +def cabinet_approach_ee_handle( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntitySelectionCfg, + cabinet_cfg: SceneEntitySelectionCfg, + threshold: float, +) -> torch.Tensor: + """Reward cabinet environments for approaching the drawer handle.""" + env_ids, ee_pos, _, handle_pos, _, _, _ = _cabinet_frames(env, robot_cfg, cabinet_cfg) + distance = torch.linalg.norm(handle_pos - ee_pos, dim=-1) + reward = torch.square(1.0 / (1.0 + distance**2)) + reward = torch.where(distance <= threshold, 2.0 * reward, reward) + return robot_cfg.scatter_to_envs(torch.nan_to_num(reward, nan=0.0, posinf=0.0, neginf=0.0)) + + +def cabinet_align_ee_handle( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntitySelectionCfg, + cabinet_cfg: SceneEntitySelectionCfg, +) -> torch.Tensor: + """Reward cabinet environments for aligning the TCP with the drawer handle.""" + env_ids, _, ee_quat, _, handle_quat, _, _ = _cabinet_frames(env, robot_cfg, cabinet_cfg) + ee_rot = math_utils.matrix_from_quat(ee_quat) + handle_rot = math_utils.matrix_from_quat(handle_quat) + align_z = torch.sum(ee_rot[..., 2] * -handle_rot[..., 0], dim=-1) + align_x = torch.sum(ee_rot[..., 0] * -handle_rot[..., 1], dim=-1) + reward = 0.5 * (torch.sign(align_z) * align_z**2 + torch.sign(align_x) * align_x**2) + return robot_cfg.scatter_to_envs(torch.nan_to_num(reward, nan=0.0, posinf=0.0, neginf=0.0)) + + +def _cabinet_grasp_alignment( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntitySelectionCfg, + cabinet_cfg: SceneEntitySelectionCfg, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Return cabinet IDs, grasp mask, handle position, and fingertip positions.""" + env_ids, _, _, handle_pos, _, left_pos, right_pos = _cabinet_frames(env, robot_cfg, cabinet_cfg) + graspable = (right_pos[:, 2] < handle_pos[:, 2]) & (left_pos[:, 2] > handle_pos[:, 2]) + return env_ids, graspable, handle_pos, torch.stack((left_pos, right_pos), dim=1) + + +def cabinet_align_grasp( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntitySelectionCfg, + cabinet_cfg: SceneEntitySelectionCfg, +) -> torch.Tensor: + """Reward cabinet environments when the fingers straddle the drawer handle.""" + env_ids, graspable, _, _ = _cabinet_grasp_alignment(env, robot_cfg, cabinet_cfg) + return robot_cfg.scatter_to_envs(graspable.float()) + + +def cabinet_approach_gripper( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntitySelectionCfg, + cabinet_cfg: SceneEntitySelectionCfg, + offset: float, +) -> torch.Tensor: + """Reward cabinet environments for placing both fingers around the handle.""" + env_ids, graspable, handle_pos, fingertips = _cabinet_grasp_alignment(env, robot_cfg, cabinet_cfg) + distances = torch.abs(fingertips[..., 2] - handle_pos[:, None, 2]) + reward = graspable * torch.sum(offset - distances, dim=-1) + return robot_cfg.scatter_to_envs(torch.nan_to_num(reward, nan=0.0, posinf=0.0, neginf=0.0)) + + +def cabinet_grasp_handle( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntitySelectionCfg, + cabinet_cfg: SceneEntitySelectionCfg, + threshold: float, + open_joint_pos: float, +) -> torch.Tensor: + """Reward closing the cabinet Franka fingers near the drawer handle.""" + env_ids, ee_pos, _, handle_pos, _, _, _ = _cabinet_frames(env, robot_cfg, cabinet_cfg) + robot: Articulation = env.scene[robot_cfg.name] + distance = torch.linalg.norm(handle_pos - ee_pos, dim=-1) + reward = (distance <= threshold) * torch.sum( + open_joint_pos - robot.data.joint_pos.torch[:, robot_cfg.joint_ids], dim=-1 + ) + return robot_cfg.scatter_to_envs(torch.nan_to_num(reward, nan=0.0, posinf=0.0, neginf=0.0)) + + +class CabinetOpenDrawerBonus(ManagerTermBase): + """Reward drawer opening and log cabinet-task success statistics.""" + + def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv) -> None: + super().__init__(cfg, env) + self.succeeded = torch.zeros(env.num_envs, dtype=torch.bool, device=env.device) + self.best_drawer_pos = torch.zeros(env.num_envs, device=env.device) + + def reset(self, env_ids: torch.Tensor) -> None: + """Log and clear episode statistics for reset cabinet environments.""" + cabinet_cfg: SceneEntitySelectionCfg = self.cfg.params["cabinet_cfg"] + selected = env_ids[cabinet_cfg.instance_ids[env_ids] >= 0] + if selected.numel() == 0: + return + log = self._env.extras.setdefault("log", {}) + log["Metrics/cabinet_success_rate"] = self.succeeded[selected].float().mean().item() + log["Metrics/cabinet_drawer_pos"] = self.best_drawer_pos[selected].mean().item() + self.succeeded[selected] = False + self.best_drawer_pos[selected] = 0.0 + + def __call__( + self, + env: ManagerBasedRLEnv, + robot_cfg: SceneEntitySelectionCfg, + cabinet_cfg: SceneEntitySelectionCfg, + success_threshold: float, + ) -> torch.Tensor: + """Compute drawer displacement with a grasp-alignment multiplier.""" + cabinet: Articulation = env.scene[cabinet_cfg.name] + drawer_pos = cabinet.data.joint_pos.torch[:, cabinet_cfg.joint_ids[0]] + drawer_limits = cabinet.data.soft_joint_pos_limits.torch[:, cabinet_cfg.joint_ids[0]] + drawer_pos = torch.nan_to_num(drawer_pos, nan=0.0, posinf=0.0, neginf=0.0) + drawer_pos = torch.maximum(torch.minimum(drawer_pos, drawer_limits[:, 1]), drawer_limits[:, 0]) + env_ids, graspable, _, _ = _cabinet_grasp_alignment(env, robot_cfg, cabinet_cfg) + self.succeeded[env_ids] |= drawer_pos > success_threshold + self.best_drawer_pos[env_ids] = torch.maximum(self.best_drawer_pos[env_ids], drawer_pos) + return robot_cfg.scatter_to_envs((graspable.float() + 1.0) * drawer_pos) + + +def cabinet_multi_stage_open( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntitySelectionCfg, + cabinet_cfg: SceneEntitySelectionCfg, +) -> torch.Tensor: + """Reward easy, medium, and hard drawer-opening milestones.""" + cabinet: Articulation = env.scene[cabinet_cfg.name] + drawer_pos = cabinet.data.joint_pos.torch[:, cabinet_cfg.joint_ids[0]] + _, graspable, _, _ = _cabinet_grasp_alignment(env, robot_cfg, cabinet_cfg) + reward = (drawer_pos > 0.01) * 0.5 + reward += (drawer_pos > 0.2) * graspable + reward += (drawer_pos > 0.3) * graspable + return cabinet_cfg.scatter_to_envs(reward) + + +def reach_position_error(env: ManagerBasedRLEnv, robot_cfg: SceneEntitySelectionCfg, command_name: str) -> torch.Tensor: + """Return UR10 end-effector position error [m] in reach environments.""" + robot: Articulation = env.scene[robot_cfg.name] + command = env.command_manager.get_command(command_name)[robot_cfg.env_ids] + goal_pos_w, _ = math_utils.combine_frame_transforms( + robot.data.root_pos_w.torch, robot.data.root_quat_w.torch, command[:, :3] + ) + error = torch.linalg.norm(robot.data.body_pos_w.torch[:, robot_cfg.body_ids[0]] - goal_pos_w, dim=-1) + return robot_cfg.scatter_to_envs(error) + + +def reach_orientation_error( + env: ManagerBasedRLEnv, robot_cfg: SceneEntitySelectionCfg, command_name: str +) -> torch.Tensor: + """Return UR10 end-effector orientation error [rad] in reach environments.""" + robot: Articulation = env.scene[robot_cfg.name] + command = env.command_manager.get_command(command_name)[robot_cfg.env_ids] + goal_quat_w = math_utils.quat_mul(robot.data.root_quat_w.torch, command[:, 3:7]) + error = math_utils.quat_error_magnitude(robot.data.body_quat_w.torch[:, robot_cfg.body_ids[0]], goal_quat_w) + return robot_cfg.scatter_to_envs(error) + + +def selected_action_l2( + env: ManagerBasedRLEnv, + task_asset_cfg: SceneEntitySelectionCfg, + action_term_names: tuple[str, ...], +) -> torch.Tensor: + """Penalize one task action head in environments containing that task asset.""" + selected_slice = _action_slice(env, action_term_names) + penalty = torch.sum(torch.square(env.action_manager.action[:, selected_slice]), dim=-1) + return penalty * (task_asset_cfg.instance_ids >= 0) + + +def selected_action_rate_l2( + env: ManagerBasedRLEnv, + task_asset_cfg: SceneEntitySelectionCfg, + action_term_names: tuple[str, ...], +) -> torch.Tensor: + """Penalize changes in one task action head where that task asset exists.""" + selected_slice = _action_slice(env, action_term_names) + delta = env.action_manager.action[:, selected_slice] - env.action_manager.prev_action[:, selected_slice] + return torch.sum(torch.square(delta), dim=-1) * (task_asset_cfg.instance_ids >= 0) + + +def selected_joint_vel_l2( + env: ManagerBasedRLEnv, + asset_cfg: SceneEntitySelectionCfg, + max_velocity: float | None = None, +) -> torch.Tensor: + """Penalize articulation joint velocities in environments containing it. + + Args: + env: Manager-based RL environment. + asset_cfg: Selection-aware articulation and joints. + max_velocity: Optional absolute velocity bound [m/s or rad/s, depending on joint type]. + """ + asset: Articulation = env.scene[asset_cfg.name] + joint_vel = torch.nan_to_num(asset.data.joint_vel.torch[:, asset_cfg.joint_ids]) + if max_velocity is not None: + joint_vel = joint_vel.clamp(-max_velocity, max_velocity) + values = torch.sum(torch.square(joint_vel), dim=-1) + return asset_cfg.scatter_to_envs(values) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/terminations.py new file mode 100644 index 000000000000..3aebab1c40f4 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/terminations.py @@ -0,0 +1,101 @@ +# 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 + +"""Selection-aware terminations for heterogeneous manipulation tasks.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from ..selection_utils import SceneEntitySelectionCfg +from .rewards import reach_orientation_error, reach_position_error + +if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject + from isaaclab.envs import ManagerBasedRLEnv + + +def reach_success( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntitySelectionCfg, + command_name: str, + position_threshold: float, + orientation_threshold: float, +) -> torch.Tensor: + """Terminate successful UR10 reach environments.""" + position_error = reach_position_error(env, robot_cfg, command_name) + orientation_error = reach_orientation_error(env, robot_cfg, command_name) + return ( + (robot_cfg.instance_ids >= 0) + & (position_error < position_threshold) + & (orientation_error < orientation_threshold) + ) + + +def lift_object_out_of_bounds( + env: ManagerBasedRLEnv, + object_cfg: SceneEntitySelectionCfg, + bounds: tuple[tuple[float, float], tuple[float, float], tuple[float, float]], +) -> torch.Tensor: + """Terminate lift environments whose object leaves configured local bounds [m].""" + object_asset: RigidObject = env.scene[object_cfg.name] + position = object_asset.data.root_pos_w.torch - env.scene.env_origins[object_cfg.env_ids] + lower = position.new_tensor([axis[0] for axis in bounds]) + upper = position.new_tensor([axis[1] for axis in bounds]) + outside = torch.any((position < lower) | (position > upper), dim=-1) + return object_cfg.scatter_to_envs(outside) + + +def lift_object_dropped( + env: ManagerBasedRLEnv, + object_cfg: SceneEntitySelectionCfg, + minimum_height: float, +) -> torch.Tensor: + """Terminate lift environments whose object falls below a local height [m].""" + object_asset: RigidObject = env.scene[object_cfg.name] + height = object_asset.data.root_pos_w.torch[:, 2] - env.scene.env_origins[object_cfg.env_ids, 2] + return object_cfg.scatter_to_envs(height < minimum_height) + + +def articulation_state_invalid( + env: ManagerBasedRLEnv, + asset_cfg: SceneEntitySelectionCfg, + max_joint_velocity: float, + joint_position_margin: float, +) -> torch.Tensor: + """Terminate selected articulations with non-finite or implausible joint state. + + Args: + env: Manager-based RL environment. + asset_cfg: Selection-aware articulation and joints. + max_joint_velocity: Maximum absolute joint velocity [m/s or rad/s, depending on joint type]. + joint_position_margin: Allowed distance beyond soft position limits [m or rad, depending on joint type]. + """ + asset: Articulation = env.scene[asset_cfg.name] + joint_pos = asset.data.joint_pos.torch[:, asset_cfg.joint_ids] + joint_vel = asset.data.joint_vel.torch[:, asset_cfg.joint_ids] + limits = asset.data.soft_joint_pos_limits.torch[:, asset_cfg.joint_ids] + invalid = ~torch.isfinite(joint_pos).all(dim=-1) | ~torch.isfinite(joint_vel).all(dim=-1) + invalid |= torch.any(torch.abs(joint_vel) > max_joint_velocity, dim=-1) + invalid |= torch.any( + (joint_pos < limits[..., 0] - joint_position_margin) | (joint_pos > limits[..., 1] + joint_position_margin), + dim=-1, + ) + return asset_cfg.scatter_to_envs(invalid) + + +def task_time_out( + env: ManagerBasedRLEnv, + task_asset_cfgs: tuple[SceneEntitySelectionCfg, ...], + episode_lengths_s: tuple[float, ...], +) -> torch.Tensor: + """Time out each task at its homogeneous episode duration [s].""" + timed_out = torch.zeros(env.num_envs, dtype=torch.bool, device=env.device) + for asset_cfg, episode_length_s in zip(task_asset_cfgs, episode_lengths_s): + max_steps = round(episode_length_s / env.step_dt) + timed_out |= (asset_cfg.instance_ids >= 0) & (env.episode_length_buf >= max_steps) + return timed_out diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/utils.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/utils.py new file mode 100644 index 000000000000..d7725d4423bc --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/mdp/utils.py @@ -0,0 +1,47 @@ +# 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 + +"""Private tensor helpers for heterogeneous manipulation MDP terms.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from isaaclab.utils import math as math_utils + +from ..selection_utils import SceneEntitySelectionCfg + +if TYPE_CHECKING: + from isaaclab.assets import Articulation + from isaaclab.envs import ManagerBasedRLEnv + + +def _offset_body_pose( + env: ManagerBasedRLEnv, + asset_cfg: SceneEntitySelectionCfg, + offset_pos: tuple[float, float, float], + offset_quat: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0), + body_index: int = 0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return selected global IDs and an offset body pose in the world frame. + + Args: + env: Manager-based RL environment. + asset_cfg: Selection-aware entity containing exactly one body. + offset_pos: Child-frame translation [m]. + offset_quat: Child-frame quaternion in ``(x, y, z, w)`` order. + body_index: Index into the resolved body selection. + + Returns: + Global environment IDs, positions [m], and quaternions for the selected view rows. + """ + asset: Articulation = env.scene[asset_cfg.name] + pose = asset.data.body_link_pose_w.torch[:, asset_cfg.body_ids[body_index]] + pos = pose.new_tensor(offset_pos).expand(len(pose), -1) + quat = math_utils.convert_quat(pose.new_tensor(offset_quat).expand(len(pose), -1), to="wxyz") + pos_w, quat_w = math_utils.combine_frame_transforms(pose[:, :3], pose[:, 3:7], pos, quat) + return asset_cfg.env_ids, pos_w, quat_w diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/multitask_env.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/multitask_env.py new file mode 100644 index 000000000000..ace525ab8eac --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/multitask_manipulation/multitask_env.py @@ -0,0 +1,61 @@ +# 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 + +"""Manager-based environment with selection-aware heterogeneous resets.""" + +from collections.abc import Sequence + +import torch + +from isaaclab.envs import ManagerBasedRLEnv + +from .selection_utils import SceneEntitySelectionCfg + + +class MultitaskManipulationEnv(ManagerBasedRLEnv): + """Manager-based manipulation environment whose assets occupy partial physics views.""" + + def _reset_idx(self, env_ids: Sequence[int]) -> None: + """Reset global environments through each asset's view-row mapping. + + Args: + env_ids: Global environment IDs to reset. + """ + global_env_ids = torch.as_tensor(env_ids, dtype=torch.long, device=self.device) + + self.curriculum_manager.compute(env_ids=global_env_ids) + for asset_name, asset in (*self.scene.articulations.items(), *self.scene.rigid_objects.items()): + rows, _ = self._get_entity_selection(asset_name).select(global_env_ids) + if rows.numel() > 0: + asset.reset(rows) + + if "reset" in self.event_manager.available_modes: + env_step_count = self._sim_step_counter // self.cfg.decimation + self.event_manager.apply(mode="reset", env_ids=global_env_ids, global_env_step_count=env_step_count) + + self.extras["log"] = {} + managers = ( + self.observation_manager, + self.action_manager, + self.reward_manager, + self.curriculum_manager, + self.command_manager, + self.event_manager, + self.termination_manager, + self.recorder_manager, + ) + for manager in managers: + self.extras["log"].update(manager.reset(global_env_ids)) + + self.episode_length_buf[global_env_ids] = 0 + self.sim.render_context.reset_scene_state_cadence() + + def _get_entity_selection(self, asset_name: str) -> SceneEntitySelectionCfg: + """Return the cached selection configuration for a complete scene asset.""" + cache: dict[str, SceneEntitySelectionCfg] = self.__dict__.setdefault("_scene_entity_selections", {}) + if asset_name not in cache: + cache[asset_name] = SceneEntitySelectionCfg(asset_name) + cache[asset_name].resolve(self.scene) + return cache[asset_name] diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/reorient/config/shadow_hand/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/reorient/config/shadow_hand/__init__.py index cf604bdc6274..25de2b02b8a6 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/reorient/config/shadow_hand/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/reorient/config/shadow_hand/__init__.py @@ -3,12 +3,11 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Contributed Shadow Hand tasks: the rendering-throughput benchmark and the OpenAI variant.""" +"""Contributed Shadow Hand tasks: the OpenAI variants.""" import gymnasium as gym from isaaclab_tasks.contrib.reorient.config.shadow_hand import agents as contrib_agents -from isaaclab_tasks.core.reorient.config.shadow_hand import agents as core_agents gym.register( id="IsaacContrib-Reorient-Cube-Shadow-OpenAI-FF-Direct", @@ -31,13 +30,3 @@ "rsl_rl_cfg_entry_point": f"{contrib_agents.__name__}.rsl_rl_ppo_cfg:ShadowHandAsymLSTMPPORunnerCfg", }, ) -gym.register( - id="IsaacContrib-Reorient-Cube-Shadow-Camera-Benchmark-Direct", - entry_point="isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_direct_camera_env:ShadowHandCameraEnv", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.shadow_hand_camera_benchmark_env_cfg:ShadowHandCameraBenchmarkEnvCfg", - "rsl_rl_cfg_entry_point": f"{core_agents.__name__}.rsl_rl_ppo_cfg:ShadowHandCameraFFPPORunnerCfg", - "rl_games_cfg_entry_point": f"{core_agents.__name__}:rl_games_ppo_camera_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/a1/rough_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/a1/rough_env_cfg.py index bad828ad1f40..6f54f541c9be 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/a1/rough_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/a1/rough_env_cfg.py @@ -41,4 +41,4 @@ def __post_init__(self): # events self.events.add_base_mass.params["asset_cfg"].body_names = "trunk" self.events.base_external_force_torque.params["asset_cfg"].body_names = "trunk" - self.events.base_com.default.params["asset_cfg"].body_names = "trunk" + self.events.base_com.params["asset_cfg"].body_names = "trunk" diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/anymal_c/rough_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/anymal_c/rough_env_cfg.py index 5badbdf012d0..4d3af315ffe7 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/anymal_c/rough_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/anymal_c/rough_env_cfg.py @@ -7,7 +7,6 @@ from isaaclab.utils import configclass from isaaclab_tasks.core.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg -from isaaclab_tasks.utils import preset ## # Pre-defined configs @@ -22,6 +21,3 @@ def __post_init__(self): # scene self.scene.robot = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.scene.robot.actuators["legs"].armature = preset( - default=0.0, newton_mjwarp=0.01, physx=0.0, isaacsim_physx=0.0 - ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/digit/agents/rsl_rl_ppo_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/digit/agents/rsl_rl_ppo_cfg.py index f22b239d7aab..fffbd52389c5 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/digit/agents/rsl_rl_ppo_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/digit/agents/rsl_rl_ppo_cfg.py @@ -30,7 +30,8 @@ class DigitRoughPPORunnerCfg(RslRlOnPolicyRunnerCfg): value_loss_coef=1.0, use_clipped_value_loss=True, clip_param=0.2, - entropy_coef=0.01, + # MJWarp needs a lower entropy floor than PhysX here (0.005 vs the usual 0.01) -- see #7520. + entropy_coef=0.005, num_learning_epochs=5, num_mini_batches=4, learning_rate=1.0e-3, diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/digit/rough_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/digit/rough_env_cfg.py index 80d7254169b1..c39dd646f0ea 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/digit/rough_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/digit/rough_env_cfg.py @@ -5,25 +5,34 @@ import math +from isaaclab_newton.sim.schemas import NewtonArticulationCfg from isaaclab_physx.physics import PhysxCfg -from isaaclab.actuators import ImplicitActuatorCfg -from isaaclab.managers import ObservationGroupCfg, ObservationTermCfg, RewardTermCfg, SceneEntityCfg, TerminationTermCfg +from isaaclab.managers import ( + ObservationGroupCfg, + ObservationTermCfg, + RewardTermCfg, + SceneEntityCfg, + TerminationTermCfg, +) from isaaclab.physics import PhysxAutoCfg from isaaclab.sim import SimulationCfg from isaaclab.utils import configclass from isaaclab.utils.noise import UniformNoiseCfg as Unoise import isaaclab_tasks.core.velocity.mdp as mdp -from isaaclab_tasks.core.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg -from isaaclab_tasks.utils import PresetCfg +from isaaclab_tasks.core.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg, RoughPhysicsCfg +from isaaclab_tasks.utils import PresetCfg, preset from isaaclab_assets.robots.agility import ARM_JOINT_NAMES, DIGIT_V4_CFG, LEG_JOINT_NAMES +_ROUGH_NEWTON_MJWARP = RoughPhysicsCfg().newton_mjwarp +"""Bound once so ``DigitPhysicsCfg`` does not construct ``RoughPhysicsCfg`` twice.""" + @configclass class DigitPhysicsCfg(PresetCfg): - """PhysX-only physics configuration for the Digit velocity environments.""" + """Physics configuration for the Digit velocity environments.""" isaacsim_physx = PhysxCfg( gpu_max_rigid_patch_count=10 * 2**15, @@ -31,6 +40,13 @@ class DigitPhysicsCfg(PresetCfg): gpu_total_aggregate_pairs_capacity=2**23, ) physx = PhysxAutoCfg(isaacsim_physx=isaacsim_physx) + # ``class_type`` is reset to ``None`` because ``NewtonCfg.__post_init__`` re-derives it from + # ``solver_cfg`` and refuses an explicit value -- ``replace()`` would otherwise carry the + # already-derived value from the source instance forward as one. + newton_mjwarp = _ROUGH_NEWTON_MJWARP.replace( + class_type=None, + solver_cfg=_ROUGH_NEWTON_MJWARP.solver_cfg.replace(njmax=5000, nconmax=2000), + ) default = isaacsim_physx @@ -244,14 +260,6 @@ def __post_init__(self): self.scene.contact_forces.history_length = self.decimation self.scene.contact_forces.update_period = self.sim.dt self.scene.height_scanner.update_period = self.decimation * self.sim.dt - # target only actuated joints explicitly — ".*" mis-indexes Digit's ball-joint DoFs - self.scene.robot.actuators = { - "legs_arms": ImplicitActuatorCfg( - joint_names_expr=LEG_JOINT_NAMES + ARM_JOINT_NAMES, - stiffness=None, - damping=None, - ), - } # commands self.commands.base_velocity.ranges.lin_vel_x = (-0.8, 0.8) self.commands.base_velocity.ranges.lin_vel_y = (-0.5, 0.5) @@ -261,8 +269,12 @@ def __post_init__(self): # events self.events.add_base_mass.params["asset_cfg"].body_names = "torso_base" self.events.base_external_force_torque.params["asset_cfg"].body_names = "torso_base" - # Digit is PhysX-only, so the inherited ``newton_mjwarp`` branch names no - # reachable backend; collapse the preset so it cannot be selected on its own. - self.events.base_com = self.events.base_com.default self.events.base_com.params["asset_cfg"].body_names = "torso_base" + # The asset authors enabledSelfCollisions=False and DIGIT_V4_CFG sets no + # articulation_props, so Newton filters every intra-articulation shape pair -- 253 of + # them, exactly C(23,2) for its 23 colliding shapes -- and the legs pass through each + # other. Which links carry colliders is left as the asset authored it. + self.scene.robot.spawn.articulation_props = preset( + default=None, newton_mjwarp=[NewtonArticulationCfg(self_collision_enabled=True)] + ) self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/go1/rough_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/go1/rough_env_cfg.py index 90d64e85db7a..678c33ceaa62 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/go1/rough_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/velocity/config/go1/rough_env_cfg.py @@ -7,7 +7,6 @@ from isaaclab.utils import configclass from isaaclab_tasks.core.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg -from isaaclab_tasks.utils import preset ## # Pre-defined configs @@ -22,7 +21,6 @@ def __post_init__(self): # scene self.scene.robot = UNITREE_GO1_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.scene.robot.actuators["base_legs"].armature = preset(default=0.0, newton_mjwarp=0.02) self.scene.height_scanner.prim_path = "{ENV_REGEX_NS}/Robot/trunk" # scale down the terrains because the robot is small self.scene.terrain.terrain_generator.sub_terrains["boxes"].grid_height_range = (0.025, 0.1) @@ -43,4 +41,4 @@ def __post_init__(self): # events self.events.add_base_mass.params["asset_cfg"].body_names = "trunk" self.events.base_external_force_torque.params["asset_cfg"].body_names = "trunk" - self.events.base_com.default.params["asset_cfg"].body_names = "trunk" + self.events.base_com.params["asset_cfg"].body_names = "trunk" diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/lift/config/franka/franka_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/lift/config/franka/franka_env_cfg.py index 26c063d78da2..b244c84dcab0 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/lift/config/franka/franka_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/lift/config/franka/franka_env_cfg.py @@ -24,6 +24,8 @@ # legacy asset so the upstream franka tasks keep their demos and baselines. FRANKA_PANDA_LIFT_CFG = FRANKA_PANDA_CFG.copy() FRANKA_PANDA_LIFT_CFG.spawn.usd_path = f"{ISAACLAB_NUCLEUS_DIR}/Robots/FrankaEmika/franka_panda.usda" +# Reset clearance was calibrated for these arm meshes; the asset's primitive colliders intersect the ground. +FRANKA_PANDA_LIFT_CFG.spawn.variants = {"Colliders": "convex_hulls"} FRANKA_PANDA_LIFT_CFG.actuators = { # Inspired by libfranka's joint_impedance_control.cpp. ``actuator_velocity_limit`` # remains the soft task-limit snapshot; ``joint_velocity_limit`` is the diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/cassie/rough_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/cassie/rough_env_cfg.py index 16170f952f12..f65b6af462a1 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/cassie/rough_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/cassie/rough_env_cfg.py @@ -13,7 +13,6 @@ LocomotionVelocityRoughEnvCfg, RewardsCfg, ) -from isaaclab_tasks.utils import preset ## # Pre-defined configs @@ -60,7 +59,6 @@ def __post_init__(self): # scene self.scene.robot = CASSIE_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.scene.robot.actuators["legs"].armature = preset(default=0.0, newton_mjwarp=0.02) self.scene.height_scanner.prim_path = "{ENV_REGEX_NS}/Robot/pelvis" # actions self.actions.joint_pos.scale = 0.5 diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/g1/agents/rsl_rl_ppo_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/g1/agents/rsl_rl_ppo_cfg.py index 197c80a59180..7b890ebdf66c 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/g1/agents/rsl_rl_ppo_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/g1/agents/rsl_rl_ppo_cfg.py @@ -7,19 +7,11 @@ from isaaclab_rl.rsl_rl import RslRlMLPModelCfg, RslRlOnPolicyRunnerCfg, RslRlPpoAlgorithmCfg -from isaaclab_tasks.utils import preset - @configclass class G1RoughPPORunnerCfg(RslRlOnPolicyRunnerCfg): num_steps_per_env = 24 - # Newton needs ~1.7x the PPO iterations to match PhysX on G1. PhysX saturates near iter 3000 - # (reward ≈ +18, ep_len ≈ 980) and does not meaningfully improve on either metric past that — - # reward oscillates +16 to +19 through iter 7500, ep_len stays flat. Newton reaches the same - # (reward, ep_len) quality at iter 5000 (+16 / 984). Comparing reward alone is misleading: - # ep_len confirms the robot is stable in both cases. The gap is sample-efficiency, not a - # ceiling — no physics or reward tuning closes it. - max_iterations = preset(default=3000, newton_mjwarp=5000) + max_iterations = 5000 save_interval = 50 experiment_name = "g1_rough" obs_groups = {"actor": ["policy"], "critic": ["policy"]} diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/go2/rough_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/go2/rough_env_cfg.py index 925318d6187f..ed7e260714f8 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/go2/rough_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/go2/rough_env_cfg.py @@ -7,7 +7,6 @@ from isaaclab.utils import configclass from isaaclab_tasks.core.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg -from isaaclab_tasks.utils import preset ## # Pre-defined configs @@ -25,7 +24,6 @@ def __post_init__(self): self.sim.use_newton_actuators = True # scene self.scene.robot = UNITREE_GO2_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.scene.robot.actuators["base_legs"].armature = preset(default=0.0, newton_mjwarp=0.02) self.scene.height_scanner.prim_path = "{ENV_REGEX_NS}/Robot/base" # scale down the terrains because the robot is small terrains = self.scene.terrain.terrain_generator.sub_terrains diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/velocity_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/velocity_env_cfg.py index 1d127d6c33e9..7aeae7f961fb 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/velocity_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/velocity_env_cfg.py @@ -36,7 +36,7 @@ from isaaclab.utils.noise import UniformNoiseCfg as Unoise import isaaclab_tasks.core.velocity.mdp as mdp -from isaaclab_tasks.utils import PresetCfg, preset +from isaaclab_tasks.utils import PresetCfg ## # Pre-defined configs @@ -223,15 +223,13 @@ class EventsCfg: }, ) - base_com = preset( - default=EventTerm( - func=mdp.randomize_rigid_body_com, - mode="startup", - params={ - "asset_cfg": SceneEntityCfg("robot", body_names="base"), - "com_range": {"x": (-0.05, 0.05), "y": (-0.05, 0.05), "z": (-0.01, 0.01)}, - }, - ), + base_com = EventTerm( + func=mdp.randomize_rigid_body_com, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names="base"), + "com_range": {"x": (-0.05, 0.05), "y": (-0.05, 0.05), "z": (-0.01, 0.01)}, + }, ) # reset diff --git a/source/isaaclab_tasks/test/benchmark/test_render_benchmark_cfg.py b/source/isaaclab_tasks/test/benchmark/test_render_benchmark_cfg.py new file mode 100644 index 000000000000..431cdfa4730e --- /dev/null +++ b/source/isaaclab_tasks/test/benchmark/test_render_benchmark_cfg.py @@ -0,0 +1,130 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Tests for the ``Isaac-RenderBenchmark-Franka-Cabinet`` task's registration and presets. + +No Kit/GPU required: these only load and resolve the config through the registry, the same way +:mod:`scripts.benchmarks.benchmark_renderer` selects a preset before launching a run. +""" + +from types import SimpleNamespace + +import gymnasium as gym +import pytest +from isaaclab_newton.physics import NewtonCfg +from isaaclab_newton.renderers import NewtonWarpRendererCfg +from isaaclab_physx.physics import PhysxCfg + +from isaaclab.physics import PhysxAutoCfg + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.benchmark.render_benchmark.render_benchmark_env import RenderBenchmarkEnv +from isaaclab_tasks.utils.hydra import collect_presets, resolve_presets +from isaaclab_tasks.utils.parse_cfg import load_cfg_from_registry + +_TASK = "Isaac-RenderBenchmark-Franka-Cabinet" + + +def _load_cfg(): + return load_cfg_from_registry(_TASK, "env_cfg_entry_point") + + +def test_task_registered_with_direct_entry_point(): + spec = gym.spec(_TASK) + + assert spec.kwargs["env_cfg_entry_point"] == ( + "isaaclab_tasks.benchmark.render_benchmark.render_benchmark_env_cfg:RenderBenchmarkFrankaCabinetEnvCfg" + ) + assert spec.disable_env_checker is True + + +def test_default_scene_and_articulations(): + cfg = _load_cfg() + + assert cfg.decimation == 2 + assert cfg.episode_length_s == pytest.approx(60.0) + assert cfg.scene.num_envs == 4 + assert set(cfg.articulations) == {"robot", "cabinet"} + assert cfg.joint_animation_amplitude == pytest.approx(0.4) + + +@pytest.mark.parametrize( + ("presets", "expected_type"), + [ + ((), NewtonCfg), + (("newton_mjwarp",), NewtonCfg), + (("physx",), PhysxAutoCfg), + (("isaacsim_physx",), PhysxCfg), + ], +) +def test_physics_presets_resolve_to_expected_backend(presets, expected_type): + cfg = resolve_presets(_load_cfg(), selected=presets) + + assert isinstance(cfg.sim.physics, expected_type) + + +_CAMERA_DATA_TYPE_PRESETS = [ + ("default", ["rgb"]), + ("rgb", ["rgb"]), + ("albedo", ["albedo"]), + ("depth", ["depth"]), + ("simple_shading_constant_diffuse", ["simple_shading_constant_diffuse"]), + ("simple_shading_diffuse_mdl", ["simple_shading_diffuse_mdl"]), + ("simple_shading_full_mdl", ["simple_shading_full_mdl"]), +] + + +@pytest.fixture(scope="module") +def render_benchmark_presets(): + """Collect every preset once, before any preset selection mutates the config in place.""" + return collect_presets(_load_cfg()) + + +@pytest.mark.parametrize("preset_name,expected_data_types", _CAMERA_DATA_TYPE_PRESETS) +def test_camera_presets_resolve_to_expected_data_types(render_benchmark_presets, preset_name, expected_data_types): + camera_presets = render_benchmark_presets["tiled_camera"] + assert preset_name in camera_presets, f"Preset '{preset_name}' not found in tiled_camera presets" + resolved = camera_presets[preset_name] + + assert resolved.data_types == expected_data_types + assert resolved.width > 0 + assert resolved.height > 0 + assert isinstance(resolved.renderer_cfg.newton_renderer, NewtonWarpRendererCfg) + + +def test_camera_resolution_defaults_to_256(): + """``BENCHMARK_RENDER_RESOLUTION`` is read once at import, so its effect is process-scoped; + this documents the default a fresh process gets when the env var is unset.""" + cfg = resolve_presets(_load_cfg(), selected=()) + + assert cfg.tiled_camera.width == 256 + assert cfg.tiled_camera.height == 256 + + +def _fake_env_for_ground_cfg(ground_size: tuple[float, float], env_spacing: float) -> SimpleNamespace: + return SimpleNamespace( + cfg=SimpleNamespace( + ground_size=ground_size, ground_thickness=0.1, ground_color=(0.5, 0.5, 0.5), ground_top_z=0.0 + ), + scene=SimpleNamespace(cfg=SimpleNamespace(env_spacing=env_spacing)), + ) + + +def test_ground_cfg_clamps_size_to_env_spacing(): + """An oversized ground tile must be clamped, or cloned environments' grounds overlap.""" + fake_env = _fake_env_for_ground_cfg(ground_size=(50.0, 50.0), env_spacing=3.0) + + ground_cfg = RenderBenchmarkEnv._ground_cfg(fake_env) + + assert ground_cfg.spawn.size[:2] == (3.0, 3.0) + + +def test_ground_cfg_keeps_size_smaller_than_env_spacing(): + """A tile already smaller than the env spacing must not be inflated.""" + fake_env = _fake_env_for_ground_cfg(ground_size=(2.0, 1.5), env_spacing=3.0) + + ground_cfg = RenderBenchmarkEnv._ground_cfg(fake_env) + + assert ground_cfg.spawn.size[:2] == (2.0, 1.5) diff --git a/source/isaaclab_tasks/test/contrib/test_contrib_environments.py b/source/isaaclab_tasks/test/contrib/test_contrib_environments.py index a1dc901834ee..3d2b961e244d 100644 --- a/source/isaaclab_tasks/test/contrib/test_contrib_environments.py +++ b/source/isaaclab_tasks/test/contrib/test_contrib_environments.py @@ -32,6 +32,11 @@ "IsaacContrib-AutoMate-Disassembly-Direct": "Requires CUDA support outside the standard environment test runner.", } _SKIPPED_TASK_SUBSTRINGS = { + # Under random actions the Kamino P-ADMM solver intermittently diverges and the whole robot state + # (root pose, joint state) turns NaN mid-episode, so the run fails nondeterministically (about 1 in 12 + # seeds locally; the sibling HoldPose task stays finite). The termination terms cannot catch a NaN state. + # Re-enable once the solver instability is resolved upstream. + "DrLegs-Walk": "Kamino solver intermittently produces NaN robot state under random actions.", "RmpFlow": "Uses SingleArticulation, which requires an update.", "Skillgen": "Requires cuRobo-specific coverage.", "Suction": "Requires CPU simulation.", @@ -67,4 +72,5 @@ def _contrib_environment_params() -> list: @pytest.mark.parametrize("task_name", _contrib_environment_params()) def test_contrib_environments(task_name): - _run_environments(task_name, device="cuda", num_envs=2) + num_envs = 3 if task_name == "IsaacContrib-Multitask-Manipulation" else 2 + _run_environments(task_name, device="cuda", num_envs=num_envs) diff --git a/source/isaaclab_tasks/test/contrib/test_franka_pour_env_cfg.py b/source/isaaclab_tasks/test/contrib/test_franka_pour_env_cfg.py index 17e936d5f5b1..37e4f5fd78dd 100644 --- a/source/isaaclab_tasks/test/contrib/test_franka_pour_env_cfg.py +++ b/source/isaaclab_tasks/test/contrib/test_franka_pour_env_cfg.py @@ -111,6 +111,7 @@ def test_reset_dataset_contract_stores_root_relative_robot_asset_path(): robot_asset = _reset_dataset_task_contract(cfg)["robot_asset"] assert robot_asset == "Robots/FrankaEmika/franka_panda.usda" assert f"{ISAACLAB_NUCLEUS_DIR}/{robot_asset}" == FRANKA_POUR_ROBOT_ASSET_ID + assert cfg.scene.robot.spawn.variants == {"Colliders": "convex_hulls"} def test_capacity_resolution_only_updates_world_dependent_solver_limits(): diff --git a/source/isaaclab_tasks/test/core/test_g1_agent_cfg.py b/source/isaaclab_tasks/test/core/test_g1_agent_cfg.py deleted file mode 100644 index 461811f05bba..000000000000 --- a/source/isaaclab_tasks/test/core/test_g1_agent_cfg.py +++ /dev/null @@ -1,14 +0,0 @@ -# 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 - -from isaaclab_tasks.core.velocity.config.g1.agents.rsl_rl_ppo_cfg import G1RoughPPORunnerCfg -from isaaclab_tasks.utils import resolve_presets - - -def test_g1_rough_newton_preset_uses_extended_training_schedule(): - """Verify MJWarp selects the tuned G1 Newton training duration.""" - agent_cfg = resolve_presets(G1RoughPPORunnerCfg(), selected=("newton_mjwarp",)) - - assert agent_cfg.max_iterations == 5000 diff --git a/source/isaaclab_tasks/test/core/test_hydra.py b/source/isaaclab_tasks/test/core/test_hydra.py index 2a77316df346..10fd5e1807c5 100644 --- a/source/isaaclab_tasks/test/core/test_hydra.py +++ b/source/isaaclab_tasks/test/core/test_hydra.py @@ -860,24 +860,15 @@ class EnvCfgFactory: # ============================================================================= -def test_go2_rough_newton_mjwarp_armature_preset(): - """Go2 rough terrain uses higher MJWarp armature without changing PhysX.""" - from isaaclab_tasks.core.velocity.config.go2.rough_env_cfg import UnitreeGo2RoughEnvCfg - - env_cfg, _ = _apply(UnitreeGo2RoughEnvCfg(), global_presets=["newton_mjwarp"]) - assert env_cfg.scene.robot.actuators["base_legs"].armature == 0.02 - - env_cfg, _ = _apply(UnitreeGo2RoughEnvCfg()) - assert env_cfg.scene.robot.actuators["base_legs"].armature == 0.0 - - def test_go2_rough_legacy_newton_alias_resolves_to_newton_mjwarp(): """Real-config alias path: ``presets=newton`` against an actual env cfg resolves to newton_mjwarp.""" + from isaaclab_newton.physics import MJWarpSolverCfg + from isaaclab_tasks.core.velocity.config.go2.rough_env_cfg import UnitreeGo2RoughEnvCfg with pytest.warns(FutureWarning, match="Preset 'newton' is deprecated"): env_cfg, _ = _apply(UnitreeGo2RoughEnvCfg(), global_presets=["newton"]) - assert env_cfg.scene.robot.actuators["base_legs"].armature == 0.02 + assert isinstance(env_cfg.sim.physics.solver_cfg, MJWarpSolverCfg) def test_velocity_events_newton_mjwarp_keeps_base_com_randomization(): diff --git a/source/isaaclab_tasks/test/core/test_lift_env_cfg.py b/source/isaaclab_tasks/test/core/test_lift_env_cfg.py index e9ab35c1dbba..eb6de255122b 100644 --- a/source/isaaclab_tasks/test/core/test_lift_env_cfg.py +++ b/source/isaaclab_tasks/test/core/test_lift_env_cfg.py @@ -10,9 +10,13 @@ import pytest import torch +from pxr import Usd + from isaaclab.managers import CommandTerm +from isaaclab.sim import select_usd_variants from isaaclab_tasks.core.lift import mdp +from isaaclab_tasks.core.lift.config.franka.franka_env_cfg import FrankaLiftEnvCfg, FrankaReorientEnvCfg from isaaclab_tasks.core.lift.config.franka_soft.franka_soft_env_cfg import FrankaSoftEnvCfg from isaaclab_tasks.core.lift.mdp.commands.pose_commands import ( CableUniformPoseCommand, @@ -58,6 +62,29 @@ def test_franka_soft_robot_physics_variant_matches_backend( assert cfg.scene.robot.spawn.variants == {"Physics": expected_physics} +@pytest.mark.parametrize("cfg_type", [FrankaLiftEnvCfg, FrankaReorientEnvCfg]) +def test_franka_rigid_tasks_select_collision_meshes_for_reset_clearance(cfg_type) -> None: + """Reset validation keeps the original arm meshes when the asset defaults to capsules.""" + cfg = cfg_type() + stage = Usd.Stage.CreateInMemory() + robot = stage.DefinePrim("/Robot", "Xform") + colliders = robot.GetVariantSets().AddVariantSet("Colliders") + for selection, prim_path, prim_type in ( + ("convex_hulls", "/Robot/link1_c/link1_c", "Mesh"), + ("primitives", "/Robot/link1_capsule", "Capsule"), + ): + colliders.AddVariant(selection) + colliders.SetVariantSelection(selection) + with colliders.GetVariantEditContext(): + stage.DefinePrim(prim_path, prim_type) + colliders.SetVariantSelection("primitives") + + select_usd_variants("/Robot", cfg.scene.robot.spawn.variants or {}, stage=stage) + + assert stage.GetPrimAtPath("/Robot/link1_c/link1_c").IsValid() + assert not stage.GetPrimAtPath("/Robot/link1_capsule").IsValid() + + def test_camera_normalization_is_stationary() -> None: """RGB and depth normalization must not depend on per-frame statistics.""" rgb = torch.tensor([0.0, 127.5, 255.0])