cmvr-es/scripts/meshcat/view_urdf_collisions.py
lgv 4e9bd398f1 feat(collision): add collision primitive generation tools
Add URDF and USD collision primitive generation, MeshCat collision visualization, configuration, documentation, and generated Isaac Sim assets for the dual-arm model.
2026-07-27 15:39:54 +08:00

722 lines
24 KiB
Python

#!/usr/bin/env python3
"""Visualize URDF collision geometry and optional visual meshes in MeshCat."""
from __future__ import annotations
import argparse
from dataclasses import dataclass
import math
import os
from pathlib import Path
import subprocess
import sys
import time
import traceback
import webbrowser
import xml.etree.ElementTree as ET
import meshcat
import meshcat.geometry as geometry
import numpy as np
@dataclass(frozen=True)
class Joint:
name: str
kind: str
parent: str
child: str
origin: np.ndarray
axis: np.ndarray
mimic: tuple[str, float, float] | None
def parse_vector(
value: str | None, size: int, default: tuple[float, ...]
) -> np.ndarray:
if value is None:
return np.asarray(default, dtype=np.float64)
parts = value.split()
if len(parts) != size:
raise ValueError(f"expected {size} values, got {value!r}")
result = np.asarray([float(part) for part in parts], dtype=np.float64)
if not np.isfinite(result).all():
raise ValueError(f"values contain NaN or infinity: {value!r}")
return result
def rpy_rotation(rpy: np.ndarray) -> np.ndarray:
roll, pitch, yaw = rpy
cr, sr = math.cos(roll), math.sin(roll)
cp, sp = math.cos(pitch), math.sin(pitch)
cy, sy = math.cos(yaw), math.sin(yaw)
return np.array(
[
[cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr],
[sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr],
[-sp, cp * sr, cp * cr],
],
dtype=np.float64,
)
def origin_transform(origin: ET.Element | None) -> np.ndarray:
transform = np.eye(4)
if origin is None:
return transform
transform[:3, :3] = rpy_rotation(
parse_vector(origin.get("rpy"), 3, (0.0, 0.0, 0.0))
)
transform[:3, 3] = parse_vector(
origin.get("xyz"), 3, (0.0, 0.0, 0.0)
)
return transform
def axis_angle_transform(axis: np.ndarray, angle: float) -> np.ndarray:
norm = float(np.linalg.norm(axis))
if norm < 1e-12:
raise ValueError("joint axis must not be zero")
x, y, z = axis / norm
c, s = math.cos(angle), math.sin(angle)
one_minus_c = 1.0 - c
transform = np.eye(4)
transform[:3, :3] = np.array(
[
[c + x * x * one_minus_c, x * y * one_minus_c - z * s, x * z * one_minus_c + y * s],
[y * x * one_minus_c + z * s, c + y * y * one_minus_c, y * z * one_minus_c - x * s],
[z * x * one_minus_c - y * s, z * y * one_minus_c + x * s, c + z * z * one_minus_c],
],
dtype=np.float64,
)
return transform
def translation_transform(offset: np.ndarray) -> np.ndarray:
transform = np.eye(4)
transform[:3, 3] = offset
return transform
def parse_joint_values(values: list[str]) -> dict[str, float]:
result: dict[str, float] = {}
for assignment in values:
name, separator, raw_value = assignment.partition("=")
if not separator or not name or not raw_value:
raise ValueError(
f"invalid --joint value {assignment!r}; expected NAME=VALUE"
)
if name in result:
raise ValueError(f"joint value specified more than once: {name}")
value = float(raw_value)
if not math.isfinite(value):
raise ValueError(f"joint value must be finite: {assignment!r}")
result[name] = value
return result
def parse_joints(root: ET.Element) -> dict[str, Joint]:
joints: dict[str, Joint] = {}
children: set[str] = set()
for element in root.findall("joint"):
name = element.get("name")
kind = element.get("type")
parent_element = element.find("parent")
child_element = element.find("child")
if not name or not kind or parent_element is None or child_element is None:
raise ValueError("every joint needs name, type, parent, and child")
parent = parent_element.get("link")
child = child_element.get("link")
if not parent or not child:
raise ValueError(f"joint {name!r} has an empty parent or child")
if name in joints:
raise ValueError(f"duplicate joint name: {name}")
if child in children:
raise ValueError(f"link {child!r} has more than one parent joint")
axis_element = element.find("axis")
axis = parse_vector(
axis_element.get("xyz") if axis_element is not None else None,
3,
(1.0, 0.0, 0.0),
)
mimic_element = element.find("mimic")
mimic = None
if mimic_element is not None:
source = mimic_element.get("joint")
if not source:
raise ValueError(f"mimic joint {name!r} has no source joint")
mimic = (
source,
float(mimic_element.get("multiplier", "1")),
float(mimic_element.get("offset", "0")),
)
joints[name] = Joint(
name=name,
kind=kind,
parent=parent,
child=child,
origin=origin_transform(element.find("origin")),
axis=axis,
mimic=mimic,
)
children.add(child)
return joints
def resolve_joint_values(
joints: dict[str, Joint], requested: dict[str, float]
) -> dict[str, float]:
unknown = set(requested) - set(joints)
if unknown:
raise ValueError(f"unknown joints: {sorted(unknown)}")
fixed = [name for name in requested if joints[name].kind == "fixed"]
if fixed:
raise ValueError(f"fixed joints cannot be assigned: {sorted(fixed)}")
resolved: dict[str, float] = {}
def resolve(name: str, stack: set[str]) -> float:
if name in resolved:
return resolved[name]
if name in stack:
raise ValueError(f"mimic joint cycle contains {name!r}")
joint = joints[name]
if name in requested:
value = requested[name]
elif joint.mimic is not None:
source, multiplier, offset = joint.mimic
if source not in joints:
raise ValueError(
f"mimic joint {name!r} references unknown joint {source!r}"
)
value = multiplier * resolve(source, stack | {name}) + offset
else:
value = 0.0
resolved[name] = value
return value
for joint_name in joints:
resolve(joint_name, set())
return resolved
def joint_motion(joint: Joint, value: float) -> np.ndarray:
if joint.kind == "fixed":
return np.eye(4)
if joint.kind in ("revolute", "continuous"):
return axis_angle_transform(joint.axis, value)
if joint.kind == "prismatic":
norm = float(np.linalg.norm(joint.axis))
if norm < 1e-12:
raise ValueError(f"joint {joint.name!r} axis must not be zero")
axis = joint.axis / norm
return translation_transform(axis * value)
raise ValueError(
f"joint {joint.name!r} uses unsupported type {joint.kind!r}; "
"supported types are fixed, revolute, continuous, and prismatic"
)
def compute_link_transforms(
root: ET.Element, joints: dict[str, Joint], values: dict[str, float]
) -> dict[str, np.ndarray]:
links = {element.get("name") for element in root.findall("link")}
if None in links:
raise ValueError("every link needs a name")
children = {joint.child for joint in joints.values()}
roots = links - children
if not roots:
raise ValueError("URDF has no root link")
transforms = {name: np.eye(4) for name in roots}
pending = list(joints.values())
while pending:
unresolved: list[Joint] = []
for joint in pending:
if joint.parent not in links or joint.child not in links:
raise ValueError(
f"joint {joint.name!r} references a missing parent or child link"
)
if joint.parent not in transforms:
unresolved.append(joint)
continue
transforms[joint.child] = (
transforms[joint.parent]
@ joint.origin
@ joint_motion(joint, values[joint.name])
)
if len(unresolved) == len(pending):
names = [joint.name for joint in unresolved]
raise ValueError(f"joint graph is cyclic or disconnected: {names}")
pending = unresolved
return transforms
def resolve_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 load_mesh(path: Path) -> geometry.Geometry:
suffix = path.suffix.lower()
if suffix == ".stl":
return geometry.StlMeshGeometry.from_file(str(path))
if suffix == ".obj":
return geometry.ObjMeshGeometry.from_file(str(path))
if suffix == ".dae":
return geometry.DaeMeshGeometry.from_file(str(path))
raise ValueError(
f"unsupported mesh format {path.suffix!r}: {path}; "
"MeshCat viewer supports STL, OBJ, and DAE"
)
def cylinder_dimensions(cylinder: ET.Element) -> tuple[float, float]:
radius = float(cylinder.get("radius", "nan"))
length = float(cylinder.get("length", "nan"))
if (
not math.isfinite(radius)
or not math.isfinite(length)
or radius <= 0.0
or length <= 0.0
):
raise ValueError(
f"cylinder radius and length must be positive: {radius}, {length}"
)
return radius, length
def cylinder_correction() -> np.ndarray:
correction = np.eye(4)
# Three.js cylinders use local Y; URDF cylinders use local Z.
correction[:3, :3] = rpy_rotation(np.array([math.pi / 2.0, 0.0, 0.0]))
return correction
def cylinder_wireframe(
radius: float,
length: float,
generator_count: int,
color: int,
opacity_value: float,
ring_segments: int = 64,
) -> geometry.LineSegments:
vertices: list[tuple[float, float, float]] = []
half_length = length / 2.0
# Smooth top and bottom rings, without cap triangulation spokes.
for y in (-half_length, half_length):
for index in range(ring_segments):
first = 2.0 * math.pi * index / ring_segments
second = 2.0 * math.pi * (index + 1) / ring_segments
vertices.extend(
[
(radius * math.cos(first), y, radius * math.sin(first)),
(radius * math.cos(second), y, radius * math.sin(second)),
]
)
# Sparse axial generator lines; six means one line every 60 degrees.
for index in range(generator_count):
angle = 2.0 * math.pi * index / generator_count
x = radius * math.cos(angle)
z = radius * math.sin(angle)
vertices.extend([(x, -half_length, z), (x, half_length, z)])
points = np.asarray(vertices, dtype=np.float32).T
material = geometry.LineBasicMaterial(
color=color,
transparent=opacity_value < 1.0,
opacity=opacity_value,
)
return geometry.LineSegments(geometry.PointsGeometry(points), material)
def geometry_object(
geometry_element: ET.Element, urdf_path: Path
) -> tuple[geometry.Geometry, np.ndarray]:
box = geometry_element.find("box")
sphere = geometry_element.find("sphere")
cylinder = geometry_element.find("cylinder")
mesh = geometry_element.find("mesh")
correction = np.eye(4)
if box is not None:
size = parse_vector(box.get("size"), 3, ())
if (size <= 0.0).any():
raise ValueError(f"box size must be positive: {size}")
return geometry.Box(size), correction
if sphere is not None:
radius = float(sphere.get("radius", "nan"))
if not math.isfinite(radius) or radius <= 0.0:
raise ValueError(f"sphere radius must be positive: {radius}")
return geometry.Sphere(radius), correction
if cylinder is not None:
radius, length = cylinder_dimensions(cylinder)
return geometry.Cylinder(length, radius), cylinder_correction()
if mesh is not None:
filename = mesh.get("filename")
if not filename:
raise ValueError("mesh geometry has no filename")
scale = parse_vector(mesh.get("scale"), 3, (1.0, 1.0, 1.0))
correction[:3, :3] = np.diag(scale)
return load_mesh(resolve_mesh_path(filename, urdf_path)), correction
raise ValueError("geometry must contain box, sphere, cylinder, or mesh")
def parse_rgba(value: str) -> tuple[float, float, float, float]:
rgba = parse_vector(value, 4, ())
if ((rgba < 0.0) | (rgba > 1.0)).any():
raise ValueError(f"RGBA values must be between 0 and 1: {value!r}")
return tuple(float(component) for component in rgba)
def rgb_integer(rgb: tuple[float, float, float]) -> int:
red, green, blue = (round(component * 255.0) for component in rgb)
return (red << 16) | (green << 8) | blue
def visual_rgba(
visual: ET.Element,
named_materials: dict[str, tuple[float, float, float, float]],
) -> tuple[float, float, float, float]:
material = visual.find("material")
if material is None:
return (0.65, 0.68, 0.72, 1.0)
color = material.find("color")
if color is not None and color.get("rgba"):
return parse_rgba(color.get("rgba", ""))
name = material.get("name")
if name and name in named_materials:
return named_materials[name]
return (0.65, 0.68, 0.72, 1.0)
def named_materials(
root: ET.Element,
) -> dict[str, tuple[float, float, float, float]]:
result: dict[str, tuple[float, float, float, float]] = {}
for material in root.findall("material"):
name = material.get("name")
color = material.find("color")
if name and color is not None and color.get("rgba"):
result[name] = parse_rgba(color.get("rgba", ""))
return result
def parse_color(value: str) -> int:
normalized = value.removeprefix("#").removeprefix("0x")
if len(normalized) != 6:
raise argparse.ArgumentTypeError("color must use RRGGBB format")
try:
result = int(normalized, 16)
except ValueError as error:
raise argparse.ArgumentTypeError("color must use RRGGBB format") from error
return result
def opacity(value: str) -> float:
result = float(value)
if not 0.0 <= result <= 1.0:
raise argparse.ArgumentTypeError("opacity must be between 0 and 1")
return result
def cylinder_lines(value: str) -> int:
result = int(value)
if result < 3:
raise argparse.ArgumentTypeError("cylinder line count must be at least 3")
return result
def safe_name(value: str) -> str:
return value.replace("/", "_")
def render_urdf(
viewer: meshcat.Visualizer,
root: ET.Element,
urdf_path: Path,
link_transforms: dict[str, np.ndarray],
*,
collision_only: bool,
visual_opacity: float,
collision_opacity: float,
collision_color: int,
wireframe: bool,
collision_cylinder_lines: int,
) -> tuple[int, int]:
viewer.delete()
materials = named_materials(root)
visual_count = 0
collision_count = 0
collision_material = geometry.MeshPhongMaterial(
color=collision_color,
transparent=collision_opacity < 1.0,
opacity=collision_opacity,
wireframe=wireframe,
)
for link in root.findall("link"):
link_name = link.get("name", "")
link_transform = link_transforms[link_name]
if not collision_only:
for index, visual in enumerate(link.findall("visual")):
geometry_element = visual.find("geometry")
if geometry_element is None:
raise ValueError(f"visual geometry missing on link {link_name!r}")
shape, correction = geometry_object(geometry_element, urdf_path)
rgba = visual_rgba(visual, materials)
alpha = visual_opacity * rgba[3]
material = geometry.MeshPhongMaterial(
color=rgb_integer(rgba[:3]),
transparent=alpha < 1.0,
opacity=alpha,
)
node = viewer[
f"robot/visual/{safe_name(link_name)}/visual_{index}"
]
node.set_object(shape, material)
node.set_transform(
link_transform
@ origin_transform(visual.find("origin"))
@ correction
)
visual_count += 1
for index, collision in enumerate(link.findall("collision")):
geometry_element = collision.find("geometry")
if geometry_element is None:
raise ValueError(f"collision geometry missing on link {link_name!r}")
cylinder = geometry_element.find("cylinder")
if wireframe and cylinder is not None:
radius, length = cylinder_dimensions(cylinder)
shape = cylinder_wireframe(
radius,
length,
collision_cylinder_lines,
collision_color,
collision_opacity,
)
correction = cylinder_correction()
custom_line_object = True
else:
shape, correction = geometry_object(geometry_element, urdf_path)
custom_line_object = False
node = viewer[
f"robot/collision/{safe_name(link_name)}/collision_{index}"
]
if custom_line_object:
node.set_object(shape)
else:
node.set_object(shape, collision_material)
node.set_transform(
link_transform
@ origin_transform(collision.find("origin"))
@ correction
)
collision_count += 1
return visual_count, collision_count
def close_viewer(viewer: meshcat.Visualizer) -> None:
"""Close MeshCat 0.3.x without relying on its broken Visualizer.close()."""
window = viewer.window
window.zmq_socket.close(linger=0)
server_process = window.server_proc
if server_process is None or server_process.poll() is not None:
return
server_process.terminate()
try:
server_process.wait(timeout=3.0)
except subprocess.TimeoutExpired:
server_process.kill()
server_process.wait(timeout=3.0)
def export_static_html(viewer: meshcat.Visualizer, output: Path) -> None:
output = output.resolve()
if output.suffix.lower() != ".html":
raise ValueError(f"MeshCat snapshot must use an .html extension: {output}")
output.parent.mkdir(parents=True, exist_ok=True)
temporary = output.with_name(f".{output.stem}.tmp{output.suffix}")
temporary.write_text(viewer.static_html(), encoding="utf-8")
os.replace(temporary, output)
def parse_args(argv: list[str]) -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--input", required=True, type=Path, help="URDF to display")
parser.add_argument(
"--joint",
action="append",
default=[],
metavar="NAME=VALUE",
help="joint position in radians, or meters for prismatic joints",
)
parser.add_argument(
"--collision-only",
action="store_true",
help="hide visual geometry and display only collisions",
)
parser.add_argument(
"--visual-opacity",
type=opacity,
default=1.0,
help="visual geometry opacity (default: 1.0)",
)
parser.add_argument(
"--collision-opacity",
type=opacity,
default=1.0,
help="collision geometry opacity (default: 1.0)",
)
parser.add_argument(
"--collision-color",
type=parse_color,
default=parse_color("00ff00"),
metavar="RRGGBB",
help="collision color in hexadecimal (default: 00ff00)",
)
parser.add_argument(
"--collision-cylinder-lines",
"--collision-cylinder-segments",
dest="collision_cylinder_lines",
type=cylinder_lines,
default=12,
metavar="COUNT",
help="axial lines on collision cylinders (default: 12, every 30 degrees)",
)
collision_style = parser.add_mutually_exclusive_group()
collision_style.add_argument(
"--wireframe",
dest="wireframe",
action="store_true",
default=True,
help="draw collision geometry as wireframe (default)",
)
collision_style.add_argument(
"--solid-collisions",
dest="wireframe",
action="store_false",
help="draw collision geometry as translucent solids",
)
parser.add_argument(
"--no-browser", action="store_true", help="do not automatically open a browser"
)
parser.add_argument(
"--export-html",
type=Path,
help="write a standalone MeshCat HTML snapshot and exit",
)
parser.add_argument(
"--exit-after-load",
action="store_true",
help=argparse.SUPPRESS,
)
parser.add_argument(
"--zmq-url", help="connect to an existing MeshCat ZMQ server"
)
return parser.parse_args(argv)
def run(args: argparse.Namespace) -> int:
urdf_path = args.input.resolve()
if not urdf_path.is_file():
raise FileNotFoundError(urdf_path)
if urdf_path.suffix.lower() != ".urdf":
raise ValueError(f"input must be a .urdf file: {urdf_path}")
root = ET.parse(urdf_path).getroot()
if root.tag != "robot":
raise ValueError(f"URDF root must be <robot>, found <{root.tag}>")
joints = parse_joints(root)
requested = parse_joint_values(args.joint)
values = resolve_joint_values(joints, requested)
link_transforms = compute_link_transforms(root, joints, values)
viewer = meshcat.Visualizer(zmq_url=args.zmq_url)
visual_count, collision_count = render_urdf(
viewer,
root,
urdf_path,
link_transforms,
collision_only=args.collision_only,
visual_opacity=args.visual_opacity,
collision_opacity=args.collision_opacity,
collision_color=args.collision_color,
wireframe=args.wireframe,
collision_cylinder_lines=args.collision_cylinder_lines,
)
if collision_count == 0:
raise RuntimeError(f"URDF contains no <collision> elements: {urdf_path}")
url = viewer.url()
print(f"Loaded: {urdf_path}")
print(f"Visual geometry: {visual_count}")
print(f"Collision geometry: {collision_count}")
print(f"MeshCat URL: {url}", flush=True)
if args.export_html is not None:
export_static_html(viewer, args.export_html)
snapshot_uri = args.export_html.resolve().as_uri()
print(f"Standalone snapshot: {snapshot_uri}", flush=True)
if not args.no_browser:
webbrowser.open(snapshot_uri, new=2)
close_viewer(viewer)
return 0
if not args.no_browser:
viewer.open()
if args.exit_after_load:
close_viewer(viewer)
return 0
print("Press Ctrl+C to stop the viewer.", flush=True)
try:
while True:
time.sleep(1.0)
except KeyboardInterrupt:
print("Stopping MeshCat viewer.")
finally:
close_viewer(viewer)
return 0
def main(argv: list[str] | None = None) -> int:
try:
return run(parse_args(sys.argv[1:] if argv is None else argv))
except Exception:
traceback.print_exc()
return 1
if __name__ == "__main__":
raise SystemExit(main())