from __future__ import annotations
import os
import shutil
import tempfile
import xml.etree.ElementTree as ET
from collections.abc import Iterator, Sequence
from pathlib import Path
from typing import Any, Literal, cast, overload
import numpy as np
from unisim.scene import resolve_scene_fragment_path
from unisim.terrain.generator import TerrainGeneratorCfg
def _enable_discardvisual(root: ET.Element) -> None:
compiler_tag = root.find("compiler")
if compiler_tag is None:
compiler_tag = ET.Element("compiler")
root.insert(0, compiler_tag)
compiler_tag.set("discardvisual", "true")
def _write_xml_root(root: ET.Element, output_path: Path) -> None:
ET.indent(root, space=" ")
output_path.write_text(ET.tostring(root, encoding="unicode"), encoding="utf-8")
[docs]
def create_discardvisual_xml(model_file: str) -> str:
tree = ET.parse(model_file)
_enable_discardvisual(tree.getroot())
return _write_temp_xml(tree, model_file)
def _iter_expanded_children(
parent: ET.Element, base_dir: Path
) -> Iterator[tuple[ET.Element, Path]]:
for child in parent:
if child.tag != "include":
yield child, base_dir
continue
include_file = child.get("file")
if not include_file:
raise ValueError(f"Invalid <include> without file attribute in {base_dir}")
include_path = (base_dir / include_file).resolve()
include_root = ET.parse(include_path).getroot()
yield from _iter_expanded_children(include_root, include_path.parent)
def _iter_named_bodies(root: ET.Element, base_dir: Path) -> Iterator[str]:
for child, child_base_dir in _iter_expanded_children(root, base_dir):
if child.tag == "body":
body_name = child.get("name")
if body_name:
yield body_name
yield from _iter_named_bodies(child, child_base_dir)
[docs]
def get_named_bodies(model_file: str) -> tuple[list[int], list[str]]:
"""List MuJoCo-style body ids and names declared in a (possibly included) XML.
Returns ``(ids, names)`` in document order, with ids starting at 1 (the
MuJoCo implicit world body 0 is not listed).
"""
model_path = Path(model_file).resolve()
names = list(_iter_named_bodies(ET.parse(model_path).getroot(), model_path.parent))
ids = list(range(1, len(names) + 1))
return ids, names
[docs]
def get_named_body_ids(model_file: str, names: Sequence[str]) -> list[int]:
"""Resolve MuJoCo-style body ids from XML without importing mujoco."""
body_ids, body_names = get_named_bodies(model_file)
body_id_by_name = dict(zip(body_names, body_ids, strict=True))
missing = [name for name in names if name not in body_id_by_name]
if missing:
missing_str = ", ".join(missing)
raise ValueError(f"Bodies not found in XML '{model_file}': {missing_str}")
return [body_id_by_name[name] for name in names]
def _mujoco_module() -> Any:
import mujoco
return cast(Any, mujoco)
def _materialize_spec_xml(spec, model_file: str) -> str:
fd, output_path = tempfile.mkstemp(
suffix=".xml", dir=os.path.dirname(os.path.abspath(model_file))
)
try:
with os.fdopen(fd, "w", encoding="utf-8") as f:
f.write(spec.to_xml())
except Exception:
os.close(fd)
raise
return output_path
def _add_w_sensors(spec, valid_bnames: list[str]) -> None:
mujoco = _mujoco_module()
for bname in valid_bnames:
spec.add_sensor(
name=f"track_pos_w_{bname}",
type=mujoco.mjtSensor.mjSENS_FRAMEPOS,
objtype=mujoco.mjtObj.mjOBJ_XBODY,
objname=bname,
)
for bname in valid_bnames:
spec.add_sensor(
name=f"track_quat_w_{bname}",
type=mujoco.mjtSensor.mjSENS_FRAMEQUAT,
objtype=mujoco.mjtObj.mjOBJ_XBODY,
objname=bname,
)
for bname in valid_bnames:
spec.add_sensor(
name=f"track_linvel_w_{bname}",
type=mujoco.mjtSensor.mjSENS_FRAMELINVEL,
objtype=mujoco.mjtObj.mjOBJ_XBODY,
objname=bname,
)
for bname in valid_bnames:
spec.add_sensor(
name=f"track_angvel_w_{bname}",
type=mujoco.mjtSensor.mjSENS_FRAMEANGVEL,
objtype=mujoco.mjtObj.mjOBJ_XBODY,
objname=bname,
)
def _add_b_sensors(spec, valid_bnames: list[str], baselink_name: str) -> None:
# Body-frame velocities are computed analytically from the world-frame
# tracking sensors (see MuJoCoBackend.get_body_*_vel_b); framelinvel /
# frameangvel sensors with a baselink reference report relative motion and
# degenerate to zero for the root body, so they are intentionally absent.
mujoco = _mujoco_module()
for bname in valid_bnames:
spec.add_sensor(
name=f"track_pos_b_{bname}",
type=mujoco.mjtSensor.mjSENS_FRAMEPOS,
objtype=mujoco.mjtObj.mjOBJ_XBODY,
objname=bname,
reftype=mujoco.mjtObj.mjOBJ_XBODY,
refname=baselink_name,
)
for bname in valid_bnames:
spec.add_sensor(
name=f"track_quat_b_{bname}",
type=mujoco.mjtSensor.mjSENS_FRAMEQUAT,
objtype=mujoco.mjtObj.mjOBJ_XBODY,
objname=bname,
reftype=mujoco.mjtObj.mjOBJ_XBODY,
refname=baselink_name,
)
def _write_temp_xml(tree: ET.ElementTree[ET.Element], model_file: str) -> str: # type: ignore[type-arg]
fd, output_path = tempfile.mkstemp(
suffix=".xml", dir=os.path.dirname(os.path.abspath(model_file))
)
os.close(fd)
tree.write(output_path)
return output_path
def _format_values(values: list[float] | tuple[float, ...]) -> str:
return " ".join(str(float(value)) for value in values)
[docs]
def materialize_scene_visual_override(
source_model_file: str,
*,
ground_texture_file: str | None = None,
ground_texrepeat: list[float] | tuple[float, float] | None = None,
skybox_rgb1: list[float] | tuple[float, float, float] | None = None,
skybox_rgb2: list[float] | tuple[float, float, float] | None = None,
) -> str:
"""Create a temporary scene XML with visual-only overrides applied."""
tree = ET.parse(source_model_file)
root = tree.getroot()
asset_tag = root.find("asset")
if asset_tag is None:
raise ValueError(f"Scene '{source_model_file}' is missing an <asset> tag.")
if skybox_rgb1 is not None or skybox_rgb2 is not None:
skybox = asset_tag.find("./texture[@type='skybox']")
if skybox is None:
raise ValueError(f"Scene '{source_model_file}' is missing a skybox texture.")
if skybox_rgb1 is not None:
skybox.set("rgb1", _format_values(tuple(skybox_rgb1)))
if skybox_rgb2 is not None:
skybox.set("rgb2", _format_values(tuple(skybox_rgb2)))
if ground_texture_file is not None:
ground_texture = asset_tag.find("./texture[@name='groundplane']")
if ground_texture is None:
raise ValueError(f"Scene '{source_model_file}' is missing the groundplane texture.")
for attr in ("builtin", "mark", "rgb1", "rgb2", "markrgb", "width", "height"):
ground_texture.attrib.pop(attr, None)
ground_texture.set("file", str(Path(ground_texture_file)))
if ground_texrepeat is not None:
ground_material = asset_tag.find("./material[@name='groundplane']")
if ground_material is None:
raise ValueError(f"Scene '{source_model_file}' is missing the groundplane material.")
ground_material.set("texrepeat", _format_values(tuple(ground_texrepeat)))
return _write_temp_xml(tree, source_model_file)
[docs]
def materialize_scene_fragments(
source_model_file: str,
*,
fragment_files: Sequence[str],
) -> str:
"""Create a temporary scene XML with task/scene fragments merged."""
tree = ET.parse(source_model_file)
root = tree.getroot()
source_path = Path(source_model_file).resolve()
for fragment_file in fragment_files:
_merge_scene_fragment(root, resolve_scene_fragment_path(fragment_file, source_path))
return _write_temp_xml(tree, source_model_file)
_ATTACH_PREFIXED_ATTRS = {
"class",
"childclass",
"name",
"material",
"texture",
"mesh",
"joint",
"site",
"geom1",
"geom2",
"body1",
"body2",
"objname",
"refname",
"hfield",
"hfieldname",
"actuator",
"target",
}
def _strip_attach_prefixes(root: ET.Element) -> None:
for elem in root.iter():
for attr, value in list(elem.attrib.items()):
if attr in _ATTACH_PREFIXED_ATTRS and value.startswith("/"):
elem.set(attr, value[1:])
def _flatten_attach_main_default(root: ET.Element) -> None:
default = root.find("default")
if default is None:
return
main = default.find("./default[@class='main']")
if main is None:
return
insert_at = list(default).index(main)
default.remove(main)
for child in list(main):
default.insert(insert_at, child)
insert_at += 1
def _ensure_child(parent: ET.Element, query: str, xml: str) -> None:
if parent.find(query) is None:
parent.append(ET.fromstring(xml))
def _merge_robot_option(root: ET.Element, robot_path: Path) -> None:
"""Preserve robot-level MuJoCo solver/contact options after MjSpec.attach."""
robot_option = ET.parse(robot_path).getroot().find("option")
if robot_option is None:
return
option = root.find("option")
if option is None:
option = ET.Element("option")
insert_at = 0
compiler = root.find("compiler")
if compiler is not None:
insert_at = list(root).index(compiler) + 1
root.insert(insert_at, option)
option.attrib.update(robot_option.attrib)
def _ensure_generated_hfield_scene_visuals(root: ET.Element, geom_name: str) -> None:
asset = root.find("asset")
if asset is None:
asset = ET.Element("asset")
root.insert(0, asset)
_ensure_child(
asset,
"./texture[@type='skybox']",
'<texture type="skybox" builtin="gradient" rgb1="0.3 0.5 0.7" '
'rgb2="0 0 0" width="512" height="3072"/>',
)
_ensure_child(
asset,
"./texture[@name='groundplane']",
'<texture type="2d" name="groundplane" builtin="checker" mark="edge" '
'rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3" markrgb="0.8 0.8 0.8" '
'width="300" height="300"/>',
)
_ensure_child(
asset,
"./material[@name='groundplane']",
'<material name="groundplane" texture="groundplane" texuniform="true" '
'texrepeat="5 5" reflectance="0.2"/>',
)
if root.find("visual") is None:
root.append(
ET.fromstring(
'<visual><headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" '
'specular="0.0 0.0 0.0"/><rgba haze="0.15 0.25 0.35 1"/>'
'<global azimuth="-130" elevation="-20"/><quality offsamples="4"/>'
'<map force="0.01"/></visual>'
)
)
terrain_geom = root.find(f".//geom[@name='{geom_name}']")
if terrain_geom is not None and terrain_geom.get("material") is None:
terrain_geom.set("material", "groundplane")
# Robot models sometimes ship with a spot- or target-mode light (e.g. G1's
# ``spotlight`` tracking the trunk). Such lights have an implicit
# ``type`` and cannot coexist with ``directional="true"``. Drop them; the
# overhead light added by the materializer plus the headlight under
# ``visual`` is sufficient for terrain visualization.
for parent in root.iter():
for light in list(parent.findall("light")):
if light.get("mode") or light.get("target") or light.get("type"):
parent.remove(light)
continue
light.set("directional", "true")
light.set("castshadow", "true")
light.set("dir", "-0.35 -0.45 -1")
def _merge_scene_fragment(root: ET.Element, fragment_file: Path) -> None:
fragment_root = ET.parse(fragment_file).getroot()
if fragment_root.tag != "mujoco":
raise ValueError(f"Scene fragment '{fragment_file}' must have a <mujoco> root.")
for child in list(fragment_root):
if child.tag in {"sensor", "keyframe", "actuator"}:
existing = root.find(child.tag)
if existing is None:
root.append(child)
else:
existing.extend(list(child))
continue
root.append(child)
def _copy_robot_asset_dir(model_file: Path, output_dir: Path) -> None:
"""Copy the robot's mesh / texture assets next to the output scene.
Honors the ``meshdir`` (and ``texturedir``) declared in the model's
``<compiler>`` tag — falling back to ``<model_file_dir>/assets`` when
the compiler tag is missing or points at a non-existent path. This
keeps materialization working for models like ``go2w.xml`` whose
meshdir is relative to a sibling robot directory.
"""
candidates: list[Path] = []
try:
root = ET.parse(model_file).getroot()
except (ET.ParseError, OSError):
root = None
if root is not None:
compiler = root.find("compiler")
if compiler is not None:
for attr in ("meshdir", "texturedir"):
value = compiler.get(attr)
if not value:
continue
path = Path(value)
if not path.is_absolute():
path = (model_file.parent / path).resolve()
if path.is_dir() and path not in candidates:
candidates.append(path)
fallback = model_file.parent / "assets"
if fallback.is_dir() and fallback not in candidates:
candidates.append(fallback)
for src in candidates:
shutil.copytree(src, output_dir / "assets", dirs_exist_ok=True)
def _collect_mujoco_assets(asset_dir: Path) -> dict[str, bytes]:
assets: dict[str, bytes] = {}
if not asset_dir.is_dir():
return assets
for path in asset_dir.rglob("*"):
if not path.is_file():
continue
rel = path.relative_to(asset_dir)
assets[str(rel)] = path.read_bytes()
assets[str(asset_dir.name / rel)] = path.read_bytes()
return assets
@overload
def materialize_mujoco_hfield_attached_scene(
*,
model_file: str,
terrain_cfg: TerrainGeneratorCfg,
output_dir: str | Path,
fragment_files: Sequence[str] = (),
hfield_name: str = "terrain_hfield",
geom_name: str = "floor",
return_surface_sampler: Literal[False] = False,
) -> tuple[Any, np.ndarray]: ...
@overload
def materialize_mujoco_hfield_attached_scene(
*,
model_file: str,
terrain_cfg: TerrainGeneratorCfg,
output_dir: str | Path,
fragment_files: Sequence[str] = (),
hfield_name: str = "terrain_hfield",
geom_name: str = "floor",
return_surface_sampler: Literal[True],
) -> tuple[Any, np.ndarray, Any]: ...
[docs]
def materialize_mujoco_hfield_attached_scene(
*,
model_file: str,
terrain_cfg: TerrainGeneratorCfg,
output_dir: str | Path,
fragment_files: Sequence[str] = (),
hfield_name: str = "terrain_hfield",
geom_name: str = "floor",
return_surface_sampler: bool = False,
) -> tuple[Any, np.ndarray] | tuple[Any, np.ndarray, Any]:
"""Build a MuJoCo model with generated hfield terrain and attached robot spec."""
import mujoco
from unisim.terrain.generator import TerrainGenerator
robot_path = Path(model_file).resolve()
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
_copy_robot_asset_dir(robot_path, output_path)
hfield_rel = Path("hfields") / "hfield.png"
generated = TerrainGenerator(terrain_cfg).write_png(output_path / hfield_rel)
spec = mujoco.MjSpec()
spec.compiler.autolimits = True
spec.compiler.meshdir = "assets"
spec.add_hfield(
name=hfield_name,
file=str((output_path / hfield_rel).resolve()),
size=list(generated.hfield_size),
)
spec.worldbody.add_light(pos=[0.0, 0.0, 8.0], dir=[0.0, 0.0, -1.0])
spec.worldbody.add_geom(
name=geom_name,
type=mujoco.mjtGeom.mjGEOM_HFIELD,
hfieldname=hfield_name,
pos=list(generated.geom_pos),
)
robot_spec = mujoco.MjSpec.from_file(str(robot_path))
frame = spec.worldbody.add_frame()
spec.attach(robot_spec, frame=frame)
root = ET.fromstring(spec.to_xml())
_strip_attach_prefixes(root)
_flatten_attach_main_default(root)
_merge_robot_option(root, robot_path)
_ensure_generated_hfield_scene_visuals(root, geom_name)
for fragment_file in fragment_files:
_merge_scene_fragment(root, resolve_scene_fragment_path(fragment_file, robot_path))
scene_xml = output_path / "scene.xml"
_write_xml_root(root, scene_xml)
physics_root = ET.fromstring(ET.tostring(root, encoding="unicode"))
_enable_discardvisual(physics_root)
model = mujoco.MjSpec.from_string(
ET.tostring(physics_root, encoding="unicode"),
assets=_collect_mujoco_assets(output_path / "assets"),
).compile()
if return_surface_sampler:
return model, generated.terrain_origins, generated.surface_sampler()
return model, generated.terrain_origins
[docs]
def inject_mujoco_tracking_sensors(
model_file: str,
baselink_name: str | None = None,
) -> tuple[str, list, list]:
"""Inject tracking sensors for the MuJoCo backend.
The generated sensors track every body in the world frame (``_w``). When
``baselink_name`` is provided, sensors in the baselink-relative frame
(``_b``) are added as well.
Returns:
(tmp_xml_path, tracked_body_ids, valid_bnames)
"""
mujoco = _mujoco_module()
tracked_body_ids, valid_bnames = get_named_bodies(model_file)
spec = mujoco.MjSpec.from_file(model_file)
_add_w_sensors(spec, valid_bnames)
if baselink_name and baselink_name in valid_bnames:
_add_b_sensors(spec, valid_bnames, baselink_name)
return _materialize_spec_xml(spec, model_file), tracked_body_ids, valid_bnames
[docs]
def processed_xml(xml_path):
xml_dir = os.path.dirname(os.path.abspath(xml_path))
tree = ET.parse(xml_path)
root = tree.getroot()
compiler = root.find("compiler")
if compiler is not None:
meshdir = compiler.get("meshdir")
if meshdir:
abs_meshdir = os.path.normpath(os.path.join(xml_dir, meshdir))
compiler.set("meshdir", abs_meshdir)
bodys = root.findall(".//body")
geom_names = []
for body in bodys:
body_name = body.get("name", "unnamed_body")
geoms = body.findall("geom")
if geoms:
filtered_geoms = []
for geom in geoms:
geom_class = geom.get("class")
if geom_class != "visual":
filtered_geoms.append(geom)
if filtered_geoms:
i = 0
for geom in filtered_geoms:
geom_name = geom.get("name", "unnamed_geom")
if geom_name == "unnamed_geom":
new_name = f"{body_name}_geom{i}"
i += 1
geom.set("name", new_name)
geom_name = new_name
geom_names.append(geom_name)
new_xml_string = ET.tostring(root, encoding="unicode")
return new_xml_string, geom_names
[docs]
def add_sensor(root, sensor_type, name, **kwargs):
"""Add a sensor child under the MuJoCo XML ``<sensor>`` node.
Args:
root: XML root node.
sensor_type: Sensor tag name, such as ``"gyro"``, ``"contact"``, or
``"framepos"``.
name: Sensor ``name`` attribute.
**kwargs: Additional XML attributes such as ``site="imu"`` or
``geom1="floor"``.
"""
# Find or create the <sensor> node.
sensor_element = root.find("sensor")
if sensor_element is None:
sensor_element = ET.SubElement(root, "sensor")
# Create the concrete sensor node.
sensor = ET.SubElement(sensor_element, sensor_type)
# Set the required name attribute.
sensor.set("name", name)
# Set any extra attributes passed by the caller.
for key, value in kwargs.items():
sensor.set(key, str(value))
return sensor