api/mjcf_check.py

"""Validate a generated MJCF scene: static safety checks, then a MuJoCo load + step test.

The load test runs in a separate process (``python -m api.mjcf_check DIR``)
with a timeout, so a pathological model cannot hang or crash the API.
"""

from __future__ import annotations

import json
import math
import os
import re
import shutil
import subprocess
import sys
import tempfile
import xml.etree.ElementTree as ET
from dataclasses import asdict, dataclass, field
from pathlib import Path

from .robots import ROBOTS, reset_to_robot_keyframe, robot_dir

MAX_MJCF_BYTES = 300_000
SCENE_FILE = "scene.xml"
# Attributes that make MuJoCo read files from disk.
FILE_ATTRS = {
    "file",
    "filename",
    "meshdir",
    "texturedir",
    "assetdir",
    "fileright",
    "fileleft",
    "fileup",
    "filedown",
    "filefront",
    "fileback",
}


@dataclass
class ValidationResult:
    valid: bool
    stage: str  # "extract" | "static" | "load" | "simulate" | "ok"
    error: str = ""
    warnings: list[str] = field(default_factory=list)
    stats: dict = field(default_factory=dict)

    def to_dict(self) -> dict:
        return asdict(self)

    def feedback(self) -> str:
        """What gets sent back to the model when asking for a repair."""
        msg = f"Validation failed at stage '{self.stage}': {self.error}"
        if self.warnings:
            msg += "\nWarnings: " + "; ".join(self.warnings)
        return msg


_XML_BLOCK = re.compile(r"```(?:xml|mjcf)?\s*\n(.*?)```", re.DOTALL | re.IGNORECASE)


def extract_mjcf(text: str) -> str | None:
    """Pull the MJCF document out of a model response (fenced block, or a bare <mujoco> element)."""
    for block in _XML_BLOCK.findall(text):
        if "<mujoco" in block:
            return block.strip()
    start, end = text.find("<mujoco"), text.rfind("</mujoco>")
    if start != -1 and end > start:
        return text[start : end + len("</mujoco>")].strip()
    return None


def static_check(xml: str, robot: str | None) -> ValidationResult | None:
    """Return a failed result if the document is unsafe or malformed; None if it may be loaded."""
    if len(xml.encode()) > MAX_MJCF_BYTES:
        return ValidationResult(False, "static", f"MJCF is larger than {MAX_MJCF_BYTES} bytes")
    if "<!DOCTYPE" in xml or "<!ENTITY" in xml:
        return ValidationResult(False, "static", "DOCTYPE/ENTITY declarations are not allowed")
    try:
        root = ET.fromstring(xml)
    except ET.ParseError as e:
        return ValidationResult(False, "static", f"XML parse error: {e}")
    if root.tag != "mujoco":
        return ValidationResult(False, "static", f"root element must be <mujoco>, got <{root.tag}>")

    allowed_includes = {ROBOTS[robot].model_file} if robot else set()
    includes = []
    for el in root.iter():
        for attr, value in el.attrib.items():
            if attr in FILE_ATTRS:
                p = value.replace("\\", "/")
                if p.startswith("/") or ".." in p.split("/") or re.match(r"^[a-zA-Z]:", p) or "://" in p:
                    return ValidationResult(
                        False, "static", f"<{el.tag} {attr}=\"{value}\">: absolute paths and '..' are not allowed"
                    )
        if el.tag == "include":
            includes.append(el.get("file", ""))
    for inc in includes:
        if inc not in allowed_includes:
            allowed = ", ".join(sorted(allowed_includes)) or "none"
            return ValidationResult(False, "static", f'<include file="{inc}"> is not allowed (allowed: {allowed})')
    if robot and ROBOTS[robot].model_file not in includes:
        return ValidationResult(False, "static", f'the scene must contain <include file="{ROBOTS[robot].model_file}"/>')
    return None


def prepare_scene_dir(xml: str, robot: str | None, menagerie_dir: Path) -> Path:
    tmp = Path(tempfile.mkdtemp(prefix="mjcf_check_"))
    if robot:
        spec = ROBOTS[robot]
        src = robot_dir(menagerie_dir, spec)
        if not (src / spec.model_file).is_file():
            shutil.rmtree(tmp)
            raise FileNotFoundError(f"robot model not found at {src / spec.model_file}; run `python -m rag.ingest`")
        (tmp / spec.model_file).symlink_to(src / spec.model_file)
        if (src / "assets").is_dir():
            (tmp / "assets").symlink_to(src / "assets", target_is_directory=True)
    (tmp / SCENE_FILE).write_text(xml)
    return tmp


