| """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", |
| "0.19 0.82 0.78 1", |
| "0.82 0.33 0.72 1", |
| "0.94 0.35 0.24 1", |
| "0.42 0.57 0.96 1", |
| ) |
| 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: |
| |
| |
| 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", |
| }) |
| |
| |
| 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() |
|
|
