From 1c597c7d94df8cfe7a751f246fe09d90226efb78 Mon Sep 17 00:00:00 2001 From: zhx06 Date: Tue, 28 Jul 2026 15:17:22 -0700 Subject: [PATCH 1/7] add joint support for robots Signed-off-by: zhx06 --- isaaclab_arena/assets/asset_cache.py | 8 + isaaclab_arena/embodiments/agibot/agibot.py | 1 + isaaclab_arena/embodiments/droid/droid.py | 1 + isaaclab_arena/embodiments/embodiment_base.py | 85 ++-- isaaclab_arena/embodiments/franka/franka.py | 1 + isaaclab_arena/embodiments/g1/g1.py | 1 + isaaclab_arena/embodiments/galbot/galbot.py | 1 + isaaclab_arena/embodiments/gr1t2/gr1t2.py | 1 + .../embodiments/kuka_allegro/kuka_allegro.py | 1 + .../embodiments/robot_on_stand_utils.py | 4 +- .../export_ready_pose_collision_meshes.py | 171 ++++++++ .../tests/test_collision_mesh_store.py | 386 +++++++++++++++++ .../tests/test_embodiment_collision_mesh.py | 32 +- isaaclab_arena/tests/test_usd_articulation.py | 409 ++++++++++++++++++ isaaclab_arena/utils/collision_mesh_store.py | 334 ++++++++++++++ isaaclab_arena/utils/isaac_sim_debug_draw.py | 34 ++ isaaclab_arena/utils/usd_articulation.py | 263 +++++++++++ isaaclab_arena/utils/usd_helpers.py | 172 +++++++- ...visualize_embodiment_placement_geometry.py | 174 ++++++++ 19 files changed, 2042 insertions(+), 37 deletions(-) create mode 100644 isaaclab_arena/scripts/export_ready_pose_collision_meshes.py create mode 100644 isaaclab_arena/tests/test_collision_mesh_store.py create mode 100644 isaaclab_arena/tests/test_usd_articulation.py create mode 100644 isaaclab_arena/utils/collision_mesh_store.py create mode 100644 isaaclab_arena/utils/usd_articulation.py create mode 100644 isaaclab_arena_examples/relations/visualize_embodiment_placement_geometry.py diff --git a/isaaclab_arena/assets/asset_cache.py b/isaaclab_arena/assets/asset_cache.py index 05628b4c56..1f93ae1c7f 100644 --- a/isaaclab_arena/assets/asset_cache.py +++ b/isaaclab_arena/assets/asset_cache.py @@ -13,3 +13,11 @@ def get_arena_asset_cache_dir() -> pathlib.Path: if not asset_cache_dir.exists(): asset_cache_dir.mkdir(parents=True, exist_ok=True) return asset_cache_dir + + +def get_arena_usd_cache_dir() -> pathlib.Path: + """Return the cache root for USDs Arena generates, such as composed and baked-geometry assets. + + The directory is not created, so callers that only compute a path do not leave one behind. + """ + return pathlib.Path.home() / ".cache" / "isaaclab_arena" / "usd" diff --git a/isaaclab_arena/embodiments/agibot/agibot.py b/isaaclab_arena/embodiments/agibot/agibot.py index 354dadfce1..9b6cb78f80 100644 --- a/isaaclab_arena/embodiments/agibot/agibot.py +++ b/isaaclab_arena/embodiments/agibot/agibot.py @@ -32,6 +32,7 @@ class AgibotEmbodiment(EmbodimentBase): name = "agibot" default_arm_mode = ArmMode.LEFT + robot_library_folder = "agibot_a2d" def __init__( self, enable_cameras: bool = False, initial_pose: Pose | None = None, arm_mode: ArmMode = ArmMode.LEFT diff --git a/isaaclab_arena/embodiments/droid/droid.py b/isaaclab_arena/embodiments/droid/droid.py index 66f6d1825e..9f88df7e3e 100644 --- a/isaaclab_arena/embodiments/droid/droid.py +++ b/isaaclab_arena/embodiments/droid/droid.py @@ -73,6 +73,7 @@ class DroidEmbodimentBase(EmbodimentBase, ABC): name = "droid" default_arm_mode = ArmMode.SINGLE_ARM + robot_library_folder = "droid" def __init__( self, diff --git a/isaaclab_arena/embodiments/embodiment_base.py b/isaaclab_arena/embodiments/embodiment_base.py index f3d0ab17d4..da0fb786f2 100644 --- a/isaaclab_arena/embodiments/embodiment_base.py +++ b/isaaclab_arena/embodiments/embodiment_base.py @@ -6,6 +6,7 @@ from __future__ import annotations from collections.abc import Mapping +from dataclasses import dataclass from typing import TYPE_CHECKING, Any from isaaclab.envs import ManagerBasedRLMimicEnv @@ -23,12 +24,36 @@ import trimesh +@dataclass(frozen=True) +class PlacementGeometrySource: + """What an embodiment's placement bounding box and collision mesh are derived from.""" + + usd_path: str + """Robot USD, as spawned.""" + + scale: tuple[float, float, float] + """Per-axis spawn scale.""" + + joint_pos: Mapping[str, float] + """Joint positions to pose the geometry at, revolute in radians, keyed by name or Isaac Lab regex.""" + + library_folder: str | None + """Robot's folder under the published robot library, or None if it publishes no collision mesh.""" + + class EmbodimentBase(PlaceableAsset): name: str | None = None tags: list[str] = ["embodiment"] default_arm_mode: ArmMode | None = None + robot_library_folder: str | None = None + """This robot's folder under the published robot library, shared by its action-space variants. + + Named separately from the USD because robots that spawn a composed asset, such as the on-stand + Droid and Franka, spawn from a per-user cache path that says nothing about where they publish. + """ + def __init__( self, enable_cameras: bool = False, @@ -56,42 +81,52 @@ def __init__( self.mimic_env: Any | None = None self.xr: Any | None = None self.termination_cfg: Any | None = None - self._collision_mesh: trimesh.Trimesh | None = None - """Lazily-extracted robot collision mesh, cached so the USD is opened once.""" + + def get_placement_geometry_source(self) -> PlacementGeometrySource: + """Return the USD, scale and joint positions the robot's placement geometry is derived from.""" + assert self.scene_config is not None, "scene_config must be populated before placement" + robot = self.scene_config.robot + assert robot is not None, "scene_config.robot must be populated before placement" + spawn = robot.spawn + assert spawn.usd_path is not None, "scene_config.robot must use a USD spawn for placement" + scale_x, scale_y, scale_z = spawn.scale or (1.0, 1.0, 1.0) + return PlacementGeometrySource( + usd_path=spawn.usd_path, + scale=(scale_x, scale_y, scale_z), + joint_pos=dict(robot.init_state.joint_pos or {}), + library_folder=self.robot_library_folder, + ) def get_bounding_box(self, prim_path: str | None = None) -> AxisAlignedBoundingBox: - """Return root-relative bounds computed from the articulation's USD geometry. + """Return root-relative bounds of the articulation posed at its configured joint positions. + + Shared and cached across callers, so treat the result as read-only. Args: prim_path: Optional sub-prim to bound (e.g. stand only). When None, bounds the full default prim. """ # Import locally because USD/pxr is available only after simulation initialization. - from isaaclab_arena.utils.usd_helpers import compute_local_bounding_box_from_usd + from isaaclab_arena.utils.usd_helpers import compute_local_bounding_box_from_usd_at_joint_pos - assert self.scene_config is not None, "scene_config must be populated before placement" - robot = self.scene_config.robot - assert robot is not None, "scene_config.robot must be populated before placement" - spawn = robot.spawn - assert spawn.usd_path is not None, "scene_config.robot must use a USD spawn for placement" - scale = tuple(spawn.scale or (1.0, 1.0, 1.0)) - # TODO(zihaox): Account for configured initial joint positions in bounds and collision meshes. - return compute_local_bounding_box_from_usd(spawn.usd_path, scale, prim_path=prim_path) + source = self.get_placement_geometry_source() + return compute_local_bounding_box_from_usd_at_joint_pos( + source.usd_path, source.joint_pos, source.scale, prim_path=prim_path + ) def get_collision_mesh(self) -> trimesh.Trimesh | None: - """Return the robot's collision mesh from its USD default prim, in the default joint pose.""" - if self._collision_mesh is None: - # Import locally because USD/pxr is available only after simulation initialization. - from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_path - - assert self.scene_config is not None, "scene_config must be populated before placement" - robot = self.scene_config.robot - assert robot is not None, "scene_config.robot must be populated before placement" - spawn = robot.spawn - assert spawn.usd_path is not None, "scene_config.robot must use a USD spawn for placement" - scale = tuple(spawn.scale or (1.0, 1.0, 1.0)) - self._collision_mesh = extract_trimesh_from_usd_path(spawn.usd_path, scale) - return self._collision_mesh + """Return the robot's collision mesh, posed at its configured initial joint positions. + + The mesh matches the robot as spawned even when its USD was authored in another joint + configuration. Shared and cached across callers, so treat the result as read-only. + """ + # Import locally because USD/pxr is available only after simulation initialization. + from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos + + source = self.get_placement_geometry_source() + return extract_trimesh_from_usd_at_joint_pos( + source.usd_path, source.joint_pos, source.scale, library_folder=source.library_folder + ) def _set_initial_pose(self, pose: Pose | PoseRange | PosePerEnv) -> None: """Store the configured pose; the construction pose is applied in ``get_scene_cfg``.""" diff --git a/isaaclab_arena/embodiments/franka/franka.py b/isaaclab_arena/embodiments/franka/franka.py index 8ec1b1d43f..dccf7a1f13 100644 --- a/isaaclab_arena/embodiments/franka/franka.py +++ b/isaaclab_arena/embodiments/franka/franka.py @@ -68,6 +68,7 @@ class FrankaEmbodimentBase(EmbodimentBase): """ default_arm_mode = ArmMode.SINGLE_ARM + robot_library_folder = "franka" def __init__( self, diff --git a/isaaclab_arena/embodiments/g1/g1.py b/isaaclab_arena/embodiments/g1/g1.py index ea91ba68a9..f2b5ff0bbd 100644 --- a/isaaclab_arena/embodiments/g1/g1.py +++ b/isaaclab_arena/embodiments/g1/g1.py @@ -47,6 +47,7 @@ class G1EmbodimentBase(EmbodimentBase): name = "g1" default_arm_mode = ArmMode.DUAL_ARM + robot_library_folder = "g1" def __init__( self, diff --git a/isaaclab_arena/embodiments/galbot/galbot.py b/isaaclab_arena/embodiments/galbot/galbot.py index 300d18e39a..ec27dcc9a5 100644 --- a/isaaclab_arena/embodiments/galbot/galbot.py +++ b/isaaclab_arena/embodiments/galbot/galbot.py @@ -38,6 +38,7 @@ class GalbotEmbodiment(EmbodimentBase): name = "galbot" default_arm_mode = ArmMode.LEFT + robot_library_folder = "galbot" def __init__(self, enable_cameras: bool = False, initial_pose: Pose | None = None, arm_mode: ArmMode | None = None): super().__init__(enable_cameras, initial_pose, arm_mode=arm_mode) diff --git a/isaaclab_arena/embodiments/gr1t2/gr1t2.py b/isaaclab_arena/embodiments/gr1t2/gr1t2.py index 152d0d5a90..ebf968d8d4 100644 --- a/isaaclab_arena/embodiments/gr1t2/gr1t2.py +++ b/isaaclab_arena/embodiments/gr1t2/gr1t2.py @@ -87,6 +87,7 @@ class GR1T2EmbodimentBase(EmbodimentBase): name = "gr1" default_arm_mode = ArmMode.RIGHT + robot_library_folder = "gr1t2" def __init__( self, diff --git a/isaaclab_arena/embodiments/kuka_allegro/kuka_allegro.py b/isaaclab_arena/embodiments/kuka_allegro/kuka_allegro.py index 8216c552ea..68dbbbd336 100644 --- a/isaaclab_arena/embodiments/kuka_allegro/kuka_allegro.py +++ b/isaaclab_arena/embodiments/kuka_allegro/kuka_allegro.py @@ -91,6 +91,7 @@ class KukaAllegroEmbodiment(EmbodimentBase): name = "kuka_allegro" default_arm_mode = ArmMode.SINGLE_ARM + robot_library_folder = "kuka" def __init__( self, diff --git a/isaaclab_arena/embodiments/robot_on_stand_utils.py b/isaaclab_arena/embodiments/robot_on_stand_utils.py index f9dd1b5b39..da577e5e73 100644 --- a/isaaclab_arena/embodiments/robot_on_stand_utils.py +++ b/isaaclab_arena/embodiments/robot_on_stand_utils.py @@ -16,7 +16,7 @@ from isaaclab.utils.assets import retrieve_file_path from pxr import Gf, Usd, UsdGeom -from isaaclab_arena.assets.asset_cache import get_arena_asset_cache_dir +from isaaclab_arena.assets.asset_cache import get_arena_usd_cache_dir _ROBOT_ON_STAND_USD_CACHE_DIR = "robot_on_stand" @@ -80,7 +80,7 @@ def compose_on_stand_usd( """ assert stand_height_m > 0.0, f"stand_height_m must be positive, got {stand_height_m}" - cache_root = get_arena_asset_cache_dir().parent / "usd" / _ROBOT_ON_STAND_USD_CACHE_DIR + cache_root = get_arena_usd_cache_dir() / _ROBOT_ON_STAND_USD_CACHE_DIR cache_root.mkdir(parents=True, exist_ok=True) out_path = cache_root / f"{output_basename}_{stand_height_m:.3f}.usd" diff --git a/isaaclab_arena/scripts/export_ready_pose_collision_meshes.py b/isaaclab_arena/scripts/export_ready_pose_collision_meshes.py new file mode 100644 index 0000000000..a0051efad8 --- /dev/null +++ b/isaaclab_arena/scripts/export_ready_pose_collision_meshes.py @@ -0,0 +1,171 @@ +# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Export every registered robot's ready-pose collision mesh, for upload to Nucleus. + +Extracting a robot's posed mesh costs 0.1-2.5 s per process. Publishing the result once means every +scene that spawns the robot reads it instead. The output holds a folder per robot and is uploaded +verbatim into ``collision_mesh_store.ARENA_ROBOT_LIBRARY_DIR``, merging with each robot's existing +folder: the relative paths are the ones the loader looks for. + +Re-run this whenever a robot's USD or configured joint positions change, since the loader validates +the pose an artifact was extracted at and falls back to extraction when it no longer matches. + +Run inside the container: + + /isaac-sim/python.sh isaaclab_arena/scripts/export_ready_pose_collision_meshes.py --headless \\ + --out_dir /tmp/ready_pose_export + +Check an export before uploading it by pointing the loader at it: + + export ISAACLAB_ARENA_ROBOT_LIBRARY_DIR=/tmp/ready_pose_export +""" + +from __future__ import annotations + +import argparse +import traceback +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser +from isaaclab_arena.utils.isaaclab_utils.simulation_app import SimulationAppContext + +if TYPE_CHECKING: + # Typing only: embodiment_base pulls in Isaac Lab, which is unavailable before sim init. + from isaaclab_arena.embodiments.embodiment_base import PlacementGeometrySource + + +def add_export_arguments(parser: argparse.ArgumentParser) -> None: + """Add the export flags.""" + group = parser.add_argument_group("Ready-Pose Mesh Export Arguments") + group.add_argument( + "--out_dir", + type=Path, + default=Path("/tmp/ready_pose_export"), + help="Directory to write per-robot folders into, merged as-is into the Arena robot library.", + ) + group.add_argument( + "--robots", + nargs="+", + default=None, + help="Registered embodiment names to export. Defaults to every registered embodiment.", + ) + + +def _embodiment_classes(names: list[str] | None) -> list[type]: + """Return the registered embodiment classes to export, in a stable order.""" + from isaaclab_arena.assets.registries import AssetRegistry + from isaaclab_arena.embodiments.embodiment_base import EmbodimentBase + + registry = AssetRegistry() + candidates = names if names is not None else registry.get_all_keys() + classes = [] + for name in sorted(candidates): + asset = registry.get_asset_by_name(name) + if isinstance(asset, type) and issubclass(asset, EmbodimentBase): + classes.append(asset) + assert classes, f"no registered embodiments among {candidates}" + return classes + + +@dataclass(frozen=True) +class PlannedArtifact: + """One artifact the export will write, and the embodiment it was planned from.""" + + embodiment_name: str + """Embodiment the mesh is extracted from, for reporting which robot an artifact came from.""" + + source: PlacementGeometrySource + """USD, scale, joint positions and library folder the artifact is written from.""" + + +def plan_artifacts(sources: Mapping[str, PlacementGeometrySource]) -> dict[str, PlannedArtifact]: + """Map each artifact's published relative path to the embodiment and source it is written from. + + Robots differing only in action space share a USD and a folder, so they collapse to one artifact. + Robots declaring no library folder have nowhere to publish and are left out. + + Args: + sources: Placement geometry source per embodiment name. + """ + from isaaclab_arena.utils.collision_mesh_store import ready_pose_artifact_name + + plan: dict[str, PlannedArtifact] = {} + for name, source in sorted(sources.items()): + if source.library_folder is None: + continue + relative_path = f"{source.library_folder}/{ready_pose_artifact_name(source.usd_path)}" + # Artifacts are validated by USD stem, so two robots sharing one within a folder would be + # served each other's mesh. Refuse to publish that rather than let placement use the wrong shape. + planned = plan.setdefault(relative_path, PlannedArtifact(name, source)) + assert planned.source.usd_path == source.usd_path, ( + f"{relative_path} would be written for both {planned.source.usd_path} and" + f" {source.usd_path}; rename one of the source USDs" + ) + return plan + + +def export_ready_pose_meshes(out_dir: Path, names: list[str] | None) -> tuple[list[str], list[str]]: + """Write each robot's ready-pose mesh into out_dir, returning the artifacts and the failures. + + Args: + out_dir: Staging directory, whose per-robot folders are uploaded as-is to the robot library. + names: Registered embodiment names to export, or None for every registered embodiment. + """ + from isaaclab_arena.utils.collision_mesh_store import export_ready_pose_mesh + from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos + + sources = {} + failed = [] + for embodiment_class in _embodiment_classes(names): + try: + sources[embodiment_class.name] = embodiment_class().get_placement_geometry_source() + except Exception as error: + failed.append(f"{embodiment_class.__name__}: {error}") + traceback.print_exc() + + skipped = sorted(name for name, source in sources.items() if source.library_folder is None) + written = [] + for relative_path, planned in plan_artifacts(sources).items(): + source = planned.source + try: + # Unit scale: one artifact serves every spawn scale, as the loader rescales on read. + mesh = extract_trimesh_from_usd_at_joint_pos(source.usd_path, source.joint_pos, (1.0, 1.0, 1.0)) + export_ready_pose_mesh(source.usd_path, source.joint_pos, mesh, out_dir, source.library_folder) + written.append(f"{relative_path} ({len(mesh.vertices)} vertices, from {planned.embodiment_name})") + except Exception as error: + failed.append(f"{planned.embodiment_name}: {error}") + traceback.print_exc() + if skipped: + print(f"\nSkipped {len(skipped)} embodiment(s) with no robot_library_folder: {', '.join(skipped)}") + return sorted(written), failed + + +def main() -> None: + args_parser = get_isaaclab_arena_cli_parser() + add_export_arguments(args_parser) + args_cli, _ = args_parser.parse_known_args() + + with SimulationAppContext(args_cli): + from isaaclab_arena.utils.collision_mesh_store import ARENA_ROBOT_LIBRARY_DIR + + written, failed = export_ready_pose_meshes(args_cli.out_dir, args_cli.robots) + + print(f"\nWrote {len(written)} artifact(s) to {args_cli.out_dir}, upload them to {ARENA_ROBOT_LIBRARY_DIR}:") + for line in written: + print(f" {line}") + if failed: + print(f"\n{len(failed)} embodiment(s) failed to export:") + for line in failed: + print(f" {line}") + # Exit non-zero so a partial export is not mistaken for a complete one and uploaded. + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/isaaclab_arena/tests/test_collision_mesh_store.py b/isaaclab_arena/tests/test_collision_mesh_store.py new file mode 100644 index 0000000000..57372c4f48 --- /dev/null +++ b/isaaclab_arena/tests/test_collision_mesh_store.py @@ -0,0 +1,386 @@ +# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for storing and reloading an articulation's collision mesh, keyed by joint pose.""" + +import contextlib +import numpy as np +import os +import shutil +import tempfile +import time +import trimesh +from pathlib import Path + +from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function + +HEADLESS = True + + +def _build_arm_usd(tmp_dir: str, joint_type: str = "revolute") -> str: + """Export the two-link fixture arm from the articulation tests to a USD file.""" + from isaaclab_arena.tests.test_usd_articulation import _build_two_link_arm + + usd_path = f"{tmp_dir}/arm.usda" + _build_two_link_arm(joint_type=joint_type).Export(usd_path) + return usd_path + + +@contextlib.contextmanager +def _published_dir(path: str): + """Read published artifacts from path, keeping tests off Nucleus and away from real artifacts.""" + from isaaclab_arena.utils.collision_mesh_store import ROBOT_LIBRARY_DIR_ENV_VAR + + previous = os.environ.get(ROBOT_LIBRARY_DIR_ENV_VAR) + os.environ[ROBOT_LIBRARY_DIR_ENV_VAR] = path + try: + yield + finally: + if previous is None: + del os.environ[ROBOT_LIBRARY_DIR_ENV_VAR] + else: + os.environ[ROBOT_LIBRARY_DIR_ENV_VAR] = previous + + +def _test_trimming_evicts_least_recently_used_artifacts(simulation_app) -> bool: + """Trimming drops artifacts oldest-first down to the budget, and leaves in-flight writes alone.""" + from isaaclab_arena.utils import collision_mesh_store as store + + with tempfile.TemporaryDirectory() as cache_dir: + # Three 1 KiB artifacts, aged a day apart, plus a staging file another process is writing. + paths = {} + for age_days, name in enumerate(["newest.usd", "middle.usd", "oldest.usd"]): + path = Path(cache_dir) / "robot" / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"x" * 1024) + os.utime(path, (0, time.time() - age_days * 86400)) + paths[name] = path + staging = Path(cache_dir) / "robot" / f"{store._STAGING_PREFIX}half_written.usd" + staging.write_bytes(b"x" * 4096) + os.utime(staging, (0, time.time() - 7 * 86400)) + + previous_budget = store.CACHE_BUDGET_BYTES + store.CACHE_BUDGET_BYTES = 2048 + try: + store._trim_cache(Path(cache_dir)) + finally: + store.CACHE_BUDGET_BYTES = previous_budget + + assert paths["newest.usd"].is_file(), "the most recently used artifact must survive" + assert paths["middle.usd"].is_file(), "the budget must be filled before evicting" + assert not paths["oldest.usd"].is_file(), "the least recently used artifact must be evicted" + # Unlinking a staging file would fail the rename of whichever process is writing it. + assert staging.is_file(), "a half-written artifact must not be evicted despite being oldest" + return True + + +def _test_pose_keys_identify_the_pose(simulation_app) -> bool: + """Every all-zero spelling shares the readable zero key; distinct poses key apart.""" + from isaaclab_arena.utils.collision_mesh_store import is_zero_pose, pose_key + + assert is_zero_pose({}) and is_zero_pose({"elbow": 0.0, "wrist": -0.0}) + assert not is_zero_pose({"elbow": 0.5}) + + assert pose_key({}) == "zero" + assert pose_key({"elbow": 0.0, "wrist": -0.0}) == "zero" + assert pose_key({"elbow": 0.5}) != "zero" + # Order must not matter, but the values must. + assert pose_key({"a": 0.5, "b": 0.25}) == pose_key({"b": 0.25, "a": 0.5}) + assert pose_key({"elbow": 0.5}) != pose_key({"elbow": 0.6}) + assert pose_key({"elbow": 0.5}) != pose_key({"wrist": 0.5}) + return True + + +def _test_distinct_poses_get_distinct_artifacts(simulation_app) -> bool: + """Posing at one configuration must never serve the mesh for another.""" + from isaaclab_arena.utils import usd_helpers + from isaaclab_arena.utils.collision_mesh_store import mesh_cache_path + + extended = {"elbow": 0.5} + with tempfile.TemporaryDirectory() as tmp_dir, contextlib.ExitStack() as published: + usd_path = _build_arm_usd(tmp_dir, joint_type="prismatic") + published.enter_context(_published_dir(tmp_dir)) + for joint_pos in ({}, extended): + mesh_cache_path(usd_path, joint_pos).unlink(missing_ok=True) + + usd_helpers._extract_trimesh_from_usd_at_joint_pos.cache_clear() + zero_mesh = usd_helpers.extract_trimesh_from_usd_at_joint_pos(usd_path, {}) + usd_helpers._extract_trimesh_from_usd_at_joint_pos.cache_clear() + posed_mesh = usd_helpers.extract_trimesh_from_usd_at_joint_pos(usd_path, extended) + + assert mesh_cache_path(usd_path, {}).is_file() and mesh_cache_path(usd_path, extended).is_file() + assert mesh_cache_path(usd_path, {}) != mesh_cache_path(usd_path, extended) + assert posed_mesh.extents[0] > zero_mesh.extents[0] + 0.4, f"prismatic pose should extend: {posed_mesh.extents}" + + # Reloading from the store must preserve that difference. + usd_helpers._extract_trimesh_from_usd_at_joint_pos.cache_clear() + reloaded_zero = usd_helpers.extract_trimesh_from_usd_at_joint_pos(usd_path, {}) + np.testing.assert_allclose(reloaded_zero.extents, zero_mesh.extents, atol=1e-6) + for joint_pos in ({}, extended): + mesh_cache_path(usd_path, joint_pos).unlink(missing_ok=True) + return True + + +def _test_stored_mesh_keeps_its_vertices_verbatim(simulation_app) -> bool: + """Reloading must not merge coincident vertices, which trimesh does by default on construction.""" + from isaaclab_arena.utils.collision_mesh_store import load_mesh, mesh_cache_path, save_mesh + + # Two coincident triangles: merging would collapse six vertices into three. + vertices = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]] * 2) + mesh = trimesh.Trimesh(vertices=vertices, faces=np.array([[0, 1, 2], [3, 4, 5]]), process=False) + + with tempfile.TemporaryDirectory() as tmp_dir, contextlib.ExitStack() as published: + source = f"{tmp_dir}/synthetic.usda" + published.enter_context(_published_dir(tmp_dir)) + save_mesh(source, {}, mesh) + loaded = load_mesh(source, {}, (1.0, 1.0, 1.0)) + assert loaded is not None + assert len(loaded.vertices) == len(vertices), f"round trip changed vertex count: {len(loaded.vertices)}" + + scaled = load_mesh(source, {}, (2.0, 2.0, 2.0)) + assert scaled is not None + assert len(scaled.vertices) == len(vertices), f"scaling changed vertex count: {len(scaled.vertices)}" + mesh_cache_path(source, {}).unlink(missing_ok=True) + return True + + +def _test_stored_mesh_is_scale_independent(simulation_app) -> bool: + """One unit-scale artifact serves every spawn scale.""" + from isaaclab_arena.utils import usd_helpers + from isaaclab_arena.utils.collision_mesh_store import mesh_cache_path + + with tempfile.TemporaryDirectory() as tmp_dir, contextlib.ExitStack() as published: + usd_path = _build_arm_usd(tmp_dir) + published.enter_context(_published_dir(tmp_dir)) + mesh_cache_path(usd_path, {}).unlink(missing_ok=True) + + usd_helpers._extract_trimesh_from_usd_at_joint_pos.cache_clear() + unit = usd_helpers.extract_trimesh_from_usd_at_joint_pos(usd_path, {}) + # Served from the store now, and must still honour the requested scale. + usd_helpers._extract_trimesh_from_usd_at_joint_pos.cache_clear() + doubled = usd_helpers.extract_trimesh_from_usd_at_joint_pos(usd_path, {}, scale=(2.0, 1.0, 1.0)) + np.testing.assert_allclose(doubled.extents[0], unit.extents[0] * 2.0, atol=1e-6) + np.testing.assert_allclose(doubled.extents[1:], unit.extents[1:], atol=1e-6) + mesh_cache_path(usd_path, {}).unlink(missing_ok=True) + return True + + +def _test_artifact_from_another_asset_or_pose_is_ignored(simulation_app) -> bool: + """A mesh recording a different source or pose is refused, so it cannot misplace a robot.""" + from isaaclab_arena.utils.collision_mesh_store import load_mesh, mesh_cache_path, save_mesh + from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos + + posed = {"elbow": 0.5} + with tempfile.TemporaryDirectory() as tmp_dir, contextlib.ExitStack() as published: + usd_path = _build_arm_usd(tmp_dir) + published.enter_context(_published_dir(tmp_dir)) + mesh = extract_trimesh_from_usd_at_joint_pos(usd_path, posed) + + # Recorded under a foreign source USD, then moved into this asset's slot. + save_mesh(f"{tmp_dir}/other_arm.usda", posed, mesh).replace(mesh_cache_path(usd_path, posed)) + assert load_mesh(usd_path, posed, (1.0, 1.0, 1.0)) is None, "foreign source must be ignored" + + # Recorded for this asset but at another pose, then moved into this pose's slot. + save_mesh(usd_path, {"elbow": 0.9}, mesh).replace(mesh_cache_path(usd_path, posed)) + assert load_mesh(usd_path, posed, (1.0, 1.0, 1.0)) is None, "foreign pose must be ignored" + mesh_cache_path(usd_path, posed).unlink(missing_ok=True) + return True + + +def _test_published_artifact_is_found_by_its_name(simulation_app) -> bool: + """An exported artifact is loaded from the published directory when the local cache is empty.""" + from isaaclab_arena.utils.collision_mesh_store import ( + export_ready_pose_mesh, + load_mesh, + mesh_cache_path, + published_ready_pose_path, + ready_pose_artifact_name, + ) + from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos + + posed = {"elbow": 0.5} + with tempfile.TemporaryDirectory() as tmp_dir, tempfile.TemporaryDirectory() as publish_dir: + usd_path = _build_arm_usd(tmp_dir) + mesh = extract_trimesh_from_usd_at_joint_pos(usd_path, posed) + artifact = export_ready_pose_mesh(usd_path, posed, mesh, Path(publish_dir), "test_arm") + + # The artifact lands in the robot's own library folder, under a name derived from the asset. + assert artifact.name == "arm_ready_pose.usd" == ready_pose_artifact_name(usd_path) + assert artifact.parent.name == "test_arm", artifact + + with _published_dir(publish_dir): + assert published_ready_pose_path(usd_path, "test_arm") == str(artifact) + # Empty the local cache so only the published copy can answer. + mesh_cache_path(usd_path, posed).unlink(missing_ok=True) + loaded = load_mesh(usd_path, posed, (1.0, 1.0, 1.0), "test_arm") + assert loaded is not None, f"published artifact {artifact} was not found" + np.testing.assert_allclose(loaded.vertices, mesh.vertices, atol=1e-6) + + # A robot that names no folder must not be served another robot's artifact. + mesh_cache_path(usd_path, posed).unlink(missing_ok=True) + assert load_mesh(usd_path, posed, (1.0, 1.0, 1.0)) is None, "published lookup needs a folder" + + # Published artifacts outlive config changes, so a repose must not be served the old shape. + assert load_mesh(usd_path, {"elbow": 0.9}, (1.0, 1.0, 1.0), "test_arm") is None, "stale pose" + mesh_cache_path(usd_path, posed).unlink(missing_ok=True) + return True + + +def _test_published_artifact_loads_for_a_relocated_source(simulation_app) -> bool: + """An artifact stays usable when the robot USD sits elsewhere, as on another machine. + + Arena composes the robot-on-stand USDs into a per-user cache directory, so validating against the + full source path would reject every artifact anyone else exported. + """ + from isaaclab_arena.utils.collision_mesh_store import export_ready_pose_mesh, load_mesh, mesh_cache_path + from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos + + posed = {"elbow": 0.5} + with ( + tempfile.TemporaryDirectory() as exporter_dir, + tempfile.TemporaryDirectory() as consumer_dir, + tempfile.TemporaryDirectory() as publish_dir, + ): + exporter_usd = _build_arm_usd(exporter_dir) + mesh = extract_trimesh_from_usd_at_joint_pos(exporter_usd, posed) + export_ready_pose_mesh(exporter_usd, posed, mesh, Path(publish_dir), "test_arm") + + # The same asset reached through another absolute path, as a second machine's cache would. + consumer_usd = f"{consumer_dir}/arm.usda" + shutil.copy(exporter_usd, consumer_usd) + + with _published_dir(publish_dir): + mesh_cache_path(consumer_usd, posed).unlink(missing_ok=True) + loaded = load_mesh(consumer_usd, posed, (1.0, 1.0, 1.0), "test_arm") + assert loaded is not None, "an artifact must load for the same asset at another path" + np.testing.assert_allclose(loaded.vertices, mesh.vertices, atol=1e-6) + for path in (exporter_usd, consumer_usd): + mesh_cache_path(path, posed).unlink(missing_ok=True) + return True + + +def _test_truncated_artifact_is_ignored(simulation_app) -> bool: + """A half-written or hand-edited artifact falls back to extraction rather than raising.""" + from pxr import Usd, UsdGeom + + from isaaclab_arena.utils.collision_mesh_store import load_mesh, save_mesh + from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos + + posed = {"elbow": 0.5} + with tempfile.TemporaryDirectory() as tmp_dir, contextlib.ExitStack() as published: + usd_path = _build_arm_usd(tmp_dir) + published.enter_context(_published_dir(tmp_dir)) + cache_path = save_mesh(usd_path, posed, extract_trimesh_from_usd_at_joint_pos(usd_path, posed)) + assert load_mesh(usd_path, posed, (1.0, 1.0, 1.0)) is not None + + # Drop the geometry but keep the provenance, as a truncated write would. + stage = Usd.Stage.Open(str(cache_path)) + mesh_prim = UsdGeom.Mesh(stage.GetPrimAtPath("/CollisionMesh")) + mesh_prim.GetPointsAttr().Clear() + mesh_prim.GetFaceVertexIndicesAttr().Clear() + stage.GetRootLayer().Save() + + assert load_mesh(usd_path, posed, (1.0, 1.0, 1.0)) is None, "an artifact without geometry must be ignored" + cache_path.unlink(missing_ok=True) + return True + + +def _test_colliding_published_names_are_refused(simulation_app) -> bool: + """Two robots sharing a stem within one folder cannot both be published, as one name serves both.""" + from isaaclab_arena.embodiments.embodiment_base import PlacementGeometrySource + from isaaclab_arena.scripts.export_ready_pose_collision_meshes import plan_artifacts + + def source(usd_path: str, library_folder: str | None) -> PlacementGeometrySource: + return PlacementGeometrySource(usd_path, (1.0, 1.0, 1.0), {}, library_folder) + + # Same stem in one folder, different assets: the published name cannot tell them apart. + colliding = { + "robot_a": source("/assets/vendor_a/robot.usd", "shared"), + "robot_b": source("/assets/vendor_b/robot.usd", "shared"), + } + try: + plan_artifacts(colliding) + except AssertionError as error: + assert "rename one of the source USDs" in str(error), error + else: + raise AssertionError("a colliding published name must be refused") + + # The same stem in separate folders is fine, which is what per-robot folders buy. + separate = { + "robot_a": source("/assets/vendor_a/robot.usd", "robot_a"), + "robot_b": source("/assets/vendor_b/robot.usd", "robot_b"), + } + assert sorted(plan_artifacts(separate)) == ["robot_a/robot_ready_pose.usd", "robot_b/robot_ready_pose.usd"] + + # Action-space variants of one robot share a USD, and so legitimately share one artifact. + shared_usd = source("/assets/vendor_a/robot.usd", "robot_a") + assert sorted(plan_artifacts({"robot_ik": shared_usd, "robot_joint_pos": shared_usd})) == [ + "robot_a/robot_ready_pose.usd" + ] + + # A robot with nowhere to publish is left out rather than written to a guessed folder. + assert plan_artifacts({"unpublished": source("/assets/vendor_a/robot.usd", None)}) == {} + return True + + +def _test_export_writes_one_named_artifact_per_robot(simulation_app) -> bool: + """The exporter deduplicates embodiments sharing a USD and reports nothing as failed.""" + from isaaclab_arena.scripts.export_ready_pose_collision_meshes import export_ready_pose_meshes + + # Droid's three action-space variants share one USD and one ready pose. + droid_variants = ["droid_abs_joint_pos", "droid_rel_joint_pos", "droid_differential_ik"] + with tempfile.TemporaryDirectory() as out_dir: + written, failed = export_ready_pose_meshes(Path(out_dir), droid_variants) + + assert not failed, f"export reported failures: {failed}" + assert len(written) == 1, f"variants of one robot must share an artifact: {written}" + # Relative paths, because the staging tree is uploaded into the robot library as it stands. + artifacts = sorted(str(path.relative_to(out_dir)) for path in Path(out_dir).rglob("*.usd")) + assert artifacts == ["droid/droid_franka_robotiq_on_stand_1.350_ready_pose.usd"], artifacts + return True + + +def test_trimming_evicts_least_recently_used_artifacts(): + assert run_simulation_app_function(_test_trimming_evicts_least_recently_used_artifacts, headless=HEADLESS) + + +def test_pose_keys_identify_the_pose(): + assert run_simulation_app_function(_test_pose_keys_identify_the_pose, headless=HEADLESS) + + +def test_truncated_artifact_is_ignored(): + assert run_simulation_app_function(_test_truncated_artifact_is_ignored, headless=HEADLESS) + + +def test_colliding_published_names_are_refused(): + assert run_simulation_app_function(_test_colliding_published_names_are_refused, headless=HEADLESS) + + +def test_export_writes_one_named_artifact_per_robot(): + assert run_simulation_app_function(_test_export_writes_one_named_artifact_per_robot, headless=HEADLESS) + + +def test_published_artifact_is_found_by_its_name(): + assert run_simulation_app_function(_test_published_artifact_is_found_by_its_name, headless=HEADLESS) + + +def test_published_artifact_loads_for_a_relocated_source(): + assert run_simulation_app_function(_test_published_artifact_loads_for_a_relocated_source, headless=HEADLESS) + + +def test_distinct_poses_get_distinct_artifacts(): + assert run_simulation_app_function(_test_distinct_poses_get_distinct_artifacts, headless=HEADLESS) + + +def test_stored_mesh_keeps_its_vertices_verbatim(): + assert run_simulation_app_function(_test_stored_mesh_keeps_its_vertices_verbatim, headless=HEADLESS) + + +def test_stored_mesh_is_scale_independent(): + assert run_simulation_app_function(_test_stored_mesh_is_scale_independent, headless=HEADLESS) + + +def test_artifact_from_another_asset_or_pose_is_ignored(): + assert run_simulation_app_function(_test_artifact_from_another_asset_or_pose_is_ignored, headless=HEADLESS) diff --git a/isaaclab_arena/tests/test_embodiment_collision_mesh.py b/isaaclab_arena/tests/test_embodiment_collision_mesh.py index c2ddb18b85..95dbe25e66 100644 --- a/isaaclab_arena/tests/test_embodiment_collision_mesh.py +++ b/isaaclab_arena/tests/test_embodiment_collision_mesh.py @@ -20,25 +20,39 @@ def _test_embodiment_provides_robot_collision_mesh(simulation_app) -> bool: assert mesh is not None, "embodiment must expose a collision mesh; None forces the loose bbox fallback" assert len(mesh.vertices) > 0 - # Mesh extraction scopes to UsdGeom.Mesh under the default prim (arm/gripper). The Droid USD also - # bakes in a stand and may reference non-mesh gprims; leaking a 50 m ground plane would blow this up. + # The spawn USD composes robot and stand, so the mesh spans the whole assembly: 1.46 x 0.91 x + # 2.10 m as measured, the tallest axis being the arm on its 1.35 m stand. A leaked 50 m ground + # plane would still blow this up by an order of magnitude. extents = mesh.extents - assert all(e < 2.0 for e in extents), f"mesh leaked non-robot geometry: extents {extents}" + assert all(e < 3.0 for e in extents), f"mesh leaked non-robot geometry: extents {extents}" - # Placement bbox comes from the full composed on-stand spawn USD (robot + stand). It should be at - # least as large as the arm mesh footprint. + # The placement bbox covers the same posed assembly plus the analytic gprims that mesh extraction + # cannot represent, so it is marginally larger but must not diverge. bbox = emb.get_bounding_box() bbox_size = (bbox.max_point - bbox.min_point)[0].tolist() for mesh_extent, box_extent in zip(extents, bbox_size): assert ( box_extent + 1e-3 >= mesh_extent ), f"placement bbox {bbox_size} should cover robot mesh extents {extents.tolist()}" - # TODO(qianl): Re-enable check for exact match when the stand with non-mesh collision geometry - # is correctly included in get_bounding_box()/extract_trimesh_from_prim() - # assert abs(mesh_extent - box_extent) < 0.2, f"mesh extents {extents} disagree withbox {bbox_size}" + assert abs(mesh_extent - box_extent) < 0.2, f"mesh extents {extents} disagree with box {bbox_size}" - # Extraction opens the USD, so the result is cached rather than recomputed per solve. + # Both derivations open the USD and pose it, so results are cached rather than recomputed per + # solve step. The cache is keyed by USD, joint positions and scale, so an identical embodiment + # shares it. assert emb.get_collision_mesh() is mesh + assert DroidAbsoluteJointPositionEmbodiment().get_collision_mesh() is mesh + assert emb.get_bounding_box() is bbox + assert DroidAbsoluteJointPositionEmbodiment().get_bounding_box() is bbox + + # Isaac Lab reaches placeable assets through EventTermCfg params and validates whatever they + # hold without tracking visited objects, so an embodiment holding its mesh would send + # validation recursing through trimesh's back-references until the stack overflows. + from isaaclab.managers import EventTermCfg + + def _noop(env, env_ids, embodiment): + pass + + EventTermCfg(func=_noop, mode="reset", params={"embodiment": emb}).validate() except Exception as e: print(f"Error: {e}") diff --git a/isaaclab_arena/tests/test_usd_articulation.py b/isaaclab_arena/tests/test_usd_articulation.py new file mode 100644 index 0000000000..7efd172425 --- /dev/null +++ b/isaaclab_arena/tests/test_usd_articulation.py @@ -0,0 +1,409 @@ +# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for posing a USD articulation from its physics joints. + +The fixture arm places the joint at the world origin and the child link one unit along +X, so a +90 degree rotation about Z must swing the child from (1, 0, 0) to (0, 1, 0). +""" + +import math +import numpy as np + +from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function + +HEADLESS = True + + +def _define_box(stage, path: str, half_extent: float = 0.1) -> None: + """Define a cube mesh centred on its own origin.""" + from pxr import Gf, UsdGeom + + mesh = UsdGeom.Mesh.Define(stage, path) + h = half_extent + mesh.GetPointsAttr().Set([ + Gf.Vec3f(-h, -h, -h), + Gf.Vec3f(h, -h, -h), + Gf.Vec3f(h, h, -h), + Gf.Vec3f(-h, h, -h), + Gf.Vec3f(-h, -h, h), + Gf.Vec3f(h, -h, h), + Gf.Vec3f(h, h, h), + Gf.Vec3f(-h, h, h), + ]) + mesh.GetFaceVertexCountsAttr().Set([4, 4, 4, 4, 4, 4]) + mesh.GetFaceVertexIndicesAttr().Set([ + 0, 1, 2, 3, + 4, 5, 6, 7, + 0, 1, 5, 4, + 2, 3, 7, 6, + 0, 3, 7, 4, + 1, 2, 6, 5, + ]) # fmt: skip + + +def _build_two_link_arm(joint_type: str = "revolute"): + """Build a base link at the origin and a child link at +X, joined at the world origin.""" + from pxr import Gf, Usd, UsdGeom, UsdPhysics + + stage = Usd.Stage.CreateInMemory() + root = UsdGeom.Xform.Define(stage, "/root") + stage.SetDefaultPrim(root.GetPrim()) + + base = UsdGeom.Xform.Define(stage, "/root/base") + UsdPhysics.RigidBodyAPI.Apply(base.GetPrim()) + _define_box(stage, "/root/base/box") + + forearm = UsdGeom.Xform.Define(stage, "/root/forearm") + forearm.AddTranslateOp().Set(Gf.Vec3d(1.0, 0.0, 0.0)) + UsdPhysics.RigidBodyAPI.Apply(forearm.GetPrim()) + _define_box(stage, "/root/forearm/box") + + schema = UsdPhysics.RevoluteJoint if joint_type == "revolute" else UsdPhysics.PrismaticJoint + joint = schema.Define(stage, "/root/elbow") + joint.CreateBody0Rel().SetTargets(["/root/base"]) + joint.CreateBody1Rel().SetTargets(["/root/forearm"]) + joint.CreateAxisAttr("Z" if joint_type == "revolute" else "X") + joint.CreateLocalPos0Attr(Gf.Vec3f(0.0, 0.0, 0.0)) + joint.CreateLocalPos1Attr(Gf.Vec3f(-1.0, 0.0, 0.0)) + return stage + + +def _posed_forearm_origin(stage, joint_pos) -> np.ndarray: + """Return where the forearm box's rest-pose origin lands once joint_pos is applied.""" + from isaaclab_arena.utils.usd_articulation import compute_posed_prim_world_deltas, resolve_prim_world_delta + + deltas = compute_posed_prim_world_deltas(stage, "/root", joint_pos) + delta = resolve_prim_world_delta("/root/forearm/box", deltas) + assert delta is not None, "forearm geometry must inherit its link's delta" + return (np.array([1.0, 0.0, 0.0, 1.0]) @ delta)[:3] + + +def _test_revolute_joint_swings_child_link(simulation_app) -> bool: + """A 90 degree revolute rotation about Z moves the child link from +X to +Y.""" + stage = _build_two_link_arm() + + posed = _posed_forearm_origin(stage, {"elbow": math.pi / 2.0}) + np.testing.assert_allclose(posed, [0.0, 1.0, 0.0], atol=1e-9) + + # Half the rotation must land on the 45 degree diagonal, ruling out a quarter-turn constant. + posed_45 = _posed_forearm_origin(stage, {"elbow": math.pi / 4.0}) + np.testing.assert_allclose(posed_45, [math.sqrt(0.5), math.sqrt(0.5), 0.0], atol=1e-9) + return True + + +def _test_zero_joint_position_preserves_authored_pose(simulation_app) -> bool: + """Posing at zero reproduces the authored transforms, so unposed callers see no change.""" + stage = _build_two_link_arm() + + np.testing.assert_allclose(_posed_forearm_origin(stage, {"elbow": 0.0}), [1.0, 0.0, 0.0], atol=1e-9) + # Omitted joints default to zero rather than to the authored pose. + np.testing.assert_allclose(_posed_forearm_origin(stage, {}), [1.0, 0.0, 0.0], atol=1e-9) + return True + + +def _test_prismatic_joint_translates_child_link(simulation_app) -> bool: + """A prismatic joint slides the child link along its axis.""" + stage = _build_two_link_arm(joint_type="prismatic") + + np.testing.assert_allclose(_posed_forearm_origin(stage, {"elbow": 0.25}), [1.25, 0.0, 0.0], atol=1e-9) + return True + + +def _test_authored_pose_away_from_joint_zero_is_corrected(simulation_app) -> bool: + """Geometry follows the joint values, not the pose the asset was authored in. + + The forearm is authored swung 90 degrees out, which no longer agrees with the joint frames. + Posing at zero must pull it back onto +X instead of trusting the authored transform. + """ + from pxr import Gf, UsdGeom, UsdPhysics + + stage = _build_two_link_arm() + forearm = UsdGeom.Xform.Define(stage, "/root/forearm") + forearm.GetPrim().RemoveProperty("xformOp:translate") + forearm.ClearXformOpOrder() + forearm.AddTranslateOp().Set(Gf.Vec3d(0.0, 1.0, 0.0)) + forearm.AddRotateZOp().Set(90.0) + UsdPhysics.RigidBodyAPI.Apply(forearm.GetPrim()) + + from isaaclab_arena.utils.usd_articulation import compute_posed_prim_world_deltas, resolve_prim_world_delta + + deltas = compute_posed_prim_world_deltas(stage, "/root", {"elbow": 0.0}) + delta = resolve_prim_world_delta("/root/forearm/box", deltas) + # The authored origin sits at (0, 1, 0); zeroing the joint returns it to (1, 0, 0). + posed = (np.array([0.0, 1.0, 0.0, 1.0]) @ delta)[:3] + np.testing.assert_allclose(posed, [1.0, 0.0, 0.0], atol=1e-9) + return True + + +def _test_unknown_joint_name_is_rejected(simulation_app) -> bool: + """A joint position naming a joint the articulation lacks is a configuration error.""" + from isaaclab_arena.utils.usd_articulation import compute_posed_prim_world_deltas + + stage = _build_two_link_arm() + try: + compute_posed_prim_world_deltas(stage, "/root", {"shoulder": 0.5}) + except AssertionError as error: + assert "shoulder" in str(error), f"assertion should name the unknown joint, got: {error}" + return True + raise AssertionError("expected an assertion for an unknown joint name") + + +def _test_joint_pos_patterns_expand_to_joint_names(simulation_app) -> bool: + """Isaac Lab regex joint keys expand to every joint they full-match, and misses are dropped.""" + from isaaclab_arena.utils.usd_articulation import resolve_joint_pos_patterns + + names = ["panda_joint1", "panda_joint2", "right_outer_knuckle_joint", "left_inner_finger_joint"] + resolved = resolve_joint_pos_patterns(names, {"panda_joint1": 0.5, "right_outer.*": 0.25}) + assert resolved == {"panda_joint1": 0.5, "right_outer_knuckle_joint": 0.25}, resolved + + # A pattern matching nothing is dropped, leaving those joints at zero rather than raising. + assert resolve_joint_pos_patterns(names, {"head_.*": 1.0}) == {} + + # A partial match must not count: Isaac Lab full-matches joint names. + assert resolve_joint_pos_patterns(names, {"panda": 1.0}) == {} + + # Later keys win where patterns overlap, matching Isaac Lab's ordering. + overlapped = resolve_joint_pos_patterns(names, {"panda_.*": 1.0, "panda_joint2": 2.0}) + assert overlapped == {"panda_joint1": 1.0, "panda_joint2": 2.0}, overlapped + return True + + +def _test_droid_geometry_tracks_configured_joint_positions(simulation_app) -> bool: + """Posing the real Droid at its configured joints reproduces its authored geometry, zero does not.""" + from pxr import Usd + + from isaaclab_arena.embodiments.droid.droid import DroidAbsoluteJointPositionEmbodiment + from isaaclab_arena.utils.usd_articulation import articulation_joint_prims + from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos, extract_trimesh_from_usd_path + + robot = DroidAbsoluteJointPositionEmbodiment().scene_config.robot + usd_path = robot.spawn.usd_path + + stage = Usd.Stage.Open(usd_path) + joint_names = set(articulation_joint_prims(stage.GetDefaultPrim())) + assert "panda_joint1" in joint_names, f"expected Franka arm joints, found {sorted(joint_names)}" + + # Passed verbatim, so the config's regex keys (e.g. "right_outer.*") must resolve to real joints. + configured = robot.init_state.joint_pos + authored = extract_trimesh_from_usd_path(usd_path) + posed = extract_trimesh_from_usd_at_joint_pos(usd_path, configured) + np.testing.assert_allclose(posed.extents, authored.extents, atol=2e-3) + + # At zero the arm stands straight up, so it is taller and narrower than the authored pose. + zero = extract_trimesh_from_usd_at_joint_pos(usd_path, {}) + assert zero.extents[2] > authored.extents[2] + 0.2, f"zero pose should stand taller: {zero.extents}" + assert zero.extents[0] < authored.extents[0] - 0.2, f"zero pose should be narrower: {zero.extents}" + return True + + +def _test_droid_posed_bounding_box_covers_all_geometry(simulation_app) -> bool: + """The posed bbox covers every gprim, not just the meshes. + + Deriving bounds from the posed mesh instead drops analytic gprims such as ``gripper_adapter``. + """ + from isaaclab_arena.embodiments.droid.droid import DroidAbsoluteJointPositionEmbodiment + from isaaclab_arena.utils.usd_helpers import ( + compute_local_bounding_box_from_usd, + compute_local_bounding_box_from_usd_at_joint_pos, + extract_trimesh_from_usd_at_joint_pos, + ) + + robot = DroidAbsoluteJointPositionEmbodiment().scene_config.robot + usd_path = robot.spawn.usd_path + configured = robot.init_state.joint_pos + + authored = compute_local_bounding_box_from_usd(usd_path) + posed = compute_local_bounding_box_from_usd_at_joint_pos(usd_path, configured) + np.testing.assert_allclose(posed.size.numpy()[0], authored.size.numpy()[0], atol=2e-3) + np.testing.assert_allclose(posed.min_point.numpy()[0], authored.min_point.numpy()[0], atol=2e-3) + + # The bbox must be at least as large as the mesh, since it also covers the non-mesh gprims. + mesh_extents = extract_trimesh_from_usd_at_joint_pos(usd_path, configured).extents + posed_size = posed.size.numpy()[0] + assert np.all(posed_size >= mesh_extents - 1e-6), f"bbox {posed_size} smaller than mesh {mesh_extents}" + # Droid's dropped gripper_adapter makes this a strict difference, so a mesh-only bbox would fail. + assert np.any(posed_size > mesh_extents + 1e-3), "expected non-mesh geometry to widen the bbox" + return True + + +def _test_offline_posing_matches_physx_link_poses(simulation_app) -> bool: + """Offline posing of the real Droid agrees with PhysX's own kinematics for every link. + + This is the ground-truth check: the simulator resolves the articulation independently, so + matching its link poses rules out frame, axis, and joint-ordering errors that a synthetic + two-link arm cannot expose. The joints PhysX settled at are used as the input, because a reset + lets the position-controlled arm sag a few milliradians off the configured values. + """ + import warp as wp + from pxr import Usd, UsdGeom, UsdPhysics + + from isaaclab_arena.cli.isaaclab_arena_cli import arena_env_builder_cfg_from_argparse, get_isaaclab_arena_cli_parser + from isaaclab_arena.embodiments.droid.droid import DroidAbsoluteJointPositionEmbodiment + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment + from isaaclab_arena.scene.scene import Scene + from isaaclab_arena.utils.usd_articulation import ( + articulation_joint_prims, + compute_posed_prim_world_deltas, + resolve_joint_pos_patterns, + ) + + embodiment = DroidAbsoluteJointPositionEmbodiment() + usd_path = embodiment.scene_config.robot.spawn.usd_path + + arena_env = IsaacLabArenaEnvironment(name="verify_articulation", embodiment=embodiment, scene=Scene(assets=[])) + args_cli = get_isaaclab_arena_cli_parser().parse_args([]) + env = ArenaEnvBuilder(arena_env, arena_env_builder_cfg_from_argparse(args_cli)).make_registered() + try: + env.reset() + robot = env.unwrapped.scene["robot"] + body_names = list(robot.body_names) + sim_pose = wp.to_torch(robot.data.body_link_pose_w)[0].double().cpu().numpy() + sim_joint_pos = wp.to_torch(robot.data.joint_pos)[0].double().cpu().numpy() + actual_joint_pos = {name: float(value) for name, value in zip(robot.joint_names, sim_joint_pos)} + finally: + env.close() + + stage = Usd.Stage.Open(usd_path) + default_prim = stage.GetDefaultPrim() + resolved = resolve_joint_pos_patterns(articulation_joint_prims(default_prim), actual_joint_pos) + deltas = compute_posed_prim_world_deltas(stage, default_prim.GetPath().pathString, resolved) + + # Gripper links nest under panda_link8, so resolve body prims by traversal rather than by layout. + body_prim_paths = { + prim.GetName(): prim.GetPath().pathString + for prim in Usd.PrimRange(default_prim) + if prim.HasAPI(UsdPhysics.RigidBodyAPI) + } + + def offline_world_position(prim_path: str) -> np.ndarray: + rest = np.array( + UsdGeom.Xformable(stage.GetPrimAtPath(prim_path)).ComputeLocalToWorldTransform(Usd.TimeCode.Default()), + dtype=np.float64, + ) + posed = rest @ deltas[prim_path] if prim_path in deltas else rest + return posed[3, :3] + + # Compare relative to the root link, which cancels the env origin and the robot's base placement. + offline_root = offline_world_position(body_prim_paths[body_names[0]]) + for index, name in enumerate(body_names): + assert name in body_prim_paths, f"no rigid-body prim named {name} under {default_prim.GetPath()}" + sim_rel = sim_pose[index, :3] - sim_pose[0, :3] + offline_rel = offline_world_position(body_prim_paths[name]) - offline_root + np.testing.assert_allclose(offline_rel, sim_rel, atol=1e-5, err_msg=f"link {name} disagrees with PhysX") + return True + + +def _test_closed_loop_articulation_poses_a_spanning_tree(simulation_app) -> bool: + """A redundant second joint to a body is dropped, leaving the spanning-tree pose untouched.""" + from pxr import Gf, UsdGeom, UsdPhysics + + stage = _build_two_link_arm() + strut = UsdGeom.Xform.Define(stage, "/root/strut") + strut.AddTranslateOp().Set(Gf.Vec3d(0.5, 0.0, 0.0)) + UsdPhysics.RigidBodyAPI.Apply(strut.GetPrim()) + _define_box(stage, "/root/strut/box") + + shoulder = UsdPhysics.RevoluteJoint.Define(stage, "/root/shoulder") + shoulder.CreateBody0Rel().SetTargets(["/root/base"]) + shoulder.CreateBody1Rel().SetTargets(["/root/strut"]) + shoulder.CreateAxisAttr("Z") + + # Closes the loop: the forearm is now reachable both directly and via the strut. + brace = UsdPhysics.RevoluteJoint.Define(stage, "/root/brace") + brace.CreateBody0Rel().SetTargets(["/root/strut"]) + brace.CreateBody1Rel().SetTargets(["/root/forearm"]) + brace.CreateAxisAttr("Z") + + # The elbow is declared first, so it keeps the forearm and the brace is the closure that drops. + np.testing.assert_allclose(_posed_forearm_origin(stage, {"elbow": math.pi / 2.0}), [0.0, 1.0, 0.0], atol=1e-9) + return True + + +def _test_instanced_geometry_is_posed_with_its_link(simulation_app) -> bool: + """Instanced geometry contributes to the posed bounds and mesh, and follows the link it hangs off.""" + import tempfile + + from pxr import Gf, UsdGeom + + from isaaclab_arena.utils.collision_mesh_store import mesh_cache_path + from isaaclab_arena.utils.usd_helpers import ( + compute_local_bounding_box_from_usd_at_joint_pos, + extract_trimesh_from_usd_at_joint_pos, + ) + + stage = _build_two_link_arm() + # The prototype lives outside the default prim, so only the instance can contribute to the bounds. + UsdGeom.Xform.Define(stage, "/prototypes/mount") + _define_box(stage, "/prototypes/mount/box", half_extent=0.2) + mount = UsdGeom.Xform.Define(stage, "/root/forearm/mount") + mount.AddTranslateOp().Set(Gf.Vec3d(0.5, 0.0, 0.0)) + mount.GetPrim().GetReferences().AddInternalReference("/prototypes/mount") + mount.GetPrim().SetInstanceable(True) + assert mount.GetPrim().IsInstance(), "fixture must exercise a real instance" + + with tempfile.TemporaryDirectory() as tmp_dir: + usd_path = f"{tmp_dir}/arm_with_instanced_mount.usda" + stage.Export(usd_path) + + # The mount reaches 1.7 along +X: forearm at 1.0, mount offset 0.5, half extent 0.2. + rest = compute_local_bounding_box_from_usd_at_joint_pos(usd_path, {}) + np.testing.assert_allclose(rest.max_point.numpy()[0][0], 1.7, atol=1e-6) + np.testing.assert_allclose(extract_trimesh_from_usd_at_joint_pos(usd_path, {}).vertices.max(axis=0)[0], 1.7) + + # Swinging the elbow must carry the instance with the forearm rather than leave it behind. + swung = compute_local_bounding_box_from_usd_at_joint_pos(usd_path, {"elbow": math.pi / 2.0}) + np.testing.assert_allclose(swung.max_point.numpy()[0][1], 1.7, atol=1e-6) + np.testing.assert_allclose(swung.max_point.numpy()[0][0], 0.2, atol=1e-6) + + # The temp USD is unique per run, so its stored meshes would otherwise pile up in the cache. + for joint_pos in ({}, {"elbow": math.pi / 2.0}): + mesh_cache_path(usd_path, joint_pos).unlink(missing_ok=True) + return True + + +def test_revolute_joint_swings_child_link(): + assert run_simulation_app_function(_test_revolute_joint_swings_child_link, headless=HEADLESS) + + +def test_zero_joint_position_preserves_authored_pose(): + assert run_simulation_app_function(_test_zero_joint_position_preserves_authored_pose, headless=HEADLESS) + + +def test_prismatic_joint_translates_child_link(): + assert run_simulation_app_function(_test_prismatic_joint_translates_child_link, headless=HEADLESS) + + +def test_authored_pose_away_from_joint_zero_is_corrected(): + assert run_simulation_app_function(_test_authored_pose_away_from_joint_zero_is_corrected, headless=HEADLESS) + + +def test_unknown_joint_name_is_rejected(): + assert run_simulation_app_function(_test_unknown_joint_name_is_rejected, headless=HEADLESS) + + +def test_joint_pos_patterns_expand_to_joint_names(): + assert run_simulation_app_function(_test_joint_pos_patterns_expand_to_joint_names, headless=HEADLESS) + + +def test_droid_geometry_tracks_configured_joint_positions(): + assert run_simulation_app_function(_test_droid_geometry_tracks_configured_joint_positions, headless=HEADLESS) + + +def test_droid_posed_bounding_box_covers_all_geometry(): + assert run_simulation_app_function(_test_droid_posed_bounding_box_covers_all_geometry, headless=HEADLESS) + + +def test_offline_posing_matches_physx_link_poses(): + assert run_simulation_app_function(_test_offline_posing_matches_physx_link_poses, headless=HEADLESS) + + +def test_closed_loop_articulation_poses_a_spanning_tree(): + assert run_simulation_app_function(_test_closed_loop_articulation_poses_a_spanning_tree, headless=HEADLESS) + + +def test_instanced_geometry_is_posed_with_its_link(): + assert run_simulation_app_function(_test_instanced_geometry_is_posed_with_its_link, headless=HEADLESS) diff --git a/isaaclab_arena/utils/collision_mesh_store.py b/isaaclab_arena/utils/collision_mesh_store.py new file mode 100644 index 0000000000..c5862ee646 --- /dev/null +++ b/isaaclab_arena/utils/collision_mesh_store.py @@ -0,0 +1,334 @@ +# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Persist an articulation's posed collision mesh as a USD, so extraction runs once per pose. + +Walking a robot USD to merge its meshes costs 0.1-2.5 s depending on the robot, which is paid again +in every new process. Artifacts are keyed by the joint pose they were extracted at, because that is +what makes a mesh reusable: embodiments spawn at their configured pose rather than at zero, so keying +on the asset alone would store a mesh nobody asks for. + +Lookup order is the local cache, then the robot's own folder under ``ARENA_ROBOT_LIBRARY_DIR`` on +Nucleus, which holds one published artifact per robot at the pose that robot spawns in, written by +``scripts/export_ready_pose_collision_meshes.py``. A miss falls back to extraction. + +Artifacts are named and validated by the source USD's stem, so a file exported on one machine loads +on another even though Arena composes the robot-on-stand USDs into a per-user cache directory. The +stem only has to be unique within a robot's folder, and a full export refuses to write a pair that +would collide there, since at a shared pose one robot would otherwise be served the other's mesh. +""" + +from __future__ import annotations + +import contextlib +import hashlib +import numpy as np +import os +import tempfile +import trimesh +from collections.abc import Mapping +from pathlib import Path + +from pxr import Usd, UsdGeom + +from isaaclab_arena.assets.asset_cache import get_arena_usd_cache_dir +from isaaclab_arena.assets.nucleus import ARENA_NUCLEUS_DIR + +_MESH_PRIM_PATH = "/CollisionMesh" +"""Prim the merged mesh is authored at, also the artifact's default prim.""" + +_ASSET_KEY = "arenaAssetKey" +"""customData key recording which asset the mesh describes, validated on load.""" + +_SOURCE_KEY = "arenaSourceUsdPath" +"""customData key recording the exporting machine's source path, kept to trace an artifact back.""" + +_POSE_KEY = "arenaJointPoseKey" +"""customData key recording which joint pose the mesh was extracted at, validated on load.""" + +_ZERO_POSE_KEY = "zero" +"""Pose key for an all-zero pose, spelled out so cached artifacts stay readable.""" + +_READY_POSE_SUFFIX = "_ready_pose.usd" +"""Suffix naming a published artifact, one per robot, at the pose that robot spawns in.""" + +ROBOT_LIBRARY_DIR_ENV_VAR = "ISAACLAB_ARENA_ROBOT_LIBRARY_DIR" +"""Environment variable redirecting which robot library published artifacts are read from.""" + +ARENA_ROBOT_LIBRARY_DIR = f"{ARENA_NUCLEUS_DIR}/Arena/assets/robot_library" +"""Nucleus robot library, holding each robot's assets in a folder of its own.""" + +_LOCAL_CACHE_DIR = "collision_mesh" +"""Local cache root under ``get_arena_usd_cache_dir()``, with one subfolder per robot when known.""" + +_STAGING_PREFIX = "_staging_" +"""Prefix marking a half-written artifact, so trimming leaves another process's write alone. + +Staging has to share the destination's directory for the rename to stay atomic, and the suffix has +to stay a USD one for ``Usd.Stage.CreateNew``, so the prefix is what keeps the two apart. +""" + +CACHE_BUDGET_BYTES = 1 << 30 +"""Disk the local cache may occupy before its least recently used artifacts are dropped. + +A robot's mesh runs 1-9 MB, so this holds a few hundred: every robot Arena ships at several poses +each, while bounding what a run against throwaway USDs can leave behind. +""" + + +def local_collision_mesh_cache_dir() -> Path: + """Return the local collision-mesh cache root (``~/.cache/.../usd/collision_mesh``).""" + return get_arena_usd_cache_dir() / _LOCAL_CACHE_DIR + + +def scaled_mesh(mesh: trimesh.Trimesh, scale: tuple[float, float, float]) -> trimesh.Trimesh: + """Return mesh with vertices scaled per axis in its own frame. + + Scaling in the root frame commutes with merging prims, so scaling a merged mesh matches scaling + each prim's vertices during extraction. + """ + if tuple(scale) == (1.0, 1.0, 1.0): + return mesh + # process=False keeps trimesh from merging vertices, which would make a scaled or reloaded mesh + # differ from the one extraction produced. + return trimesh.Trimesh( + vertices=mesh.vertices * np.asarray(scale, dtype=np.float64), faces=mesh.faces, process=False + ) + + +def is_zero_pose(joint_pos: Mapping[str, float]) -> bool: + """Whether joint_pos poses every joint at zero, including by naming no joints at all. + + Omitted joints are posed at zero, so an empty mapping is the zero pose rather than the asset's + authored configuration. + """ + return all(float(value) == 0.0 for value in joint_pos.values()) + + +def pose_key(joint_pos: Mapping[str, float]) -> str: + """Return a filename-safe key identifying a joint pose. + + Keys are derived from the joint names and values as given, so two spellings of one pose, a regex + and the names it expands to, key differently and are extracted separately. + """ + if is_zero_pose(joint_pos): + return _ZERO_POSE_KEY + canonical = ";".join(f"{name}={float(value) + 0.0:.9g}" for name, value in sorted(joint_pos.items())) + return hashlib.sha1(canonical.encode()).hexdigest()[:12] + + +def asset_key(source_usd_path: str) -> str: + """Return the identity a published artifact is named and validated by: the source USD's stem.""" + # The stem rather than the full path, so an artifact exported from a per-user cache directory + # still matches elsewhere. Stems already spell out the variant: droid_franka_robotiq_on_stand_1.350. + return Path(str(source_usd_path)).stem + + +def ready_pose_artifact_name(source_usd_path: str) -> str: + """Return the filename a robot USD's ready-pose mesh is published under.""" + return f"{asset_key(source_usd_path)}{_READY_POSE_SUFFIX}" + + +def published_ready_pose_dir() -> str: + """Return the robot library that published ready-pose artifacts are read from. + + Defaults to the Arena Nucleus robot library and is redirected by + ``ISAACLAB_ARENA_ROBOT_LIBRARY_DIR``, which lets an export be checked locally before upload. + """ + return os.environ.get(ROBOT_LIBRARY_DIR_ENV_VAR) or ARENA_ROBOT_LIBRARY_DIR + + +def published_ready_pose_path(source_usd_path: str, library_folder: str) -> str: + """Return the full path a robot USD's ready-pose mesh is published at. + + The artifact sits in the robot's own library folder rather than beside its source USD, because a + robot that spawns a composed asset sources from a per-user cache path that is nobody else's. + """ + return f"{published_ready_pose_dir().rstrip('/')}/{library_folder}/{ready_pose_artifact_name(source_usd_path)}" + + +def mesh_cache_path(source_usd_path: str, joint_pos: Mapping[str, float], library_folder: str | None = None) -> Path: + """Return the local cache path for a robot USD's mesh at joint_pos. + + Layout is ``collision_mesh/{library_folder}/{filename}`` when ``library_folder`` is set (the + embodiment's ``robot_library_folder``), otherwise ``collision_mesh/{filename}`` for throwaway + fixtures that have no published robot folder. The full source path is hashed into the filename + so assets sharing a stem stay distinct on one machine even though they would share a published + name. + """ + source_digest = hashlib.sha1(str(source_usd_path).encode()).hexdigest()[:12] + name = f"{asset_key(source_usd_path)}_{source_digest}_collision_{pose_key(joint_pos)}_pose.usd" + cache_root = local_collision_mesh_cache_dir() + if library_folder: + return cache_root / library_folder / name + return cache_root / name + + +def load_mesh( + source_usd_path: str, + joint_pos: Mapping[str, float], + scale: tuple[float, float, float], + library_folder: str | None = None, +) -> trimesh.Trimesh | None: + """Return the stored mesh for a robot USD at joint_pos scaled to scale, or None if unavailable. + + Reads the local cache first and the published robot library second, copying a published hit + into the local cache on the way out. An artifact recording another asset or another pose is + ignored rather than trusted. + + Args: + source_usd_path: Robot USD the mesh should describe. + joint_pos: Joint positions the mesh should be posed at. + scale: Per-axis scale to apply to the stored vertices. + library_folder: Robot's folder in the published library (and local cache), or None to + skip the published lookup and use the flat local-cache path. + """ + cached = mesh_cache_path(source_usd_path, joint_pos, library_folder) + mesh = _read_mesh_usd(str(cached), expected_asset=asset_key(source_usd_path), expected_pose=pose_key(joint_pos)) + if mesh is not None: + # Mark the artifact used, so trimming evicts by last use rather than by write time. + with contextlib.suppress(OSError): + os.utime(cached) + return scaled_mesh(mesh, scale) + + if library_folder is None: + return None + published = published_ready_pose_path(source_usd_path, library_folder) + mesh = _read_mesh_usd(published, expected_asset=asset_key(source_usd_path), expected_pose=pose_key(joint_pos)) + if mesh is None: + return None + # Copy the published artifact locally, so only the first process pays the Nucleus round trip. + with contextlib.suppress(OSError): + save_mesh(source_usd_path, joint_pos, mesh, library_folder) + return scaled_mesh(mesh, scale) + + +def save_mesh( + source_usd_path: str, + joint_pos: Mapping[str, float], + mesh: trimesh.Trimesh, + library_folder: str | None = None, +) -> Path: + """Write a robot USD's unscaled mesh at joint_pos to the local cache and return its path. + + Least recently used artifacts are dropped to keep the cache within ``CACHE_BUDGET_BYTES``. + + Args: + source_usd_path: Robot USD the mesh was extracted from, recorded for validation on load. + joint_pos: Joint positions the mesh was posed at, recorded for validation on load. + mesh: Merged mesh at unit scale, in the robot's default-prim frame. + library_folder: Robot's folder under the local cache, or None for a flat cache path. + """ + out_path = _write_mesh_usd( + mesh_cache_path(source_usd_path, joint_pos, library_folder), source_usd_path, joint_pos, mesh + ) + _trim_cache(local_collision_mesh_cache_dir()) + return out_path + + +def _trim_cache(cache_dir: Path) -> None: + """Delete cached artifacts, least recently used first, until the cache fits its budget. + + Artifacts another process is still staging are left alone, since unlinking one would fail that + process's rename. + """ + if not cache_dir.is_dir(): + return + artifacts = (path for path in cache_dir.rglob("*.usd") if not path.name.startswith(_STAGING_PREFIX)) + cumulative_bytes = 0 + for path in sorted(artifacts, key=_last_used, reverse=True): + with contextlib.suppress(OSError): + # Counts what has been walked rather than what survives, so eviction keeps a strict + # most-recently-used prefix instead of backfilling with whatever happens to fit. + cumulative_bytes += path.stat().st_size + if cumulative_bytes > CACHE_BUDGET_BYTES: + path.unlink() + + +def _last_used(path: Path) -> float: + """Return when an artifact was last read or written, or 0.0 if it has since gone.""" + try: + return path.stat().st_mtime + except OSError: + return 0.0 + + +def export_ready_pose_mesh( + source_usd_path: str, joint_pos: Mapping[str, float], mesh: trimesh.Trimesh, out_dir: Path, library_folder: str +) -> Path: + """Write a robot's ready-pose mesh into out_dir under its published name, for upload. + + Args: + source_usd_path: Robot USD the mesh was extracted from. + joint_pos: The robot's configured joint positions, recorded for validation on load. + mesh: Merged mesh at unit scale, in the robot's default-prim frame. + out_dir: Staging directory whose layout mirrors ``ARENA_ROBOT_LIBRARY_DIR`` for upload as-is. + library_folder: Robot's folder in the published library, created under out_dir. + """ + artifact_path = out_dir / library_folder / ready_pose_artifact_name(source_usd_path) + artifact_path.parent.mkdir(parents=True, exist_ok=True) + return _write_mesh_usd(artifact_path, source_usd_path, joint_pos, mesh) + + +def _write_mesh_usd( + out_path: Path, source_usd_path: str, joint_pos: Mapping[str, float], mesh: trimesh.Trimesh +) -> Path: + """Author a mesh artifact at out_path, staged and renamed so no reader sees a partial file.""" + out_path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + prefix=_STAGING_PREFIX, suffix=".usd", dir=out_path.parent, delete=False + ) as tmp_file: + tmp_path = Path(tmp_file.name) + try: + stage = Usd.Stage.CreateNew(str(tmp_path)) + mesh_prim = UsdGeom.Mesh.Define(stage, _MESH_PRIM_PATH) + stage.SetDefaultPrim(mesh_prim.GetPrim()) + faces = np.asarray(mesh.faces, dtype=np.int32) + mesh_prim.GetPointsAttr().Set(np.asarray(mesh.vertices, dtype=np.float32)) + mesh_prim.GetFaceVertexCountsAttr().Set([3] * len(faces)) + mesh_prim.GetFaceVertexIndicesAttr().Set(faces.reshape(-1)) + mesh_prim.GetPrim().SetCustomDataByKey(_ASSET_KEY, asset_key(source_usd_path)) + mesh_prim.GetPrim().SetCustomDataByKey(_POSE_KEY, pose_key(joint_pos)) + mesh_prim.GetPrim().SetCustomDataByKey(_SOURCE_KEY, str(source_usd_path)) + assert stage.GetRootLayer().Save(), f"failed to save collision mesh to {tmp_path}" + os.replace(tmp_path, out_path) + except Exception: + tmp_path.unlink(missing_ok=True) + raise + return out_path + + +def _read_mesh_usd(usd_path: str, expected_asset: str, expected_pose: str) -> trimesh.Trimesh | None: + """Read a stored artifact, returning None when it is absent or describes another asset or pose.""" + from isaaclab.utils.assets import check_file_path + + # Stage.Open raises rather than returning None on a missing remote path, and a published artifact + # is missing for every robot nobody has exported yet. + if check_file_path(usd_path) == 0: + return None + stage = Usd.Stage.Open(usd_path) + if stage is None: + return None + mesh_prim = UsdGeom.Mesh(stage.GetPrimAtPath(_MESH_PRIM_PATH)) + if not mesh_prim: + return None + + recorded = ( + mesh_prim.GetPrim().GetCustomDataByKey(_ASSET_KEY), + mesh_prim.GetPrim().GetCustomDataByKey(_POSE_KEY), + ) + if recorded != (expected_asset, expected_pose): + print(f"Ignoring collision mesh {usd_path}: recorded {recorded} != {(expected_asset, expected_pose)}") + return None + + points = mesh_prim.GetPointsAttr().Get() + indices = mesh_prim.GetFaceVertexIndicesAttr().Get() + if points is None or indices is None: + return None + return trimesh.Trimesh( + vertices=np.asarray(points, dtype=np.float64), + faces=np.asarray(indices, dtype=np.int32).reshape(-1, 3), + process=False, + ) diff --git a/isaaclab_arena/utils/isaac_sim_debug_draw.py b/isaaclab_arena/utils/isaac_sim_debug_draw.py index 95d808a5c2..85b44ea08b 100644 --- a/isaaclab_arena/utils/isaac_sim_debug_draw.py +++ b/isaaclab_arena/utils/isaac_sim_debug_draw.py @@ -79,6 +79,40 @@ def draw_object_bboxes( else: print(f"Skipping {obj.name}: no bbox coordinates") + def draw_line_segments( + self, + start_points: list[tuple[float, float, float]], + end_points: list[tuple[float, float, float]], + color: tuple[float, float, float, float] = DEFAULT_COLOR, + thickness: float = 3.0, + ) -> None: + """Draw one line per (start, end) pair, for wireframes ``draw_bbox`` cannot express. + + Args: + start_points: Segment start positions in world coordinates. + end_points: Segment end positions, matching start_points element-wise. + color: RGBA colour applied to every segment. + thickness: Line thickness in pixels. + """ + assert len(start_points) == len(end_points), "each start point needs a matching end point" + count = len(start_points) + self._draw.draw_lines(start_points, end_points, [color] * count, [thickness] * count) + + def draw_points( + self, + points: list[tuple[float, float, float]], + color: tuple[float, float, float, float] = DEFAULT_COLOR, + size: float = 5.0, + ) -> None: + """Draw a point cloud, for geometry such as mesh vertices. + + Args: + points: Positions in world coordinates. + color: RGBA colour applied to every point. + size: Point size in pixels. + """ + self._draw.draw_points(points, [color] * len(points), [size] * len(points)) + def clear(self) -> None: """Clear all debug drawings.""" self._draw.clear_lines() diff --git a/isaaclab_arena/utils/usd_articulation.py b/isaaclab_arena/utils/usd_articulation.py new file mode 100644 index 0000000000..6283440cf5 --- /dev/null +++ b/isaaclab_arena/utils/usd_articulation.py @@ -0,0 +1,263 @@ +# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Forward kinematics over a USD articulation's physics joints. + +Authored USD transforms only describe the one joint configuration an asset happens to be saved in, +so posing an articulation is a prerequisite for placement geometry that matches what is spawned. +Matrices here follow USD's row-vector convention: ``point_parent = point_local @ matrix``. +""" + +from __future__ import annotations + +import numpy as np +import re +from collections.abc import Iterable, Mapping +from dataclasses import dataclass + +from pxr import Gf, Sdf, Usd, UsdGeom, UsdPhysics + +_AXIS_VECTORS: dict[str, Gf.Vec3d] = { + "X": Gf.Vec3d(1.0, 0.0, 0.0), + "Y": Gf.Vec3d(0.0, 1.0, 0.0), + "Z": Gf.Vec3d(0.0, 0.0, 1.0), +} + + +def articulation_joint_prims(root_prim: Usd.Prim) -> dict[str, Usd.Prim]: + """Return movable joint prims under root_prim, keyed by joint name.""" + joint_prims: dict[str, Usd.Prim] = {} + for prim in Usd.PrimRange(root_prim): + if not prim.IsA(UsdPhysics.RevoluteJoint) and not prim.IsA(UsdPhysics.PrismaticJoint): + continue + name = prim.GetName() + assert name not in joint_prims, ( + f"Duplicate joint name '{name}' under {root_prim.GetPath()}:" + f" {joint_prims[name].GetPath()} and {prim.GetPath()}." + ) + joint_prims[name] = prim + return joint_prims + + +def resolve_joint_pos_patterns( + joint_names: Iterable[str], + joint_pos: Mapping[str, float], +) -> dict[str, float]: + """Expand Isaac Lab-style joint keys against the joint names an articulation actually has. + + Isaac Lab full-matches ``init_state.joint_pos`` keys as regexes, so one key such as + ``"right_outer.*"`` sets several joints at once. Keys matching no joint are dropped rather than + rejected, because a config may name joints only some variants of an asset carry; those joints + stay at zero. + + Args: + joint_names: Joint names the articulation has. + joint_pos: Joint positions keyed by exact joint name or by regex. + + Returns: + Joint positions keyed by exact joint name. Later keys win where patterns overlap. + """ + names = list(joint_names) + resolved: dict[str, float] = {} + for pattern, value in joint_pos.items(): + for name in names: + if name == pattern or re.fullmatch(pattern, name): + resolved[name] = float(value) + return resolved + + +def compute_posed_prim_world_deltas( + stage: Usd.Stage, + root_prim_path: str, + joint_pos: Mapping[str, float], +) -> dict[str, np.ndarray]: + """Return the world-transform delta that posing at joint_pos applies to each articulated body. + + A point already in stage world space moves to its posed location via ``point @ delta`` for the + delta of the body it belongs to. Bodies absent from the result do not move, so a caller resolves + a geometry prim's delta from its nearest ancestor present in the mapping. + + Joints not named in joint_pos are posed at zero. Revolute values are radians, matching Isaac + Lab rather than USD's degrees. Prismatic values are in stage linear units, which coincide with + metres for robot USDs authored at ``metersPerUnit = 1``. + + Args: + stage: Stage holding the articulation. + root_prim_path: Prim path to search for joints and bodies under. + joint_pos: Joint positions keyed by joint name. + + Returns: + Row-vector 4x4 world deltas keyed by body prim path. + """ + root_prim = stage.GetPrimAtPath(root_prim_path) + assert root_prim, f"No prim found at path {root_prim_path}" + + joint_prims = articulation_joint_prims(root_prim) + unknown_joints = set(joint_pos) - set(joint_prims) + assert not unknown_joints, ( + f"Joint positions name joints absent under {root_prim_path}: {sorted(unknown_joints)}." + f" Available joints: {sorted(joint_prims)}." + ) + + edges = _joint_edges(stage, root_prim, joint_prims, joint_pos) + if not edges: + return {} + + rest_transforms = { + path: _local_to_world(stage, path) + for path in {path for edge in edges for path in (edge.parent, edge.child) if path} + } + posed_transforms = _propagate_joint_motion(edges, rest_transforms) + return {path: np.linalg.inv(rest_transforms[path]) @ posed for path, posed in posed_transforms.items()} + + +def resolve_prim_world_delta(prim_path: str, body_deltas: Mapping[str, np.ndarray]) -> np.ndarray | None: + """Return the delta of the nearest ancestor of prim_path in body_deltas, or None if unposed.""" + candidate = prim_path + while candidate and candidate != "/": + if candidate in body_deltas: + return body_deltas[candidate] + candidate = candidate.rsplit("/", 1)[0] + return None + + +@dataclass(frozen=True, slots=True) +class _JointEdge: + """A parent-to-child articulation link with the motion its joint applies.""" + + parent: str + """Prim path of the joint's body0, empty when the joint attaches to the world.""" + + child: str + """Prim path of the joint's body1.""" + + motion: np.ndarray + """Row-vector 4x4 mapping the child joint frame into the parent joint frame.""" + + local_0: np.ndarray + """Row-vector 4x4 mapping the joint frame into body0.""" + + local_1: np.ndarray + """Row-vector 4x4 mapping the joint frame into body1.""" + + +def _joint_edges( + stage: Usd.Stage, + root_prim: Usd.Prim, + joint_prims: Mapping[str, Usd.Prim], + joint_pos: Mapping[str, float], +) -> list[_JointEdge]: + """Build the articulation's parent-to-child edges, including fixed joints that carry no motion.""" + root_path = root_prim.GetPath().pathString + movable_paths = {prim.GetPath().pathString for prim in joint_prims.values()} + values_by_path = {joint_prims[name].GetPath().pathString: value for name, value in joint_pos.items()} + + edges: list[_JointEdge] = [] + for prim in Usd.PrimRange(root_prim): + if not prim.IsA(UsdPhysics.Joint): + continue + joint_path = prim.GetPath().pathString + parent, child = _joint_body_paths(prim) + # Compare path components: a plain prefix test would also match a sibling root's children. + if child is None or not Sdf.Path(child).HasPrefix(Sdf.Path(root_path)): + continue + if joint_path in movable_paths: + motion = _joint_motion(prim, values_by_path.get(joint_path, 0.0)) + else: + motion = np.eye(4) + edges.append( + _JointEdge( + parent=parent if parent is not None else "", + child=child, + motion=motion, + local_0=_joint_local_frame(prim, index=0), + local_1=_joint_local_frame(prim, index=1), + ) + ) + return edges + + +def _propagate_joint_motion( + edges: list[_JointEdge], + rest_transforms: Mapping[str, np.ndarray], +) -> dict[str, np.ndarray]: + """Walk the articulation outward from its roots, composing each joint's motion.""" + edges_by_parent: dict[str, list[_JointEdge]] = {} + for edge in edges: + edges_by_parent.setdefault(edge.parent, []).append(edge) + + child_paths = {edge.child for edge in edges} + root_paths = [edge.parent for edge in edges if edge.parent not in child_paths] + assert root_paths, "articulation has no root body: every body is a joint child, so it is fully cyclic" + + posed: dict[str, np.ndarray] = {} + queue: list[tuple[str, np.ndarray]] = [] + for path in dict.fromkeys(root_paths): + # A world-attached joint has no body0 prim, so its parent frame is the stage origin. + parent_world = rest_transforms[path] if path else np.eye(4) + if path: + posed[path] = parent_world + queue.append((path, parent_world)) + + while queue: + parent_path, parent_world = queue.pop() + for edge in edges_by_parent.get(parent_path, ()): + # Closed loops reach a body through a second joint: Agibot's gripper four-bars, Galbot's + # fixed suction-cup mount. Keep the first path to the body and drop the loop closure, as + # PhysX also poses a spanning tree and enforces the remaining joint as a constraint. + if edge.child in posed: + continue + child_world = np.linalg.inv(edge.local_1) @ edge.motion @ edge.local_0 @ parent_world + posed[edge.child] = child_world + queue.append((edge.child, child_world)) + return posed + + +def _joint_body_paths(joint_prim: Usd.Prim) -> tuple[str | None, str | None]: + """Return the (body0, body1) prim paths a joint connects, using None for a world attachment.""" + joint = UsdPhysics.Joint(joint_prim) + body_paths: list[str | None] = [] + for relationship in (joint.GetBody0Rel(), joint.GetBody1Rel()): + targets = relationship.GetTargets() if relationship else [] + body_paths.append(targets[0].pathString if targets else None) + return body_paths[0], body_paths[1] + + +def _joint_local_frame(joint_prim: Usd.Prim, index: int) -> np.ndarray: + """Return the joint frame relative to the joint's body at index (0 or 1).""" + joint = UsdPhysics.Joint(joint_prim) + position_attr = joint.GetLocalPos0Attr() if index == 0 else joint.GetLocalPos1Attr() + rotation_attr = joint.GetLocalRot0Attr() if index == 0 else joint.GetLocalRot1Attr() + position = position_attr.Get() or Gf.Vec3f(0.0, 0.0, 0.0) + rotation = rotation_attr.Get() or Gf.Quatf(1.0, 0.0, 0.0, 0.0) + quaternion = Gf.Quatd(float(rotation.GetReal()), Gf.Vec3d(*rotation.GetImaginary())) + matrix = Gf.Matrix4d() + matrix.SetTransform(Gf.Rotation(quaternion), Gf.Vec3d(*position)) + return np.array(matrix, dtype=np.float64) + + +def _joint_motion(joint_prim: Usd.Prim, value: float) -> np.ndarray: + """Return the motion a movable joint applies at value, in its own joint frame.""" + axis_token = ( + UsdPhysics.RevoluteJoint(joint_prim).GetAxisAttr().Get() + if joint_prim.IsA(UsdPhysics.RevoluteJoint) + else UsdPhysics.PrismaticJoint(joint_prim).GetAxisAttr().Get() + ) + axis = _AXIS_VECTORS.get(axis_token or "X") + assert axis is not None, f"Joint {joint_prim.GetPath()} has unsupported axis '{axis_token}'." + + matrix = Gf.Matrix4d() + if joint_prim.IsA(UsdPhysics.RevoluteJoint): + matrix.SetRotate(Gf.Rotation(axis, np.degrees(value))) + else: + matrix.SetTranslate(axis * value) + return np.array(matrix, dtype=np.float64) + + +def _local_to_world(stage: Usd.Stage, prim_path: str) -> np.ndarray: + """Return a prim's authored local-to-world transform.""" + prim = stage.GetPrimAtPath(prim_path) + assert prim, f"Joint references a missing prim: {prim_path}" + return np.array(UsdGeom.Xformable(prim).ComputeLocalToWorldTransform(Usd.TimeCode.Default()), dtype=np.float64) diff --git a/isaaclab_arena/utils/usd_helpers.py b/isaaclab_arena/utils/usd_helpers.py index 3e25076188..4e388ab2fb 100644 --- a/isaaclab_arena/utils/usd_helpers.py +++ b/isaaclab_arena/utils/usd_helpers.py @@ -5,14 +5,30 @@ from __future__ import annotations +import functools import numpy as np import trimesh +from collections.abc import Mapping from contextlib import contextmanager from pxr import Gf, Usd, UsdGeom, UsdLux, UsdPhysics from isaaclab_arena.assets.object_type import ObjectType from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox +from isaaclab_arena.utils.collision_mesh_store import load_mesh, save_mesh, scaled_mesh +from isaaclab_arena.utils.usd_articulation import ( + articulation_joint_prims, + compute_posed_prim_world_deltas, + resolve_joint_pos_patterns, + resolve_prim_world_delta, +) + +_POSED_GEOMETRY_CACHE_SIZE = 16 +"""Distinct (USD, joint positions, scale) combinations to keep posed geometry for. + +An environment poses a handful of articulations, so a small cache spares repeated USD opens +without pinning many multi-megabyte meshes. +""" class NoCollisionMeshError(ValueError): @@ -358,10 +374,19 @@ def extract_trimesh_from_prim( stage: Usd.Stage, prim_path: str, scale: tuple[float, float, float] = (1.0, 1.0, 1.0), + prim_world_deltas: Mapping[str, np.ndarray] | None = None, ) -> trimesh.Trimesh: """Extract UsdGeom.Mesh geometry under a prim into the prim's local frame. Other Gprim geometry is rejected, not silently dropped. + + Args: + stage: Stage containing the prim. + prim_path: Root prim to gather mesh geometry under. + scale: Per-axis scale applied in the root frame. + prim_world_deltas: Optional world-space transform per prim path, applied before the root + frame conversion so callers can relocate geometry (e.g. posing an articulation). A mesh + inherits the delta of its nearest ancestor present in the mapping. """ assert all( s > 0 for s in scale @@ -382,7 +407,8 @@ def extract_trimesh_from_prim( skipped_gprims: list[str] = [] offset = 0 - for prim in Usd.PrimRange(root_prim): + # Instance proxies must be traversed explicitly, or an instanceable stand contributes no vertices. + for prim in Usd.PrimRange(root_prim, Usd.TraverseInstanceProxies()): if not prim.IsA(UsdGeom.Mesh): if prim.IsA(UsdGeom.Gprim): skipped_gprims.append(str(prim.GetPath())) @@ -395,6 +421,10 @@ def extract_trimesh_from_prim( continue prim_world_tf = np.array(UsdGeom.Xformable(prim).ComputeLocalToWorldTransform(Usd.TimeCode.Default())) + if prim_world_deltas is not None: + delta = resolve_prim_world_delta(str(prim.GetPath()), prim_world_deltas) + if delta is not None: + prim_world_tf = prim_world_tf @ delta prim_to_root_tf = prim_world_tf @ root_world_tf_inv verts = np.asarray(points, dtype=np.float64) verts_h = np.hstack([verts, np.ones((len(verts), 1))]) @@ -437,3 +467,143 @@ def extract_trimesh_from_usd_path( assert stage is not None, f"could not open USD: {usd_path}" default_prim = stage.GetDefaultPrim() or stage.GetPseudoRoot() return extract_trimesh_from_prim(stage, default_prim.GetPath().pathString, scale) + + +def extract_trimesh_from_usd_at_joint_pos( + usd_path: str, + joint_pos: Mapping[str, float], + scale: tuple[float, float, float] = (1.0, 1.0, 1.0), + library_folder: str | None = None, +) -> trimesh.Trimesh: + """Extract an articulation's mesh posed at joint_pos, in its default prim's local frame. + + Joints the articulation has but joint_pos omits are posed at zero, so the result depends only on + joint_pos and not on the configuration the asset happens to be authored in. Cached in process and + on disk, and shared with every other caller, so treat the result as read-only. + + Args: + usd_path: Path to the articulation's .usd/.usda/.usdc file. + joint_pos: Joint positions keyed by exact joint name or Isaac Lab regex, revolute in radians. + scale: Spawn-time scale passed to ``UsdFileCfg``. + library_folder: Robot's folder in the published library, or None to skip the published lookup. + + Returns: + Combined trimesh in the scaled default-prim frame. + """ + return _extract_trimesh_from_usd_at_joint_pos( + usd_path, tuple(sorted(joint_pos.items())), tuple(scale), library_folder + ) + + +# NOTE(zihaox, 2026-07-28): Cache here rather than on the asset. Isaac Lab reaches assets through +# EventTermCfg params, and configclass's validation walk tracks no visited set, so a trimesh held by +# an asset sends it recursing through trimesh's internal back-references until the stack overflows. +@functools.lru_cache(maxsize=_POSED_GEOMETRY_CACHE_SIZE) +def _extract_trimesh_from_usd_at_joint_pos( + usd_path: str, + joint_pos_items: tuple[tuple[str, float], ...], + scale: tuple[float, float, float], + library_folder: str | None, +) -> trimesh.Trimesh: + """Cacheable body of ``extract_trimesh_from_usd_at_joint_pos``, keyed by hashable arguments.""" + joint_pos = dict(joint_pos_items) + stored = load_mesh(usd_path, joint_pos, scale, library_folder) + if stored is not None: + return stored + + stage = Usd.Stage.Open(usd_path) + assert stage is not None, f"could not open USD: {usd_path}" + default_prim = stage.GetDefaultPrim() or stage.GetPseudoRoot() + default_prim_path = default_prim.GetPath().pathString + resolved = resolve_joint_pos_patterns(articulation_joint_prims(default_prim), joint_pos) + deltas = compute_posed_prim_world_deltas(stage, default_prim_path, resolved) + + # Store at unit scale so one artifact serves every spawn scale of the asset. + unscaled = extract_trimesh_from_prim(stage, default_prim_path, (1.0, 1.0, 1.0), prim_world_deltas=deltas) + save_mesh(usd_path, joint_pos, unscaled, library_folder=library_folder) + return scaled_mesh(unscaled, scale) + + +def compute_local_bounding_box_from_usd_at_joint_pos( + usd_path: str, + joint_pos: Mapping[str, float], + scale: tuple[float, float, float] = (1.0, 1.0, 1.0), + prim_path: str | None = None, +) -> AxisAlignedBoundingBox: + """Compute posed bounds under a prim in the articulation's default-prim-local frame. + + Every ``UsdGeom.Gprim`` contributes, matching the geometry the unposed + ``compute_local_bounding_box_from_usd`` covers. Bounds must not come from the posed mesh alone: + mesh extraction drops analytic gprims (Droid's ``gripper_adapter``), which would understate the + robot's footprint by centimetres. + + The result is cached on the arguments and shared with every other caller, so treat it as + read-only. Relation losses ask for bounds every optimisation step, which would otherwise reopen + the USD and redo the posing per step. + + Args: + usd_path: Path to the articulation's .usd/.usda/.usdc file. + joint_pos: Joint positions keyed by exact joint name or Isaac Lab regex, revolute in radians. + scale: Spawn-time scale passed to ``UsdFileCfg``. + prim_path: Optional sub-prim to bound. When None, bounds the full default prim. + + Returns: + AxisAlignedBoundingBox containing the posed local bounds. + """ + return _compute_local_bounding_box_from_usd_at_joint_pos( + usd_path, tuple(sorted(joint_pos.items())), tuple(scale), prim_path + ) + + +@functools.lru_cache(maxsize=_POSED_GEOMETRY_CACHE_SIZE) +def _compute_local_bounding_box_from_usd_at_joint_pos( + usd_path: str, + joint_pos_items: tuple[tuple[str, float], ...], + scale: tuple[float, float, float], + prim_path: str | None, +) -> AxisAlignedBoundingBox: + """Cacheable body of ``compute_local_bounding_box_from_usd_at_joint_pos``, keyed by hashable args.""" + joint_pos = dict(joint_pos_items) + stage = Usd.Stage.Open(usd_path) + assert stage is not None, f"could not open USD: {usd_path}" + default_prim = stage.GetDefaultPrim() or stage.GetPseudoRoot() + root_path = default_prim.GetPath().pathString + bound_prim = stage.GetPrimAtPath(prim_path) if prim_path is not None else default_prim + assert bound_prim.IsValid(), f"prim not found: {prim_path} in {usd_path}" + resolved = resolve_joint_pos_patterns(articulation_joint_prims(default_prim), joint_pos) + deltas = compute_posed_prim_world_deltas(stage, root_path, resolved) + + # Expressing corners relative to the root cancels its own transform, including the root scale + # Isaac Lab's spawner ignores, so only the caller's spawn scale applies. + bbox_cache = UsdGeom.BBoxCache(Usd.TimeCode.Default(), includedPurposes=[UsdGeom.Tokens.default_]) + root_world_tf_inv = np.linalg.inv( + np.array(UsdGeom.Xformable(default_prim).ComputeLocalToWorldTransform(Usd.TimeCode.Default())) + ) + scale_np = np.asarray(scale, dtype=np.float64) + + corners: list[np.ndarray] = [] + # Instance proxies must be traversed explicitly: robot-on-stand assemblies mark the stand (and for + # Franka, every visual) instanceable, so a default traversal stops at the instance and reports bounds + # covering the arm alone. + for prim in Usd.PrimRange(bound_prim, Usd.TraverseInstanceProxies()): + if not prim.IsA(UsdGeom.Gprim): + continue + local_range = bbox_cache.ComputeUntransformedBound(prim).ComputeAlignedRange() + if local_range.IsEmpty(): + continue + low, high = local_range.GetMin(), local_range.GetMax() + box_corners = np.array( + [[x, y, z, 1.0] for x in (low[0], high[0]) for y in (low[1], high[1]) for z in (low[2], high[2])] + ) + prim_world_tf = np.array(UsdGeom.Xformable(prim).ComputeLocalToWorldTransform(Usd.TimeCode.Default())) + delta = resolve_prim_world_delta(str(prim.GetPath()), deltas) + if delta is not None: + prim_world_tf = prim_world_tf @ delta + corners.append((box_corners @ (prim_world_tf @ root_world_tf_inv))[:, :3] * scale_np) + + assert corners, f"no bounded geometry found under {root_path} in {usd_path}" + stacked = np.vstack(corners) + return AxisAlignedBoundingBox( + min_point=tuple(stacked.min(axis=0)), + max_point=tuple(stacked.max(axis=0)), + ) diff --git a/isaaclab_arena_examples/relations/visualize_embodiment_placement_geometry.py b/isaaclab_arena_examples/relations/visualize_embodiment_placement_geometry.py new file mode 100644 index 0000000000..e045b39689 --- /dev/null +++ b/isaaclab_arena_examples/relations/visualize_embodiment_placement_geometry.py @@ -0,0 +1,174 @@ +# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Overlay an embodiment's placement bounding box and collision mesh on the robot as spawned. + +Run with the GUI and inspect by eye: the green wireframe is ``EmbodimentBase.get_bounding_box()`` and +the magenta points are ``get_collision_mesh()``, both posed at the robot's configured initial joint +positions and drawn through the spawned robot prim's transform. At ``num_steps=0`` the arm has not yet +sagged under gravity, so the overlay should hug the robot; stepping lets the joints settle a few +milliradians away and the overlay drifts by millimetres. + +Passing ``spec_yaml`` builds a full env graph instead of an empty ground plane, which shows the same +geometry where the relation solver actually placed the robot — the overlay then also reveals whether a +yaw relation rotated the placement geometry along with the robot. The spec YAML may also be given as +the first command-line argument, alongside the usual Isaac Lab flags such as ``--headless``. +""" + +# %% +from __future__ import annotations + +from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser +from isaaclab_arena.utils.isaaclab_utils.simulation_app import get_app_launcher + +# pyright: reportArgumentType=false, reportCallIssue=false, reportAttributeAccessIssue=false + + +_args, _positional_args = get_isaaclab_arena_cli_parser().parse_known_args() +if not _args.headless: + # A bare AppLauncher(headless=False) starts Kit without a window: Isaac Lab 3.0 only builds the + # viewport when a visualizer is requested, so ask for "kit" the way Arena's test harness does. + _args.visualizer = ["kit"] +print(f"Launching simulation app (headless={_args.headless})") +simulation_app = get_app_launcher(_args).app + +# %% + +MESH_POINT_BUDGET = 7000 +"""Mesh vertices are subsampled to this many debug points to keep the viewport responsive.""" + +BBOX_COLOR = (0.0, 1.0, 0.2, 1.0) +MESH_COLOR = (1.0, 0.2, 0.9, 0.9) + + +def _build_environment(embodiment_name: str, spec_yaml: str | None): + """Return an environment and its embodiment, either from an env graph spec or an empty scene.""" + if spec_yaml is not None: + from isaaclab_arena.environment_spec.arena_env_graph_conversion_utils import build_arena_env_from_graph_spec + from isaaclab_arena.environment_spec.arena_env_graph_spec import ArenaEnvGraphSpec + + environment = build_arena_env_from_graph_spec(ArenaEnvGraphSpec.from_yaml(spec_yaml)) + return environment, environment.embodiment + + from isaaclab_arena.assets.registries import AssetRegistry + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment + from isaaclab_arena.scene.scene import Scene + + asset_registry = AssetRegistry() + embodiment = asset_registry.get_asset_by_name(embodiment_name)() + scene_assets = [asset_registry.get_asset_by_name(name)() for name in ("ground_plane", "light")] + environment = IsaacLabArenaEnvironment( + name="visualize_embodiment_placement_geometry", + embodiment=embodiment, + scene=Scene(assets=scene_assets), + ) + return environment, embodiment + + +def visualize_embodiment_placement_geometry( + embodiment_name: str = "droid_abs_joint_pos", + spec_yaml: str | None = None, + num_steps: int = 0, + mesh_point_budget: int = MESH_POINT_BUDGET, + hold_open: bool = True, +): + """Spawn an embodiment and draw its placement bbox and collision mesh over it. + + Args: + embodiment_name: Registry name of the embodiment to inspect. Ignored when ``spec_yaml`` is set, + which carries its own embodiment. + spec_yaml: Path to an env graph spec YAML to build instead of an empty ground-plane scene. + num_steps: Simulation steps to run before drawing. Zero keeps the joints exactly at their + configured positions, which is what the placement geometry is posed at. + mesh_point_budget: Approximate number of mesh vertices to draw. + hold_open: Keep rendering until the window closes. Set False to exit after drawing, which is + what a headless check of the spec wants. + """ + import numpy as np + import torch + + import omni.usd + from pxr import Usd, UsdGeom + + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder, ArenaEnvBuilderCfg + from isaaclab_arena.utils.isaac_sim_debug_draw import IsaacSimDebugDraw + + environment, embodiment = _build_environment(embodiment_name, spec_yaml) + + # Placement geometry, in the embodiment's local (default-prim) frame. + bbox = embodiment.get_bounding_box() + mesh = embodiment.get_collision_mesh() + local_min = np.asarray(bbox.min_point.numpy()[0], dtype=np.float64) + local_max = np.asarray(bbox.max_point.numpy()[0], dtype=np.float64) + print(f"local bbox min = {np.round(local_min, 4)}") + print(f"local bbox max = {np.round(local_max, 4)}") + print(f"local bbox size = {np.round(local_max - local_min, 4)}") + print(f"local mesh size = {np.round(mesh.extents, 4)} ({len(mesh.vertices)} verts)") + # A positive difference is the analytic (non-mesh) geometry the bbox covers but the mesh cannot. + print(f"bbox minus mesh = {np.round((local_max - local_min) - mesh.extents, 4)}") + + env = ArenaEnvBuilder(environment, ArenaEnvBuilderCfg()).make_registered() + env.reset() + + for _ in range(num_steps): + with torch.inference_mode(): + env.step(torch.zeros(env.action_space.shape, device=env.unwrapped.device)) + + # The local frame is the spawned robot prim's frame, so draw through its world transform. + stage = omni.usd.get_context().get_stage() + robot_prim_path = env.unwrapped.scene["robot"].cfg.prim_path.replace("env_.*", "env_0") + robot_prim = stage.GetPrimAtPath(robot_prim_path) + assert robot_prim, f"no robot prim at {robot_prim_path}" + robot_world = np.array( + UsdGeom.Xformable(robot_prim).ComputeLocalToWorldTransform(Usd.TimeCode.Default()), dtype=np.float64 + ) + + def to_world(points: np.ndarray) -> np.ndarray: + return (np.hstack([points, np.ones((len(points), 1))]) @ robot_world)[:, :3] + + corners = to_world( + np.array([ + [x, y, z] + for x in (local_min[0], local_max[0]) + for y in (local_min[1], local_max[1]) + for z in (local_min[2], local_max[2]) + ]) + ) + # Corner order above is (x, y, z) bit-major, so edges join indices differing in exactly one bit. + edges = [(i, i ^ bit) for i in range(8) for bit in (1, 2, 4) if i < (i ^ bit)] + + debug_draw = IsaacSimDebugDraw() + debug_draw.clear() + debug_draw.draw_line_segments( + [tuple(corners[i]) for i, _ in edges], + [tuple(corners[j]) for _, j in edges], + color=BBOX_COLOR, + thickness=5.0, + ) + + vertices = np.asarray(mesh.vertices, dtype=np.float64) + if len(vertices) > mesh_point_budget: + vertices = vertices[:: len(vertices) // mesh_point_budget + 1] + world_vertices = to_world(vertices) + debug_draw.draw_points([tuple(point) for point in world_vertices], color=MESH_COLOR) + print(f"\ndrew {len(edges)} bbox edges (green) and {len(world_vertices)} mesh points (magenta)") + print(f"robot world position = {np.round(robot_world[3, :3], 4)}") + if not hold_open: + return + print("Inspect the overlay in the viewport. Close the window to exit.") + + # Render without stepping so the viewport stays interactive while the joints stay exactly where + # the placement geometry was posed. Debug draws persist across frames, so one draw is enough. + while simulation_app.is_running(): + env.unwrapped.sim.render() + + +# %% +visualize_embodiment_placement_geometry( + spec_yaml=_positional_args[0] if _positional_args else None, + hold_open=not _args.headless, +) + +# %% From 8c39b91e5f2e1939fdd63d37ca3e48cd85d99845 Mon Sep 17 00:00:00 2001 From: zhx06 Date: Tue, 28 Jul 2026 17:18:04 -0700 Subject: [PATCH 2/7] simplify codebase Signed-off-by: zhx06 --- isaaclab_arena/embodiments/agibot/agibot.py | 1 - isaaclab_arena/embodiments/droid/droid.py | 5 +- isaaclab_arena/embodiments/embodiment_base.py | 15 +- isaaclab_arena/embodiments/franka/franka.py | 22 +- isaaclab_arena/embodiments/g1/g1.py | 1 - isaaclab_arena/embodiments/galbot/galbot.py | 1 - isaaclab_arena/embodiments/gr1t2/gr1t2.py | 1 - .../embodiments/kuka_allegro/kuka_allegro.py | 1 - .../export_ready_pose_collision_meshes.py | 24 +- .../tests/test_collision_mesh_store.py | 152 +++++----- .../tests/test_embodiment_collision_mesh.py | 59 ++++ isaaclab_arena/tests/test_usd_articulation.py | 4 +- isaaclab_arena/utils/collision_mesh_store.py | 281 +++++++++--------- isaaclab_arena/utils/isaac_sim_debug_draw.py | 34 --- isaaclab_arena/utils/usd_helpers.py | 11 +- ...visualize_embodiment_placement_geometry.py | 174 ----------- 16 files changed, 313 insertions(+), 473 deletions(-) delete mode 100644 isaaclab_arena_examples/relations/visualize_embodiment_placement_geometry.py diff --git a/isaaclab_arena/embodiments/agibot/agibot.py b/isaaclab_arena/embodiments/agibot/agibot.py index 9b6cb78f80..354dadfce1 100644 --- a/isaaclab_arena/embodiments/agibot/agibot.py +++ b/isaaclab_arena/embodiments/agibot/agibot.py @@ -32,7 +32,6 @@ class AgibotEmbodiment(EmbodimentBase): name = "agibot" default_arm_mode = ArmMode.LEFT - robot_library_folder = "agibot_a2d" def __init__( self, enable_cameras: bool = False, initial_pose: Pose | None = None, arm_mode: ArmMode = ArmMode.LEFT diff --git a/isaaclab_arena/embodiments/droid/droid.py b/isaaclab_arena/embodiments/droid/droid.py index 9f88df7e3e..a0b1f87479 100644 --- a/isaaclab_arena/embodiments/droid/droid.py +++ b/isaaclab_arena/embodiments/droid/droid.py @@ -73,7 +73,6 @@ class DroidEmbodimentBase(EmbodimentBase, ABC): name = "droid" default_arm_mode = ArmMode.SINGLE_ARM - robot_library_folder = "droid" def __init__( self, @@ -424,10 +423,10 @@ class DroidEventCfg: 0.0, # panda_joint7 0.0, # finger_joint 0.0, # right_outer_knuckle_joint + 0.0, # left_inner_finger_joint 0.0, # right_inner_finger_joint - 0.0, # right_inner_finger_knuckle_joint 0.0, # left_inner_finger_knuckle_joint - 0.0, # left_inner_finger_joint + 0.0, # right_inner_finger_knuckle_joint ], }, ) diff --git a/isaaclab_arena/embodiments/embodiment_base.py b/isaaclab_arena/embodiments/embodiment_base.py index da0fb786f2..06f04244f9 100644 --- a/isaaclab_arena/embodiments/embodiment_base.py +++ b/isaaclab_arena/embodiments/embodiment_base.py @@ -37,9 +37,6 @@ class PlacementGeometrySource: joint_pos: Mapping[str, float] """Joint positions to pose the geometry at, revolute in radians, keyed by name or Isaac Lab regex.""" - library_folder: str | None - """Robot's folder under the published robot library, or None if it publishes no collision mesh.""" - class EmbodimentBase(PlaceableAsset): @@ -47,13 +44,6 @@ class EmbodimentBase(PlaceableAsset): tags: list[str] = ["embodiment"] default_arm_mode: ArmMode | None = None - robot_library_folder: str | None = None - """This robot's folder under the published robot library, shared by its action-space variants. - - Named separately from the USD because robots that spawn a composed asset, such as the on-stand - Droid and Franka, spawn from a per-user cache path that says nothing about where they publish. - """ - def __init__( self, enable_cameras: bool = False, @@ -94,7 +84,6 @@ def get_placement_geometry_source(self) -> PlacementGeometrySource: usd_path=spawn.usd_path, scale=(scale_x, scale_y, scale_z), joint_pos=dict(robot.init_state.joint_pos or {}), - library_folder=self.robot_library_folder, ) def get_bounding_box(self, prim_path: str | None = None) -> AxisAlignedBoundingBox: @@ -124,9 +113,7 @@ def get_collision_mesh(self) -> trimesh.Trimesh | None: from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos source = self.get_placement_geometry_source() - return extract_trimesh_from_usd_at_joint_pos( - source.usd_path, source.joint_pos, source.scale, library_folder=source.library_folder - ) + return extract_trimesh_from_usd_at_joint_pos(source.usd_path, source.joint_pos, source.scale) def _set_initial_pose(self, pose: Pose | PoseRange | PosePerEnv) -> None: """Store the configured pose; the construction pose is applied in ``get_scene_cfg``.""" diff --git a/isaaclab_arena/embodiments/franka/franka.py b/isaaclab_arena/embodiments/franka/franka.py index dccf7a1f13..8ff60fd89d 100644 --- a/isaaclab_arena/embodiments/franka/franka.py +++ b/isaaclab_arena/embodiments/franka/franka.py @@ -68,7 +68,6 @@ class FrankaEmbodimentBase(EmbodimentBase): """ default_arm_mode = ArmMode.SINGLE_ARM - robot_library_folder = "franka" def __init__( self, @@ -284,6 +283,23 @@ def __post_init__(self): policy: PolicyCfg = PolicyCfg() +_FRANKA_READY_POSE = { + "panda_joint1": 0.0, + "panda_joint2": -0.785, + "panda_joint3": -0.1107, + "panda_joint4": -1.1775, + "panda_joint5": 0.0, + "panda_joint6": 0.785, + "panda_joint7": 0.785, + "panda_finger_joint.*": 0.0400, +} +"""The arm pose the Franka spawns and resets in, spelled by name for the spawn state. + +``init_franka_arm_pose`` below repeats it positionally, as Isaac Lab's reset event assigns joints by +index. ``test_spawn_pose_matches_the_reset_pose`` holds the two together. +""" + + @configclass class FrankaEventCfg: """Configuration for Franka.""" @@ -460,6 +476,10 @@ def get_object_poses(self, env_ids: Sequence[int] | None = None): def _franka_robot_cfg_on_stand(robot_cfg: ArticulationCfg) -> ArticulationCfg: """Copy ``robot_cfg`` onto ``{ENV_REGEX_NS}/Robot`` with the composed on-stand USD.""" cfg = robot_cfg.replace(prim_path="{ENV_REGEX_NS}/Robot") + # Spawn where ``init_franka_arm_pose`` resets the arm to, rather than at Isaac Lab's own pose. + # The two described different arm configurations, which left placement geometry, derived from + # the spawn state, describing an arm the scene never contains. + cfg.init_state = cfg.init_state.replace(joint_pos=_FRANKA_READY_POSE) cfg.spawn.usd_path = compose_on_stand_usd( _FRANKA_ROBOT_PRIM, _FRANKA_STAND_PRIM, diff --git a/isaaclab_arena/embodiments/g1/g1.py b/isaaclab_arena/embodiments/g1/g1.py index f2b5ff0bbd..ea91ba68a9 100644 --- a/isaaclab_arena/embodiments/g1/g1.py +++ b/isaaclab_arena/embodiments/g1/g1.py @@ -47,7 +47,6 @@ class G1EmbodimentBase(EmbodimentBase): name = "g1" default_arm_mode = ArmMode.DUAL_ARM - robot_library_folder = "g1" def __init__( self, diff --git a/isaaclab_arena/embodiments/galbot/galbot.py b/isaaclab_arena/embodiments/galbot/galbot.py index ec27dcc9a5..300d18e39a 100644 --- a/isaaclab_arena/embodiments/galbot/galbot.py +++ b/isaaclab_arena/embodiments/galbot/galbot.py @@ -38,7 +38,6 @@ class GalbotEmbodiment(EmbodimentBase): name = "galbot" default_arm_mode = ArmMode.LEFT - robot_library_folder = "galbot" def __init__(self, enable_cameras: bool = False, initial_pose: Pose | None = None, arm_mode: ArmMode | None = None): super().__init__(enable_cameras, initial_pose, arm_mode=arm_mode) diff --git a/isaaclab_arena/embodiments/gr1t2/gr1t2.py b/isaaclab_arena/embodiments/gr1t2/gr1t2.py index ebf968d8d4..152d0d5a90 100644 --- a/isaaclab_arena/embodiments/gr1t2/gr1t2.py +++ b/isaaclab_arena/embodiments/gr1t2/gr1t2.py @@ -87,7 +87,6 @@ class GR1T2EmbodimentBase(EmbodimentBase): name = "gr1" default_arm_mode = ArmMode.RIGHT - robot_library_folder = "gr1t2" def __init__( self, diff --git a/isaaclab_arena/embodiments/kuka_allegro/kuka_allegro.py b/isaaclab_arena/embodiments/kuka_allegro/kuka_allegro.py index 68dbbbd336..8216c552ea 100644 --- a/isaaclab_arena/embodiments/kuka_allegro/kuka_allegro.py +++ b/isaaclab_arena/embodiments/kuka_allegro/kuka_allegro.py @@ -91,7 +91,6 @@ class KukaAllegroEmbodiment(EmbodimentBase): name = "kuka_allegro" default_arm_mode = ArmMode.SINGLE_ARM - robot_library_folder = "kuka" def __init__( self, diff --git a/isaaclab_arena/scripts/export_ready_pose_collision_meshes.py b/isaaclab_arena/scripts/export_ready_pose_collision_meshes.py index a0051efad8..0169ca0311 100644 --- a/isaaclab_arena/scripts/export_ready_pose_collision_meshes.py +++ b/isaaclab_arena/scripts/export_ready_pose_collision_meshes.py @@ -81,27 +81,27 @@ class PlannedArtifact: """Embodiment the mesh is extracted from, for reporting which robot an artifact came from.""" source: PlacementGeometrySource - """USD, scale, joint positions and library folder the artifact is written from.""" + """USD, scale and joint positions the artifact is written from.""" def plan_artifacts(sources: Mapping[str, PlacementGeometrySource]) -> dict[str, PlannedArtifact]: """Map each artifact's published relative path to the embodiment and source it is written from. - Robots differing only in action space share a USD and a folder, so they collapse to one artifact. - Robots declaring no library folder have nowhere to publish and are left out. + Robots differing only in action space share a USD, so they collapse to one artifact. Robots + missing from ``ROBOT_LIBRARY_FOLDERS`` have nowhere to publish and are left out. Args: sources: Placement geometry source per embodiment name. """ - from isaaclab_arena.utils.collision_mesh_store import ready_pose_artifact_name + from isaaclab_arena.utils.collision_mesh_store import published_relative_path plan: dict[str, PlannedArtifact] = {} for name, source in sorted(sources.items()): - if source.library_folder is None: + relative_path = published_relative_path(source.usd_path) + if relative_path is None: continue - relative_path = f"{source.library_folder}/{ready_pose_artifact_name(source.usd_path)}" - # Artifacts are validated by USD stem, so two robots sharing one within a folder would be - # served each other's mesh. Refuse to publish that rather than let placement use the wrong shape. + # Artifacts are validated by USD stem, so two robots sharing one would be served each other's + # mesh. Refuse to publish that rather than let placement use the wrong shape. planned = plan.setdefault(relative_path, PlannedArtifact(name, source)) assert planned.source.usd_path == source.usd_path, ( f"{relative_path} would be written for both {planned.source.usd_path} and" @@ -117,7 +117,7 @@ def export_ready_pose_meshes(out_dir: Path, names: list[str] | None) -> tuple[li out_dir: Staging directory, whose per-robot folders are uploaded as-is to the robot library. names: Registered embodiment names to export, or None for every registered embodiment. """ - from isaaclab_arena.utils.collision_mesh_store import export_ready_pose_mesh + from isaaclab_arena.utils.collision_mesh_store import export_ready_pose_mesh, published_relative_path from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos sources = {} @@ -129,20 +129,20 @@ def export_ready_pose_meshes(out_dir: Path, names: list[str] | None) -> tuple[li failed.append(f"{embodiment_class.__name__}: {error}") traceback.print_exc() - skipped = sorted(name for name, source in sources.items() if source.library_folder is None) + skipped = sorted(name for name, source in sources.items() if published_relative_path(source.usd_path) is None) written = [] for relative_path, planned in plan_artifacts(sources).items(): source = planned.source try: # Unit scale: one artifact serves every spawn scale, as the loader rescales on read. mesh = extract_trimesh_from_usd_at_joint_pos(source.usd_path, source.joint_pos, (1.0, 1.0, 1.0)) - export_ready_pose_mesh(source.usd_path, source.joint_pos, mesh, out_dir, source.library_folder) + export_ready_pose_mesh(source.usd_path, source.joint_pos, mesh, out_dir) written.append(f"{relative_path} ({len(mesh.vertices)} vertices, from {planned.embodiment_name})") except Exception as error: failed.append(f"{planned.embodiment_name}: {error}") traceback.print_exc() if skipped: - print(f"\nSkipped {len(skipped)} embodiment(s) with no robot_library_folder: {', '.join(skipped)}") + print(f"\nSkipped {len(skipped)} embodiment(s) missing from ROBOT_LIBRARY_FOLDERS: {', '.join(skipped)}") return sorted(written), failed diff --git a/isaaclab_arena/tests/test_collision_mesh_store.py b/isaaclab_arena/tests/test_collision_mesh_store.py index 57372c4f48..477b16d889 100644 --- a/isaaclab_arena/tests/test_collision_mesh_store.py +++ b/isaaclab_arena/tests/test_collision_mesh_store.py @@ -44,6 +44,18 @@ def _published_dir(path: str): os.environ[ROBOT_LIBRARY_DIR_ENV_VAR] = previous +@contextlib.contextmanager +def _publishes_as(stem: str, folder: str): + """Give a fixture asset a robot-library folder, as a shipped robot has.""" + from isaaclab_arena.utils import collision_mesh_store as store + + store.ROBOT_LIBRARY_FOLDERS[stem] = folder + try: + yield + finally: + store.ROBOT_LIBRARY_FOLDERS.pop(stem, None) + + def _test_trimming_evicts_least_recently_used_artifacts(simulation_app) -> bool: """Trimming drops artifacts oldest-first down to the budget, and leaves in-flight writes alone.""" from isaaclab_arena.utils import collision_mesh_store as store @@ -78,40 +90,38 @@ def _test_trimming_evicts_least_recently_used_artifacts(simulation_app) -> bool: def _test_pose_keys_identify_the_pose(simulation_app) -> bool: """Every all-zero spelling shares the readable zero key; distinct poses key apart.""" - from isaaclab_arena.utils.collision_mesh_store import is_zero_pose, pose_key - - assert is_zero_pose({}) and is_zero_pose({"elbow": 0.0, "wrist": -0.0}) - assert not is_zero_pose({"elbow": 0.5}) + from isaaclab_arena.utils.collision_mesh_store import _pose_key - assert pose_key({}) == "zero" - assert pose_key({"elbow": 0.0, "wrist": -0.0}) == "zero" - assert pose_key({"elbow": 0.5}) != "zero" + # Omitted joints are posed at zero, so an empty mapping is the zero pose. + assert _pose_key({}) == "zero" + assert _pose_key({"elbow": 0.0, "wrist": -0.0}) == "zero" + assert _pose_key({"elbow": 0.5}) != "zero" # Order must not matter, but the values must. - assert pose_key({"a": 0.5, "b": 0.25}) == pose_key({"b": 0.25, "a": 0.5}) - assert pose_key({"elbow": 0.5}) != pose_key({"elbow": 0.6}) - assert pose_key({"elbow": 0.5}) != pose_key({"wrist": 0.5}) + assert _pose_key({"a": 0.5, "b": 0.25}) == _pose_key({"b": 0.25, "a": 0.5}) + assert _pose_key({"elbow": 0.5}) != _pose_key({"elbow": 0.6}) + assert _pose_key({"elbow": 0.5}) != _pose_key({"wrist": 0.5}) return True def _test_distinct_poses_get_distinct_artifacts(simulation_app) -> bool: """Posing at one configuration must never serve the mesh for another.""" from isaaclab_arena.utils import usd_helpers - from isaaclab_arena.utils.collision_mesh_store import mesh_cache_path + from isaaclab_arena.utils.collision_mesh_store import _cache_path extended = {"elbow": 0.5} with tempfile.TemporaryDirectory() as tmp_dir, contextlib.ExitStack() as published: usd_path = _build_arm_usd(tmp_dir, joint_type="prismatic") published.enter_context(_published_dir(tmp_dir)) for joint_pos in ({}, extended): - mesh_cache_path(usd_path, joint_pos).unlink(missing_ok=True) + _cache_path(usd_path, joint_pos).unlink(missing_ok=True) usd_helpers._extract_trimesh_from_usd_at_joint_pos.cache_clear() zero_mesh = usd_helpers.extract_trimesh_from_usd_at_joint_pos(usd_path, {}) usd_helpers._extract_trimesh_from_usd_at_joint_pos.cache_clear() posed_mesh = usd_helpers.extract_trimesh_from_usd_at_joint_pos(usd_path, extended) - assert mesh_cache_path(usd_path, {}).is_file() and mesh_cache_path(usd_path, extended).is_file() - assert mesh_cache_path(usd_path, {}) != mesh_cache_path(usd_path, extended) + assert _cache_path(usd_path, {}).is_file() and _cache_path(usd_path, extended).is_file() + assert _cache_path(usd_path, {}) != _cache_path(usd_path, extended) assert posed_mesh.extents[0] > zero_mesh.extents[0] + 0.4, f"prismatic pose should extend: {posed_mesh.extents}" # Reloading from the store must preserve that difference. @@ -119,13 +129,13 @@ def _test_distinct_poses_get_distinct_artifacts(simulation_app) -> bool: reloaded_zero = usd_helpers.extract_trimesh_from_usd_at_joint_pos(usd_path, {}) np.testing.assert_allclose(reloaded_zero.extents, zero_mesh.extents, atol=1e-6) for joint_pos in ({}, extended): - mesh_cache_path(usd_path, joint_pos).unlink(missing_ok=True) + _cache_path(usd_path, joint_pos).unlink(missing_ok=True) return True def _test_stored_mesh_keeps_its_vertices_verbatim(simulation_app) -> bool: """Reloading must not merge coincident vertices, which trimesh does by default on construction.""" - from isaaclab_arena.utils.collision_mesh_store import load_mesh, mesh_cache_path, save_mesh + from isaaclab_arena.utils.collision_mesh_store import _cache_path, load_mesh, save_mesh # Two coincident triangles: merging would collapse six vertices into three. vertices = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]] * 2) @@ -142,19 +152,19 @@ def _test_stored_mesh_keeps_its_vertices_verbatim(simulation_app) -> bool: scaled = load_mesh(source, {}, (2.0, 2.0, 2.0)) assert scaled is not None assert len(scaled.vertices) == len(vertices), f"scaling changed vertex count: {len(scaled.vertices)}" - mesh_cache_path(source, {}).unlink(missing_ok=True) + _cache_path(source, {}).unlink(missing_ok=True) return True def _test_stored_mesh_is_scale_independent(simulation_app) -> bool: """One unit-scale artifact serves every spawn scale.""" from isaaclab_arena.utils import usd_helpers - from isaaclab_arena.utils.collision_mesh_store import mesh_cache_path + from isaaclab_arena.utils.collision_mesh_store import _cache_path with tempfile.TemporaryDirectory() as tmp_dir, contextlib.ExitStack() as published: usd_path = _build_arm_usd(tmp_dir) published.enter_context(_published_dir(tmp_dir)) - mesh_cache_path(usd_path, {}).unlink(missing_ok=True) + _cache_path(usd_path, {}).unlink(missing_ok=True) usd_helpers._extract_trimesh_from_usd_at_joint_pos.cache_clear() unit = usd_helpers.extract_trimesh_from_usd_at_joint_pos(usd_path, {}) @@ -163,13 +173,13 @@ def _test_stored_mesh_is_scale_independent(simulation_app) -> bool: doubled = usd_helpers.extract_trimesh_from_usd_at_joint_pos(usd_path, {}, scale=(2.0, 1.0, 1.0)) np.testing.assert_allclose(doubled.extents[0], unit.extents[0] * 2.0, atol=1e-6) np.testing.assert_allclose(doubled.extents[1:], unit.extents[1:], atol=1e-6) - mesh_cache_path(usd_path, {}).unlink(missing_ok=True) + _cache_path(usd_path, {}).unlink(missing_ok=True) return True def _test_artifact_from_another_asset_or_pose_is_ignored(simulation_app) -> bool: """A mesh recording a different source or pose is refused, so it cannot misplace a robot.""" - from isaaclab_arena.utils.collision_mesh_store import load_mesh, mesh_cache_path, save_mesh + from isaaclab_arena.utils.collision_mesh_store import _cache_path, load_mesh, save_mesh from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos posed = {"elbow": 0.5} @@ -179,52 +189,52 @@ def _test_artifact_from_another_asset_or_pose_is_ignored(simulation_app) -> bool mesh = extract_trimesh_from_usd_at_joint_pos(usd_path, posed) # Recorded under a foreign source USD, then moved into this asset's slot. - save_mesh(f"{tmp_dir}/other_arm.usda", posed, mesh).replace(mesh_cache_path(usd_path, posed)) + save_mesh(f"{tmp_dir}/other_arm.usda", posed, mesh).replace(_cache_path(usd_path, posed)) assert load_mesh(usd_path, posed, (1.0, 1.0, 1.0)) is None, "foreign source must be ignored" # Recorded for this asset but at another pose, then moved into this pose's slot. - save_mesh(usd_path, {"elbow": 0.9}, mesh).replace(mesh_cache_path(usd_path, posed)) + save_mesh(usd_path, {"elbow": 0.9}, mesh).replace(_cache_path(usd_path, posed)) assert load_mesh(usd_path, posed, (1.0, 1.0, 1.0)) is None, "foreign pose must be ignored" - mesh_cache_path(usd_path, posed).unlink(missing_ok=True) + _cache_path(usd_path, posed).unlink(missing_ok=True) return True def _test_published_artifact_is_found_by_its_name(simulation_app) -> bool: """An exported artifact is loaded from the published directory when the local cache is empty.""" from isaaclab_arena.utils.collision_mesh_store import ( + _cache_path, + _published_path, export_ready_pose_mesh, load_mesh, - mesh_cache_path, - published_ready_pose_path, - ready_pose_artifact_name, + published_relative_path, ) from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos posed = {"elbow": 0.5} - with tempfile.TemporaryDirectory() as tmp_dir, tempfile.TemporaryDirectory() as publish_dir: + with ( + tempfile.TemporaryDirectory() as tmp_dir, + tempfile.TemporaryDirectory() as publish_dir, + _publishes_as("arm", "test_arm"), + ): usd_path = _build_arm_usd(tmp_dir) mesh = extract_trimesh_from_usd_at_joint_pos(usd_path, posed) - artifact = export_ready_pose_mesh(usd_path, posed, mesh, Path(publish_dir), "test_arm") + artifact = export_ready_pose_mesh(usd_path, posed, mesh, Path(publish_dir)) # The artifact lands in the robot's own library folder, under a name derived from the asset. - assert artifact.name == "arm_ready_pose.usd" == ready_pose_artifact_name(usd_path) - assert artifact.parent.name == "test_arm", artifact + assert published_relative_path(usd_path) == "test_arm/arm_ready_pose.usd" + assert artifact.name == "arm_ready_pose.usd" and artifact.parent.name == "test_arm", artifact with _published_dir(publish_dir): - assert published_ready_pose_path(usd_path, "test_arm") == str(artifact) + assert _published_path(usd_path) == str(artifact) # Empty the local cache so only the published copy can answer. - mesh_cache_path(usd_path, posed).unlink(missing_ok=True) - loaded = load_mesh(usd_path, posed, (1.0, 1.0, 1.0), "test_arm") + _cache_path(usd_path, posed).unlink(missing_ok=True) + loaded = load_mesh(usd_path, posed, (1.0, 1.0, 1.0)) assert loaded is not None, f"published artifact {artifact} was not found" np.testing.assert_allclose(loaded.vertices, mesh.vertices, atol=1e-6) - # A robot that names no folder must not be served another robot's artifact. - mesh_cache_path(usd_path, posed).unlink(missing_ok=True) - assert load_mesh(usd_path, posed, (1.0, 1.0, 1.0)) is None, "published lookup needs a folder" - # Published artifacts outlive config changes, so a repose must not be served the old shape. - assert load_mesh(usd_path, {"elbow": 0.9}, (1.0, 1.0, 1.0), "test_arm") is None, "stale pose" - mesh_cache_path(usd_path, posed).unlink(missing_ok=True) + assert load_mesh(usd_path, {"elbow": 0.9}, (1.0, 1.0, 1.0)) is None, "stale pose" + _cache_path(usd_path, posed).unlink(missing_ok=True) return True @@ -234,7 +244,7 @@ def _test_published_artifact_loads_for_a_relocated_source(simulation_app) -> boo Arena composes the robot-on-stand USDs into a per-user cache directory, so validating against the full source path would reject every artifact anyone else exported. """ - from isaaclab_arena.utils.collision_mesh_store import export_ready_pose_mesh, load_mesh, mesh_cache_path + from isaaclab_arena.utils.collision_mesh_store import _cache_path, export_ready_pose_mesh, load_mesh from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos posed = {"elbow": 0.5} @@ -242,22 +252,23 @@ def _test_published_artifact_loads_for_a_relocated_source(simulation_app) -> boo tempfile.TemporaryDirectory() as exporter_dir, tempfile.TemporaryDirectory() as consumer_dir, tempfile.TemporaryDirectory() as publish_dir, + _publishes_as("arm", "test_arm"), ): exporter_usd = _build_arm_usd(exporter_dir) mesh = extract_trimesh_from_usd_at_joint_pos(exporter_usd, posed) - export_ready_pose_mesh(exporter_usd, posed, mesh, Path(publish_dir), "test_arm") + export_ready_pose_mesh(exporter_usd, posed, mesh, Path(publish_dir)) # The same asset reached through another absolute path, as a second machine's cache would. consumer_usd = f"{consumer_dir}/arm.usda" shutil.copy(exporter_usd, consumer_usd) with _published_dir(publish_dir): - mesh_cache_path(consumer_usd, posed).unlink(missing_ok=True) - loaded = load_mesh(consumer_usd, posed, (1.0, 1.0, 1.0), "test_arm") + _cache_path(consumer_usd, posed).unlink(missing_ok=True) + loaded = load_mesh(consumer_usd, posed, (1.0, 1.0, 1.0)) assert loaded is not None, "an artifact must load for the same asset at another path" np.testing.assert_allclose(loaded.vertices, mesh.vertices, atol=1e-6) for path in (exporter_usd, consumer_usd): - mesh_cache_path(path, posed).unlink(missing_ok=True) + _cache_path(path, posed).unlink(missing_ok=True) return True @@ -288,40 +299,31 @@ def _test_truncated_artifact_is_ignored(simulation_app) -> bool: def _test_colliding_published_names_are_refused(simulation_app) -> bool: - """Two robots sharing a stem within one folder cannot both be published, as one name serves both.""" + """Two robots sharing a USD stem cannot both be published, as one artifact name serves both.""" from isaaclab_arena.embodiments.embodiment_base import PlacementGeometrySource from isaaclab_arena.scripts.export_ready_pose_collision_meshes import plan_artifacts - def source(usd_path: str, library_folder: str | None) -> PlacementGeometrySource: - return PlacementGeometrySource(usd_path, (1.0, 1.0, 1.0), {}, library_folder) + def source(usd_path: str) -> PlacementGeometrySource: + return PlacementGeometrySource(usd_path, (1.0, 1.0, 1.0), {}) - # Same stem in one folder, different assets: the published name cannot tell them apart. - colliding = { - "robot_a": source("/assets/vendor_a/robot.usd", "shared"), - "robot_b": source("/assets/vendor_b/robot.usd", "shared"), - } - try: - plan_artifacts(colliding) - except AssertionError as error: - assert "rename one of the source USDs" in str(error), error - else: - raise AssertionError("a colliding published name must be refused") - - # The same stem in separate folders is fine, which is what per-robot folders buy. - separate = { - "robot_a": source("/assets/vendor_a/robot.usd", "robot_a"), - "robot_b": source("/assets/vendor_b/robot.usd", "robot_b"), - } - assert sorted(plan_artifacts(separate)) == ["robot_a/robot_ready_pose.usd", "robot_b/robot_ready_pose.usd"] - - # Action-space variants of one robot share a USD, and so legitimately share one artifact. - shared_usd = source("/assets/vendor_a/robot.usd", "robot_a") - assert sorted(plan_artifacts({"robot_ik": shared_usd, "robot_joint_pos": shared_usd})) == [ - "robot_a/robot_ready_pose.usd" - ] - - # A robot with nowhere to publish is left out rather than written to a guessed folder. - assert plan_artifacts({"unpublished": source("/assets/vendor_a/robot.usd", None)}) == {} + with _publishes_as("robot", "shared"): + # Same stem, different assets: the published name cannot tell them apart. + colliding = {"robot_a": source("/assets/vendor_a/robot.usd"), "robot_b": source("/assets/vendor_b/robot.usd")} + try: + plan_artifacts(colliding) + except AssertionError as error: + assert "rename one of the source USDs" in str(error), error + else: + raise AssertionError("a colliding published name must be refused") + + # Action-space variants of one robot share a USD, and so legitimately share one artifact. + shared_usd = source("/assets/vendor_a/robot.usd") + assert sorted(plan_artifacts({"robot_ik": shared_usd, "robot_joint_pos": shared_usd})) == [ + "shared/robot_ready_pose.usd" + ] + + # A robot missing from ROBOT_LIBRARY_FOLDERS is left out rather than written to a guessed folder. + assert plan_artifacts({"unpublished": source("/assets/vendor_a/robot.usd")}) == {} return True diff --git a/isaaclab_arena/tests/test_embodiment_collision_mesh.py b/isaaclab_arena/tests/test_embodiment_collision_mesh.py index 95dbe25e66..5d0d60ca9d 100644 --- a/isaaclab_arena/tests/test_embodiment_collision_mesh.py +++ b/isaaclab_arena/tests/test_embodiment_collision_mesh.py @@ -62,11 +62,70 @@ def _noop(env, env_ids, embodiment): return True +def _test_spawn_pose_matches_the_reset_pose(simulation_app) -> bool: + """The spawn joint positions placement geometry is posed at are the ones a reset drives to. + + Both robots also reach their arm pose through an event that assigns joints positionally, which + cannot be read off the config; the articulation supplies the joint order here instead. + """ + + import torch + + from isaaclab_arena.cli.isaaclab_arena_cli import arena_env_builder_cfg_from_argparse, get_isaaclab_arena_cli_parser + from isaaclab_arena.embodiments.droid.droid import DroidAbsoluteJointPositionEmbodiment + from isaaclab_arena.embodiments.franka.franka import FrankaIKEmbodiment + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment + from isaaclab_arena.scene.scene import Scene + from isaaclab_arena.utils.usd_articulation import resolve_joint_pos_patterns + + for embodiment_class in (FrankaIKEmbodiment, DroidAbsoluteJointPositionEmbodiment): + env = None + try: + embodiment = embodiment_class() + environment = IsaacLabArenaEnvironment( + name=f"spawn_pose_{embodiment_class.__name__}", embodiment=embodiment, scene=Scene(assets=[]) + ) + builder_cfg = arena_env_builder_cfg_from_argparse(get_isaaclab_arena_cli_parser().parse_args([])) + env = ArenaEnvBuilder(environment, builder_cfg).make_registered() + env.reset() + + robot = env.unwrapped.scene["robot"] + joint_names = list(robot.joint_names) + spawn_pose = resolve_joint_pos_patterns(joint_names, embodiment.get_placement_geometry_source().joint_pos) + # Compare against the defaults a reset restores rather than measured positions, which + # carry gravity sag and the joint randomisation event on top. + reset_to = torch.as_tensor(robot.data.default_joint_pos)[0].cpu() + for index, joint_name in enumerate(joint_names): + assert abs(spawn_pose.get(joint_name, 0.0) - float(reset_to[index])) < 1e-6, ( + f"{embodiment_class.__name__} poses placement geometry with {joint_name} at" + f" {spawn_pose.get(joint_name, 0.0)}, but resets it to {float(reset_to[index])}" + ) + + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + + finally: + if env is not None: + env.close() + + return True + + def test_embodiment_provides_robot_collision_mesh(): """Pytest entry point for the embodiment collision-mesh test.""" result = run_function_with_persistent_simulation_app(_test_embodiment_provides_robot_collision_mesh, headless=True) assert result, f"Test {test_embodiment_provides_robot_collision_mesh.__name__} failed" +def test_spawn_pose_matches_the_reset_pose(): + """Pytest entry point for the spawn-versus-reset joint-position test.""" + result = run_simulation_app_function(_test_spawn_pose_matches_the_reset_pose, headless=True) + assert result, f"Test {test_spawn_pose_matches_the_reset_pose.__name__} failed" + + if __name__ == "__main__": test_embodiment_provides_robot_collision_mesh() + test_spawn_pose_matches_the_reset_pose() diff --git a/isaaclab_arena/tests/test_usd_articulation.py b/isaaclab_arena/tests/test_usd_articulation.py index 7efd172425..0b240587c5 100644 --- a/isaaclab_arena/tests/test_usd_articulation.py +++ b/isaaclab_arena/tests/test_usd_articulation.py @@ -329,7 +329,7 @@ def _test_instanced_geometry_is_posed_with_its_link(simulation_app) -> bool: from pxr import Gf, UsdGeom - from isaaclab_arena.utils.collision_mesh_store import mesh_cache_path + from isaaclab_arena.utils.collision_mesh_store import _cache_path from isaaclab_arena.utils.usd_helpers import ( compute_local_bounding_box_from_usd_at_joint_pos, extract_trimesh_from_usd_at_joint_pos, @@ -361,7 +361,7 @@ def _test_instanced_geometry_is_posed_with_its_link(simulation_app) -> bool: # The temp USD is unique per run, so its stored meshes would otherwise pile up in the cache. for joint_pos in ({}, {"elbow": math.pi / 2.0}): - mesh_cache_path(usd_path, joint_pos).unlink(missing_ok=True) + _cache_path(usd_path, joint_pos).unlink(missing_ok=True) return True diff --git a/isaaclab_arena/utils/collision_mesh_store.py b/isaaclab_arena/utils/collision_mesh_store.py index c5862ee646..4625911e03 100644 --- a/isaaclab_arena/utils/collision_mesh_store.py +++ b/isaaclab_arena/utils/collision_mesh_store.py @@ -15,9 +15,8 @@ ``scripts/export_ready_pose_collision_meshes.py``. A miss falls back to extraction. Artifacts are named and validated by the source USD's stem, so a file exported on one machine loads -on another even though Arena composes the robot-on-stand USDs into a per-user cache directory. The -stem only has to be unique within a robot's folder, and a full export refuses to write a pair that -would collide there, since at a shared pose one robot would otherwise be served the other's mesh. +on another even though Arena composes the robot-on-stand USDs into a per-user cache directory. +``ROBOT_LIBRARY_FOLDERS`` maps that stem to the robot's published folder. """ from __future__ import annotations @@ -36,6 +35,35 @@ from isaaclab_arena.assets.asset_cache import get_arena_usd_cache_dir from isaaclab_arena.assets.nucleus import ARENA_NUCLEUS_DIR +ROBOT_LIBRARY_FOLDERS = { + "A2D_physics": "agibot_a2d", + "droid_franka_robotiq_on_stand_1.350": "droid", + "franka_panda_on_stand_0.875": "franka", + "g1_29dof_with_hand_rev_1_0": "g1", + "galbot_one_charlie": "galbot", + "GR1T2_fourier_hand_6dof": "gr1t2", + "kuka": "kuka", +} +"""Folder each robot publishes its collision mesh to, keyed by the stem of the USD Arena spawns. + +Robots differing only in action space share a USD and so one artifact. The robots Arena composes +onto a stand key by the composed stem, stand height included, because a robot at another stand +height is different geometry that nothing has published; it falls back to extraction. +""" + +ARENA_ROBOT_LIBRARY_DIR = f"{ARENA_NUCLEUS_DIR}/Arena/assets/robot_library" +"""Nucleus robot library, holding each robot's assets in a folder of its own.""" + +ROBOT_LIBRARY_DIR_ENV_VAR = "ISAACLAB_ARENA_ROBOT_LIBRARY_DIR" +"""Environment variable redirecting which robot library published artifacts are read from.""" + +CACHE_BUDGET_BYTES = 1 << 30 +"""Disk the local cache may occupy before its least recently used artifacts are dropped. + +A robot's mesh runs 1-9 MB, so this holds a few hundred: every robot Arena ships at several poses +each, while bounding what a run against throwaway USDs can leave behind. +""" + _MESH_PRIM_PATH = "/CollisionMesh" """Prim the merged mesh is authored at, also the artifact's default prim.""" @@ -54,12 +82,6 @@ _READY_POSE_SUFFIX = "_ready_pose.usd" """Suffix naming a published artifact, one per robot, at the pose that robot spawns in.""" -ROBOT_LIBRARY_DIR_ENV_VAR = "ISAACLAB_ARENA_ROBOT_LIBRARY_DIR" -"""Environment variable redirecting which robot library published artifacts are read from.""" - -ARENA_ROBOT_LIBRARY_DIR = f"{ARENA_NUCLEUS_DIR}/Arena/assets/robot_library" -"""Nucleus robot library, holding each robot's assets in a folder of its own.""" - _LOCAL_CACHE_DIR = "collision_mesh" """Local cache root under ``get_arena_usd_cache_dir()``, with one subfolder per robot when known.""" @@ -70,162 +92,148 @@ to stay a USD one for ``Usd.Stage.CreateNew``, so the prefix is what keeps the two apart. """ -CACHE_BUDGET_BYTES = 1 << 30 -"""Disk the local cache may occupy before its least recently used artifacts are dropped. -A robot's mesh runs 1-9 MB, so this holds a few hundred: every robot Arena ships at several poses -each, while bounding what a run against throwaway USDs can leave behind. -""" +def load_mesh( + source_usd_path: str, joint_pos: Mapping[str, float], scale: tuple[float, float, float] +) -> trimesh.Trimesh | None: + """Return the stored mesh for a robot USD at joint_pos scaled to scale, or None if unavailable. + Reads the local cache first and the published robot library second, copying a published hit into + the local cache on the way out. An artifact recording another asset or another pose is ignored + rather than trusted. -def local_collision_mesh_cache_dir() -> Path: - """Return the local collision-mesh cache root (``~/.cache/.../usd/collision_mesh``).""" - return get_arena_usd_cache_dir() / _LOCAL_CACHE_DIR + Args: + source_usd_path: Robot USD the mesh should describe. + joint_pos: Joint positions the mesh should be posed at. + scale: Per-axis scale to apply to the stored vertices. + """ + expected = {"expected_asset": _asset_key(source_usd_path), "expected_pose": _pose_key(joint_pos)} + cached = _cache_path(source_usd_path, joint_pos) + mesh = _read_mesh_usd(str(cached), **expected) + if mesh is not None: + # Mark the artifact used, so trimming evicts by last use rather than by write time. + with contextlib.suppress(OSError): + os.utime(cached) + return scaled_mesh(mesh, scale) -def scaled_mesh(mesh: trimesh.Trimesh, scale: tuple[float, float, float]) -> trimesh.Trimesh: - """Return mesh with vertices scaled per axis in its own frame. + published = _published_path(source_usd_path) + if published is None: + return None + mesh = _read_mesh_usd(published, **expected) + if mesh is None: + return None + # Copy the published artifact locally, so only the first process pays the Nucleus round trip. + with contextlib.suppress(OSError): + save_mesh(source_usd_path, joint_pos, mesh) + return scaled_mesh(mesh, scale) - Scaling in the root frame commutes with merging prims, so scaling a merged mesh matches scaling - each prim's vertices during extraction. - """ - if tuple(scale) == (1.0, 1.0, 1.0): - return mesh - # process=False keeps trimesh from merging vertices, which would make a scaled or reloaded mesh - # differ from the one extraction produced. - return trimesh.Trimesh( - vertices=mesh.vertices * np.asarray(scale, dtype=np.float64), faces=mesh.faces, process=False - ) +def save_mesh(source_usd_path: str, joint_pos: Mapping[str, float], mesh: trimesh.Trimesh) -> Path: + """Write a robot USD's unscaled mesh at joint_pos to the local cache and return its path. -def is_zero_pose(joint_pos: Mapping[str, float]) -> bool: - """Whether joint_pos poses every joint at zero, including by naming no joints at all. + Least recently used artifacts are dropped to keep the cache within ``CACHE_BUDGET_BYTES``. - Omitted joints are posed at zero, so an empty mapping is the zero pose rather than the asset's - authored configuration. + Args: + source_usd_path: Robot USD the mesh was extracted from, recorded for validation on load. + joint_pos: Joint positions the mesh was posed at, recorded for validation on load. + mesh: Merged mesh at unit scale, in the robot's default-prim frame. """ - return all(float(value) == 0.0 for value in joint_pos.values()) + out_path = _write_mesh_usd(_cache_path(source_usd_path, joint_pos), source_usd_path, joint_pos, mesh) + _trim_cache(_local_cache_dir()) + return out_path -def pose_key(joint_pos: Mapping[str, float]) -> str: - """Return a filename-safe key identifying a joint pose. +def export_ready_pose_mesh( + source_usd_path: str, joint_pos: Mapping[str, float], mesh: trimesh.Trimesh, out_dir: Path +) -> Path: + """Write a robot's ready-pose mesh into out_dir under its published name, for upload. - Keys are derived from the joint names and values as given, so two spellings of one pose, a regex - and the names it expands to, key differently and are extracted separately. + Args: + source_usd_path: Robot USD the mesh was extracted from, which must be a published robot. + joint_pos: The robot's configured joint positions, recorded for validation on load. + mesh: Merged mesh at unit scale, in the robot's default-prim frame. + out_dir: Staging directory whose layout mirrors ``ARENA_ROBOT_LIBRARY_DIR`` for upload as-is. """ - if is_zero_pose(joint_pos): - return _ZERO_POSE_KEY - canonical = ";".join(f"{name}={float(value) + 0.0:.9g}" for name, value in sorted(joint_pos.items())) - return hashlib.sha1(canonical.encode()).hexdigest()[:12] - - -def asset_key(source_usd_path: str) -> str: - """Return the identity a published artifact is named and validated by: the source USD's stem.""" - # The stem rather than the full path, so an artifact exported from a per-user cache directory - # still matches elsewhere. Stems already spell out the variant: droid_franka_robotiq_on_stand_1.350. - return Path(str(source_usd_path)).stem + relative_path = published_relative_path(source_usd_path) + assert ( + relative_path is not None + ), f"{_asset_key(source_usd_path)} has no entry in ROBOT_LIBRARY_FOLDERS, so there is no folder to publish it to" + return _write_mesh_usd(out_dir / relative_path, source_usd_path, joint_pos, mesh) -def ready_pose_artifact_name(source_usd_path: str) -> str: - """Return the filename a robot USD's ready-pose mesh is published under.""" - return f"{asset_key(source_usd_path)}{_READY_POSE_SUFFIX}" +def published_relative_path(source_usd_path: str) -> str | None: + """Return a robot USD's artifact path within the robot library, or None if it publishes none.""" + folder = ROBOT_LIBRARY_FOLDERS.get(_asset_key(source_usd_path)) + if folder is None: + return None + return f"{folder}/{_asset_key(source_usd_path)}{_READY_POSE_SUFFIX}" -def published_ready_pose_dir() -> str: - """Return the robot library that published ready-pose artifacts are read from. +def _published_path(source_usd_path: str) -> str | None: + """Return the full path a robot USD's ready-pose mesh is read from, or None if it publishes none. Defaults to the Arena Nucleus robot library and is redirected by ``ISAACLAB_ARENA_ROBOT_LIBRARY_DIR``, which lets an export be checked locally before upload. """ - return os.environ.get(ROBOT_LIBRARY_DIR_ENV_VAR) or ARENA_ROBOT_LIBRARY_DIR - - -def published_ready_pose_path(source_usd_path: str, library_folder: str) -> str: - """Return the full path a robot USD's ready-pose mesh is published at. - - The artifact sits in the robot's own library folder rather than beside its source USD, because a - robot that spawns a composed asset sources from a per-user cache path that is nobody else's. - """ - return f"{published_ready_pose_dir().rstrip('/')}/{library_folder}/{ready_pose_artifact_name(source_usd_path)}" + relative_path = published_relative_path(source_usd_path) + if relative_path is None: + return None + library_dir = os.environ.get(ROBOT_LIBRARY_DIR_ENV_VAR) or ARENA_ROBOT_LIBRARY_DIR + return f"{library_dir.rstrip('/')}/{relative_path}" -def mesh_cache_path(source_usd_path: str, joint_pos: Mapping[str, float], library_folder: str | None = None) -> Path: +def _cache_path(source_usd_path: str, joint_pos: Mapping[str, float]) -> Path: """Return the local cache path for a robot USD's mesh at joint_pos. - Layout is ``collision_mesh/{library_folder}/{filename}`` when ``library_folder`` is set (the - embodiment's ``robot_library_folder``), otherwise ``collision_mesh/{filename}`` for throwaway - fixtures that have no published robot folder. The full source path is hashed into the filename - so assets sharing a stem stay distinct on one machine even though they would share a published - name. + Published robots mirror their library folder, and everything else, such as a throwaway test + fixture, sits at the cache root. The full source path is hashed into the filename so assets + sharing a stem stay distinct on one machine even though they would share a published name. """ source_digest = hashlib.sha1(str(source_usd_path).encode()).hexdigest()[:12] - name = f"{asset_key(source_usd_path)}_{source_digest}_collision_{pose_key(joint_pos)}_pose.usd" - cache_root = local_collision_mesh_cache_dir() - if library_folder: - return cache_root / library_folder / name - return cache_root / name + name = f"{_asset_key(source_usd_path)}_{source_digest}_collision_{_pose_key(joint_pos)}_pose.usd" + folder = ROBOT_LIBRARY_FOLDERS.get(_asset_key(source_usd_path)) + return _local_cache_dir() / folder / name if folder else _local_cache_dir() / name -def load_mesh( - source_usd_path: str, - joint_pos: Mapping[str, float], - scale: tuple[float, float, float], - library_folder: str | None = None, -) -> trimesh.Trimesh | None: - """Return the stored mesh for a robot USD at joint_pos scaled to scale, or None if unavailable. +def _local_cache_dir() -> Path: + """Return the local collision-mesh cache root (``~/.cache/.../usd/collision_mesh``).""" + return get_arena_usd_cache_dir() / _LOCAL_CACHE_DIR - Reads the local cache first and the published robot library second, copying a published hit - into the local cache on the way out. An artifact recording another asset or another pose is - ignored rather than trusted. - Args: - source_usd_path: Robot USD the mesh should describe. - joint_pos: Joint positions the mesh should be posed at. - scale: Per-axis scale to apply to the stored vertices. - library_folder: Robot's folder in the published library (and local cache), or None to - skip the published lookup and use the flat local-cache path. - """ - cached = mesh_cache_path(source_usd_path, joint_pos, library_folder) - mesh = _read_mesh_usd(str(cached), expected_asset=asset_key(source_usd_path), expected_pose=pose_key(joint_pos)) - if mesh is not None: - # Mark the artifact used, so trimming evicts by last use rather than by write time. - with contextlib.suppress(OSError): - os.utime(cached) - return scaled_mesh(mesh, scale) +def _asset_key(source_usd_path: str) -> str: + """Return the identity a published artifact is named and validated by: the source USD's stem.""" + # The stem rather than the full path, so an artifact exported from a per-user cache directory + # still matches elsewhere. Stems already spell out the variant: droid_franka_robotiq_on_stand_1.350. + return Path(str(source_usd_path)).stem - if library_folder is None: - return None - published = published_ready_pose_path(source_usd_path, library_folder) - mesh = _read_mesh_usd(published, expected_asset=asset_key(source_usd_path), expected_pose=pose_key(joint_pos)) - if mesh is None: - return None - # Copy the published artifact locally, so only the first process pays the Nucleus round trip. - with contextlib.suppress(OSError): - save_mesh(source_usd_path, joint_pos, mesh, library_folder) - return scaled_mesh(mesh, scale) +def _pose_key(joint_pos: Mapping[str, float]) -> str: + """Return a filename-safe key identifying a joint pose. -def save_mesh( - source_usd_path: str, - joint_pos: Mapping[str, float], - mesh: trimesh.Trimesh, - library_folder: str | None = None, -) -> Path: - """Write a robot USD's unscaled mesh at joint_pos to the local cache and return its path. + Keys are derived from the joint names and values as given, so two spellings of one pose, a regex + and the names it expands to, key differently and are extracted separately. + """ + # Omitted joints are posed at zero, so an empty mapping is the zero pose, not the authored one. + if all(float(value) == 0.0 for value in joint_pos.values()): + return _ZERO_POSE_KEY + canonical = ";".join(f"{name}={float(value) + 0.0:.9g}" for name, value in sorted(joint_pos.items())) + return hashlib.sha1(canonical.encode()).hexdigest()[:12] - Least recently used artifacts are dropped to keep the cache within ``CACHE_BUDGET_BYTES``. - Args: - source_usd_path: Robot USD the mesh was extracted from, recorded for validation on load. - joint_pos: Joint positions the mesh was posed at, recorded for validation on load. - mesh: Merged mesh at unit scale, in the robot's default-prim frame. - library_folder: Robot's folder under the local cache, or None for a flat cache path. +def scaled_mesh(mesh: trimesh.Trimesh, scale: tuple[float, float, float]) -> trimesh.Trimesh: + """Return mesh with vertices scaled per axis in its own frame. + + Scaling in the root frame commutes with merging prims, so scaling a merged mesh matches scaling + each prim's vertices during extraction. """ - out_path = _write_mesh_usd( - mesh_cache_path(source_usd_path, joint_pos, library_folder), source_usd_path, joint_pos, mesh + if tuple(scale) == (1.0, 1.0, 1.0): + return mesh + # process=False keeps trimesh from merging vertices, which would make a scaled or reloaded mesh + # differ from the one extraction produced. + return trimesh.Trimesh( + vertices=mesh.vertices * np.asarray(scale, dtype=np.float64), faces=mesh.faces, process=False ) - _trim_cache(local_collision_mesh_cache_dir()) - return out_path def _trim_cache(cache_dir: Path) -> None: @@ -255,23 +263,6 @@ def _last_used(path: Path) -> float: return 0.0 -def export_ready_pose_mesh( - source_usd_path: str, joint_pos: Mapping[str, float], mesh: trimesh.Trimesh, out_dir: Path, library_folder: str -) -> Path: - """Write a robot's ready-pose mesh into out_dir under its published name, for upload. - - Args: - source_usd_path: Robot USD the mesh was extracted from. - joint_pos: The robot's configured joint positions, recorded for validation on load. - mesh: Merged mesh at unit scale, in the robot's default-prim frame. - out_dir: Staging directory whose layout mirrors ``ARENA_ROBOT_LIBRARY_DIR`` for upload as-is. - library_folder: Robot's folder in the published library, created under out_dir. - """ - artifact_path = out_dir / library_folder / ready_pose_artifact_name(source_usd_path) - artifact_path.parent.mkdir(parents=True, exist_ok=True) - return _write_mesh_usd(artifact_path, source_usd_path, joint_pos, mesh) - - def _write_mesh_usd( out_path: Path, source_usd_path: str, joint_pos: Mapping[str, float], mesh: trimesh.Trimesh ) -> Path: @@ -289,8 +280,8 @@ def _write_mesh_usd( mesh_prim.GetPointsAttr().Set(np.asarray(mesh.vertices, dtype=np.float32)) mesh_prim.GetFaceVertexCountsAttr().Set([3] * len(faces)) mesh_prim.GetFaceVertexIndicesAttr().Set(faces.reshape(-1)) - mesh_prim.GetPrim().SetCustomDataByKey(_ASSET_KEY, asset_key(source_usd_path)) - mesh_prim.GetPrim().SetCustomDataByKey(_POSE_KEY, pose_key(joint_pos)) + mesh_prim.GetPrim().SetCustomDataByKey(_ASSET_KEY, _asset_key(source_usd_path)) + mesh_prim.GetPrim().SetCustomDataByKey(_POSE_KEY, _pose_key(joint_pos)) mesh_prim.GetPrim().SetCustomDataByKey(_SOURCE_KEY, str(source_usd_path)) assert stage.GetRootLayer().Save(), f"failed to save collision mesh to {tmp_path}" os.replace(tmp_path, out_path) diff --git a/isaaclab_arena/utils/isaac_sim_debug_draw.py b/isaaclab_arena/utils/isaac_sim_debug_draw.py index 85b44ea08b..95d808a5c2 100644 --- a/isaaclab_arena/utils/isaac_sim_debug_draw.py +++ b/isaaclab_arena/utils/isaac_sim_debug_draw.py @@ -79,40 +79,6 @@ def draw_object_bboxes( else: print(f"Skipping {obj.name}: no bbox coordinates") - def draw_line_segments( - self, - start_points: list[tuple[float, float, float]], - end_points: list[tuple[float, float, float]], - color: tuple[float, float, float, float] = DEFAULT_COLOR, - thickness: float = 3.0, - ) -> None: - """Draw one line per (start, end) pair, for wireframes ``draw_bbox`` cannot express. - - Args: - start_points: Segment start positions in world coordinates. - end_points: Segment end positions, matching start_points element-wise. - color: RGBA colour applied to every segment. - thickness: Line thickness in pixels. - """ - assert len(start_points) == len(end_points), "each start point needs a matching end point" - count = len(start_points) - self._draw.draw_lines(start_points, end_points, [color] * count, [thickness] * count) - - def draw_points( - self, - points: list[tuple[float, float, float]], - color: tuple[float, float, float, float] = DEFAULT_COLOR, - size: float = 5.0, - ) -> None: - """Draw a point cloud, for geometry such as mesh vertices. - - Args: - points: Positions in world coordinates. - color: RGBA colour applied to every point. - size: Point size in pixels. - """ - self._draw.draw_points(points, [color] * len(points), [size] * len(points)) - def clear(self) -> None: """Clear all debug drawings.""" self._draw.clear_lines() diff --git a/isaaclab_arena/utils/usd_helpers.py b/isaaclab_arena/utils/usd_helpers.py index 4e388ab2fb..8349e6219c 100644 --- a/isaaclab_arena/utils/usd_helpers.py +++ b/isaaclab_arena/utils/usd_helpers.py @@ -473,7 +473,6 @@ def extract_trimesh_from_usd_at_joint_pos( usd_path: str, joint_pos: Mapping[str, float], scale: tuple[float, float, float] = (1.0, 1.0, 1.0), - library_folder: str | None = None, ) -> trimesh.Trimesh: """Extract an articulation's mesh posed at joint_pos, in its default prim's local frame. @@ -485,14 +484,11 @@ def extract_trimesh_from_usd_at_joint_pos( usd_path: Path to the articulation's .usd/.usda/.usdc file. joint_pos: Joint positions keyed by exact joint name or Isaac Lab regex, revolute in radians. scale: Spawn-time scale passed to ``UsdFileCfg``. - library_folder: Robot's folder in the published library, or None to skip the published lookup. Returns: Combined trimesh in the scaled default-prim frame. """ - return _extract_trimesh_from_usd_at_joint_pos( - usd_path, tuple(sorted(joint_pos.items())), tuple(scale), library_folder - ) + return _extract_trimesh_from_usd_at_joint_pos(usd_path, tuple(sorted(joint_pos.items())), tuple(scale)) # NOTE(zihaox, 2026-07-28): Cache here rather than on the asset. Isaac Lab reaches assets through @@ -503,11 +499,10 @@ def _extract_trimesh_from_usd_at_joint_pos( usd_path: str, joint_pos_items: tuple[tuple[str, float], ...], scale: tuple[float, float, float], - library_folder: str | None, ) -> trimesh.Trimesh: """Cacheable body of ``extract_trimesh_from_usd_at_joint_pos``, keyed by hashable arguments.""" joint_pos = dict(joint_pos_items) - stored = load_mesh(usd_path, joint_pos, scale, library_folder) + stored = load_mesh(usd_path, joint_pos, scale) if stored is not None: return stored @@ -520,7 +515,7 @@ def _extract_trimesh_from_usd_at_joint_pos( # Store at unit scale so one artifact serves every spawn scale of the asset. unscaled = extract_trimesh_from_prim(stage, default_prim_path, (1.0, 1.0, 1.0), prim_world_deltas=deltas) - save_mesh(usd_path, joint_pos, unscaled, library_folder=library_folder) + save_mesh(usd_path, joint_pos, unscaled) return scaled_mesh(unscaled, scale) diff --git a/isaaclab_arena_examples/relations/visualize_embodiment_placement_geometry.py b/isaaclab_arena_examples/relations/visualize_embodiment_placement_geometry.py deleted file mode 100644 index e045b39689..0000000000 --- a/isaaclab_arena_examples/relations/visualize_embodiment_placement_geometry.py +++ /dev/null @@ -1,174 +0,0 @@ -# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Overlay an embodiment's placement bounding box and collision mesh on the robot as spawned. - -Run with the GUI and inspect by eye: the green wireframe is ``EmbodimentBase.get_bounding_box()`` and -the magenta points are ``get_collision_mesh()``, both posed at the robot's configured initial joint -positions and drawn through the spawned robot prim's transform. At ``num_steps=0`` the arm has not yet -sagged under gravity, so the overlay should hug the robot; stepping lets the joints settle a few -milliradians away and the overlay drifts by millimetres. - -Passing ``spec_yaml`` builds a full env graph instead of an empty ground plane, which shows the same -geometry where the relation solver actually placed the robot — the overlay then also reveals whether a -yaw relation rotated the placement geometry along with the robot. The spec YAML may also be given as -the first command-line argument, alongside the usual Isaac Lab flags such as ``--headless``. -""" - -# %% -from __future__ import annotations - -from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser -from isaaclab_arena.utils.isaaclab_utils.simulation_app import get_app_launcher - -# pyright: reportArgumentType=false, reportCallIssue=false, reportAttributeAccessIssue=false - - -_args, _positional_args = get_isaaclab_arena_cli_parser().parse_known_args() -if not _args.headless: - # A bare AppLauncher(headless=False) starts Kit without a window: Isaac Lab 3.0 only builds the - # viewport when a visualizer is requested, so ask for "kit" the way Arena's test harness does. - _args.visualizer = ["kit"] -print(f"Launching simulation app (headless={_args.headless})") -simulation_app = get_app_launcher(_args).app - -# %% - -MESH_POINT_BUDGET = 7000 -"""Mesh vertices are subsampled to this many debug points to keep the viewport responsive.""" - -BBOX_COLOR = (0.0, 1.0, 0.2, 1.0) -MESH_COLOR = (1.0, 0.2, 0.9, 0.9) - - -def _build_environment(embodiment_name: str, spec_yaml: str | None): - """Return an environment and its embodiment, either from an env graph spec or an empty scene.""" - if spec_yaml is not None: - from isaaclab_arena.environment_spec.arena_env_graph_conversion_utils import build_arena_env_from_graph_spec - from isaaclab_arena.environment_spec.arena_env_graph_spec import ArenaEnvGraphSpec - - environment = build_arena_env_from_graph_spec(ArenaEnvGraphSpec.from_yaml(spec_yaml)) - return environment, environment.embodiment - - from isaaclab_arena.assets.registries import AssetRegistry - from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment - from isaaclab_arena.scene.scene import Scene - - asset_registry = AssetRegistry() - embodiment = asset_registry.get_asset_by_name(embodiment_name)() - scene_assets = [asset_registry.get_asset_by_name(name)() for name in ("ground_plane", "light")] - environment = IsaacLabArenaEnvironment( - name="visualize_embodiment_placement_geometry", - embodiment=embodiment, - scene=Scene(assets=scene_assets), - ) - return environment, embodiment - - -def visualize_embodiment_placement_geometry( - embodiment_name: str = "droid_abs_joint_pos", - spec_yaml: str | None = None, - num_steps: int = 0, - mesh_point_budget: int = MESH_POINT_BUDGET, - hold_open: bool = True, -): - """Spawn an embodiment and draw its placement bbox and collision mesh over it. - - Args: - embodiment_name: Registry name of the embodiment to inspect. Ignored when ``spec_yaml`` is set, - which carries its own embodiment. - spec_yaml: Path to an env graph spec YAML to build instead of an empty ground-plane scene. - num_steps: Simulation steps to run before drawing. Zero keeps the joints exactly at their - configured positions, which is what the placement geometry is posed at. - mesh_point_budget: Approximate number of mesh vertices to draw. - hold_open: Keep rendering until the window closes. Set False to exit after drawing, which is - what a headless check of the spec wants. - """ - import numpy as np - import torch - - import omni.usd - from pxr import Usd, UsdGeom - - from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder, ArenaEnvBuilderCfg - from isaaclab_arena.utils.isaac_sim_debug_draw import IsaacSimDebugDraw - - environment, embodiment = _build_environment(embodiment_name, spec_yaml) - - # Placement geometry, in the embodiment's local (default-prim) frame. - bbox = embodiment.get_bounding_box() - mesh = embodiment.get_collision_mesh() - local_min = np.asarray(bbox.min_point.numpy()[0], dtype=np.float64) - local_max = np.asarray(bbox.max_point.numpy()[0], dtype=np.float64) - print(f"local bbox min = {np.round(local_min, 4)}") - print(f"local bbox max = {np.round(local_max, 4)}") - print(f"local bbox size = {np.round(local_max - local_min, 4)}") - print(f"local mesh size = {np.round(mesh.extents, 4)} ({len(mesh.vertices)} verts)") - # A positive difference is the analytic (non-mesh) geometry the bbox covers but the mesh cannot. - print(f"bbox minus mesh = {np.round((local_max - local_min) - mesh.extents, 4)}") - - env = ArenaEnvBuilder(environment, ArenaEnvBuilderCfg()).make_registered() - env.reset() - - for _ in range(num_steps): - with torch.inference_mode(): - env.step(torch.zeros(env.action_space.shape, device=env.unwrapped.device)) - - # The local frame is the spawned robot prim's frame, so draw through its world transform. - stage = omni.usd.get_context().get_stage() - robot_prim_path = env.unwrapped.scene["robot"].cfg.prim_path.replace("env_.*", "env_0") - robot_prim = stage.GetPrimAtPath(robot_prim_path) - assert robot_prim, f"no robot prim at {robot_prim_path}" - robot_world = np.array( - UsdGeom.Xformable(robot_prim).ComputeLocalToWorldTransform(Usd.TimeCode.Default()), dtype=np.float64 - ) - - def to_world(points: np.ndarray) -> np.ndarray: - return (np.hstack([points, np.ones((len(points), 1))]) @ robot_world)[:, :3] - - corners = to_world( - np.array([ - [x, y, z] - for x in (local_min[0], local_max[0]) - for y in (local_min[1], local_max[1]) - for z in (local_min[2], local_max[2]) - ]) - ) - # Corner order above is (x, y, z) bit-major, so edges join indices differing in exactly one bit. - edges = [(i, i ^ bit) for i in range(8) for bit in (1, 2, 4) if i < (i ^ bit)] - - debug_draw = IsaacSimDebugDraw() - debug_draw.clear() - debug_draw.draw_line_segments( - [tuple(corners[i]) for i, _ in edges], - [tuple(corners[j]) for _, j in edges], - color=BBOX_COLOR, - thickness=5.0, - ) - - vertices = np.asarray(mesh.vertices, dtype=np.float64) - if len(vertices) > mesh_point_budget: - vertices = vertices[:: len(vertices) // mesh_point_budget + 1] - world_vertices = to_world(vertices) - debug_draw.draw_points([tuple(point) for point in world_vertices], color=MESH_COLOR) - print(f"\ndrew {len(edges)} bbox edges (green) and {len(world_vertices)} mesh points (magenta)") - print(f"robot world position = {np.round(robot_world[3, :3], 4)}") - if not hold_open: - return - print("Inspect the overlay in the viewport. Close the window to exit.") - - # Render without stepping so the viewport stays interactive while the joints stay exactly where - # the placement geometry was posed. Debug draws persist across frames, so one draw is enough. - while simulation_app.is_running(): - env.unwrapped.sim.render() - - -# %% -visualize_embodiment_placement_geometry( - spec_yaml=_positional_args[0] if _positional_args else None, - hold_open=not _args.headless, -) - -# %% From d964de691c3b3547be0d325c2a6ae2558052793c Mon Sep 17 00:00:00 2001 From: zhx06 Date: Tue, 28 Jul 2026 23:29:01 -0700 Subject: [PATCH 3/7] fix franka Signed-off-by: zhx06 --- isaaclab_arena/embodiments/droid/droid.py | 42 +++------- isaaclab_arena/embodiments/franka/franka.py | 49 ++++++----- .../tests/test_embodiment_collision_mesh.py | 81 +++++++++++++++++-- .../cube_goal_pose_environment.py | 14 +++- .../franka_put_and_close_door_environment.py | 11 ++- .../sorting_environment.py | 14 +++- 6 files changed, 138 insertions(+), 73 deletions(-) diff --git a/isaaclab_arena/embodiments/droid/droid.py b/isaaclab_arena/embodiments/droid/droid.py index a0b1f87479..f08747617d 100644 --- a/isaaclab_arena/embodiments/droid/droid.py +++ b/isaaclab_arena/embodiments/droid/droid.py @@ -6,6 +6,7 @@ import torch from abc import ABC +from collections.abc import Mapping import isaaclab.envs.mdp as mdp_isaac_lab import isaaclab.sim as sim_utils @@ -78,7 +79,6 @@ def __init__( self, enable_cameras: bool = False, initial_pose: Pose | None = None, - initial_joint_pose: list[float] | None = None, concatenate_observation_terms: bool = False, arm_mode: ArmMode | None = None, stand_height_m: float = _DROID_STAND_PRIM.stand_default_height, @@ -98,8 +98,6 @@ def __init__( self.camera_config = DroidCameraCfg() self.observation_config = DroidObservationsCfg() self.event_config = DroidEventCfg() - if initial_joint_pose is not None: - self.set_initial_joint_pose(initial_joint_pose) self.reward_config = None self.mimic_env = None self.add_camera_variations(self.camera_config) @@ -113,8 +111,15 @@ def get_bounding_box(self) -> AxisAlignedBoundingBox: prim_path = _DROID_ROBOT_PRIM.stand_prim_path if self.placement_bbox_stand_only else None return super().get_bounding_box(prim_path=prim_path) - def set_initial_joint_pose(self, initial_joint_pose: list[float]) -> None: - self.event_config.init_franka_arm_pose.params["default_pose"] = initial_joint_pose + def set_initial_joint_pose(self, initial_joint_pose: Mapping[str, float]) -> None: + """Spawn and reset the arm at ``initial_joint_pose``, keyed by joint name or Isaac Lab regex. + + Placement geometry is derived from the same spawn state, so it follows the new pose. + """ + assert self.scene_config is not None, "scene_config must be populated before setting the joint pose" + robot = self.scene_config.robot + assert robot is not None, "scene_config.robot must be populated before setting the joint pose" + robot.init_state = robot.init_state.replace(joint_pos=dict(initial_joint_pose)) def get_ee_frame_name(self, arm_mode: ArmMode) -> str: return "ee_frame" @@ -134,7 +139,6 @@ def __init__( self, enable_cameras: bool = False, initial_pose: Pose | None = None, - initial_joint_pose: list[float] | None = None, concatenate_observation_terms: bool = False, arm_mode: ArmMode | None = None, stand_height_m: float = _DROID_STAND_PRIM.stand_default_height, @@ -143,7 +147,6 @@ def __init__( super().__init__( enable_cameras, initial_pose, - initial_joint_pose, concatenate_observation_terms, arm_mode, stand_height_m, @@ -163,7 +166,6 @@ def __init__( self, enable_cameras: bool = False, initial_pose: Pose | None = None, - initial_joint_pose: list[float] | None = None, concatenate_observation_terms: bool = False, arm_mode: ArmMode | None = None, stand_height_m: float = _DROID_STAND_PRIM.stand_default_height, @@ -172,7 +174,6 @@ def __init__( super().__init__( enable_cameras, initial_pose, - initial_joint_pose, concatenate_observation_terms, arm_mode, stand_height_m, @@ -193,7 +194,6 @@ def __init__( self, enable_cameras: bool = False, initial_pose: Pose | None = None, - initial_joint_pose: list[float] | None = None, concatenate_observation_terms: bool = False, arm_mode: ArmMode | None = None, stand_height_m: float = _DROID_STAND_PRIM.stand_default_height, @@ -202,7 +202,6 @@ def __init__( super().__init__( enable_cameras, initial_pose, - initial_joint_pose, concatenate_observation_terms, arm_mode, stand_height_m, @@ -409,27 +408,6 @@ def __post_init__(self): class DroidEventCfg: """Configuration for Franka.""" - init_franka_arm_pose = EventTerm( - func=franka_stack_events.set_default_joint_pose, - mode="reset", - params={ - "default_pose": [ - 0.0, # panda_joint1 - -1 / 5 * torch.pi, # panda_joint2 - 0.0, # panda_joint3 - -4 / 5 * torch.pi, # panda_joint4 - 0.0, # panda_joint5 - 3 / 5 * torch.pi, # panda_joint6 - 0.0, # panda_joint7 - 0.0, # finger_joint - 0.0, # right_outer_knuckle_joint - 0.0, # left_inner_finger_joint - 0.0, # right_inner_finger_joint - 0.0, # left_inner_finger_knuckle_joint - 0.0, # right_inner_finger_knuckle_joint - ], - }, - ) randomize_franka_joint_state = EventTerm( func=franka_stack_events.randomize_joint_by_gaussian_offset, mode="reset", diff --git a/isaaclab_arena/embodiments/franka/franka.py b/isaaclab_arena/embodiments/franka/franka.py index 8ff60fd89d..3f8f084015 100644 --- a/isaaclab_arena/embodiments/franka/franka.py +++ b/isaaclab_arena/embodiments/franka/franka.py @@ -5,7 +5,7 @@ import torch -from collections.abc import Sequence +from collections.abc import Mapping, Sequence import isaaclab.envs.mdp as mdp_isaac_lab import isaaclab.sim as sim_utils @@ -73,14 +73,11 @@ def __init__( self, enable_cameras: bool = False, initial_pose: Pose | None = None, - initial_joint_pose: list[float] | None = None, concatenate_observation_terms: bool = False, arm_mode: ArmMode | None = None, ): super().__init__(enable_cameras, initial_pose, concatenate_observation_terms, arm_mode) self.event_config = FrankaEventCfg() - if initial_joint_pose is not None: - self.set_initial_joint_pose(initial_joint_pose) self.reward_config = FrankaRewardsCfg() self.mimic_env = FrankaMimicEnv self.camera_config = FrankaCameraCfg() @@ -89,8 +86,16 @@ def __init__( self.observation_config.policy.concatenate_terms = self.concatenate_observation_terms self.add_camera_variations(self.camera_config) - def set_initial_joint_pose(self, initial_joint_pose: list[float]) -> None: - self.event_config.init_franka_arm_pose.params["default_pose"] = initial_joint_pose + def set_initial_joint_pose(self, initial_joint_pose: Mapping[str, float]) -> None: + """Spawn and reset the arm at ``initial_joint_pose``, keyed by joint name or Isaac Lab regex. + + Call after construction, once :attr:`scene_config` holds the robot. Placement geometry is + derived from the same spawn state, so it follows the new pose. + """ + assert self.scene_config is not None, "scene_config must be populated before setting the joint pose" + robot = self.scene_config.robot + assert robot is not None, "scene_config.robot must be populated before setting the joint pose" + robot.init_state = robot.init_state.replace(joint_pos=dict(initial_joint_pose)) def get_ee_frame_name(self, arm_mode: ArmMode) -> str: return "ee_frame" @@ -107,14 +112,12 @@ def __init__( self, enable_cameras: bool = False, initial_pose: Pose | None = None, - initial_joint_pose: list[float] | None = None, concatenate_observation_terms: bool = False, arm_mode: ArmMode | None = None, ): super().__init__( enable_cameras=enable_cameras, initial_pose=initial_pose, - initial_joint_pose=initial_joint_pose, concatenate_observation_terms=concatenate_observation_terms, arm_mode=arm_mode, ) @@ -160,14 +163,12 @@ def __init__( self, enable_cameras: bool = False, initial_pose: Pose | None = None, - initial_joint_pose: list[float] | None = None, concatenate_observation_terms: bool = False, arm_mode: ArmMode | None = None, ): super().__init__( enable_cameras=enable_cameras, initial_pose=initial_pose, - initial_joint_pose=initial_joint_pose, concatenate_observation_terms=concatenate_observation_terms, arm_mode=arm_mode, ) @@ -186,7 +187,15 @@ class FrankaJointPosActionsCfg: asset_name="robot", joint_names=["panda_joint.*"], scale=0.5, - use_default_offset=True, + # Actions are displacements from Isaac Lab's default Franka pose, which is the zero point + # trained policies were fitted against. Stated rather than left to ``use_default_offset``, + # which reads the spawn state and so would move the zero point whenever that pose changes. + use_default_offset=False, + offset={ + name: value + for name, value in FRANKA_PANDA_CFG.init_state.joint_pos.items() + if name.startswith("panda_joint") + }, ) gripper_action: ActionTermCfg = BinaryJointPositionActionCfg( @@ -293,24 +302,13 @@ def __post_init__(self): "panda_joint7": 0.785, "panda_finger_joint.*": 0.0400, } -"""The arm pose the Franka spawns and resets in, spelled by name for the spawn state. - -``init_franka_arm_pose`` below repeats it positionally, as Isaac Lab's reset event assigns joints by -index. ``test_spawn_pose_matches_the_reset_pose`` holds the two together. -""" +"""The arm pose the Franka spawns and resets in, overridable via ``set_initial_joint_pose``.""" @configclass class FrankaEventCfg: """Configuration for Franka.""" - init_franka_arm_pose = EventTerm( - func=franka_stack_events.set_default_joint_pose, - mode="reset", - params={ - "default_pose": [0.0, -0.785, -0.1107, -1.1775, 0.0, 0.785, 0.785, 0.0400, 0.0400], - }, - ) randomize_franka_joint_state = EventTerm( func=franka_stack_events.randomize_joint_by_gaussian_offset, mode="reset", @@ -476,9 +474,8 @@ def get_object_poses(self, env_ids: Sequence[int] | None = None): def _franka_robot_cfg_on_stand(robot_cfg: ArticulationCfg) -> ArticulationCfg: """Copy ``robot_cfg`` onto ``{ENV_REGEX_NS}/Robot`` with the composed on-stand USD.""" cfg = robot_cfg.replace(prim_path="{ENV_REGEX_NS}/Robot") - # Spawn where ``init_franka_arm_pose`` resets the arm to, rather than at Isaac Lab's own pose. - # The two described different arm configurations, which left placement geometry, derived from - # the spawn state, describing an arm the scene never contains. + # Arena reaches for objects on a table, so it spawns at its own ready pose rather than at the + # pose Isaac Lab ships, which folds the elbow back. cfg.init_state = cfg.init_state.replace(joint_pos=_FRANKA_READY_POSE) cfg.spawn.usd_path = compose_on_stand_usd( _FRANKA_ROBOT_PRIM, diff --git a/isaaclab_arena/tests/test_embodiment_collision_mesh.py b/isaaclab_arena/tests/test_embodiment_collision_mesh.py index 5d0d60ca9d..fdacf018f6 100644 --- a/isaaclab_arena/tests/test_embodiment_collision_mesh.py +++ b/isaaclab_arena/tests/test_embodiment_collision_mesh.py @@ -65,8 +65,8 @@ def _noop(env, env_ids, embodiment): def _test_spawn_pose_matches_the_reset_pose(simulation_app) -> bool: """The spawn joint positions placement geometry is posed at are the ones a reset drives to. - Both robots also reach their arm pose through an event that assigns joints positionally, which - cannot be read off the config; the articulation supplies the joint order here instead. + Also covers ``set_initial_joint_pose``, which environments use to pose the arm for their scene + and which has to move the spawn state placement geometry reads, not just the reset. """ import torch @@ -79,12 +79,32 @@ def _test_spawn_pose_matches_the_reset_pose(simulation_app) -> bool: from isaaclab_arena.scene.scene import Scene from isaaclab_arena.utils.usd_articulation import resolve_joint_pos_patterns - for embodiment_class in (FrankaIKEmbodiment, DroidAbsoluteJointPositionEmbodiment): + # The pose isaaclab_arena_environments/cube_goal_pose_environment.py reaches from. + overridden_pose = { + "panda_joint1": 0.0444, + "panda_joint2": -0.1894, + "panda_joint3": -0.1107, + "panda_joint4": -2.5148, + "panda_joint5": 0.0044, + "panda_joint6": 2.3775, + "panda_joint7": 0.6952, + "panda_finger_joint.*": 0.0400, + } + + for embodiment_class, override in ( + (FrankaIKEmbodiment, None), + (DroidAbsoluteJointPositionEmbodiment, None), + (FrankaIKEmbodiment, overridden_pose), + ): env = None try: embodiment = embodiment_class() + if override is not None: + embodiment.set_initial_joint_pose(override) environment = IsaacLabArenaEnvironment( - name=f"spawn_pose_{embodiment_class.__name__}", embodiment=embodiment, scene=Scene(assets=[]) + name=f"spawn_pose_{embodiment_class.__name__}_{override is not None}", + embodiment=embodiment, + scene=Scene(assets=[]), ) builder_cfg = arena_env_builder_cfg_from_argparse(get_isaaclab_arena_cli_parser().parse_args([])) env = ArenaEnvBuilder(environment, builder_cfg).make_registered() @@ -93,13 +113,21 @@ def _test_spawn_pose_matches_the_reset_pose(simulation_app) -> bool: robot = env.unwrapped.scene["robot"] joint_names = list(robot.joint_names) spawn_pose = resolve_joint_pos_patterns(joint_names, embodiment.get_placement_geometry_source().joint_pos) + # Both sides below are read back from the spawn state, so pin the override against what + # was asked for; otherwise a setter that quietly dropped it would still agree with itself. + if override is not None: + assert spawn_pose == resolve_joint_pos_patterns(joint_names, override), ( + "set_initial_joint_pose did not reach the spawn state placement geometry reads:" + f" {spawn_pose} != {resolve_joint_pos_patterns(joint_names, override)}" + ) # Compare against the defaults a reset restores rather than measured positions, which # carry gravity sag and the joint randomisation event on top. reset_to = torch.as_tensor(robot.data.default_joint_pos)[0].cpu() for index, joint_name in enumerate(joint_names): assert abs(spawn_pose.get(joint_name, 0.0) - float(reset_to[index])) < 1e-6, ( - f"{embodiment_class.__name__} poses placement geometry with {joint_name} at" - f" {spawn_pose.get(joint_name, 0.0)}, but resets it to {float(reset_to[index])}" + f"{embodiment_class.__name__}" + f"{' with an overridden pose' if override is not None else ''} poses placement geometry with" + f" {joint_name} at {spawn_pose.get(joint_name, 0.0)}, but resets it to {float(reset_to[index])}" ) except Exception as e: @@ -114,6 +142,40 @@ def _test_spawn_pose_matches_the_reset_pose(simulation_app) -> bool: return True +def _test_joint_position_action_offset_is_pinned(simulation_app) -> bool: + """Reposing the arm leaves the joint-position action's zero point where policies were fitted. + + The zero point is a displacement origin trained policies depend on, so it is spelled out here + rather than read from the config: this fails if the spawn pose starts feeding it again, or if an + Isaac Lab bump moves the default pose it is derived from. + """ + from isaaclab_arena.embodiments.franka.franka import FrankaJointPosEmbodiment + + trained_against = { + "panda_joint1": 0.0, + "panda_joint2": -0.569, + "panda_joint3": 0.0, + "panda_joint4": -2.810, + "panda_joint5": 0.0, + "panda_joint6": 3.037, + "panda_joint7": 0.741, + } + + try: + embodiment = FrankaJointPosEmbodiment() + embodiment.set_initial_joint_pose({"panda_joint.*": 0.0, "panda_finger_joint.*": 0.0}) + arm_action = embodiment.action_config.arm_action + assert not arm_action.use_default_offset, "the action zero point would follow the spawn pose" + assert arm_action.offset == trained_against, f"action zero point moved to {arm_action.offset}" + + except Exception as e: + print(f"Error: {e}") + traceback.print_exc() + return False + + return True + + def test_embodiment_provides_robot_collision_mesh(): """Pytest entry point for the embodiment collision-mesh test.""" result = run_function_with_persistent_simulation_app(_test_embodiment_provides_robot_collision_mesh, headless=True) @@ -126,6 +188,13 @@ def test_spawn_pose_matches_the_reset_pose(): assert result, f"Test {test_spawn_pose_matches_the_reset_pose.__name__} failed" +def test_joint_position_action_offset_is_pinned(): + """Pytest entry point for the joint-position action offset test.""" + result = run_simulation_app_function(_test_joint_position_action_offset_is_pinned, headless=True) + assert result, f"Test {test_joint_position_action_offset_is_pinned.__name__} failed" + + if __name__ == "__main__": test_embodiment_provides_robot_collision_mesh() test_spawn_pose_matches_the_reset_pose() + test_joint_position_action_offset_is_pinned() diff --git a/isaaclab_arena_environments/cube_goal_pose_environment.py b/isaaclab_arena_environments/cube_goal_pose_environment.py index 2648e7c602..427e484f17 100644 --- a/isaaclab_arena_environments/cube_goal_pose_environment.py +++ b/isaaclab_arena_environments/cube_goal_pose_environment.py @@ -59,10 +59,16 @@ def build(self, cfg: CubeGoalPoseEnvironmentCfg) -> IsaacLabArenaEnvironment: rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) - # order: [panda_joint1, panda_joint2, panda_joint3, panda_joint4, panda_joint5, panda_joint6, panda_joint7, panda_finger_joint1, panda_finger_joint2] - embodiment.set_initial_joint_pose( - initial_joint_pose=[0.0444, -0.1894, -0.1107, -2.5148, 0.0044, 2.3775, 0.6952, 0.0400, 0.0400] - ) + embodiment.set_initial_joint_pose({ + "panda_joint1": 0.0444, + "panda_joint2": -0.1894, + "panda_joint3": -0.1107, + "panda_joint4": -2.5148, + "panda_joint5": 0.0044, + "panda_joint6": 2.3775, + "panda_joint7": 0.6952, + "panda_finger_joint.*": 0.0400, + }) if cfg.teleop_device is not None: teleop_device = self.device_registry.get_device_by_name(cfg.teleop_device)() diff --git a/isaaclab_arena_environments/franka_put_and_close_door_environment.py b/isaaclab_arena_environments/franka_put_and_close_door_environment.py index c717a963b1..825df002c7 100644 --- a/isaaclab_arena_environments/franka_put_and_close_door_environment.py +++ b/isaaclab_arena_environments/franka_put_and_close_door_environment.py @@ -86,7 +86,16 @@ def build(self, cfg: FrankaPutAndCloseDoorEnvironmentCfg) -> IsaacLabArenaEnviro if cfg.embodiment == "franka_ik": # Set Franka arm pose for kitchen setup - embodiment.set_initial_joint_pose([0.0, -1.309, 0.0, -2.793, 0.0, 3.037, 0.740, 0.04, 0.04]) + embodiment.set_initial_joint_pose({ + "panda_joint1": 0.0, + "panda_joint2": -1.309, + "panda_joint3": 0.0, + "panda_joint4": -2.793, + "panda_joint5": 0.0, + "panda_joint6": 3.037, + "panda_joint7": 0.740, + "panda_finger_joint.*": 0.04, + }) # Create destination reference destination_ref = ObjectReference( diff --git a/isaaclab_arena_environments/sorting_environment.py b/isaaclab_arena_environments/sorting_environment.py index fd264e699a..c12c096195 100644 --- a/isaaclab_arena_environments/sorting_environment.py +++ b/isaaclab_arena_environments/sorting_environment.py @@ -67,10 +67,16 @@ def build(self, cfg: TableTopSortCubesEnvironmentCfg) -> IsaacLabArenaEnvironmen ) ) - # order: [panda_joint1, panda_joint2, panda_joint3, panda_joint4, panda_joint5, panda_joint6, panda_joint7, panda_finger_joint1, panda_finger_joint2] - embodiment.set_initial_joint_pose( - initial_joint_pose=[0.0444, -0.1894, -0.1107, -2.5148, 0.0044, 2.3775, 0.6952, 0.0400, 0.0400] - ) + embodiment.set_initial_joint_pose({ + "panda_joint1": 0.0444, + "panda_joint2": -0.1894, + "panda_joint3": -0.1107, + "panda_joint4": -2.5148, + "panda_joint5": 0.0044, + "panda_joint6": 2.3775, + "panda_joint7": 0.6952, + "panda_finger_joint.*": 0.0400, + }) else: raise NotImplementedError From e7f1cab08d1d2d0ddedcfff684456b1ac350a763 Mon Sep 17 00:00:00 2001 From: zhx06 Date: Thu, 30 Jul 2026 10:52:47 -0700 Subject: [PATCH 4/7] remove caching, keep joint_position bbox/meshes Signed-off-by: zhx06 --- isaaclab_arena/assets/asset_cache.py | 8 - .../embodiments/robot_on_stand_utils.py | 4 +- .../export_ready_pose_collision_meshes.py | 171 -------- .../tests/test_collision_mesh_store.py | 388 ------------------ isaaclab_arena/tests/test_usd_articulation.py | 5 - isaaclab_arena/utils/collision_mesh_store.py | 325 --------------- isaaclab_arena/utils/usd_helpers.py | 15 +- 7 files changed, 5 insertions(+), 911 deletions(-) delete mode 100644 isaaclab_arena/scripts/export_ready_pose_collision_meshes.py delete mode 100644 isaaclab_arena/tests/test_collision_mesh_store.py delete mode 100644 isaaclab_arena/utils/collision_mesh_store.py diff --git a/isaaclab_arena/assets/asset_cache.py b/isaaclab_arena/assets/asset_cache.py index 1f93ae1c7f..05628b4c56 100644 --- a/isaaclab_arena/assets/asset_cache.py +++ b/isaaclab_arena/assets/asset_cache.py @@ -13,11 +13,3 @@ def get_arena_asset_cache_dir() -> pathlib.Path: if not asset_cache_dir.exists(): asset_cache_dir.mkdir(parents=True, exist_ok=True) return asset_cache_dir - - -def get_arena_usd_cache_dir() -> pathlib.Path: - """Return the cache root for USDs Arena generates, such as composed and baked-geometry assets. - - The directory is not created, so callers that only compute a path do not leave one behind. - """ - return pathlib.Path.home() / ".cache" / "isaaclab_arena" / "usd" diff --git a/isaaclab_arena/embodiments/robot_on_stand_utils.py b/isaaclab_arena/embodiments/robot_on_stand_utils.py index da577e5e73..f9dd1b5b39 100644 --- a/isaaclab_arena/embodiments/robot_on_stand_utils.py +++ b/isaaclab_arena/embodiments/robot_on_stand_utils.py @@ -16,7 +16,7 @@ from isaaclab.utils.assets import retrieve_file_path from pxr import Gf, Usd, UsdGeom -from isaaclab_arena.assets.asset_cache import get_arena_usd_cache_dir +from isaaclab_arena.assets.asset_cache import get_arena_asset_cache_dir _ROBOT_ON_STAND_USD_CACHE_DIR = "robot_on_stand" @@ -80,7 +80,7 @@ def compose_on_stand_usd( """ assert stand_height_m > 0.0, f"stand_height_m must be positive, got {stand_height_m}" - cache_root = get_arena_usd_cache_dir() / _ROBOT_ON_STAND_USD_CACHE_DIR + cache_root = get_arena_asset_cache_dir().parent / "usd" / _ROBOT_ON_STAND_USD_CACHE_DIR cache_root.mkdir(parents=True, exist_ok=True) out_path = cache_root / f"{output_basename}_{stand_height_m:.3f}.usd" diff --git a/isaaclab_arena/scripts/export_ready_pose_collision_meshes.py b/isaaclab_arena/scripts/export_ready_pose_collision_meshes.py deleted file mode 100644 index 0169ca0311..0000000000 --- a/isaaclab_arena/scripts/export_ready_pose_collision_meshes.py +++ /dev/null @@ -1,171 +0,0 @@ -# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Export every registered robot's ready-pose collision mesh, for upload to Nucleus. - -Extracting a robot's posed mesh costs 0.1-2.5 s per process. Publishing the result once means every -scene that spawns the robot reads it instead. The output holds a folder per robot and is uploaded -verbatim into ``collision_mesh_store.ARENA_ROBOT_LIBRARY_DIR``, merging with each robot's existing -folder: the relative paths are the ones the loader looks for. - -Re-run this whenever a robot's USD or configured joint positions change, since the loader validates -the pose an artifact was extracted at and falls back to extraction when it no longer matches. - -Run inside the container: - - /isaac-sim/python.sh isaaclab_arena/scripts/export_ready_pose_collision_meshes.py --headless \\ - --out_dir /tmp/ready_pose_export - -Check an export before uploading it by pointing the loader at it: - - export ISAACLAB_ARENA_ROBOT_LIBRARY_DIR=/tmp/ready_pose_export -""" - -from __future__ import annotations - -import argparse -import traceback -from collections.abc import Mapping -from dataclasses import dataclass -from pathlib import Path -from typing import TYPE_CHECKING - -from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser -from isaaclab_arena.utils.isaaclab_utils.simulation_app import SimulationAppContext - -if TYPE_CHECKING: - # Typing only: embodiment_base pulls in Isaac Lab, which is unavailable before sim init. - from isaaclab_arena.embodiments.embodiment_base import PlacementGeometrySource - - -def add_export_arguments(parser: argparse.ArgumentParser) -> None: - """Add the export flags.""" - group = parser.add_argument_group("Ready-Pose Mesh Export Arguments") - group.add_argument( - "--out_dir", - type=Path, - default=Path("/tmp/ready_pose_export"), - help="Directory to write per-robot folders into, merged as-is into the Arena robot library.", - ) - group.add_argument( - "--robots", - nargs="+", - default=None, - help="Registered embodiment names to export. Defaults to every registered embodiment.", - ) - - -def _embodiment_classes(names: list[str] | None) -> list[type]: - """Return the registered embodiment classes to export, in a stable order.""" - from isaaclab_arena.assets.registries import AssetRegistry - from isaaclab_arena.embodiments.embodiment_base import EmbodimentBase - - registry = AssetRegistry() - candidates = names if names is not None else registry.get_all_keys() - classes = [] - for name in sorted(candidates): - asset = registry.get_asset_by_name(name) - if isinstance(asset, type) and issubclass(asset, EmbodimentBase): - classes.append(asset) - assert classes, f"no registered embodiments among {candidates}" - return classes - - -@dataclass(frozen=True) -class PlannedArtifact: - """One artifact the export will write, and the embodiment it was planned from.""" - - embodiment_name: str - """Embodiment the mesh is extracted from, for reporting which robot an artifact came from.""" - - source: PlacementGeometrySource - """USD, scale and joint positions the artifact is written from.""" - - -def plan_artifacts(sources: Mapping[str, PlacementGeometrySource]) -> dict[str, PlannedArtifact]: - """Map each artifact's published relative path to the embodiment and source it is written from. - - Robots differing only in action space share a USD, so they collapse to one artifact. Robots - missing from ``ROBOT_LIBRARY_FOLDERS`` have nowhere to publish and are left out. - - Args: - sources: Placement geometry source per embodiment name. - """ - from isaaclab_arena.utils.collision_mesh_store import published_relative_path - - plan: dict[str, PlannedArtifact] = {} - for name, source in sorted(sources.items()): - relative_path = published_relative_path(source.usd_path) - if relative_path is None: - continue - # Artifacts are validated by USD stem, so two robots sharing one would be served each other's - # mesh. Refuse to publish that rather than let placement use the wrong shape. - planned = plan.setdefault(relative_path, PlannedArtifact(name, source)) - assert planned.source.usd_path == source.usd_path, ( - f"{relative_path} would be written for both {planned.source.usd_path} and" - f" {source.usd_path}; rename one of the source USDs" - ) - return plan - - -def export_ready_pose_meshes(out_dir: Path, names: list[str] | None) -> tuple[list[str], list[str]]: - """Write each robot's ready-pose mesh into out_dir, returning the artifacts and the failures. - - Args: - out_dir: Staging directory, whose per-robot folders are uploaded as-is to the robot library. - names: Registered embodiment names to export, or None for every registered embodiment. - """ - from isaaclab_arena.utils.collision_mesh_store import export_ready_pose_mesh, published_relative_path - from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos - - sources = {} - failed = [] - for embodiment_class in _embodiment_classes(names): - try: - sources[embodiment_class.name] = embodiment_class().get_placement_geometry_source() - except Exception as error: - failed.append(f"{embodiment_class.__name__}: {error}") - traceback.print_exc() - - skipped = sorted(name for name, source in sources.items() if published_relative_path(source.usd_path) is None) - written = [] - for relative_path, planned in plan_artifacts(sources).items(): - source = planned.source - try: - # Unit scale: one artifact serves every spawn scale, as the loader rescales on read. - mesh = extract_trimesh_from_usd_at_joint_pos(source.usd_path, source.joint_pos, (1.0, 1.0, 1.0)) - export_ready_pose_mesh(source.usd_path, source.joint_pos, mesh, out_dir) - written.append(f"{relative_path} ({len(mesh.vertices)} vertices, from {planned.embodiment_name})") - except Exception as error: - failed.append(f"{planned.embodiment_name}: {error}") - traceback.print_exc() - if skipped: - print(f"\nSkipped {len(skipped)} embodiment(s) missing from ROBOT_LIBRARY_FOLDERS: {', '.join(skipped)}") - return sorted(written), failed - - -def main() -> None: - args_parser = get_isaaclab_arena_cli_parser() - add_export_arguments(args_parser) - args_cli, _ = args_parser.parse_known_args() - - with SimulationAppContext(args_cli): - from isaaclab_arena.utils.collision_mesh_store import ARENA_ROBOT_LIBRARY_DIR - - written, failed = export_ready_pose_meshes(args_cli.out_dir, args_cli.robots) - - print(f"\nWrote {len(written)} artifact(s) to {args_cli.out_dir}, upload them to {ARENA_ROBOT_LIBRARY_DIR}:") - for line in written: - print(f" {line}") - if failed: - print(f"\n{len(failed)} embodiment(s) failed to export:") - for line in failed: - print(f" {line}") - # Exit non-zero so a partial export is not mistaken for a complete one and uploaded. - raise SystemExit(1) - - -if __name__ == "__main__": - main() diff --git a/isaaclab_arena/tests/test_collision_mesh_store.py b/isaaclab_arena/tests/test_collision_mesh_store.py deleted file mode 100644 index 477b16d889..0000000000 --- a/isaaclab_arena/tests/test_collision_mesh_store.py +++ /dev/null @@ -1,388 +0,0 @@ -# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for storing and reloading an articulation's collision mesh, keyed by joint pose.""" - -import contextlib -import numpy as np -import os -import shutil -import tempfile -import time -import trimesh -from pathlib import Path - -from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function - -HEADLESS = True - - -def _build_arm_usd(tmp_dir: str, joint_type: str = "revolute") -> str: - """Export the two-link fixture arm from the articulation tests to a USD file.""" - from isaaclab_arena.tests.test_usd_articulation import _build_two_link_arm - - usd_path = f"{tmp_dir}/arm.usda" - _build_two_link_arm(joint_type=joint_type).Export(usd_path) - return usd_path - - -@contextlib.contextmanager -def _published_dir(path: str): - """Read published artifacts from path, keeping tests off Nucleus and away from real artifacts.""" - from isaaclab_arena.utils.collision_mesh_store import ROBOT_LIBRARY_DIR_ENV_VAR - - previous = os.environ.get(ROBOT_LIBRARY_DIR_ENV_VAR) - os.environ[ROBOT_LIBRARY_DIR_ENV_VAR] = path - try: - yield - finally: - if previous is None: - del os.environ[ROBOT_LIBRARY_DIR_ENV_VAR] - else: - os.environ[ROBOT_LIBRARY_DIR_ENV_VAR] = previous - - -@contextlib.contextmanager -def _publishes_as(stem: str, folder: str): - """Give a fixture asset a robot-library folder, as a shipped robot has.""" - from isaaclab_arena.utils import collision_mesh_store as store - - store.ROBOT_LIBRARY_FOLDERS[stem] = folder - try: - yield - finally: - store.ROBOT_LIBRARY_FOLDERS.pop(stem, None) - - -def _test_trimming_evicts_least_recently_used_artifacts(simulation_app) -> bool: - """Trimming drops artifacts oldest-first down to the budget, and leaves in-flight writes alone.""" - from isaaclab_arena.utils import collision_mesh_store as store - - with tempfile.TemporaryDirectory() as cache_dir: - # Three 1 KiB artifacts, aged a day apart, plus a staging file another process is writing. - paths = {} - for age_days, name in enumerate(["newest.usd", "middle.usd", "oldest.usd"]): - path = Path(cache_dir) / "robot" / name - path.parent.mkdir(parents=True, exist_ok=True) - path.write_bytes(b"x" * 1024) - os.utime(path, (0, time.time() - age_days * 86400)) - paths[name] = path - staging = Path(cache_dir) / "robot" / f"{store._STAGING_PREFIX}half_written.usd" - staging.write_bytes(b"x" * 4096) - os.utime(staging, (0, time.time() - 7 * 86400)) - - previous_budget = store.CACHE_BUDGET_BYTES - store.CACHE_BUDGET_BYTES = 2048 - try: - store._trim_cache(Path(cache_dir)) - finally: - store.CACHE_BUDGET_BYTES = previous_budget - - assert paths["newest.usd"].is_file(), "the most recently used artifact must survive" - assert paths["middle.usd"].is_file(), "the budget must be filled before evicting" - assert not paths["oldest.usd"].is_file(), "the least recently used artifact must be evicted" - # Unlinking a staging file would fail the rename of whichever process is writing it. - assert staging.is_file(), "a half-written artifact must not be evicted despite being oldest" - return True - - -def _test_pose_keys_identify_the_pose(simulation_app) -> bool: - """Every all-zero spelling shares the readable zero key; distinct poses key apart.""" - from isaaclab_arena.utils.collision_mesh_store import _pose_key - - # Omitted joints are posed at zero, so an empty mapping is the zero pose. - assert _pose_key({}) == "zero" - assert _pose_key({"elbow": 0.0, "wrist": -0.0}) == "zero" - assert _pose_key({"elbow": 0.5}) != "zero" - # Order must not matter, but the values must. - assert _pose_key({"a": 0.5, "b": 0.25}) == _pose_key({"b": 0.25, "a": 0.5}) - assert _pose_key({"elbow": 0.5}) != _pose_key({"elbow": 0.6}) - assert _pose_key({"elbow": 0.5}) != _pose_key({"wrist": 0.5}) - return True - - -def _test_distinct_poses_get_distinct_artifacts(simulation_app) -> bool: - """Posing at one configuration must never serve the mesh for another.""" - from isaaclab_arena.utils import usd_helpers - from isaaclab_arena.utils.collision_mesh_store import _cache_path - - extended = {"elbow": 0.5} - with tempfile.TemporaryDirectory() as tmp_dir, contextlib.ExitStack() as published: - usd_path = _build_arm_usd(tmp_dir, joint_type="prismatic") - published.enter_context(_published_dir(tmp_dir)) - for joint_pos in ({}, extended): - _cache_path(usd_path, joint_pos).unlink(missing_ok=True) - - usd_helpers._extract_trimesh_from_usd_at_joint_pos.cache_clear() - zero_mesh = usd_helpers.extract_trimesh_from_usd_at_joint_pos(usd_path, {}) - usd_helpers._extract_trimesh_from_usd_at_joint_pos.cache_clear() - posed_mesh = usd_helpers.extract_trimesh_from_usd_at_joint_pos(usd_path, extended) - - assert _cache_path(usd_path, {}).is_file() and _cache_path(usd_path, extended).is_file() - assert _cache_path(usd_path, {}) != _cache_path(usd_path, extended) - assert posed_mesh.extents[0] > zero_mesh.extents[0] + 0.4, f"prismatic pose should extend: {posed_mesh.extents}" - - # Reloading from the store must preserve that difference. - usd_helpers._extract_trimesh_from_usd_at_joint_pos.cache_clear() - reloaded_zero = usd_helpers.extract_trimesh_from_usd_at_joint_pos(usd_path, {}) - np.testing.assert_allclose(reloaded_zero.extents, zero_mesh.extents, atol=1e-6) - for joint_pos in ({}, extended): - _cache_path(usd_path, joint_pos).unlink(missing_ok=True) - return True - - -def _test_stored_mesh_keeps_its_vertices_verbatim(simulation_app) -> bool: - """Reloading must not merge coincident vertices, which trimesh does by default on construction.""" - from isaaclab_arena.utils.collision_mesh_store import _cache_path, load_mesh, save_mesh - - # Two coincident triangles: merging would collapse six vertices into three. - vertices = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]] * 2) - mesh = trimesh.Trimesh(vertices=vertices, faces=np.array([[0, 1, 2], [3, 4, 5]]), process=False) - - with tempfile.TemporaryDirectory() as tmp_dir, contextlib.ExitStack() as published: - source = f"{tmp_dir}/synthetic.usda" - published.enter_context(_published_dir(tmp_dir)) - save_mesh(source, {}, mesh) - loaded = load_mesh(source, {}, (1.0, 1.0, 1.0)) - assert loaded is not None - assert len(loaded.vertices) == len(vertices), f"round trip changed vertex count: {len(loaded.vertices)}" - - scaled = load_mesh(source, {}, (2.0, 2.0, 2.0)) - assert scaled is not None - assert len(scaled.vertices) == len(vertices), f"scaling changed vertex count: {len(scaled.vertices)}" - _cache_path(source, {}).unlink(missing_ok=True) - return True - - -def _test_stored_mesh_is_scale_independent(simulation_app) -> bool: - """One unit-scale artifact serves every spawn scale.""" - from isaaclab_arena.utils import usd_helpers - from isaaclab_arena.utils.collision_mesh_store import _cache_path - - with tempfile.TemporaryDirectory() as tmp_dir, contextlib.ExitStack() as published: - usd_path = _build_arm_usd(tmp_dir) - published.enter_context(_published_dir(tmp_dir)) - _cache_path(usd_path, {}).unlink(missing_ok=True) - - usd_helpers._extract_trimesh_from_usd_at_joint_pos.cache_clear() - unit = usd_helpers.extract_trimesh_from_usd_at_joint_pos(usd_path, {}) - # Served from the store now, and must still honour the requested scale. - usd_helpers._extract_trimesh_from_usd_at_joint_pos.cache_clear() - doubled = usd_helpers.extract_trimesh_from_usd_at_joint_pos(usd_path, {}, scale=(2.0, 1.0, 1.0)) - np.testing.assert_allclose(doubled.extents[0], unit.extents[0] * 2.0, atol=1e-6) - np.testing.assert_allclose(doubled.extents[1:], unit.extents[1:], atol=1e-6) - _cache_path(usd_path, {}).unlink(missing_ok=True) - return True - - -def _test_artifact_from_another_asset_or_pose_is_ignored(simulation_app) -> bool: - """A mesh recording a different source or pose is refused, so it cannot misplace a robot.""" - from isaaclab_arena.utils.collision_mesh_store import _cache_path, load_mesh, save_mesh - from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos - - posed = {"elbow": 0.5} - with tempfile.TemporaryDirectory() as tmp_dir, contextlib.ExitStack() as published: - usd_path = _build_arm_usd(tmp_dir) - published.enter_context(_published_dir(tmp_dir)) - mesh = extract_trimesh_from_usd_at_joint_pos(usd_path, posed) - - # Recorded under a foreign source USD, then moved into this asset's slot. - save_mesh(f"{tmp_dir}/other_arm.usda", posed, mesh).replace(_cache_path(usd_path, posed)) - assert load_mesh(usd_path, posed, (1.0, 1.0, 1.0)) is None, "foreign source must be ignored" - - # Recorded for this asset but at another pose, then moved into this pose's slot. - save_mesh(usd_path, {"elbow": 0.9}, mesh).replace(_cache_path(usd_path, posed)) - assert load_mesh(usd_path, posed, (1.0, 1.0, 1.0)) is None, "foreign pose must be ignored" - _cache_path(usd_path, posed).unlink(missing_ok=True) - return True - - -def _test_published_artifact_is_found_by_its_name(simulation_app) -> bool: - """An exported artifact is loaded from the published directory when the local cache is empty.""" - from isaaclab_arena.utils.collision_mesh_store import ( - _cache_path, - _published_path, - export_ready_pose_mesh, - load_mesh, - published_relative_path, - ) - from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos - - posed = {"elbow": 0.5} - with ( - tempfile.TemporaryDirectory() as tmp_dir, - tempfile.TemporaryDirectory() as publish_dir, - _publishes_as("arm", "test_arm"), - ): - usd_path = _build_arm_usd(tmp_dir) - mesh = extract_trimesh_from_usd_at_joint_pos(usd_path, posed) - artifact = export_ready_pose_mesh(usd_path, posed, mesh, Path(publish_dir)) - - # The artifact lands in the robot's own library folder, under a name derived from the asset. - assert published_relative_path(usd_path) == "test_arm/arm_ready_pose.usd" - assert artifact.name == "arm_ready_pose.usd" and artifact.parent.name == "test_arm", artifact - - with _published_dir(publish_dir): - assert _published_path(usd_path) == str(artifact) - # Empty the local cache so only the published copy can answer. - _cache_path(usd_path, posed).unlink(missing_ok=True) - loaded = load_mesh(usd_path, posed, (1.0, 1.0, 1.0)) - assert loaded is not None, f"published artifact {artifact} was not found" - np.testing.assert_allclose(loaded.vertices, mesh.vertices, atol=1e-6) - - # Published artifacts outlive config changes, so a repose must not be served the old shape. - assert load_mesh(usd_path, {"elbow": 0.9}, (1.0, 1.0, 1.0)) is None, "stale pose" - _cache_path(usd_path, posed).unlink(missing_ok=True) - return True - - -def _test_published_artifact_loads_for_a_relocated_source(simulation_app) -> bool: - """An artifact stays usable when the robot USD sits elsewhere, as on another machine. - - Arena composes the robot-on-stand USDs into a per-user cache directory, so validating against the - full source path would reject every artifact anyone else exported. - """ - from isaaclab_arena.utils.collision_mesh_store import _cache_path, export_ready_pose_mesh, load_mesh - from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos - - posed = {"elbow": 0.5} - with ( - tempfile.TemporaryDirectory() as exporter_dir, - tempfile.TemporaryDirectory() as consumer_dir, - tempfile.TemporaryDirectory() as publish_dir, - _publishes_as("arm", "test_arm"), - ): - exporter_usd = _build_arm_usd(exporter_dir) - mesh = extract_trimesh_from_usd_at_joint_pos(exporter_usd, posed) - export_ready_pose_mesh(exporter_usd, posed, mesh, Path(publish_dir)) - - # The same asset reached through another absolute path, as a second machine's cache would. - consumer_usd = f"{consumer_dir}/arm.usda" - shutil.copy(exporter_usd, consumer_usd) - - with _published_dir(publish_dir): - _cache_path(consumer_usd, posed).unlink(missing_ok=True) - loaded = load_mesh(consumer_usd, posed, (1.0, 1.0, 1.0)) - assert loaded is not None, "an artifact must load for the same asset at another path" - np.testing.assert_allclose(loaded.vertices, mesh.vertices, atol=1e-6) - for path in (exporter_usd, consumer_usd): - _cache_path(path, posed).unlink(missing_ok=True) - return True - - -def _test_truncated_artifact_is_ignored(simulation_app) -> bool: - """A half-written or hand-edited artifact falls back to extraction rather than raising.""" - from pxr import Usd, UsdGeom - - from isaaclab_arena.utils.collision_mesh_store import load_mesh, save_mesh - from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos - - posed = {"elbow": 0.5} - with tempfile.TemporaryDirectory() as tmp_dir, contextlib.ExitStack() as published: - usd_path = _build_arm_usd(tmp_dir) - published.enter_context(_published_dir(tmp_dir)) - cache_path = save_mesh(usd_path, posed, extract_trimesh_from_usd_at_joint_pos(usd_path, posed)) - assert load_mesh(usd_path, posed, (1.0, 1.0, 1.0)) is not None - - # Drop the geometry but keep the provenance, as a truncated write would. - stage = Usd.Stage.Open(str(cache_path)) - mesh_prim = UsdGeom.Mesh(stage.GetPrimAtPath("/CollisionMesh")) - mesh_prim.GetPointsAttr().Clear() - mesh_prim.GetFaceVertexIndicesAttr().Clear() - stage.GetRootLayer().Save() - - assert load_mesh(usd_path, posed, (1.0, 1.0, 1.0)) is None, "an artifact without geometry must be ignored" - cache_path.unlink(missing_ok=True) - return True - - -def _test_colliding_published_names_are_refused(simulation_app) -> bool: - """Two robots sharing a USD stem cannot both be published, as one artifact name serves both.""" - from isaaclab_arena.embodiments.embodiment_base import PlacementGeometrySource - from isaaclab_arena.scripts.export_ready_pose_collision_meshes import plan_artifacts - - def source(usd_path: str) -> PlacementGeometrySource: - return PlacementGeometrySource(usd_path, (1.0, 1.0, 1.0), {}) - - with _publishes_as("robot", "shared"): - # Same stem, different assets: the published name cannot tell them apart. - colliding = {"robot_a": source("/assets/vendor_a/robot.usd"), "robot_b": source("/assets/vendor_b/robot.usd")} - try: - plan_artifacts(colliding) - except AssertionError as error: - assert "rename one of the source USDs" in str(error), error - else: - raise AssertionError("a colliding published name must be refused") - - # Action-space variants of one robot share a USD, and so legitimately share one artifact. - shared_usd = source("/assets/vendor_a/robot.usd") - assert sorted(plan_artifacts({"robot_ik": shared_usd, "robot_joint_pos": shared_usd})) == [ - "shared/robot_ready_pose.usd" - ] - - # A robot missing from ROBOT_LIBRARY_FOLDERS is left out rather than written to a guessed folder. - assert plan_artifacts({"unpublished": source("/assets/vendor_a/robot.usd")}) == {} - return True - - -def _test_export_writes_one_named_artifact_per_robot(simulation_app) -> bool: - """The exporter deduplicates embodiments sharing a USD and reports nothing as failed.""" - from isaaclab_arena.scripts.export_ready_pose_collision_meshes import export_ready_pose_meshes - - # Droid's three action-space variants share one USD and one ready pose. - droid_variants = ["droid_abs_joint_pos", "droid_rel_joint_pos", "droid_differential_ik"] - with tempfile.TemporaryDirectory() as out_dir: - written, failed = export_ready_pose_meshes(Path(out_dir), droid_variants) - - assert not failed, f"export reported failures: {failed}" - assert len(written) == 1, f"variants of one robot must share an artifact: {written}" - # Relative paths, because the staging tree is uploaded into the robot library as it stands. - artifacts = sorted(str(path.relative_to(out_dir)) for path in Path(out_dir).rglob("*.usd")) - assert artifacts == ["droid/droid_franka_robotiq_on_stand_1.350_ready_pose.usd"], artifacts - return True - - -def test_trimming_evicts_least_recently_used_artifacts(): - assert run_simulation_app_function(_test_trimming_evicts_least_recently_used_artifacts, headless=HEADLESS) - - -def test_pose_keys_identify_the_pose(): - assert run_simulation_app_function(_test_pose_keys_identify_the_pose, headless=HEADLESS) - - -def test_truncated_artifact_is_ignored(): - assert run_simulation_app_function(_test_truncated_artifact_is_ignored, headless=HEADLESS) - - -def test_colliding_published_names_are_refused(): - assert run_simulation_app_function(_test_colliding_published_names_are_refused, headless=HEADLESS) - - -def test_export_writes_one_named_artifact_per_robot(): - assert run_simulation_app_function(_test_export_writes_one_named_artifact_per_robot, headless=HEADLESS) - - -def test_published_artifact_is_found_by_its_name(): - assert run_simulation_app_function(_test_published_artifact_is_found_by_its_name, headless=HEADLESS) - - -def test_published_artifact_loads_for_a_relocated_source(): - assert run_simulation_app_function(_test_published_artifact_loads_for_a_relocated_source, headless=HEADLESS) - - -def test_distinct_poses_get_distinct_artifacts(): - assert run_simulation_app_function(_test_distinct_poses_get_distinct_artifacts, headless=HEADLESS) - - -def test_stored_mesh_keeps_its_vertices_verbatim(): - assert run_simulation_app_function(_test_stored_mesh_keeps_its_vertices_verbatim, headless=HEADLESS) - - -def test_stored_mesh_is_scale_independent(): - assert run_simulation_app_function(_test_stored_mesh_is_scale_independent, headless=HEADLESS) - - -def test_artifact_from_another_asset_or_pose_is_ignored(): - assert run_simulation_app_function(_test_artifact_from_another_asset_or_pose_is_ignored, headless=HEADLESS) diff --git a/isaaclab_arena/tests/test_usd_articulation.py b/isaaclab_arena/tests/test_usd_articulation.py index 0b240587c5..059a386864 100644 --- a/isaaclab_arena/tests/test_usd_articulation.py +++ b/isaaclab_arena/tests/test_usd_articulation.py @@ -329,7 +329,6 @@ def _test_instanced_geometry_is_posed_with_its_link(simulation_app) -> bool: from pxr import Gf, UsdGeom - from isaaclab_arena.utils.collision_mesh_store import _cache_path from isaaclab_arena.utils.usd_helpers import ( compute_local_bounding_box_from_usd_at_joint_pos, extract_trimesh_from_usd_at_joint_pos, @@ -358,10 +357,6 @@ def _test_instanced_geometry_is_posed_with_its_link(simulation_app) -> bool: swung = compute_local_bounding_box_from_usd_at_joint_pos(usd_path, {"elbow": math.pi / 2.0}) np.testing.assert_allclose(swung.max_point.numpy()[0][1], 1.7, atol=1e-6) np.testing.assert_allclose(swung.max_point.numpy()[0][0], 0.2, atol=1e-6) - - # The temp USD is unique per run, so its stored meshes would otherwise pile up in the cache. - for joint_pos in ({}, {"elbow": math.pi / 2.0}): - _cache_path(usd_path, joint_pos).unlink(missing_ok=True) return True diff --git a/isaaclab_arena/utils/collision_mesh_store.py b/isaaclab_arena/utils/collision_mesh_store.py deleted file mode 100644 index 4625911e03..0000000000 --- a/isaaclab_arena/utils/collision_mesh_store.py +++ /dev/null @@ -1,325 +0,0 @@ -# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Persist an articulation's posed collision mesh as a USD, so extraction runs once per pose. - -Walking a robot USD to merge its meshes costs 0.1-2.5 s depending on the robot, which is paid again -in every new process. Artifacts are keyed by the joint pose they were extracted at, because that is -what makes a mesh reusable: embodiments spawn at their configured pose rather than at zero, so keying -on the asset alone would store a mesh nobody asks for. - -Lookup order is the local cache, then the robot's own folder under ``ARENA_ROBOT_LIBRARY_DIR`` on -Nucleus, which holds one published artifact per robot at the pose that robot spawns in, written by -``scripts/export_ready_pose_collision_meshes.py``. A miss falls back to extraction. - -Artifacts are named and validated by the source USD's stem, so a file exported on one machine loads -on another even though Arena composes the robot-on-stand USDs into a per-user cache directory. -``ROBOT_LIBRARY_FOLDERS`` maps that stem to the robot's published folder. -""" - -from __future__ import annotations - -import contextlib -import hashlib -import numpy as np -import os -import tempfile -import trimesh -from collections.abc import Mapping -from pathlib import Path - -from pxr import Usd, UsdGeom - -from isaaclab_arena.assets.asset_cache import get_arena_usd_cache_dir -from isaaclab_arena.assets.nucleus import ARENA_NUCLEUS_DIR - -ROBOT_LIBRARY_FOLDERS = { - "A2D_physics": "agibot_a2d", - "droid_franka_robotiq_on_stand_1.350": "droid", - "franka_panda_on_stand_0.875": "franka", - "g1_29dof_with_hand_rev_1_0": "g1", - "galbot_one_charlie": "galbot", - "GR1T2_fourier_hand_6dof": "gr1t2", - "kuka": "kuka", -} -"""Folder each robot publishes its collision mesh to, keyed by the stem of the USD Arena spawns. - -Robots differing only in action space share a USD and so one artifact. The robots Arena composes -onto a stand key by the composed stem, stand height included, because a robot at another stand -height is different geometry that nothing has published; it falls back to extraction. -""" - -ARENA_ROBOT_LIBRARY_DIR = f"{ARENA_NUCLEUS_DIR}/Arena/assets/robot_library" -"""Nucleus robot library, holding each robot's assets in a folder of its own.""" - -ROBOT_LIBRARY_DIR_ENV_VAR = "ISAACLAB_ARENA_ROBOT_LIBRARY_DIR" -"""Environment variable redirecting which robot library published artifacts are read from.""" - -CACHE_BUDGET_BYTES = 1 << 30 -"""Disk the local cache may occupy before its least recently used artifacts are dropped. - -A robot's mesh runs 1-9 MB, so this holds a few hundred: every robot Arena ships at several poses -each, while bounding what a run against throwaway USDs can leave behind. -""" - -_MESH_PRIM_PATH = "/CollisionMesh" -"""Prim the merged mesh is authored at, also the artifact's default prim.""" - -_ASSET_KEY = "arenaAssetKey" -"""customData key recording which asset the mesh describes, validated on load.""" - -_SOURCE_KEY = "arenaSourceUsdPath" -"""customData key recording the exporting machine's source path, kept to trace an artifact back.""" - -_POSE_KEY = "arenaJointPoseKey" -"""customData key recording which joint pose the mesh was extracted at, validated on load.""" - -_ZERO_POSE_KEY = "zero" -"""Pose key for an all-zero pose, spelled out so cached artifacts stay readable.""" - -_READY_POSE_SUFFIX = "_ready_pose.usd" -"""Suffix naming a published artifact, one per robot, at the pose that robot spawns in.""" - -_LOCAL_CACHE_DIR = "collision_mesh" -"""Local cache root under ``get_arena_usd_cache_dir()``, with one subfolder per robot when known.""" - -_STAGING_PREFIX = "_staging_" -"""Prefix marking a half-written artifact, so trimming leaves another process's write alone. - -Staging has to share the destination's directory for the rename to stay atomic, and the suffix has -to stay a USD one for ``Usd.Stage.CreateNew``, so the prefix is what keeps the two apart. -""" - - -def load_mesh( - source_usd_path: str, joint_pos: Mapping[str, float], scale: tuple[float, float, float] -) -> trimesh.Trimesh | None: - """Return the stored mesh for a robot USD at joint_pos scaled to scale, or None if unavailable. - - Reads the local cache first and the published robot library second, copying a published hit into - the local cache on the way out. An artifact recording another asset or another pose is ignored - rather than trusted. - - Args: - source_usd_path: Robot USD the mesh should describe. - joint_pos: Joint positions the mesh should be posed at. - scale: Per-axis scale to apply to the stored vertices. - """ - expected = {"expected_asset": _asset_key(source_usd_path), "expected_pose": _pose_key(joint_pos)} - - cached = _cache_path(source_usd_path, joint_pos) - mesh = _read_mesh_usd(str(cached), **expected) - if mesh is not None: - # Mark the artifact used, so trimming evicts by last use rather than by write time. - with contextlib.suppress(OSError): - os.utime(cached) - return scaled_mesh(mesh, scale) - - published = _published_path(source_usd_path) - if published is None: - return None - mesh = _read_mesh_usd(published, **expected) - if mesh is None: - return None - # Copy the published artifact locally, so only the first process pays the Nucleus round trip. - with contextlib.suppress(OSError): - save_mesh(source_usd_path, joint_pos, mesh) - return scaled_mesh(mesh, scale) - - -def save_mesh(source_usd_path: str, joint_pos: Mapping[str, float], mesh: trimesh.Trimesh) -> Path: - """Write a robot USD's unscaled mesh at joint_pos to the local cache and return its path. - - Least recently used artifacts are dropped to keep the cache within ``CACHE_BUDGET_BYTES``. - - Args: - source_usd_path: Robot USD the mesh was extracted from, recorded for validation on load. - joint_pos: Joint positions the mesh was posed at, recorded for validation on load. - mesh: Merged mesh at unit scale, in the robot's default-prim frame. - """ - out_path = _write_mesh_usd(_cache_path(source_usd_path, joint_pos), source_usd_path, joint_pos, mesh) - _trim_cache(_local_cache_dir()) - return out_path - - -def export_ready_pose_mesh( - source_usd_path: str, joint_pos: Mapping[str, float], mesh: trimesh.Trimesh, out_dir: Path -) -> Path: - """Write a robot's ready-pose mesh into out_dir under its published name, for upload. - - Args: - source_usd_path: Robot USD the mesh was extracted from, which must be a published robot. - joint_pos: The robot's configured joint positions, recorded for validation on load. - mesh: Merged mesh at unit scale, in the robot's default-prim frame. - out_dir: Staging directory whose layout mirrors ``ARENA_ROBOT_LIBRARY_DIR`` for upload as-is. - """ - relative_path = published_relative_path(source_usd_path) - assert ( - relative_path is not None - ), f"{_asset_key(source_usd_path)} has no entry in ROBOT_LIBRARY_FOLDERS, so there is no folder to publish it to" - return _write_mesh_usd(out_dir / relative_path, source_usd_path, joint_pos, mesh) - - -def published_relative_path(source_usd_path: str) -> str | None: - """Return a robot USD's artifact path within the robot library, or None if it publishes none.""" - folder = ROBOT_LIBRARY_FOLDERS.get(_asset_key(source_usd_path)) - if folder is None: - return None - return f"{folder}/{_asset_key(source_usd_path)}{_READY_POSE_SUFFIX}" - - -def _published_path(source_usd_path: str) -> str | None: - """Return the full path a robot USD's ready-pose mesh is read from, or None if it publishes none. - - Defaults to the Arena Nucleus robot library and is redirected by - ``ISAACLAB_ARENA_ROBOT_LIBRARY_DIR``, which lets an export be checked locally before upload. - """ - relative_path = published_relative_path(source_usd_path) - if relative_path is None: - return None - library_dir = os.environ.get(ROBOT_LIBRARY_DIR_ENV_VAR) or ARENA_ROBOT_LIBRARY_DIR - return f"{library_dir.rstrip('/')}/{relative_path}" - - -def _cache_path(source_usd_path: str, joint_pos: Mapping[str, float]) -> Path: - """Return the local cache path for a robot USD's mesh at joint_pos. - - Published robots mirror their library folder, and everything else, such as a throwaway test - fixture, sits at the cache root. The full source path is hashed into the filename so assets - sharing a stem stay distinct on one machine even though they would share a published name. - """ - source_digest = hashlib.sha1(str(source_usd_path).encode()).hexdigest()[:12] - name = f"{_asset_key(source_usd_path)}_{source_digest}_collision_{_pose_key(joint_pos)}_pose.usd" - folder = ROBOT_LIBRARY_FOLDERS.get(_asset_key(source_usd_path)) - return _local_cache_dir() / folder / name if folder else _local_cache_dir() / name - - -def _local_cache_dir() -> Path: - """Return the local collision-mesh cache root (``~/.cache/.../usd/collision_mesh``).""" - return get_arena_usd_cache_dir() / _LOCAL_CACHE_DIR - - -def _asset_key(source_usd_path: str) -> str: - """Return the identity a published artifact is named and validated by: the source USD's stem.""" - # The stem rather than the full path, so an artifact exported from a per-user cache directory - # still matches elsewhere. Stems already spell out the variant: droid_franka_robotiq_on_stand_1.350. - return Path(str(source_usd_path)).stem - - -def _pose_key(joint_pos: Mapping[str, float]) -> str: - """Return a filename-safe key identifying a joint pose. - - Keys are derived from the joint names and values as given, so two spellings of one pose, a regex - and the names it expands to, key differently and are extracted separately. - """ - # Omitted joints are posed at zero, so an empty mapping is the zero pose, not the authored one. - if all(float(value) == 0.0 for value in joint_pos.values()): - return _ZERO_POSE_KEY - canonical = ";".join(f"{name}={float(value) + 0.0:.9g}" for name, value in sorted(joint_pos.items())) - return hashlib.sha1(canonical.encode()).hexdigest()[:12] - - -def scaled_mesh(mesh: trimesh.Trimesh, scale: tuple[float, float, float]) -> trimesh.Trimesh: - """Return mesh with vertices scaled per axis in its own frame. - - Scaling in the root frame commutes with merging prims, so scaling a merged mesh matches scaling - each prim's vertices during extraction. - """ - if tuple(scale) == (1.0, 1.0, 1.0): - return mesh - # process=False keeps trimesh from merging vertices, which would make a scaled or reloaded mesh - # differ from the one extraction produced. - return trimesh.Trimesh( - vertices=mesh.vertices * np.asarray(scale, dtype=np.float64), faces=mesh.faces, process=False - ) - - -def _trim_cache(cache_dir: Path) -> None: - """Delete cached artifacts, least recently used first, until the cache fits its budget. - - Artifacts another process is still staging are left alone, since unlinking one would fail that - process's rename. - """ - if not cache_dir.is_dir(): - return - artifacts = (path for path in cache_dir.rglob("*.usd") if not path.name.startswith(_STAGING_PREFIX)) - cumulative_bytes = 0 - for path in sorted(artifacts, key=_last_used, reverse=True): - with contextlib.suppress(OSError): - # Counts what has been walked rather than what survives, so eviction keeps a strict - # most-recently-used prefix instead of backfilling with whatever happens to fit. - cumulative_bytes += path.stat().st_size - if cumulative_bytes > CACHE_BUDGET_BYTES: - path.unlink() - - -def _last_used(path: Path) -> float: - """Return when an artifact was last read or written, or 0.0 if it has since gone.""" - try: - return path.stat().st_mtime - except OSError: - return 0.0 - - -def _write_mesh_usd( - out_path: Path, source_usd_path: str, joint_pos: Mapping[str, float], mesh: trimesh.Trimesh -) -> Path: - """Author a mesh artifact at out_path, staged and renamed so no reader sees a partial file.""" - out_path.parent.mkdir(parents=True, exist_ok=True) - with tempfile.NamedTemporaryFile( - prefix=_STAGING_PREFIX, suffix=".usd", dir=out_path.parent, delete=False - ) as tmp_file: - tmp_path = Path(tmp_file.name) - try: - stage = Usd.Stage.CreateNew(str(tmp_path)) - mesh_prim = UsdGeom.Mesh.Define(stage, _MESH_PRIM_PATH) - stage.SetDefaultPrim(mesh_prim.GetPrim()) - faces = np.asarray(mesh.faces, dtype=np.int32) - mesh_prim.GetPointsAttr().Set(np.asarray(mesh.vertices, dtype=np.float32)) - mesh_prim.GetFaceVertexCountsAttr().Set([3] * len(faces)) - mesh_prim.GetFaceVertexIndicesAttr().Set(faces.reshape(-1)) - mesh_prim.GetPrim().SetCustomDataByKey(_ASSET_KEY, _asset_key(source_usd_path)) - mesh_prim.GetPrim().SetCustomDataByKey(_POSE_KEY, _pose_key(joint_pos)) - mesh_prim.GetPrim().SetCustomDataByKey(_SOURCE_KEY, str(source_usd_path)) - assert stage.GetRootLayer().Save(), f"failed to save collision mesh to {tmp_path}" - os.replace(tmp_path, out_path) - except Exception: - tmp_path.unlink(missing_ok=True) - raise - return out_path - - -def _read_mesh_usd(usd_path: str, expected_asset: str, expected_pose: str) -> trimesh.Trimesh | None: - """Read a stored artifact, returning None when it is absent or describes another asset or pose.""" - from isaaclab.utils.assets import check_file_path - - # Stage.Open raises rather than returning None on a missing remote path, and a published artifact - # is missing for every robot nobody has exported yet. - if check_file_path(usd_path) == 0: - return None - stage = Usd.Stage.Open(usd_path) - if stage is None: - return None - mesh_prim = UsdGeom.Mesh(stage.GetPrimAtPath(_MESH_PRIM_PATH)) - if not mesh_prim: - return None - - recorded = ( - mesh_prim.GetPrim().GetCustomDataByKey(_ASSET_KEY), - mesh_prim.GetPrim().GetCustomDataByKey(_POSE_KEY), - ) - if recorded != (expected_asset, expected_pose): - print(f"Ignoring collision mesh {usd_path}: recorded {recorded} != {(expected_asset, expected_pose)}") - return None - - points = mesh_prim.GetPointsAttr().Get() - indices = mesh_prim.GetFaceVertexIndicesAttr().Get() - if points is None or indices is None: - return None - return trimesh.Trimesh( - vertices=np.asarray(points, dtype=np.float64), - faces=np.asarray(indices, dtype=np.int32).reshape(-1, 3), - process=False, - ) diff --git a/isaaclab_arena/utils/usd_helpers.py b/isaaclab_arena/utils/usd_helpers.py index 8349e6219c..39e407ddc1 100644 --- a/isaaclab_arena/utils/usd_helpers.py +++ b/isaaclab_arena/utils/usd_helpers.py @@ -15,7 +15,6 @@ from isaaclab_arena.assets.object_type import ObjectType from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox -from isaaclab_arena.utils.collision_mesh_store import load_mesh, save_mesh, scaled_mesh from isaaclab_arena.utils.usd_articulation import ( articulation_joint_prims, compute_posed_prim_world_deltas, @@ -477,8 +476,8 @@ def extract_trimesh_from_usd_at_joint_pos( """Extract an articulation's mesh posed at joint_pos, in its default prim's local frame. Joints the articulation has but joint_pos omits are posed at zero, so the result depends only on - joint_pos and not on the configuration the asset happens to be authored in. Cached in process and - on disk, and shared with every other caller, so treat the result as read-only. + joint_pos and not on the configuration the asset happens to be authored in. Cached in process + and shared with every other caller, so treat the result as read-only. Args: usd_path: Path to the articulation's .usd/.usda/.usdc file. @@ -502,21 +501,13 @@ def _extract_trimesh_from_usd_at_joint_pos( ) -> trimesh.Trimesh: """Cacheable body of ``extract_trimesh_from_usd_at_joint_pos``, keyed by hashable arguments.""" joint_pos = dict(joint_pos_items) - stored = load_mesh(usd_path, joint_pos, scale) - if stored is not None: - return stored - stage = Usd.Stage.Open(usd_path) assert stage is not None, f"could not open USD: {usd_path}" default_prim = stage.GetDefaultPrim() or stage.GetPseudoRoot() default_prim_path = default_prim.GetPath().pathString resolved = resolve_joint_pos_patterns(articulation_joint_prims(default_prim), joint_pos) deltas = compute_posed_prim_world_deltas(stage, default_prim_path, resolved) - - # Store at unit scale so one artifact serves every spawn scale of the asset. - unscaled = extract_trimesh_from_prim(stage, default_prim_path, (1.0, 1.0, 1.0), prim_world_deltas=deltas) - save_mesh(usd_path, joint_pos, unscaled) - return scaled_mesh(unscaled, scale) + return extract_trimesh_from_prim(stage, default_prim_path, scale, prim_world_deltas=deltas) def compute_local_bounding_box_from_usd_at_joint_pos( From 6974e96470b28bf81176e6789f381366a1fe178f Mon Sep 17 00:00:00 2001 From: zhx06 Date: Mon, 3 Aug 2026 08:54:37 -0700 Subject: [PATCH 5/7] address review comments Signed-off-by: zhx06 --- isaaclab_arena/embodiments/droid/droid.py | 100 +++++++++--- isaaclab_arena/embodiments/embodiment_base.py | 16 +- isaaclab_arena/embodiments/franka/franka.py | 43 +++++- .../tests/test_embodiment_collision_mesh.py | 61 ++++---- isaaclab_arena/tests/test_usd_articulation.py | 79 ++++++---- isaaclab_arena/utils/usd_helpers.py | 144 +++++++++++++++--- .../cube_goal_pose_environment.py | 11 +- .../franka_put_and_close_door_environment.py | 11 +- ...oid_pick_and_place_lightwheel_kitchen.yaml | 1 + .../sorting_environment.py | 13 +- 10 files changed, 330 insertions(+), 149 deletions(-) diff --git a/isaaclab_arena/embodiments/droid/droid.py b/isaaclab_arena/embodiments/droid/droid.py index f08747617d..42085976e5 100644 --- a/isaaclab_arena/embodiments/droid/droid.py +++ b/isaaclab_arena/embodiments/droid/droid.py @@ -3,10 +3,11 @@ # # SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations import torch from abc import ABC -from collections.abc import Mapping +from typing import TYPE_CHECKING import isaaclab.envs.mdp as mdp_isaac_lab import isaaclab.sim as sim_utils @@ -37,10 +38,14 @@ from isaaclab_arena.embodiments.embodiment_base import EmbodimentBase from isaaclab_arena.embodiments.franka.franka import franka_stack_events from isaaclab_arena.embodiments.robot_on_stand_utils import RobotPrimSpec, StandPrimSpec, compose_on_stand_usd +from isaaclab_arena.relations.collision_mode import CollisionMode from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.cameras import ArenaCameraCfg from isaaclab_arena.utils.pose import Pose +if TYPE_CHECKING: + import trimesh + _DROID_ROBOT_PRIM = RobotPrimSpec( robot_usd_path=f"{ARENA_NUCLEUS_DIR}/Arena/assets/robot_library/droid/franka_robotiq_2f_85_flattened.usd", root_prim_path="/panda", @@ -55,6 +60,21 @@ footprint_scale_xy=(1.2, 1.2), stand_default_height=1.35, ) +_DROID_JOINT_NAMES = ( + "panda_joint1", + "panda_joint2", + "panda_joint3", + "panda_joint4", + "panda_joint5", + "panda_joint6", + "panda_joint7", + "finger_joint", + "right_outer_knuckle_joint", + "right_inner_finger_joint", + "right_inner_finger_knuckle_joint", + "left_inner_finger_knuckle_joint", + "left_inner_finger_joint", +) class DroidEmbodimentBase(EmbodimentBase, ABC): @@ -79,12 +99,20 @@ def __init__( self, enable_cameras: bool = False, initial_pose: Pose | None = None, + initial_joint_pose: list[float] | None = None, concatenate_observation_terms: bool = False, arm_mode: ArmMode | None = None, stand_height_m: float = _DROID_STAND_PRIM.stand_default_height, placement_bbox_stand_only: bool = False, + collision_mode: CollisionMode | str | None = None, ): - super().__init__(enable_cameras, initial_pose, concatenate_observation_terms, arm_mode) + super().__init__( + enable_cameras=enable_cameras, + initial_pose=initial_pose, + concatenate_observation_terms=concatenate_observation_terms, + arm_mode=arm_mode, + collision_mode=collision_mode, + ) self.stand_height_m = stand_height_m self.placement_bbox_stand_only = placement_bbox_stand_only self.scene_config = DroidSceneCfg() @@ -98,6 +126,8 @@ def __init__( self.camera_config = DroidCameraCfg() self.observation_config = DroidObservationsCfg() self.event_config = DroidEventCfg() + if initial_joint_pose is not None: + self.set_initial_joint_pose(initial_joint_pose) self.reward_config = None self.mimic_env = None self.add_camera_variations(self.camera_config) @@ -111,15 +141,23 @@ def get_bounding_box(self) -> AxisAlignedBoundingBox: prim_path = _DROID_ROBOT_PRIM.stand_prim_path if self.placement_bbox_stand_only else None return super().get_bounding_box(prim_path=prim_path) - def set_initial_joint_pose(self, initial_joint_pose: Mapping[str, float]) -> None: - """Spawn and reset the arm at ``initial_joint_pose``, keyed by joint name or Isaac Lab regex. + def get_collision_mesh(self) -> trimesh.Trimesh: + """Return one posed box mesh for the robot and stand.""" + from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos - Placement geometry is derived from the same spawn state, so it follows the new pose. - """ + source = self.get_placement_geometry_source() + return extract_trimesh_from_usd_at_joint_pos(source.usd_path, source.joint_pos, source.scale) + + def set_initial_joint_pose(self, initial_joint_pose: list[float]) -> None: + """Set the spawn and reset joint positions in articulation order.""" + expected_joint_count = len(_DROID_JOINT_NAMES) + assert ( + len(initial_joint_pose) == expected_joint_count + ), f"expected {expected_joint_count} joint positions, got {len(initial_joint_pose)}" assert self.scene_config is not None, "scene_config must be populated before setting the joint pose" robot = self.scene_config.robot assert robot is not None, "scene_config.robot must be populated before setting the joint pose" - robot.init_state = robot.init_state.replace(joint_pos=dict(initial_joint_pose)) + robot.init_state = robot.init_state.replace(joint_pos=dict(zip(_DROID_JOINT_NAMES, initial_joint_pose))) def get_ee_frame_name(self, arm_mode: ArmMode) -> str: return "ee_frame" @@ -139,18 +177,22 @@ def __init__( self, enable_cameras: bool = False, initial_pose: Pose | None = None, + initial_joint_pose: list[float] | None = None, concatenate_observation_terms: bool = False, arm_mode: ArmMode | None = None, stand_height_m: float = _DROID_STAND_PRIM.stand_default_height, placement_bbox_stand_only: bool = False, + collision_mode: CollisionMode | str | None = None, ): super().__init__( - enable_cameras, - initial_pose, - concatenate_observation_terms, - arm_mode, - stand_height_m, - placement_bbox_stand_only, + enable_cameras=enable_cameras, + initial_pose=initial_pose, + initial_joint_pose=initial_joint_pose, + concatenate_observation_terms=concatenate_observation_terms, + arm_mode=arm_mode, + stand_height_m=stand_height_m, + placement_bbox_stand_only=placement_bbox_stand_only, + collision_mode=collision_mode, ) self.action_config = DroidDifferentialIKActionsCfg() @@ -166,18 +208,22 @@ def __init__( self, enable_cameras: bool = False, initial_pose: Pose | None = None, + initial_joint_pose: list[float] | None = None, concatenate_observation_terms: bool = False, arm_mode: ArmMode | None = None, stand_height_m: float = _DROID_STAND_PRIM.stand_default_height, placement_bbox_stand_only: bool = False, + collision_mode: CollisionMode | str | None = None, ): super().__init__( - enable_cameras, - initial_pose, - concatenate_observation_terms, - arm_mode, - stand_height_m, - placement_bbox_stand_only, + enable_cameras=enable_cameras, + initial_pose=initial_pose, + initial_joint_pose=initial_joint_pose, + concatenate_observation_terms=concatenate_observation_terms, + arm_mode=arm_mode, + stand_height_m=stand_height_m, + placement_bbox_stand_only=placement_bbox_stand_only, + collision_mode=collision_mode, ) self.action_config = DroidRelativeJointPositionActionsCfg() @@ -194,18 +240,22 @@ def __init__( self, enable_cameras: bool = False, initial_pose: Pose | None = None, + initial_joint_pose: list[float] | None = None, concatenate_observation_terms: bool = False, arm_mode: ArmMode | None = None, stand_height_m: float = _DROID_STAND_PRIM.stand_default_height, placement_bbox_stand_only: bool = False, + collision_mode: CollisionMode | str | None = None, ): super().__init__( - enable_cameras, - initial_pose, - concatenate_observation_terms, - arm_mode, - stand_height_m, - placement_bbox_stand_only, + enable_cameras=enable_cameras, + initial_pose=initial_pose, + initial_joint_pose=initial_joint_pose, + concatenate_observation_terms=concatenate_observation_terms, + arm_mode=arm_mode, + stand_height_m=stand_height_m, + placement_bbox_stand_only=placement_bbox_stand_only, + collision_mode=collision_mode, ) self.action_config = DroidAbsoluteJointPositionActionsCfg() diff --git a/isaaclab_arena/embodiments/embodiment_base.py b/isaaclab_arena/embodiments/embodiment_base.py index 06f04244f9..6c767668a6 100644 --- a/isaaclab_arena/embodiments/embodiment_base.py +++ b/isaaclab_arena/embodiments/embodiment_base.py @@ -14,6 +14,7 @@ from isaaclab.managers.recorder_manager import RecorderManagerBaseCfg from isaaclab_arena.embodiments.common.arm_mode import ArmMode +from isaaclab_arena.relations.collision_mode import CollisionMode from isaaclab_arena.relations.placement_asset import PlaceableAsset from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox from isaaclab_arena.utils.cameras import ArenaCameraCfg, make_camera_observation_cfg @@ -50,9 +51,10 @@ def __init__( initial_pose: Pose | None = None, concatenate_observation_terms: bool = False, arm_mode: ArmMode | None = None, + collision_mode: CollisionMode | str | None = None, ): assert self.name is not None, "Embodiment name is required" - super().__init__(name=self.name, tags=self.tags) + super().__init__(name=self.name, tags=self.tags, collision_mode=collision_mode) if "embodiment" not in self.tags: self.tags.append("embodiment") self.enable_cameras = enable_cameras @@ -89,8 +91,6 @@ def get_placement_geometry_source(self) -> PlacementGeometrySource: def get_bounding_box(self, prim_path: str | None = None) -> AxisAlignedBoundingBox: """Return root-relative bounds of the articulation posed at its configured joint positions. - Shared and cached across callers, so treat the result as read-only. - Args: prim_path: Optional sub-prim to bound (e.g. stand only). When None, bounds the full default prim. @@ -104,16 +104,12 @@ def get_bounding_box(self, prim_path: str | None = None) -> AxisAlignedBoundingB ) def get_collision_mesh(self) -> trimesh.Trimesh | None: - """Return the robot's collision mesh, posed at its configured initial joint positions. - - The mesh matches the robot as spawned even when its USD was authored in another joint - configuration. Shared and cached across callers, so treat the result as read-only. - """ + """Return the robot mesh from its USD default prim.""" # Import locally because USD/pxr is available only after simulation initialization. - from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos + from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_path source = self.get_placement_geometry_source() - return extract_trimesh_from_usd_at_joint_pos(source.usd_path, source.joint_pos, source.scale) + return extract_trimesh_from_usd_path(source.usd_path, source.scale) def _set_initial_pose(self, pose: Pose | PoseRange | PosePerEnv) -> None: """Store the configured pose; the construction pose is applied in ``get_scene_cfg``.""" diff --git a/isaaclab_arena/embodiments/franka/franka.py b/isaaclab_arena/embodiments/franka/franka.py index 3f8f084015..917c5ad226 100644 --- a/isaaclab_arena/embodiments/franka/franka.py +++ b/isaaclab_arena/embodiments/franka/franka.py @@ -3,9 +3,11 @@ # # SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations import torch -from collections.abc import Mapping, Sequence +from collections.abc import Sequence +from typing import TYPE_CHECKING import isaaclab.envs.mdp as mdp_isaac_lab import isaaclab.sim as sim_utils @@ -41,6 +43,9 @@ from isaaclab_arena.utils.cameras import ArenaCameraCfg from isaaclab_arena.utils.pose import Pose +if TYPE_CHECKING: + import trimesh + _DEFAULT_CAMERA_OFFSET = Pose(position_xyz=(0.11, -0.031, -0.074), rotation_xyzw=(0.0, 0.0, 0.70711, 0.70711)) _FRANKA_ROBOT_PRIM = RobotPrimSpec( @@ -58,6 +63,17 @@ footprint_scale_xy=(1.2, 1.2), stand_default_height=0.8755, ) +_FRANKA_JOINT_NAMES = ( + "panda_joint1", + "panda_joint2", + "panda_joint3", + "panda_joint4", + "panda_joint5", + "panda_joint6", + "panda_joint7", + "panda_finger_joint1", + "panda_finger_joint2", +) class FrankaEmbodimentBase(EmbodimentBase): @@ -86,16 +102,23 @@ def __init__( self.observation_config.policy.concatenate_terms = self.concatenate_observation_terms self.add_camera_variations(self.camera_config) - def set_initial_joint_pose(self, initial_joint_pose: Mapping[str, float]) -> None: - """Spawn and reset the arm at ``initial_joint_pose``, keyed by joint name or Isaac Lab regex. + def get_collision_mesh(self) -> trimesh.Trimesh: + """Return one posed box mesh for the robot and stand.""" + from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos - Call after construction, once :attr:`scene_config` holds the robot. Placement geometry is - derived from the same spawn state, so it follows the new pose. - """ + source = self.get_placement_geometry_source() + return extract_trimesh_from_usd_at_joint_pos(source.usd_path, source.joint_pos, source.scale) + + def set_initial_joint_pose(self, initial_joint_pose: list[float]) -> None: + """Set the spawn and reset joint positions in articulation order.""" + expected_joint_count = len(_FRANKA_JOINT_NAMES) + assert ( + len(initial_joint_pose) == expected_joint_count + ), f"expected {expected_joint_count} joint positions, got {len(initial_joint_pose)}" assert self.scene_config is not None, "scene_config must be populated before setting the joint pose" robot = self.scene_config.robot assert robot is not None, "scene_config.robot must be populated before setting the joint pose" - robot.init_state = robot.init_state.replace(joint_pos=dict(initial_joint_pose)) + robot.init_state = robot.init_state.replace(joint_pos=dict(zip(_FRANKA_JOINT_NAMES, initial_joint_pose))) def get_ee_frame_name(self, arm_mode: ArmMode) -> str: return "ee_frame" @@ -112,6 +135,7 @@ def __init__( self, enable_cameras: bool = False, initial_pose: Pose | None = None, + initial_joint_pose: list[float] | None = None, concatenate_observation_terms: bool = False, arm_mode: ArmMode | None = None, ): @@ -122,6 +146,8 @@ def __init__( arm_mode=arm_mode, ) self.scene_config.robot = _franka_robot_cfg_on_stand(FRANKA_PANDA_HIGH_PD_CFG.copy()) + if initial_joint_pose is not None: + self.set_initial_joint_pose(initial_joint_pose) self.action_config = FrankaIKActionCfg() def get_command_body_name(self) -> str: @@ -163,6 +189,7 @@ def __init__( self, enable_cameras: bool = False, initial_pose: Pose | None = None, + initial_joint_pose: list[float] | None = None, concatenate_observation_terms: bool = False, arm_mode: ArmMode | None = None, ): @@ -174,6 +201,8 @@ def __init__( ) self.action_config = FrankaJointPosActionsCfg() self.scene_config.robot = _franka_robot_cfg_on_stand(FRANKA_PANDA_CFG.copy()) + if initial_joint_pose is not None: + self.set_initial_joint_pose(initial_joint_pose) def get_command_body_name(self) -> str: return "panda_hand" diff --git a/isaaclab_arena/tests/test_embodiment_collision_mesh.py b/isaaclab_arena/tests/test_embodiment_collision_mesh.py index fdacf018f6..1488f86fd2 100644 --- a/isaaclab_arena/tests/test_embodiment_collision_mesh.py +++ b/isaaclab_arena/tests/test_embodiment_collision_mesh.py @@ -26,23 +26,27 @@ def _test_embodiment_provides_robot_collision_mesh(simulation_app) -> bool: extents = mesh.extents assert all(e < 3.0 for e in extents), f"mesh leaked non-robot geometry: extents {extents}" - # The placement bbox covers the same posed assembly plus the analytic gprims that mesh extraction - # cannot represent, so it is marginally larger but must not diverge. + # Link boxes conservatively cover the same posed Gprims as the placement bbox. bbox = emb.get_bounding_box() bbox_size = (bbox.max_point - bbox.min_point)[0].tolist() for mesh_extent, box_extent in zip(extents, bbox_size): assert ( - box_extent + 1e-3 >= mesh_extent - ), f"placement bbox {bbox_size} should cover robot mesh extents {extents.tolist()}" + mesh_extent + 1e-3 >= box_extent + ), f"robot mesh extents {extents.tolist()} should cover placement bbox {bbox_size}" assert abs(mesh_extent - box_extent) < 0.2, f"mesh extents {extents} disagree with box {bbox_size}" - # Both derivations open the USD and pose it, so results are cached rather than recomputed per - # solve step. The cache is keyed by USD, joint positions and scale, so an identical embodiment - # shares it. - assert emb.get_collision_mesh() is mesh - assert DroidAbsoluteJointPositionEmbodiment().get_collision_mesh() is mesh - assert emb.get_bounding_box() is bbox - assert DroidAbsoluteJointPositionEmbodiment().get_bounding_box() is bbox + # Cached geometry is copied before return so caller mutation cannot poison later queries. + another_mesh = emb.get_collision_mesh() + assert another_mesh is not mesh + original_vertex = another_mesh.vertices[0].copy() + mesh.vertices[0] += 100.0 + assert (emb.get_collision_mesh().vertices[0] == original_vertex).all() + + another_bbox = emb.get_bounding_box() + assert another_bbox is not bbox + original_min = another_bbox.min_point.clone() + bbox.min_point.add_(100.0) + assert emb.get_bounding_box().min_point.equal(original_min) # Isaac Lab reaches placeable assets through EventTermCfg params and validates whatever they # hold without tracking visited objects, so an embodiment holding its mesh would send @@ -65,8 +69,8 @@ def _noop(env, env_ids, embodiment): def _test_spawn_pose_matches_the_reset_pose(simulation_app) -> bool: """The spawn joint positions placement geometry is posed at are the ones a reset drives to. - Also covers ``set_initial_joint_pose``, which environments use to pose the arm for their scene - and which has to move the spawn state placement geometry reads, not just the reset. + Also covers constructor joint-pose overrides, which must update the spawn state placement + geometry reads. """ import torch @@ -79,28 +83,18 @@ def _test_spawn_pose_matches_the_reset_pose(simulation_app) -> bool: from isaaclab_arena.scene.scene import Scene from isaaclab_arena.utils.usd_articulation import resolve_joint_pos_patterns - # The pose isaaclab_arena_environments/cube_goal_pose_environment.py reaches from. - overridden_pose = { - "panda_joint1": 0.0444, - "panda_joint2": -0.1894, - "panda_joint3": -0.1107, - "panda_joint4": -2.5148, - "panda_joint5": 0.0044, - "panda_joint6": 2.3775, - "panda_joint7": 0.6952, - "panda_finger_joint.*": 0.0400, - } + franka_override = [0.0444, -0.1894, -0.1107, -2.5148, 0.0044, 2.3775, 0.6952, 0.0400, 0.0400] + droid_override = [0.0444, -0.1894, -0.1107, -2.5148, 0.0044, 2.3775, 0.6952, 0.0400] + [0.0] * 5 for embodiment_class, override in ( (FrankaIKEmbodiment, None), (DroidAbsoluteJointPositionEmbodiment, None), - (FrankaIKEmbodiment, overridden_pose), + (FrankaIKEmbodiment, franka_override), + (DroidAbsoluteJointPositionEmbodiment, droid_override), ): env = None try: - embodiment = embodiment_class() - if override is not None: - embodiment.set_initial_joint_pose(override) + embodiment = embodiment_class(initial_joint_pose=override) environment = IsaacLabArenaEnvironment( name=f"spawn_pose_{embodiment_class.__name__}_{override is not None}", embodiment=embodiment, @@ -116,9 +110,10 @@ def _test_spawn_pose_matches_the_reset_pose(simulation_app) -> bool: # Both sides below are read back from the spawn state, so pin the override against what # was asked for; otherwise a setter that quietly dropped it would still agree with itself. if override is not None: - assert spawn_pose == resolve_joint_pos_patterns(joint_names, override), ( + expected_override = dict(zip(joint_names, override)) + assert spawn_pose == expected_override, ( "set_initial_joint_pose did not reach the spawn state placement geometry reads:" - f" {spawn_pose} != {resolve_joint_pos_patterns(joint_names, override)}" + f" {spawn_pose} != {expected_override}" ) # Compare against the defaults a reset restores rather than measured positions, which # carry gravity sag and the joint randomisation event on top. @@ -163,7 +158,7 @@ def _test_joint_position_action_offset_is_pinned(simulation_app) -> bool: try: embodiment = FrankaJointPosEmbodiment() - embodiment.set_initial_joint_pose({"panda_joint.*": 0.0, "panda_finger_joint.*": 0.0}) + embodiment.set_initial_joint_pose([0.0] * 9) arm_action = embodiment.action_config.arm_action assert not arm_action.use_default_offset, "the action zero point would follow the spawn pose" assert arm_action.offset == trained_against, f"action zero point moved to {arm_action.offset}" @@ -184,13 +179,13 @@ def test_embodiment_provides_robot_collision_mesh(): def test_spawn_pose_matches_the_reset_pose(): """Pytest entry point for the spawn-versus-reset joint-position test.""" - result = run_simulation_app_function(_test_spawn_pose_matches_the_reset_pose, headless=True) + result = run_function_with_persistent_simulation_app(_test_spawn_pose_matches_the_reset_pose, headless=True) assert result, f"Test {test_spawn_pose_matches_the_reset_pose.__name__} failed" def test_joint_position_action_offset_is_pinned(): """Pytest entry point for the joint-position action offset test.""" - result = run_simulation_app_function(_test_joint_position_action_offset_is_pinned, headless=True) + result = run_function_with_persistent_simulation_app(_test_joint_position_action_offset_is_pinned, headless=True) assert result, f"Test {test_joint_position_action_offset_is_pinned.__name__} failed" diff --git a/isaaclab_arena/tests/test_usd_articulation.py b/isaaclab_arena/tests/test_usd_articulation.py index 059a386864..ebca92a134 100644 --- a/isaaclab_arena/tests/test_usd_articulation.py +++ b/isaaclab_arena/tests/test_usd_articulation.py @@ -12,7 +12,7 @@ import math import numpy as np -from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function +from isaaclab_arena.tests.utils.persistent_simulation_app import run_function_with_persistent_simulation_app HEADLESS = True @@ -172,12 +172,15 @@ def _test_joint_pos_patterns_expand_to_joint_names(simulation_app) -> bool: def _test_droid_geometry_tracks_configured_joint_positions(simulation_app) -> bool: - """Posing the real Droid at its configured joints reproduces its authored geometry, zero does not.""" + """Droid's configured and zero joint positions produce different conservative link-box proxies.""" from pxr import Usd from isaaclab_arena.embodiments.droid.droid import DroidAbsoluteJointPositionEmbodiment from isaaclab_arena.utils.usd_articulation import articulation_joint_prims - from isaaclab_arena.utils.usd_helpers import extract_trimesh_from_usd_at_joint_pos, extract_trimesh_from_usd_path + from isaaclab_arena.utils.usd_helpers import ( + compute_local_bounding_box_from_usd_at_joint_pos, + extract_trimesh_from_usd_at_joint_pos, + ) robot = DroidAbsoluteJointPositionEmbodiment().scene_config.robot usd_path = robot.spawn.usd_path @@ -188,22 +191,19 @@ def _test_droid_geometry_tracks_configured_joint_positions(simulation_app) -> bo # Passed verbatim, so the config's regex keys (e.g. "right_outer.*") must resolve to real joints. configured = robot.init_state.joint_pos - authored = extract_trimesh_from_usd_path(usd_path) posed = extract_trimesh_from_usd_at_joint_pos(usd_path, configured) - np.testing.assert_allclose(posed.extents, authored.extents, atol=2e-3) + posed_bbox = compute_local_bounding_box_from_usd_at_joint_pos(usd_path, configured) + np.testing.assert_allclose(posed.extents, posed_bbox.size.numpy()[0], atol=2e-3) - # At zero the arm stands straight up, so it is taller and narrower than the authored pose. + # At zero the arm stands straight up, so it is taller and narrower than the configured pose. zero = extract_trimesh_from_usd_at_joint_pos(usd_path, {}) - assert zero.extents[2] > authored.extents[2] + 0.2, f"zero pose should stand taller: {zero.extents}" - assert zero.extents[0] < authored.extents[0] - 0.2, f"zero pose should be narrower: {zero.extents}" + assert zero.extents[2] > posed.extents[2] + 0.2, f"zero pose should stand taller: {zero.extents}" + assert zero.extents[0] < posed.extents[0] - 0.2, f"zero pose should be narrower: {zero.extents}" return True def _test_droid_posed_bounding_box_covers_all_geometry(simulation_app) -> bool: - """The posed bbox covers every gprim, not just the meshes. - - Deriving bounds from the posed mesh instead drops analytic gprims such as ``gripper_adapter``. - """ + """The link-box proxy covers every posed Gprim, including analytic geometry.""" from isaaclab_arena.embodiments.droid.droid import DroidAbsoluteJointPositionEmbodiment from isaaclab_arena.utils.usd_helpers import ( compute_local_bounding_box_from_usd, @@ -220,12 +220,12 @@ def _test_droid_posed_bounding_box_covers_all_geometry(simulation_app) -> bool: np.testing.assert_allclose(posed.size.numpy()[0], authored.size.numpy()[0], atol=2e-3) np.testing.assert_allclose(posed.min_point.numpy()[0], authored.min_point.numpy()[0], atol=2e-3) - # The bbox must be at least as large as the mesh, since it also covers the non-mesh gprims. + # Both paths include every Gprim. Link-local boxes are conservative, so their aggregate can be + # larger than the exact posed bound but must never be smaller. mesh_extents = extract_trimesh_from_usd_at_joint_pos(usd_path, configured).extents posed_size = posed.size.numpy()[0] - assert np.all(posed_size >= mesh_extents - 1e-6), f"bbox {posed_size} smaller than mesh {mesh_extents}" - # Droid's dropped gripper_adapter makes this a strict difference, so a mesh-only bbox would fail. - assert np.any(posed_size > mesh_extents + 1e-3), "expected non-mesh geometry to widen the bbox" + assert np.all(mesh_extents >= posed_size - 1e-6), f"proxy {mesh_extents} does not cover bbox {posed_size}" + assert np.all(mesh_extents - posed_size < 0.2), f"link-local proxy is unexpectedly loose: {mesh_extents}" return True @@ -331,6 +331,7 @@ def _test_instanced_geometry_is_posed_with_its_link(simulation_app) -> bool: from isaaclab_arena.utils.usd_helpers import ( compute_local_bounding_box_from_usd_at_joint_pos, + extract_link_bbox_meshes_from_usd_at_joint_pos, extract_trimesh_from_usd_at_joint_pos, ) @@ -348,6 +349,16 @@ def _test_instanced_geometry_is_posed_with_its_link(simulation_app) -> bool: usd_path = f"{tmp_dir}/arm_with_instanced_mount.usda" stage.Export(usd_path) + components = extract_link_bbox_meshes_from_usd_at_joint_pos(usd_path, {}) + assert len(components) == 2, "the base and forearm must each produce exactly one collision component" + assert all(component.is_watertight for component in components) + try: + extract_link_bbox_meshes_from_usd_at_joint_pos(usd_path, {}, scale=(-1.0, 1.0, 1.0)) + except AssertionError as error: + assert "positive" in str(error) + else: + raise AssertionError("negative spawn scale must be rejected before it flips SDF signs") + # The mount reaches 1.7 along +X: forearm at 1.0, mount offset 0.5, half extent 0.2. rest = compute_local_bounding_box_from_usd_at_joint_pos(usd_path, {}) np.testing.assert_allclose(rest.max_point.numpy()[0][0], 1.7, atol=1e-6) @@ -361,44 +372,58 @@ def _test_instanced_geometry_is_posed_with_its_link(simulation_app) -> bool: def test_revolute_joint_swings_child_link(): - assert run_simulation_app_function(_test_revolute_joint_swings_child_link, headless=HEADLESS) + assert run_function_with_persistent_simulation_app(_test_revolute_joint_swings_child_link, headless=HEADLESS) def test_zero_joint_position_preserves_authored_pose(): - assert run_simulation_app_function(_test_zero_joint_position_preserves_authored_pose, headless=HEADLESS) + assert run_function_with_persistent_simulation_app( + _test_zero_joint_position_preserves_authored_pose, headless=HEADLESS + ) def test_prismatic_joint_translates_child_link(): - assert run_simulation_app_function(_test_prismatic_joint_translates_child_link, headless=HEADLESS) + assert run_function_with_persistent_simulation_app(_test_prismatic_joint_translates_child_link, headless=HEADLESS) def test_authored_pose_away_from_joint_zero_is_corrected(): - assert run_simulation_app_function(_test_authored_pose_away_from_joint_zero_is_corrected, headless=HEADLESS) + assert run_function_with_persistent_simulation_app( + _test_authored_pose_away_from_joint_zero_is_corrected, headless=HEADLESS + ) def test_unknown_joint_name_is_rejected(): - assert run_simulation_app_function(_test_unknown_joint_name_is_rejected, headless=HEADLESS) + assert run_function_with_persistent_simulation_app(_test_unknown_joint_name_is_rejected, headless=HEADLESS) def test_joint_pos_patterns_expand_to_joint_names(): - assert run_simulation_app_function(_test_joint_pos_patterns_expand_to_joint_names, headless=HEADLESS) + assert run_function_with_persistent_simulation_app( + _test_joint_pos_patterns_expand_to_joint_names, headless=HEADLESS + ) def test_droid_geometry_tracks_configured_joint_positions(): - assert run_simulation_app_function(_test_droid_geometry_tracks_configured_joint_positions, headless=HEADLESS) + assert run_function_with_persistent_simulation_app( + _test_droid_geometry_tracks_configured_joint_positions, headless=HEADLESS + ) def test_droid_posed_bounding_box_covers_all_geometry(): - assert run_simulation_app_function(_test_droid_posed_bounding_box_covers_all_geometry, headless=HEADLESS) + assert run_function_with_persistent_simulation_app( + _test_droid_posed_bounding_box_covers_all_geometry, headless=HEADLESS + ) def test_offline_posing_matches_physx_link_poses(): - assert run_simulation_app_function(_test_offline_posing_matches_physx_link_poses, headless=HEADLESS) + assert run_function_with_persistent_simulation_app(_test_offline_posing_matches_physx_link_poses, headless=HEADLESS) def test_closed_loop_articulation_poses_a_spanning_tree(): - assert run_simulation_app_function(_test_closed_loop_articulation_poses_a_spanning_tree, headless=HEADLESS) + assert run_function_with_persistent_simulation_app( + _test_closed_loop_articulation_poses_a_spanning_tree, headless=HEADLESS + ) def test_instanced_geometry_is_posed_with_its_link(): - assert run_simulation_app_function(_test_instanced_geometry_is_posed_with_its_link, headless=HEADLESS) + assert run_function_with_persistent_simulation_app( + _test_instanced_geometry_is_posed_with_its_link, headless=HEADLESS + ) diff --git a/isaaclab_arena/utils/usd_helpers.py b/isaaclab_arena/utils/usd_helpers.py index 39e407ddc1..6a6667fc1d 100644 --- a/isaaclab_arena/utils/usd_helpers.py +++ b/isaaclab_arena/utils/usd_helpers.py @@ -468,16 +468,19 @@ def extract_trimesh_from_usd_path( return extract_trimesh_from_prim(stage, default_prim.GetPath().pathString, scale) -def extract_trimesh_from_usd_at_joint_pos( +def extract_link_bbox_meshes_from_usd_at_joint_pos( usd_path: str, joint_pos: Mapping[str, float], scale: tuple[float, float, float] = (1.0, 1.0, 1.0), -) -> trimesh.Trimesh: - """Extract an articulation's mesh posed at joint_pos, in its default prim's local frame. +) -> tuple[trimesh.Trimesh, ...]: + """Return posed box meshes grouped by rigid body, in the default prim's local frame. Joints the articulation has but joint_pos omits are posed at zero, so the result depends only on - joint_pos and not on the configuration the asset happens to be authored in. Cached in process - and shared with every other caller, so treat the result as read-only. + joint_pos and not on the configuration the asset happens to be authored in. Each box encloses all + default-purpose geometry attached to one rigid body in that body's local frame, then follows the + body's posed transform. Geometry not attached to a rigid body is represented by one root box. + + Results are cached in process and copied before return. Args: usd_path: Path to the articulation's .usd/.usda/.usdc file. @@ -485,29 +488,136 @@ def extract_trimesh_from_usd_at_joint_pos( scale: Spawn-time scale passed to ``UsdFileCfg``. Returns: - Combined trimesh in the scaled default-prim frame. + Watertight link-box meshes in the scaled default-prim frame. """ - return _extract_trimesh_from_usd_at_joint_pos(usd_path, tuple(sorted(joint_pos.items())), tuple(scale)) + meshes = _extract_link_bbox_meshes_from_usd_at_joint_pos(usd_path, tuple(sorted(joint_pos.items())), tuple(scale)) + return tuple(mesh.copy() for mesh in meshes) # NOTE(zihaox, 2026-07-28): Cache here rather than on the asset. Isaac Lab reaches assets through # EventTermCfg params, and configclass's validation walk tracks no visited set, so a trimesh held by # an asset sends it recursing through trimesh's internal back-references until the stack overflows. @functools.lru_cache(maxsize=_POSED_GEOMETRY_CACHE_SIZE) -def _extract_trimesh_from_usd_at_joint_pos( +def _extract_link_bbox_meshes_from_usd_at_joint_pos( usd_path: str, joint_pos_items: tuple[tuple[str, float], ...], scale: tuple[float, float, float], -) -> trimesh.Trimesh: - """Cacheable body of ``extract_trimesh_from_usd_at_joint_pos``, keyed by hashable arguments.""" +) -> tuple[trimesh.Trimesh, ...]: + """Cacheable body of ``extract_link_bbox_meshes_from_usd_at_joint_pos``.""" + assert all( + component > 0 for component in scale + ), f"All scale components must be positive (negative scale flips winding/SDF sign), got {scale}" joint_pos = dict(joint_pos_items) stage = Usd.Stage.Open(usd_path) assert stage is not None, f"could not open USD: {usd_path}" default_prim = stage.GetDefaultPrim() or stage.GetPseudoRoot() - default_prim_path = default_prim.GetPath().pathString + root_path = default_prim.GetPath().pathString resolved = resolve_joint_pos_patterns(articulation_joint_prims(default_prim), joint_pos) - deltas = compute_posed_prim_world_deltas(stage, default_prim_path, resolved) - return extract_trimesh_from_prim(stage, default_prim_path, scale, prim_world_deltas=deltas) + deltas = compute_posed_prim_world_deltas(stage, root_path, resolved) + return _posed_link_bbox_meshes(stage, default_prim, deltas, scale) + + +@functools.lru_cache(maxsize=_POSED_GEOMETRY_CACHE_SIZE) +def _extract_trimesh_from_usd_at_joint_pos( + usd_path: str, + joint_pos_items: tuple[tuple[str, float], ...], + scale: tuple[float, float, float], +) -> trimesh.Trimesh: + """Return a cached mesh containing all of an articulation's link boxes.""" + meshes = _extract_link_bbox_meshes_from_usd_at_joint_pos(usd_path, joint_pos_items, scale) + return trimesh.util.concatenate(meshes) + + +def extract_trimesh_from_usd_at_joint_pos( + usd_path: str, + joint_pos: Mapping[str, float], + scale: tuple[float, float, float] = (1.0, 1.0, 1.0), +) -> trimesh.Trimesh: + """Return one mesh containing all of an articulation's posed link boxes.""" + mesh = _extract_trimesh_from_usd_at_joint_pos(usd_path, tuple(sorted(joint_pos.items())), tuple(scale)) + return mesh.copy() + + +def _nearest_rigid_body_ancestor(prim: Usd.Prim, root_prim: Usd.Prim) -> Usd.Prim | None: + """Return the nearest rigid-body ancestor at or below root_prim.""" + candidate = prim + root_path = root_prim.GetPath() + while candidate and candidate.IsValid() and candidate.GetPath().HasPrefix(root_path): + if candidate.HasAPI(UsdPhysics.RigidBodyAPI): + return candidate + if candidate == root_prim: + break + candidate = candidate.GetParent() + return None + + +def _untransformed_gprim_corners( + prim: Usd.Prim, + bbox_cache: UsdGeom.BBoxCache, +) -> np.ndarray | None: + """Return a Gprim's local homogeneous bound corners, or None for an empty bound.""" + local_range = bbox_cache.ComputeUntransformedBound(prim).ComputeAlignedRange() + if local_range.IsEmpty(): + return None + low, high = local_range.GetMin(), local_range.GetMax() + return np.array( + [[x, y, z, 1.0] for x in (low[0], high[0]) for y in (low[1], high[1]) for z in (low[2], high[2])], + dtype=np.float64, + ) + + +def _posed_link_bbox_meshes( + stage: Usd.Stage, + default_prim: Usd.Prim, + body_deltas: Mapping[str, np.ndarray], + scale: tuple[float, float, float], +) -> tuple[trimesh.Trimesh, ...]: + """Build and pose one local bounding box per rigid body or static root.""" + time = Usd.TimeCode.Default() + bbox_cache = UsdGeom.BBoxCache(time, includedPurposes=[UsdGeom.Tokens.default_]) + root_world_tf = np.array(UsdGeom.Xformable(default_prim).ComputeLocalToWorldTransform(time)) + root_world_tf_inv = np.linalg.inv(root_world_tf) + scale_np = np.asarray(scale, dtype=np.float64) + + frames: dict[str, Usd.Prim] = {} + corners_by_frame: dict[str, list[np.ndarray]] = {} + for prim in Usd.PrimRange(default_prim, Usd.TraverseInstanceProxies()): + if not prim.IsA(UsdGeom.Gprim): + continue + local_corners = _untransformed_gprim_corners(prim, bbox_cache) + if local_corners is None: + continue + body_prim = _nearest_rigid_body_ancestor(prim, default_prim) + frame_prim = body_prim or default_prim + frame_path = frame_prim.GetPath().pathString + frames.setdefault(frame_path, frame_prim) + frame_world_tf = np.array(UsdGeom.Xformable(frame_prim).ComputeLocalToWorldTransform(time)) + prim_world_tf = np.array(UsdGeom.Xformable(prim).ComputeLocalToWorldTransform(time)) + corners_by_frame.setdefault(frame_path, []).append( + local_corners @ (prim_world_tf @ np.linalg.inv(frame_world_tf)) + ) + + meshes: list[trimesh.Trimesh] = [] + for frame_path, local_corners in corners_by_frame.items(): + stacked = np.vstack(local_corners)[:, :3] + low, high = stacked.min(axis=0), stacked.max(axis=0) + extents = high - low + assert np.all(extents > 0.0), f"degenerate geometry under {frame_path}: extents={extents}" + box = trimesh.creation.box(extents=extents) + box.apply_translation((low + high) / 2.0) + + frame_prim = frames[frame_path] + posed_frame_world_tf = np.array(UsdGeom.Xformable(frame_prim).ComputeLocalToWorldTransform(time)) + delta = body_deltas.get(frame_path) + if delta is not None: + posed_frame_world_tf = posed_frame_world_tf @ delta + frame_to_root = posed_frame_world_tf @ root_world_tf_inv + vertices_h = np.column_stack([box.vertices, np.ones(len(box.vertices))]) + vertices = (vertices_h @ frame_to_root)[:, :3] * scale_np + meshes.append(trimesh.Trimesh(vertices=vertices, faces=box.faces, process=False)) + + assert meshes, f"no bounded geometry found under {default_prim.GetPath()}" + return tuple(meshes) def compute_local_bounding_box_from_usd_at_joint_pos( @@ -523,9 +633,8 @@ def compute_local_bounding_box_from_usd_at_joint_pos( mesh extraction drops analytic gprims (Droid's ``gripper_adapter``), which would understate the robot's footprint by centimetres. - The result is cached on the arguments and shared with every other caller, so treat it as - read-only. Relation losses ask for bounds every optimisation step, which would otherwise reopen - the USD and redo the posing per step. + The result is cached on the arguments and copied before return. Relation losses ask for bounds + every optimisation step, which would otherwise reopen the USD and redo the posing per step. Args: usd_path: Path to the articulation's .usd/.usda/.usdc file. @@ -536,9 +645,10 @@ def compute_local_bounding_box_from_usd_at_joint_pos( Returns: AxisAlignedBoundingBox containing the posed local bounds. """ - return _compute_local_bounding_box_from_usd_at_joint_pos( + bbox = _compute_local_bounding_box_from_usd_at_joint_pos( usd_path, tuple(sorted(joint_pos.items())), tuple(scale), prim_path ) + return AxisAlignedBoundingBox(bbox.min_point.clone(), bbox.max_point.clone()) @functools.lru_cache(maxsize=_POSED_GEOMETRY_CACHE_SIZE) diff --git a/isaaclab_arena_environments/cube_goal_pose_environment.py b/isaaclab_arena_environments/cube_goal_pose_environment.py index 427e484f17..828df5c49a 100644 --- a/isaaclab_arena_environments/cube_goal_pose_environment.py +++ b/isaaclab_arena_environments/cube_goal_pose_environment.py @@ -59,16 +59,7 @@ def build(self, cfg: CubeGoalPoseEnvironmentCfg) -> IsaacLabArenaEnvironment: rotation_xyzw=(0.0, 0.0, 0.0, 1.0), ) ) - embodiment.set_initial_joint_pose({ - "panda_joint1": 0.0444, - "panda_joint2": -0.1894, - "panda_joint3": -0.1107, - "panda_joint4": -2.5148, - "panda_joint5": 0.0044, - "panda_joint6": 2.3775, - "panda_joint7": 0.6952, - "panda_finger_joint.*": 0.0400, - }) + embodiment.set_initial_joint_pose([0.0444, -0.1894, -0.1107, -2.5148, 0.0044, 2.3775, 0.6952, 0.0400, 0.0400]) if cfg.teleop_device is not None: teleop_device = self.device_registry.get_device_by_name(cfg.teleop_device)() diff --git a/isaaclab_arena_environments/franka_put_and_close_door_environment.py b/isaaclab_arena_environments/franka_put_and_close_door_environment.py index 825df002c7..c717a963b1 100644 --- a/isaaclab_arena_environments/franka_put_and_close_door_environment.py +++ b/isaaclab_arena_environments/franka_put_and_close_door_environment.py @@ -86,16 +86,7 @@ def build(self, cfg: FrankaPutAndCloseDoorEnvironmentCfg) -> IsaacLabArenaEnviro if cfg.embodiment == "franka_ik": # Set Franka arm pose for kitchen setup - embodiment.set_initial_joint_pose({ - "panda_joint1": 0.0, - "panda_joint2": -1.309, - "panda_joint3": 0.0, - "panda_joint4": -2.793, - "panda_joint5": 0.0, - "panda_joint6": 3.037, - "panda_joint7": 0.740, - "panda_finger_joint.*": 0.04, - }) + embodiment.set_initial_joint_pose([0.0, -1.309, 0.0, -2.793, 0.0, 3.037, 0.740, 0.04, 0.04]) # Create destination reference destination_ref = ObjectReference( diff --git a/isaaclab_arena_environments/kitchen_bench/droid_pick_and_place_lightwheel_kitchen.yaml b/isaaclab_arena_environments/kitchen_bench/droid_pick_and_place_lightwheel_kitchen.yaml index 15d0836214..d2bac53943 100644 --- a/isaaclab_arena_environments/kitchen_bench/droid_pick_and_place_lightwheel_kitchen.yaml +++ b/isaaclab_arena_environments/kitchen_bench/droid_pick_and_place_lightwheel_kitchen.yaml @@ -10,6 +10,7 @@ embodiment: params: stand_height_m: 0.8 placement_bbox_stand_only: true + collision_mode: mesh background: id: kitchen registry_name: lightwheel_robocasa_kitchen diff --git a/isaaclab_arena_environments/sorting_environment.py b/isaaclab_arena_environments/sorting_environment.py index c12c096195..a47fbf104a 100644 --- a/isaaclab_arena_environments/sorting_environment.py +++ b/isaaclab_arena_environments/sorting_environment.py @@ -67,16 +67,9 @@ def build(self, cfg: TableTopSortCubesEnvironmentCfg) -> IsaacLabArenaEnvironmen ) ) - embodiment.set_initial_joint_pose({ - "panda_joint1": 0.0444, - "panda_joint2": -0.1894, - "panda_joint3": -0.1107, - "panda_joint4": -2.5148, - "panda_joint5": 0.0044, - "panda_joint6": 2.3775, - "panda_joint7": 0.6952, - "panda_finger_joint.*": 0.0400, - }) + embodiment.set_initial_joint_pose( + [0.0444, -0.1894, -0.1107, -2.5148, 0.0044, 2.3775, 0.6952, 0.0400, 0.0400] + ) else: raise NotImplementedError From 33861a92fdaa1b765f7197db774f08ed39247196 Mon Sep 17 00:00:00 2001 From: zhx06 Date: Mon, 3 Aug 2026 11:18:54 -0700 Subject: [PATCH 6/7] utils cleanup Signed-off-by: zhx06 --- isaaclab_arena/embodiments/embodiment_base.py | 10 +-- isaaclab_arena/tests/test_usd_articulation.py | 10 ++- isaaclab_arena/utils/usd_articulation.py | 40 +++++------- isaaclab_arena/utils/usd_helpers.py | 64 ++++++------------- 4 files changed, 47 insertions(+), 77 deletions(-) diff --git a/isaaclab_arena/embodiments/embodiment_base.py b/isaaclab_arena/embodiments/embodiment_base.py index 6c767668a6..8c8625c2f0 100644 --- a/isaaclab_arena/embodiments/embodiment_base.py +++ b/isaaclab_arena/embodiments/embodiment_base.py @@ -26,8 +26,8 @@ @dataclass(frozen=True) -class PlacementGeometrySource: - """What an embodiment's placement bounding box and collision mesh are derived from.""" +class ArticulationGeometrySpec: + """USD articulation state used to compute embodiment geometry.""" usd_path: str """Robot USD, as spawned.""" @@ -74,15 +74,15 @@ def __init__( self.xr: Any | None = None self.termination_cfg: Any | None = None - def get_placement_geometry_source(self) -> PlacementGeometrySource: - """Return the USD, scale and joint positions the robot's placement geometry is derived from.""" + def get_placement_geometry_source(self) -> ArticulationGeometrySpec: + """Return the USD articulation state used to compute embodiment geometry.""" assert self.scene_config is not None, "scene_config must be populated before placement" robot = self.scene_config.robot assert robot is not None, "scene_config.robot must be populated before placement" spawn = robot.spawn assert spawn.usd_path is not None, "scene_config.robot must use a USD spawn for placement" scale_x, scale_y, scale_z = spawn.scale or (1.0, 1.0, 1.0) - return PlacementGeometrySource( + return ArticulationGeometrySpec( usd_path=spawn.usd_path, scale=(scale_x, scale_y, scale_z), joint_pos=dict(robot.init_state.joint_pos or {}), diff --git a/isaaclab_arena/tests/test_usd_articulation.py b/isaaclab_arena/tests/test_usd_articulation.py index ebca92a134..794b9d2d84 100644 --- a/isaaclab_arena/tests/test_usd_articulation.py +++ b/isaaclab_arena/tests/test_usd_articulation.py @@ -331,7 +331,6 @@ def _test_instanced_geometry_is_posed_with_its_link(simulation_app) -> bool: from isaaclab_arena.utils.usd_helpers import ( compute_local_bounding_box_from_usd_at_joint_pos, - extract_link_bbox_meshes_from_usd_at_joint_pos, extract_trimesh_from_usd_at_joint_pos, ) @@ -349,11 +348,10 @@ def _test_instanced_geometry_is_posed_with_its_link(simulation_app) -> bool: usd_path = f"{tmp_dir}/arm_with_instanced_mount.usda" stage.Export(usd_path) - components = extract_link_bbox_meshes_from_usd_at_joint_pos(usd_path, {}) - assert len(components) == 2, "the base and forearm must each produce exactly one collision component" - assert all(component.is_watertight for component in components) + rest_mesh = extract_trimesh_from_usd_at_joint_pos(usd_path, {}) + assert rest_mesh.is_watertight try: - extract_link_bbox_meshes_from_usd_at_joint_pos(usd_path, {}, scale=(-1.0, 1.0, 1.0)) + extract_trimesh_from_usd_at_joint_pos(usd_path, {}, scale=(-1.0, 1.0, 1.0)) except AssertionError as error: assert "positive" in str(error) else: @@ -362,7 +360,7 @@ def _test_instanced_geometry_is_posed_with_its_link(simulation_app) -> bool: # The mount reaches 1.7 along +X: forearm at 1.0, mount offset 0.5, half extent 0.2. rest = compute_local_bounding_box_from_usd_at_joint_pos(usd_path, {}) np.testing.assert_allclose(rest.max_point.numpy()[0][0], 1.7, atol=1e-6) - np.testing.assert_allclose(extract_trimesh_from_usd_at_joint_pos(usd_path, {}).vertices.max(axis=0)[0], 1.7) + np.testing.assert_allclose(rest_mesh.vertices.max(axis=0)[0], 1.7) # Swinging the elbow must carry the instance with the forearm rather than leave it behind. swung = compute_local_bounding_box_from_usd_at_joint_pos(usd_path, {"elbow": math.pi / 2.0}) diff --git a/isaaclab_arena/utils/usd_articulation.py b/isaaclab_arena/utils/usd_articulation.py index 6283440cf5..26e2561056 100644 --- a/isaaclab_arena/utils/usd_articulation.py +++ b/isaaclab_arena/utils/usd_articulation.py @@ -3,11 +3,12 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Forward kinematics over a USD articulation's physics joints. +"""Pose USD articulation geometry without spawning a simulation. -Authored USD transforms only describe the one joint configuration an asset happens to be saved in, -so posing an articulation is a prerequisite for placement geometry that matches what is spawned. -Matrices here follow USD's row-vector convention: ``point_parent = point_local @ matrix``. +Placement geometry is computed before an Isaac Lab articulation exists, so runtime link poses are +unavailable. This USD-joint implementation also avoids requiring a separate CuRobo kinematics +configuration for every embodiment. Matrices follow USD's row-vector convention: +``point_parent = point_local @ matrix``. """ from __future__ import annotations @@ -101,14 +102,16 @@ def compute_posed_prim_world_deltas( f" Available joints: {sorted(joint_prims)}." ) - edges = _joint_edges(stage, root_prim, joint_prims, joint_pos) + edges = _joint_edges(root_prim, joint_prims, joint_pos) if not edges: return {} - rest_transforms = { - path: _local_to_world(stage, path) - for path in {path for edge in edges for path in (edge.parent, edge.child) if path} - } + xform_cache = UsdGeom.XformCache(Usd.TimeCode.Default()) + rest_transforms: dict[str, np.ndarray] = {} + for path in {path for edge in edges for path in (edge.parent, edge.child) if path}: + prim = stage.GetPrimAtPath(path) + assert prim, f"Joint references a missing prim: {path}" + rest_transforms[path] = np.array(xform_cache.GetLocalToWorldTransform(prim), dtype=np.float64) posed_transforms = _propagate_joint_motion(edges, rest_transforms) return {path: np.linalg.inv(rest_transforms[path]) @ posed for path, posed in posed_transforms.items()} @@ -144,15 +147,13 @@ class _JointEdge: def _joint_edges( - stage: Usd.Stage, root_prim: Usd.Prim, joint_prims: Mapping[str, Usd.Prim], joint_pos: Mapping[str, float], ) -> list[_JointEdge]: """Build the articulation's parent-to-child edges, including fixed joints that carry no motion.""" - root_path = root_prim.GetPath().pathString - movable_paths = {prim.GetPath().pathString for prim in joint_prims.values()} - values_by_path = {joint_prims[name].GetPath().pathString: value for name, value in joint_pos.items()} + root_sdf_path = root_prim.GetPath() + joint_pos_by_path = {prim.GetPath().pathString: joint_pos.get(name, 0.0) for name, prim in joint_prims.items()} edges: list[_JointEdge] = [] for prim in Usd.PrimRange(root_prim): @@ -161,10 +162,10 @@ def _joint_edges( joint_path = prim.GetPath().pathString parent, child = _joint_body_paths(prim) # Compare path components: a plain prefix test would also match a sibling root's children. - if child is None or not Sdf.Path(child).HasPrefix(Sdf.Path(root_path)): + if child is None or not Sdf.Path(child).HasPrefix(root_sdf_path): continue - if joint_path in movable_paths: - motion = _joint_motion(prim, values_by_path.get(joint_path, 0.0)) + if joint_path in joint_pos_by_path: + motion = _joint_motion(prim, joint_pos_by_path[joint_path]) else: motion = np.eye(4) edges.append( @@ -254,10 +255,3 @@ def _joint_motion(joint_prim: Usd.Prim, value: float) -> np.ndarray: else: matrix.SetTranslate(axis * value) return np.array(matrix, dtype=np.float64) - - -def _local_to_world(stage: Usd.Stage, prim_path: str) -> np.ndarray: - """Return a prim's authored local-to-world transform.""" - prim = stage.GetPrimAtPath(prim_path) - assert prim, f"Joint references a missing prim: {prim_path}" - return np.array(UsdGeom.Xformable(prim).ComputeLocalToWorldTransform(Usd.TimeCode.Default()), dtype=np.float64) diff --git a/isaaclab_arena/utils/usd_helpers.py b/isaaclab_arena/utils/usd_helpers.py index 6a6667fc1d..6a09378ab5 100644 --- a/isaaclab_arena/utils/usd_helpers.py +++ b/isaaclab_arena/utils/usd_helpers.py @@ -468,42 +468,41 @@ def extract_trimesh_from_usd_path( return extract_trimesh_from_prim(stage, default_prim.GetPath().pathString, scale) -def extract_link_bbox_meshes_from_usd_at_joint_pos( +# ----------------------------------------------------------------------------- +# Joint-posed articulation geometry helpers +# ----------------------------------------------------------------------------- + + +def extract_trimesh_from_usd_at_joint_pos( usd_path: str, joint_pos: Mapping[str, float], scale: tuple[float, float, float] = (1.0, 1.0, 1.0), -) -> tuple[trimesh.Trimesh, ...]: - """Return posed box meshes grouped by rigid body, in the default prim's local frame. - - Joints the articulation has but joint_pos omits are posed at zero, so the result depends only on - joint_pos and not on the configuration the asset happens to be authored in. Each box encloses all - default-purpose geometry attached to one rigid body in that body's local frame, then follows the - body's posed transform. Geometry not attached to a rigid body is represented by one root box. - - Results are cached in process and copied before return. - - Args: - usd_path: Path to the articulation's .usd/.usda/.usdc file. - joint_pos: Joint positions keyed by exact joint name or Isaac Lab regex, revolute in radians. - scale: Spawn-time scale passed to ``UsdFileCfg``. - - Returns: - Watertight link-box meshes in the scaled default-prim frame. - """ - meshes = _extract_link_bbox_meshes_from_usd_at_joint_pos(usd_path, tuple(sorted(joint_pos.items())), tuple(scale)) - return tuple(mesh.copy() for mesh in meshes) +) -> trimesh.Trimesh: + """Return one mesh containing all of an articulation's posed link boxes.""" + mesh = _extract_trimesh_from_usd_at_joint_pos(usd_path, tuple(sorted(joint_pos.items())), tuple(scale)) + return mesh.copy() # NOTE(zihaox, 2026-07-28): Cache here rather than on the asset. Isaac Lab reaches assets through # EventTermCfg params, and configclass's validation walk tracks no visited set, so a trimesh held by # an asset sends it recursing through trimesh's internal back-references until the stack overflows. @functools.lru_cache(maxsize=_POSED_GEOMETRY_CACHE_SIZE) +def _extract_trimesh_from_usd_at_joint_pos( + usd_path: str, + joint_pos_items: tuple[tuple[str, float], ...], + scale: tuple[float, float, float], +) -> trimesh.Trimesh: + """Return a cached mesh containing all of an articulation's link boxes.""" + meshes = _extract_link_bbox_meshes_from_usd_at_joint_pos(usd_path, joint_pos_items, scale) + return trimesh.util.concatenate(meshes) + + def _extract_link_bbox_meshes_from_usd_at_joint_pos( usd_path: str, joint_pos_items: tuple[tuple[str, float], ...], scale: tuple[float, float, float], ) -> tuple[trimesh.Trimesh, ...]: - """Cacheable body of ``extract_link_bbox_meshes_from_usd_at_joint_pos``.""" + """Build posed box meshes grouped by rigid body.""" assert all( component > 0 for component in scale ), f"All scale components must be positive (negative scale flips winding/SDF sign), got {scale}" @@ -517,27 +516,6 @@ def _extract_link_bbox_meshes_from_usd_at_joint_pos( return _posed_link_bbox_meshes(stage, default_prim, deltas, scale) -@functools.lru_cache(maxsize=_POSED_GEOMETRY_CACHE_SIZE) -def _extract_trimesh_from_usd_at_joint_pos( - usd_path: str, - joint_pos_items: tuple[tuple[str, float], ...], - scale: tuple[float, float, float], -) -> trimesh.Trimesh: - """Return a cached mesh containing all of an articulation's link boxes.""" - meshes = _extract_link_bbox_meshes_from_usd_at_joint_pos(usd_path, joint_pos_items, scale) - return trimesh.util.concatenate(meshes) - - -def extract_trimesh_from_usd_at_joint_pos( - usd_path: str, - joint_pos: Mapping[str, float], - scale: tuple[float, float, float] = (1.0, 1.0, 1.0), -) -> trimesh.Trimesh: - """Return one mesh containing all of an articulation's posed link boxes.""" - mesh = _extract_trimesh_from_usd_at_joint_pos(usd_path, tuple(sorted(joint_pos.items())), tuple(scale)) - return mesh.copy() - - def _nearest_rigid_body_ancestor(prim: Usd.Prim, root_prim: Usd.Prim) -> Usd.Prim | None: """Return the nearest rigid-body ancestor at or below root_prim.""" candidate = prim From d0dad620fdcc02066668867d6b0f08357cebf637 Mon Sep 17 00:00:00 2001 From: zhx06 Date: Mon, 3 Aug 2026 21:05:12 -0700 Subject: [PATCH 7/7] inline helper function Signed-off-by: zhx06 --- isaaclab_arena/utils/usd_helpers.py | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/isaaclab_arena/utils/usd_helpers.py b/isaaclab_arena/utils/usd_helpers.py index 6a09378ab5..945cefdc22 100644 --- a/isaaclab_arena/utils/usd_helpers.py +++ b/isaaclab_arena/utils/usd_helpers.py @@ -493,16 +493,6 @@ def _extract_trimesh_from_usd_at_joint_pos( scale: tuple[float, float, float], ) -> trimesh.Trimesh: """Return a cached mesh containing all of an articulation's link boxes.""" - meshes = _extract_link_bbox_meshes_from_usd_at_joint_pos(usd_path, joint_pos_items, scale) - return trimesh.util.concatenate(meshes) - - -def _extract_link_bbox_meshes_from_usd_at_joint_pos( - usd_path: str, - joint_pos_items: tuple[tuple[str, float], ...], - scale: tuple[float, float, float], -) -> tuple[trimesh.Trimesh, ...]: - """Build posed box meshes grouped by rigid body.""" assert all( component > 0 for component in scale ), f"All scale components must be positive (negative scale flips winding/SDF sign), got {scale}" @@ -513,7 +503,7 @@ def _extract_link_bbox_meshes_from_usd_at_joint_pos( root_path = default_prim.GetPath().pathString resolved = resolve_joint_pos_patterns(articulation_joint_prims(default_prim), joint_pos) deltas = compute_posed_prim_world_deltas(stage, root_path, resolved) - return _posed_link_bbox_meshes(stage, default_prim, deltas, scale) + return trimesh.util.concatenate(_posed_link_bbox_meshes(stage, default_prim, deltas, scale)) def _nearest_rigid_body_ancestor(prim: Usd.Prim, root_prim: Usd.Prim) -> Usd.Prim | None: