#!/usr/bin/env python3 """Fit primitive colliders to a robot's visual meshes and write URDF or USD. The input and output formats are selected from their file extensions. URDF output is a new complete document; USD output is an overlay of the input USD. The source asset is never modified. """ from __future__ import annotations import argparse import os from dataclasses import dataclass from pathlib import Path import sys import tomllib import traceback import xml.etree.ElementTree as ET from typing import Any import numpy as np import trimesh GENERATOR_TAG = "trimesh-primitives-v1" GENERATED_PREFIX = "AUTO_COLLISION_" PRIMITIVE_TYPES = ("box", "sphere", "cylinder") @dataclass(frozen=True) class PrimitiveFit: kind: str transform: np.ndarray dimensions: tuple[float, ...] volume: float def _points_array(points: np.ndarray) -> np.ndarray: result = np.asarray(points, dtype=np.float64) if result.ndim != 2 or result.shape[1] != 3 or len(result) < 4: raise ValueError("at least four 3D points are required") if not np.isfinite(result).all(): raise ValueError("points contain NaN or infinity") return result def fit_box(points: np.ndarray, padding: float = 0.0, scale: float = 1.0) -> PrimitiveFit: points = _points_array(points) to_box, extents = trimesh.bounds.oriented_bounds(points) extents = np.asarray(extents, dtype=np.float64) * scale + 2.0 * padding transform = np.linalg.inv(np.asarray(to_box, dtype=np.float64)) return PrimitiveFit("box", transform, tuple(extents), float(np.prod(extents))) def fit_link_aligned_box( points: np.ndarray, padding: float = 0.0, scale: float = 1.0 ) -> PrimitiveFit: points = _points_array(points) lower = points.min(axis=0) upper = points.max(axis=0) extents = (upper - lower) * scale + 2.0 * padding transform = np.eye(4) transform[:3, 3] = (lower + upper) / 2.0 return PrimitiveFit("box", transform, tuple(extents), float(np.prod(extents))) def fit_sphere(points: np.ndarray, padding: float = 0.0, scale: float = 1.0) -> PrimitiveFit: points = _points_array(points) center, radius = trimesh.nsphere.minimum_nsphere(points) radius = float(radius) * scale + padding transform = np.eye(4) transform[:3, 3] = center return PrimitiveFit("sphere", transform, (radius,), float(4.0 * np.pi * radius**3 / 3.0)) def fit_cylinder( points: np.ndarray, padding: float = 0.0, scale: float = 1.0, sample_count: int = 6, angle_tol: float = 0.001, ) -> PrimitiveFit: points = _points_array(points) result = trimesh.bounds.minimum_cylinder( points, sample_count=sample_count, angle_tol=angle_tol ) radius = float(result["radius"]) * scale + padding height = float(result["height"]) * scale + 2.0 * padding transform = np.asarray(result["transform"], dtype=np.float64) volume = float(np.pi * radius**2 * height) return PrimitiveFit("cylinder", transform, (radius, height), volume) def fit_axis_aligned_cylinder( points: np.ndarray, axis: str, padding: float = 0.0, scale: float = 1.0, ) -> PrimitiveFit: points = _points_array(points) axis = axis.lower() if axis not in "xyz": raise ValueError(f"cylinder axis must be x, y, or z: {axis}") axis_index = "xyz".index(axis) radial_indices = [index for index in range(3) if index != axis_index] radial_center, radius = trimesh.nsphere.minimum_nsphere( points[:, radial_indices] ) axial_min = float(points[:, axis_index].min()) axial_max = float(points[:, axis_index].max()) center = np.zeros(3) center[axis_index] = (axial_min + axial_max) / 2.0 center[radial_indices] = radial_center radius = float(radius) * scale + padding height = (axial_max - axial_min) * scale + 2.0 * padding transform = np.eye(4) if axis == "x": transform[:3, :3] = np.array( [[0.0, 0.0, 1.0], [0.0, 1.0, 0.0], [-1.0, 0.0, 0.0]] ) elif axis == "y": transform[:3, :3] = np.array( [[1.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, -1.0, 0.0]] ) transform[:3, 3] = center volume = float(np.pi * radius**2 * height) return PrimitiveFit("cylinder", transform, (radius, height), volume) def fit_primitive( points: np.ndarray, kind: str, *, allowed: list[str], padding: float, scale: float, cylinder_sample_count: int, cylinder_angle_tol: float, alignment: str = "oriented", axis: str | None = None, ) -> PrimitiveFit: def fit(candidate: str) -> PrimitiveFit: if candidate == "box": if alignment == "link": return fit_link_aligned_box(points, padding, scale) return fit_box(points, padding, scale) if candidate == "sphere": return fit_sphere(points, padding, scale) if candidate == "cylinder": if axis: return fit_axis_aligned_cylinder(points, axis, padding, scale) return fit_cylinder( points, padding, scale, cylinder_sample_count, cylinder_angle_tol, ) raise ValueError(f"unsupported primitive type: {candidate}") if kind != "auto": return fit(kind) candidates: list[PrimitiveFit] = [] failures: list[str] = [] for candidate in allowed: try: candidates.append(fit(candidate)) except Exception as exc: # A degenerate mesh may fail one fitter only. failures.append(f"{candidate}: {exc}") if not candidates: raise RuntimeError("all primitive fits failed: " + "; ".join(failures)) return min(candidates, key=lambda candidate: candidate.volume) def load_config(path: Path | None) -> dict[str, Any]: if path is None: return {} with path.open("rb") as stream: config = tomllib.load(stream) config["_config_dir"] = str(path.parent.resolve()) return config def link_settings(config: dict[str, Any], link_name: str) -> dict[str, Any]: settings = dict(config.get("defaults", {})) settings.update(config.get("links", {}).get(link_name, {})) return settings def _find_robot_root(stage: Any, requested_path: str | None) -> Any: if requested_path: prim = stage.GetPrimAtPath(requested_path) if not prim: raise ValueError(f"robot root does not exist: {requested_path}") return prim default_prim = stage.GetDefaultPrim() if default_prim and default_prim.GetRelationship("isaac:physics:robotLinks").IsValid(): return default_prim for prim in stage.Traverse(): if prim.GetRelationship("isaac:physics:robotLinks").IsValid(): return prim raise RuntimeError("could not find an Isaac robotLinks relationship; pass --robot-root") def find_robot_links(stage: Any, robot_root_path: str | None) -> list[Any]: root = _find_robot_root(stage, robot_root_path) targets = root.GetRelationship("isaac:physics:robotLinks").GetTargets() links = [stage.GetPrimAtPath(path) for path in targets] links = [prim for prim in links if prim] if not links: raise RuntimeError(f"robot has no resolved links: {root.GetPath()}") return links def _computed_purpose(prim: Any, UsdGeom: Any) -> str: imageable = UsdGeom.Imageable(prim) if not imageable: return "" return str(imageable.ComputePurpose()) def collect_visual_points(link: Any, link_paths: set[str], Usd: Any, UsdGeom: Any, UsdPhysics: Any) -> tuple[np.ndarray, set[str]]: """Return visual vertices in link coordinates and direct mesh-collider roots.""" cache = UsdGeom.XformCache() link_to_world = cache.GetLocalToWorldTransform(link) world_to_link = np.asarray(link_to_world.GetInverse(), dtype=np.float64) point_sets: list[np.ndarray] = [] mesh_collision_roots: set[str] = set() for child in link.GetChildren(): if str(child.GetPath()) in link_paths or child.GetName().startswith(GENERATED_PREFIX): continue child_has_mesh_collision = False for prim in Usd.PrimRange(child, Usd.TraverseInstanceProxies()): if not prim.IsA(UsdGeom.Mesh): continue purpose = _computed_purpose(prim, UsdGeom) is_collision = prim.HasAPI(UsdPhysics.CollisionAPI) or purpose == str(UsdGeom.Tokens.guide) if is_collision: child_has_mesh_collision = True continue if purpose not in ("", str(UsdGeom.Tokens.default_), str(UsdGeom.Tokens.render)): continue points = np.asarray(UsdGeom.Mesh(prim).GetPointsAttr().Get(), dtype=np.float64) if not len(points): continue mesh_to_world = np.asarray(cache.GetLocalToWorldTransform(prim), dtype=np.float64) mesh_to_link = mesh_to_world @ world_to_link local_points = points @ mesh_to_link[:3, :3] + mesh_to_link[3, :3] point_sets.append(local_points) if child_has_mesh_collision: mesh_collision_roots.add(str(child.GetPath())) if not point_sets: raise RuntimeError(f"no visual mesh vertices found below {link.GetPath()}") return np.concatenate(point_sets), mesh_collision_roots def collect_mesh_file_points(path: Path) -> np.ndarray: """Load a URDF visual mesh whose vertices are already link-local.""" if not path.is_file(): raise FileNotFoundError(path) loaded = trimesh.load(path, force="mesh") if isinstance(loaded, trimesh.Scene): meshes = tuple(loaded.geometry.values()) if not meshes: raise RuntimeError(f"mesh file has no geometry: {path}") loaded = trimesh.util.concatenate(meshes) return _points_array(np.asarray(loaded.vertices, dtype=np.float64)) def _set_transform(prim: Any, transform: np.ndarray, scale: tuple[float, float, float] | None, Gf: Any, UsdGeom: Any) -> None: xformable = UsdGeom.Xformable(prim) translation = transform[:3, 3] quaternion = trimesh.transformations.quaternion_from_matrix(transform) xformable.AddTranslateOp().Set(Gf.Vec3d(*translation.tolist())) xformable.AddOrientOp(UsdGeom.XformOp.PrecisionDouble).Set( Gf.Quatd(float(quaternion[0]), Gf.Vec3d(*quaternion[1:4].tolist())) ) if scale is not None: xformable.AddScaleOp().Set(Gf.Vec3d(*scale)) def prepare_authoring_links(stage: Any, link_paths: list[str]) -> None: """De-instance only branches that contain a link requiring a new child.""" instance_roots: set[str] = set() for link_path in link_paths: link = stage.GetPrimAtPath(link_path) if not link: continue if link.IsInstance(): instance_roots.add(link_path) continue if not link.IsInstanceProxy(): continue instance_root = link while instance_root.IsInstanceProxy(): instance_root = instance_root.GetParent() if not instance_root or not instance_root.IsInstance(): raise RuntimeError(f"could not find an instance root for collision link: {link_path}") instance_roots.add(str(instance_root.GetPath())) if not instance_roots: return for instance_root in instance_roots: stage.OverridePrim(instance_root).SetInstanceable(False) stage.GetRootLayer().Save() stage.Reload() still_proxies = [ link_path for link_path in link_paths if stage.GetPrimAtPath(link_path).IsInstanceProxy() ] if still_proxies: raise RuntimeError(f"links remained instance proxies: {still_proxies}") def author_primitive(stage: Any, link_path: str, fit: PrimitiveFit, Gf: Any, Sdf: Any, UsdGeom: Any, UsdPhysics: Any) -> str: prim_path = f"{link_path}/{GENERATED_PREFIX}{fit.kind.upper()}" if fit.kind == "box": shape = UsdGeom.Cube.Define(stage, prim_path) shape.CreateSizeAttr(1.0) scale = tuple(float(value) for value in fit.dimensions) elif fit.kind == "sphere": shape = UsdGeom.Sphere.Define(stage, prim_path) shape.CreateRadiusAttr(float(fit.dimensions[0])) scale = None elif fit.kind == "cylinder": shape = UsdGeom.Cylinder.Define(stage, prim_path) shape.CreateAxisAttr(UsdGeom.Tokens.z) shape.CreateRadiusAttr(float(fit.dimensions[0])) shape.CreateHeightAttr(float(fit.dimensions[1])) scale = None else: raise AssertionError(fit.kind) prim = shape.GetPrim() _set_transform(prim, fit.transform, scale, Gf, UsdGeom) UsdPhysics.CollisionAPI.Apply(prim).CreateCollisionEnabledAttr(True) UsdGeom.Imageable(prim).CreatePurposeAttr(UsdGeom.Tokens.guide) prim.SetCustomDataByKey("collisionGenerator", GENERATOR_TAG) prim.CreateAttribute("collision:primitiveType", Sdf.ValueTypeNames.Token, custom=True).Set(fit.kind) return prim_path def create_overlay_stage( input_path: Path, temporary_output: Path, default_prim_path: str, source_stage: Any, Usd: Any, ) -> Any: stage = Usd.Stage.CreateNew(str(temporary_output)) relative_input = os.path.relpath(input_path, temporary_output.parent) stage.GetRootLayer().subLayerPaths = [relative_input] for metadata_key in ("upAxis", "metersPerUnit", "kilogramsPerUnit"): metadata_value = source_stage.GetMetadata(metadata_key) if metadata_value is not None: stage.SetMetadata(metadata_key, metadata_value) source_default = stage.GetPrimAtPath(default_prim_path) if not source_default: raise RuntimeError(f"default prim did not compose into overlay: {default_prim_path}") stage.SetDefaultPrim(source_default) return stage def validate_usd_output(path: Path, expected_count: int, Usd: Any, UsdGeom: Any, UsdPhysics: Any) -> None: stage = Usd.Stage.Open(str(path)) if not stage.GetDefaultPrim(): raise RuntimeError("output USD has no default prim") generated = [ prim for prim in stage.Traverse() if prim.GetCustomDataByKey("collisionGenerator") == GENERATOR_TAG ] if len(generated) != expected_count: raise RuntimeError(f"expected {expected_count} generated colliders, found {len(generated)}") for prim in generated: if prim.GetTypeName() not in ("Cube", "Sphere", "Cylinder"): raise RuntimeError(f"generated collider is not a primitive: {prim.GetPath()}") if not prim.HasAPI(UsdPhysics.CollisionAPI): raise RuntimeError(f"CollisionAPI missing: {prim.GetPath()}") if _computed_purpose(prim, UsdGeom) != str(UsdGeom.Tokens.guide): raise RuntimeError(f"guide purpose missing: {prim.GetPath()}") def run_usd(args: argparse.Namespace, Usd: Any, UsdGeom: Any, UsdPhysics: Any, Gf: Any, Sdf: Any) -> int: input_path = args.input.resolve() if not input_path.is_file(): raise FileNotFoundError(input_path) if not args.dry_run and args.output is None: raise ValueError("--output is required unless --dry-run is used") config = load_config(args.config.resolve() if args.config else None) defaults = config.get("defaults", {}) allowed = list(defaults.get("allowed_primitives", PRIMITIVE_TYPES)) invalid = set(allowed) - set(PRIMITIVE_TYPES) if invalid: raise ValueError(f"invalid allowed_primitives: {sorted(invalid)}") source_stage = Usd.Stage.Open(str(input_path)) if not source_stage: raise RuntimeError(f"could not open USD: {input_path}") links = find_robot_links(source_stage, args.robot_root) source_default = source_stage.GetDefaultPrim() if not source_default: raise RuntimeError("input USD has no default prim") link_paths = {str(link.GetPath()) for link in links} selected = set(args.only) results: list[tuple[str, PrimitiveFit, set[str]]] = [] for link in links: name = link.GetName() if selected and name not in selected: continue settings = link_settings(config, name) if not settings.get("enabled", True): print(f"SKIP {name}: disabled by configuration") continue try: points, collision_roots = collect_visual_points( link, link_paths, Usd, UsdGeom, UsdPhysics ) except RuntimeError as error: mesh_file = settings.get("mesh_file") if not mesh_file: raise RuntimeError( f"{error}; no mesh_file configured for link name {name!r}" ) from error mesh_path = Path(mesh_file) if not mesh_path.is_absolute(): mesh_path = Path(config["_config_dir"]) / mesh_path points = collect_mesh_file_points(mesh_path.resolve()) collision_roots = set() print(f"FALLBACK {name}: loaded {mesh_path}") kind = str(settings.get("primitive", "auto")) if kind not in (*PRIMITIVE_TYPES, "auto"): raise ValueError(f"invalid primitive for {name}: {kind}") fit = fit_primitive( points, kind, allowed=list(settings.get("allowed_primitives", allowed)), padding=float(settings.get("padding", 0.0)), scale=float(settings.get("scale", 1.0)), cylinder_sample_count=int(settings.get("cylinder_sample_count", defaults.get("cylinder_sample_count", 6))), cylinder_angle_tol=float(settings.get("cylinder_angle_tol", defaults.get("cylinder_angle_tol", 0.001))), alignment=str(settings.get("alignment", "oriented")), axis=str(settings["axis"]) if "axis" in settings else None, ) results.append((str(link.GetPath()), fit, collision_roots)) dimensions = ", ".join(f"{value:.6f}" for value in fit.dimensions) print(f"FIT {name}: {fit.kind} ({dimensions}), vertices={len(points)}, volume={fit.volume:.8f}") if selected: found = {Path(path).name for path, _, _ in results} missing = selected - found if missing: raise ValueError(f"selected links were not generated: {sorted(missing)}") if args.dry_run: print(f"Dry run complete: {len(results)} collider(s) fitted") return 0 output_path = args.output.resolve() if output_path == input_path: raise ValueError("input and output must be different files") if output_path.exists() and not args.replace: raise FileExistsError(f"output exists; pass --replace: {output_path}") output_path.parent.mkdir(parents=True, exist_ok=True) temporary = output_path.with_name(f".{output_path.stem}.tmp{output_path.suffix}") if temporary.exists(): temporary.unlink() print(f"CREATE overlay: {temporary}", flush=True) output_stage = create_overlay_stage( input_path, temporary, str(source_default.GetPath()), source_stage, Usd, ) print("CREATE overlay: composed", flush=True) prepare_authoring_links(output_stage, [link_path for link_path, _, _ in results]) disable_meshes = bool(defaults.get("disable_existing_mesh_collisions", True)) for link_path, fit, collision_roots in results: if disable_meshes: for collision_root in collision_roots: output_stage.OverridePrim(collision_root).SetActive(False) authored = author_primitive( output_stage, link_path, fit, Gf, Sdf, UsdGeom, UsdPhysics ) print(f"WRITE {authored}") output_stage.GetRootLayer().Save() del output_stage os.replace(temporary, output_path) if args.validate: validate_usd_output(output_path, len(results), Usd, UsdGeom, UsdPhysics) print(f"Validated {len(results)} generated collider(s)") print(f"Output: {output_path}") return 0 def _parse_vector( value: str | None, size: int, default: tuple[float, ...] ) -> np.ndarray: if value is None: return np.asarray(default, dtype=np.float64) result = np.fromstring(value, sep=" ", dtype=np.float64) if len(result) != size or not np.isfinite(result).all(): raise ValueError(f"expected {size} finite values, got {value!r}") return result def _urdf_origin_transform(origin: ET.Element | None) -> np.ndarray: if origin is None: return np.eye(4) xyz = _parse_vector(origin.get("xyz"), 3, (0.0, 0.0, 0.0)) rpy = _parse_vector(origin.get("rpy"), 3, (0.0, 0.0, 0.0)) transform = trimesh.transformations.euler_matrix(*rpy, axes="sxyz") transform[:3, 3] = xyz return transform def _resolve_urdf_mesh_path(filename: str, urdf_path: Path) -> Path: if filename.startswith("file://"): path = Path(filename.removeprefix("file://")) elif filename.startswith("package://"): package_path = Path(filename.removeprefix("package://")) if len(package_path.parts) < 2: raise ValueError(f"invalid package URI: {filename}") package_name, relative_parts = package_path.parts[0], package_path.parts[1:] candidates = [ parent / package_name / Path(*relative_parts) for parent in (urdf_path.parent, *urdf_path.parents) ] candidates.extend( parent / Path(*relative_parts) for parent in urdf_path.parents if parent.name == package_name ) for candidate in candidates: if candidate.is_file(): return candidate.resolve() raise FileNotFoundError( f"could not resolve {filename!r} relative to {urdf_path}" ) else: path = Path(filename) if not path.is_absolute(): path = urdf_path.parent / path path = path.resolve() if not path.is_file(): raise FileNotFoundError(path) return path def collect_urdf_visual_points(link: ET.Element, urdf_path: Path) -> np.ndarray: """Collect all visual mesh vertices in the link-local coordinate frame.""" point_sets: list[np.ndarray] = [] for visual in link.findall("visual"): visual_transform = _urdf_origin_transform(visual.find("origin")) geometry = visual.find("geometry") mesh = geometry.find("mesh") if geometry is not None else None if mesh is None: continue filename = mesh.get("filename") if not filename: raise ValueError( f"visual mesh has no filename in link {link.get('name')!r}" ) vertices = collect_mesh_file_points( _resolve_urdf_mesh_path(filename, urdf_path) ) mesh_scale = _parse_vector(mesh.get("scale"), 3, (1.0, 1.0, 1.0)) vertices = trimesh.transform_points(vertices * mesh_scale, visual_transform) point_sets.append(vertices) if not point_sets: raise RuntimeError(f"no visual mesh found in link {link.get('name')!r}") return np.concatenate(point_sets) def _format_number(value: float) -> str: if abs(value) < 5e-13: value = 0.0 return f"{value:.12g}" def _format_vector(values: np.ndarray | tuple[float, ...]) -> str: return " ".join(_format_number(float(value)) for value in values) def create_urdf_collision(fit: PrimitiveFit) -> ET.Element: collision = ET.Element( "collision", {"name": f"{GENERATED_PREFIX}{fit.kind.upper()}"} ) translation = fit.transform[:3, 3] rpy = trimesh.transformations.euler_from_matrix(fit.transform, axes="sxyz") ET.SubElement( collision, "origin", {"xyz": _format_vector(translation), "rpy": _format_vector(rpy)}, ) geometry = ET.SubElement(collision, "geometry") if fit.kind == "box": ET.SubElement(geometry, "box", {"size": _format_vector(fit.dimensions)}) elif fit.kind == "sphere": ET.SubElement( geometry, "sphere", {"radius": _format_number(fit.dimensions[0])} ) elif fit.kind == "cylinder": ET.SubElement( geometry, "cylinder", { "radius": _format_number(fit.dimensions[0]), "length": _format_number(fit.dimensions[1]), }, ) else: raise AssertionError(fit.kind) return collision def _fit_urdf_link( link: ET.Element, urdf_path: Path, config: dict[str, Any], defaults: dict[str, Any], allowed: list[str], ) -> tuple[PrimitiveFit, int]: name = link.get("name", "") settings = link_settings(config, name) try: points = collect_urdf_visual_points(link, urdf_path) except RuntimeError as error: mesh_file = settings.get("mesh_file") if not mesh_file: raise RuntimeError(str(error)) from error mesh_path = Path(mesh_file) if not mesh_path.is_absolute(): mesh_path = Path(config["_config_dir"]) / mesh_path points = collect_mesh_file_points(mesh_path.resolve()) print(f"FALLBACK {name}: loaded {mesh_path}") kind = str(settings.get("primitive", "auto")) if kind not in (*PRIMITIVE_TYPES, "auto"): raise ValueError(f"invalid primitive for {name}: {kind}") link_allowed = list(settings.get("allowed_primitives", allowed)) invalid = set(link_allowed) - set(PRIMITIVE_TYPES) if invalid: raise ValueError(f"invalid allowed_primitives for {name}: {sorted(invalid)}") fit = fit_primitive( points, kind, allowed=link_allowed, padding=float(settings.get("padding", 0.0)), scale=float(settings.get("scale", 1.0)), cylinder_sample_count=int( settings.get( "cylinder_sample_count", defaults.get("cylinder_sample_count", 6) ) ), cylinder_angle_tol=float( settings.get( "cylinder_angle_tol", defaults.get("cylinder_angle_tol", 0.001) ) ), alignment=str(settings.get("alignment", "oriented")), axis=str(settings["axis"]) if "axis" in settings else None, ) return fit, len(points) def _parse_urdf(path: Path) -> ET.ElementTree: parser = ET.XMLParser(target=ET.TreeBuilder(insert_comments=True)) tree = ET.parse(path, parser=parser) root = tree.getroot() if root.tag != "robot": raise ValueError(f"URDF root must be , found <{root.tag}>") return tree def validate_urdf_output(path: Path, expected: dict[str, PrimitiveFit]) -> None: root = _parse_urdf(path).getroot() links = {link.get("name", ""): link for link in root.findall("link")} missing = set(expected) - set(links) if missing: raise RuntimeError(f"output URDF is missing links: {sorted(missing)}") for name, fit in expected.items(): generated = [ collision for collision in links[name].findall("collision") if collision.get("name", "").startswith(GENERATED_PREFIX) ] if len(generated) != 1: raise RuntimeError( f"expected one generated collider on {name}, found {len(generated)}" ) geometry = generated[0].find("geometry") primitive_count = ( sum(geometry.find(kind) is not None for kind in PRIMITIVE_TYPES) if geometry is not None else 0 ) if primitive_count != 1: raise RuntimeError(f"invalid generated collision geometry on {name}") written_transform = _urdf_origin_transform(generated[0].find("origin")) if not np.allclose(written_transform, fit.transform, atol=1e-9, rtol=1e-9): error = float(np.max(np.abs(written_transform - fit.transform))) raise RuntimeError( f"generated collision transform changed on {name}: max error {error}" ) primitive = geometry.find(fit.kind) if primitive is None: raise RuntimeError(f"expected {fit.kind} collision geometry on {name}") if fit.kind == "box": dimensions = _parse_vector(primitive.get("size"), 3, ()) elif fit.kind == "sphere": dimensions = np.asarray([float(primitive.get("radius", "nan"))]) else: dimensions = np.asarray( [ float(primitive.get("radius", "nan")), float(primitive.get("length", "nan")), ] ) expected_dimensions = np.asarray(fit.dimensions) if ( not np.isfinite(dimensions).all() or (dimensions <= 0.0).any() or not np.allclose( dimensions, expected_dimensions, atol=1e-9, rtol=1e-9 ) ): raise RuntimeError( f"generated collision dimensions changed on {name}: " f"expected {expected_dimensions}, found {dimensions}" ) def run_urdf(args: argparse.Namespace) -> int: input_path = args.input.resolve() if not input_path.is_file(): raise FileNotFoundError(input_path) if not args.dry_run and args.output is None: raise ValueError("--output is required unless --dry-run is used") config = load_config(args.config.resolve() if args.config else None) defaults = config.get("defaults", {}) allowed = list(defaults.get("allowed_primitives", PRIMITIVE_TYPES)) invalid = set(allowed) - set(PRIMITIVE_TYPES) if invalid: raise ValueError(f"invalid allowed_primitives: {sorted(invalid)}") tree = _parse_urdf(input_path) root = tree.getroot() selected = set(args.only) known_names = {link.get("name", "") for link in root.findall("link")} unknown = selected - known_names if unknown: raise ValueError(f"selected links do not exist: {sorted(unknown)}") results: list[tuple[ET.Element, PrimitiveFit]] = [] for link in root.findall("link"): name = link.get("name", "") if selected and name not in selected: continue settings = link_settings(config, name) if not settings.get("enabled", True): print(f"SKIP {name}: disabled by configuration") continue fit, vertex_count = _fit_urdf_link( link, input_path, config, defaults, allowed ) results.append((link, fit)) dimensions = ", ".join(f"{value:.6f}" for value in fit.dimensions) print( f"FIT {name}: {fit.kind} ({dimensions}), " f"vertices={vertex_count}, volume={fit.volume:.8f}" ) if args.dry_run: print(f"Dry run complete: {len(results)} collider(s) fitted") return 0 output_path = args.output.resolve() if output_path == input_path: raise ValueError("input and output must be different files") if output_path.exists() and not args.replace: raise FileExistsError(f"output exists; pass --replace: {output_path}") expected: dict[str, PrimitiveFit] = {} for link, fit in results: if not args.keep_existing: for collision in list(link.findall("collision")): link.remove(collision) link.append(create_urdf_collision(fit)) expected[link.get("name", "")] = fit output_path.parent.mkdir(parents=True, exist_ok=True) temporary = output_path.with_name(f".{output_path.stem}.tmp{output_path.suffix}") if temporary.exists(): temporary.unlink() ET.indent(tree, space=" ") tree.write(temporary, encoding="utf-8", xml_declaration=True) os.replace(temporary, output_path) if args.validate: validate_urdf_output(output_path, expected) print(f"Validated {len(expected)} generated collider(s)") print(f"Output: {output_path}") return 0 def parse_args(argv: list[str]) -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--input", required=True, type=Path, help="source URDF or USD") parser.add_argument("--output", type=Path, help="generated URDF or overlay USD") parser.add_argument("--config", type=Path, help="TOML fitting configuration") parser.add_argument( "--robot-root", help="USD robot root prim path, if auto-detection fails" ) parser.add_argument( "--only", action="append", default=[], metavar="LINK", help="generate only selected link (repeatable)", ) parser.add_argument( "--keep-existing", action="store_true", help="URDF only: append instead of replacing collisions on generated links", ) parser.add_argument( "--dry-run", action="store_true", help="fit and print without writing output" ) parser.add_argument( "--replace", action="store_true", help="atomically replace an existing output" ) parser.add_argument( "--validate", action="store_true", help="reopen and validate generated output" ) return parser.parse_args(argv) def _asset_format(path: Path) -> str: suffix = path.suffix.lower() if suffix == ".urdf": return "urdf" if suffix in (".usd", ".usda", ".usdc"): return "usd" raise ValueError( f"unsupported file extension {path.suffix!r}; expected .urdf, .usd, .usda, or .usdc" ) def _validate_format_options(args: argparse.Namespace, asset_format: str) -> None: if args.output is not None and _asset_format(args.output) != asset_format: raise ValueError("input and output formats must match") if asset_format == "urdf" and args.robot_root: raise ValueError("--robot-root is only valid for USD input") if asset_format == "usd" and args.keep_existing: raise ValueError("--keep-existing is only valid for URDF input") def main(argv: list[str] | None = None) -> int: try: args = parse_args(sys.argv[1:] if argv is None else argv) asset_format = _asset_format(args.input) _validate_format_options(args, asset_format) if asset_format == "urdf": return run_urdf(args) from isaacsim import SimulationApp simulation_app = SimulationApp({"headless": True}) try: from pxr import Gf, Sdf, Usd, UsdGeom, UsdPhysics return run_usd(args, Usd, UsdGeom, UsdPhysics, Gf, Sdf) finally: simulation_app.close() except Exception: traceback.print_exc() sys.stderr.flush() return 1 if __name__ == "__main__": raise SystemExit(main())