"""Generate a collision-safe 5x5 Microduck MJCF scene from the official model.""" from __future__ import annotations import argparse from copy import deepcopy from pathlib import Path import xml.etree.ElementTree as ET from choreography import STAND_POSE, formation_positions PROJECT_ROOT = Path(__file__).resolve().parents[1] SOURCE_MODEL = PROJECT_ROOT / "src/mjlab_microduck/robot/microduck/robot_allcollisions.xml" DEFAULT_OUTPUT = Path(__file__).resolve().parent / "generated/microduck_25_scene.xml" REFERENCE_ATTRIBUTES = { "body", "body1", "body2", "child", "geom", "geom1", "geom2", "joint", "joint1", "joint2", "objname", "site", "site1", "site2", "target", "tendon", } PALETTE = ( "0.98 0.72 0.04 1", # candy gold "0.19 0.82 0.78 1", # teal "0.82 0.33 0.72 1", # magenta "0.94 0.35 0.24 1", # red-orange "0.42 0.57 0.96 1", # electric blue ) ACCENT_MESHES = { "left_shell", "right_shell", "top_head_shell", "bottom_head_shell", "foot_left", "foot_right", "ankle_left", "ankle_right", "jaw", } START_YAW = 0.5 * 3.141592653589793 START_QUATERNION = ( 0.7071067811865476, 0.0, 0.0, 0.7071067811865475, ) def _prefix_instance(element: ET.Element, prefix: str, color: str) -> None: for node in element.iter(): if "name" in node.attrib: node.set("name", prefix + node.get("name", "")) for attribute in REFERENCE_ATTRIBUTES: if attribute in node.attrib: node.set(attribute, prefix + node.get(attribute, "")) if node.tag == "geom" and node.get("mesh") in ACCENT_MESHES and node.get("class") == "visual": node.attrib.pop("material", None) node.set("rgba", color) def _add_stage(worldbody: ET.Element, extent: float) -> None: # Bright, soft studio lighting keeps the duck silhouettes and footwork # legible. Colored rim lights add energy without darkening the background. ET.SubElement(worldbody, "light", { "pos": "0 0 5.2", "dir": "0 0 -1", "directional": "true", "diffuse": "0.72 0.74 0.80", "ambient": "0.34 0.36 0.40", }) ET.SubElement(worldbody, "light", { "pos": "-3 -3 2.8", "dir": "0.65 0.65 -0.8", "diffuse": "0.9 0.40 0.66", "specular": "0.4 0.2 0.3", }) ET.SubElement(worldbody, "light", { "pos": "3 3 2.8", "dir": "-0.65 -0.65 -0.8", "diffuse": "0.28 0.72 0.94", "specular": "0.2 0.35 0.4", }) ET.SubElement(worldbody, "geom", { "name": "dance_floor", "type": "plane", "size": f"{extent:.2f} {extent:.2f} 0.05", "material": "dance_floor_material", "condim": "3", "friction": "0.9 0.02 0.002", }) # Explicit non-colliding tiles make the checkerboard visible even when # bright lights or video compression flatten the procedural texture. tile_count = 14 tile_half_size = extent / tile_count for row in range(tile_count): for col in range(tile_count): x = -extent + tile_half_size + col * tile_half_size * 2 y = -extent + tile_half_size + row * tile_half_size * 2 ET.SubElement(worldbody, "geom", { "name": f"checker_{row:02d}_{col:02d}", "type": "box", "size": f"{tile_half_size:.6f} {tile_half_size:.6f} 0.001", "pos": f"{x:.6f} {y:.6f} 0.001", "material": "tile_light" if (row + col) % 2 == 0 else "tile_slate", "contype": "0", "conaffinity": "0", }) def build_scene(output: Path = DEFAULT_OUTPUT, grid: int = 5, spacing: float = 1.5) -> Path: source = ET.parse(SOURCE_MODEL).getroot() root = ET.Element("mujoco", {"model": f"microduck_electric_slide_{grid}x{grid}"}) ET.SubElement(root, "compiler", { "angle": "radian", "meshdir": str(SOURCE_MODEL.parent.joinpath("assets").resolve()), "autolimits": "true", }) ET.SubElement(root, "option", {"timestep": "0.005", "integrator": "implicitfast"}) for defaults in source.findall("default"): root.append(deepcopy(defaults)) visual = ET.SubElement(root, "visual") ET.SubElement(visual, "headlight", { "diffuse": "0.68 0.68 0.72", "ambient": "0.34 0.34 0.38", "specular": "0.15 0.15 0.15", }) ET.SubElement(visual, "rgba", {"haze": "0.88 0.92 0.98 1"}) ET.SubElement(visual, "global", { "azimuth": "135", "elevation": "-32", "offwidth": "1280", "offheight": "720", }) asset = deepcopy(source.find("asset")) if asset is None: raise ValueError(f"No asset section in {SOURCE_MODEL}") ET.SubElement(asset, "texture", { "type": "skybox", "builtin": "gradient", "rgb1": "0.98 0.99 1", "rgb2": "0.72 0.82 0.94", "width": "512", "height": "3072", }) ET.SubElement(asset, "texture", { "type": "2d", "name": "dance_floor", "builtin": "checker", "mark": "edge", "rgb1": "0.86 0.88 0.92", "rgb2": "0.32 0.38 0.47", "markrgb": "0.20 0.54 0.72", "width": "512", "height": "512", }) ET.SubElement(asset, "material", { "name": "dance_floor_material", "texture": "dance_floor", "texuniform": "true", "texrepeat": "14 14", "reflectance": "0.08", }) ET.SubElement(asset, "material", { "name": "tile_light", "rgba": "0.82 0.84 0.88 1", "reflectance": "0.04", }) ET.SubElement(asset, "material", { "name": "tile_slate", "rgba": "0.24 0.30 0.38 1", "reflectance": "0.04", }) root.append(asset) source_worldbody = source.find("worldbody") if source_worldbody is None or source_worldbody.find("body") is None: raise ValueError(f"No robot body in {SOURCE_MODEL}") source_body = source_worldbody.find("body") worldbody = ET.SubElement(root, "worldbody") _add_stage(worldbody, extent=max(3.0, spacing * grid * 1.4)) positions = formation_positions(grid, spacing) for index, (x, y) in enumerate(positions): body = deepcopy(source_body) prefix = f"duck_{index:02d}_" _prefix_instance(body, prefix, PALETTE[index % len(PALETTE)]) body.set("pos", f"{x:.6f} {y:.6f} 0.12") worldbody.append(body) actuator = ET.SubElement(root, "actuator") source_actuator = source.find("actuator") if source_actuator is None: raise ValueError(f"No actuator section in {SOURCE_MODEL}") for index in range(grid * grid): prefix = f"duck_{index:02d}_" for item in source_actuator: copied = deepcopy(item) _prefix_instance(copied, prefix, PALETTE[index % len(PALETTE)]) actuator.append(copied) source_sensor = source.find("sensor") if source_sensor is not None: sensor = ET.SubElement(root, "sensor") for index in range(grid * grid): prefix = f"duck_{index:02d}_" for item in source_sensor: copied = deepcopy(item) _prefix_instance(copied, prefix, PALETTE[index % len(PALETTE)]) sensor.append(copied) keyframe = ET.SubElement(root, "keyframe") qpos: list[float] = [] ctrl: list[float] = [] for x, y in positions: qpos.extend((x, y, 0.12, *START_QUATERNION, *STAND_POSE)) ctrl.extend(STAND_POSE) ET.SubElement(keyframe, "key", { "name": "STAND_5X5", "qpos": " ".join(f"{value:.9g}" for value in qpos), "ctrl": " ".join(f"{value:.9g}" for value in ctrl), }) output.parent.mkdir(parents=True, exist_ok=True) tree = ET.ElementTree(root) ET.indent(tree, space=" ") tree.write(output, encoding="utf-8", xml_declaration=True) return output def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT) parser.add_argument("--grid", type=int, default=5) parser.add_argument("--spacing", type=float, default=1.5) args = parser.parse_args() path = build_scene(args.output.resolve(), args.grid, args.spacing) print(path) if __name__ == "__main__": main()