def validate_mjcf(
    xml: str, robot: str | None, menagerie_dir: Path, sim_seconds: float = 2.0, timeout: float = 60.0
) -> ValidationResult:
    failed = static_check(xml, robot)
    if failed:
        return failed
    tmp = prepare_scene_dir(xml, robot, menagerie_dir)
    try:
        args = [sys.executable, "-m", "api.mjcf_check", str(tmp), str(sim_seconds), robot or ""]
        # Loading and stepping never renders: drop rendering settings so a GL backend that is
        # unavailable here (e.g. MUJOCO_GL=egl without a GPU) cannot break `import mujoco`.
        env = {k: v for k, v in os.environ.items() if k not in ("MUJOCO_GL", "PYOPENGL_PLATFORM")}
        env["PYTHONPATH"] = str(Path(__file__).resolve().parent.parent)
        try:
            proc = subprocess.run(args, capture_output=True, text=True, timeout=timeout, env=env)
        except subprocess.TimeoutExpired:
            return ValidationResult(False, "simulate", f"load/simulation took longer than {timeout:.0f}s")
        lines = [ln for ln in proc.stdout.splitlines() if ln.startswith("{")]
        if not lines:
            return ValidationResult(False, "load", f"validator crashed (exit {proc.returncode}): {proc.stderr[-500:]}")
        return ValidationResult(**json.loads(lines[-1]))
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


# --------------------------------------------------------------------------- subprocess side


def _load_and_step(directory: Path, sim_seconds: float, robot: str | None) -> ValidationResult:
    import mujoco
    import numpy as np

    try:
        model = mujoco.MjModel.from_xml_path(str(directory / SCENE_FILE))
    except Exception as e:  # mujoco raises ValueError with the compiler message
        return ValidationResult(False, "load", str(e).strip())

    def count(obj, n):
        return sum(1 for i in range(n) if mujoco.mj_id2name(model, obj, i))

    stats = {
        "nbody": int(model.nbody),
        "ngeom": int(model.ngeom),
        "njnt": int(model.njnt),
        "nq": int(model.nq),
        "nu": int(model.nu),
        "nlight": int(model.nlight),
        "named_geoms": count(mujoco.mjtObj.mjOBJ_GEOM, model.ngeom),
        "timestep": float(model.opt.timestep),
    }
    if robot and mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, ROBOTS[robot].root_body) < 0:
        return ValidationResult(False, "load", f"robot root body '{ROBOTS[robot].root_body}' missing", stats=stats)

    data = mujoco.MjData(model)
    spec = ROBOTS[robot] if robot else None
    reset_to_robot_keyframe(model, data, spec.root_body if spec else None, spec.home_qpos if spec else None)
    steps = max(1, int(math.ceil(sim_seconds / model.opt.timestep)))
    for _ in range(steps):
        mujoco.mj_step(model, data)
        if not (np.all(np.isfinite(data.qpos)) and np.all(np.isfinite(data.qvel))):
            return ValidationResult(
                False, "simulate", f"state became NaN/inf at t={data.time:.3f}s (unstable model)", stats=stats
            )
    warnings = []
    for i in range(len(data.warning)):
        n = int(data.warning[i].number)
        if n:
            warnings.append(f"{mujoco.mjtWarning(i).name} x{n}")
    stats["simulated_seconds"] = round(float(data.time), 3)
    unstable = [w for w in warnings if w.startswith(("mjWARN_BADQACC", "mjWARN_BADQPOS", "mjWARN_BADQVEL"))]
    if unstable:
        return ValidationResult(False, "simulate", "simulation is unstable: " + ", ".join(unstable), warnings, stats)
    return ValidationResult(True, "ok", "", warnings, stats)


def main(argv: list[str]) -> int:
    directory, sim_seconds, robot = Path(argv[0]), float(argv[1]), (argv[2] or None) if len(argv) > 2 else None
    result = _load_and_step(directory, sim_seconds, robot)
    print(json.dumps(result.to_dict()))
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))