Spaces:
Paused
Paused
| """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", | |
| } | |
| 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:])) | |