Spaces:
Paused
Paused
| """Robustness sweep: the same scene many times with randomized physics, in parallel on the GPU. | |
| Each copy ("world") gets its own | |
| - ground/object friction scale (geom_friction) | |
| - mass scale of every free object (body_mass + body_inertia, objects only, not the robot) | |
| - small joint noise on the robot's starting pose | |
| - one horizontal push on the robot's root body, in a random direction, at `push_at` seconds | |
| and runs `duration_s` with the same controller as /simulate/run. The result is how many worlds end with the | |
| robot fallen, and which parameters those were. Engine: MuJoCo Warp (NVIDIA Warp, one GPU kernel launch per | |
| step for all worlds) when CUDA is available, else plain MuJoCo on the CPU, one world after another. Both apply | |
| exactly the same randomization, so results agree up to floating point. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import math | |
| import time | |
| from dataclasses import asdict, dataclass | |
| from pathlib import Path | |
| import mujoco | |
| import numpy as np | |
| from api.robots import ROBOTS, reset_to_robot_keyframe | |
| from worker.sim_runner import _hold_controller, config_for | |
| MAX_WORLDS = 256 | |
| MAX_DURATION_S = 10.0 | |
| class SweepConfig: | |
| n_worlds: int = 64 | |
| duration_s: float = 5.0 | |
| controller: str = "hold" | |
| friction: tuple = (0.5, 1.5) # scale range | |
| object_mass: tuple = (0.5, 2.0) # scale range for free objects | |
| push_n: tuple = (0.0, 80.0) # push force range (N), applied for push_s | |
| push_at: float = 1.0 | |
| push_s: float = 0.2 | |
| joint_noise: float = 0.02 # rad, uniform | |
| seed: int = 0 | |
| class World: | |
| friction: float | |
| object_mass: float | |
| push_n: float | |
| push_dir_deg: float | |
| fell: bool = False | |
| fell_at: float | None = None | |
| min_height: float | None = None | |
| final_height: float | None = None | |
| def sample_worlds(cfg: SweepConfig) -> list[World]: | |
| rng = np.random.default_rng(cfg.seed) | |
| return [ | |
| World( | |
| friction=round(float(rng.uniform(*cfg.friction)), 3), | |
| object_mass=round(float(rng.uniform(*cfg.object_mass)), 3), | |
| push_n=round(float(rng.uniform(*cfg.push_n)), 1), | |
| push_dir_deg=round(float(rng.uniform(0, 360)), 1), | |
| ) | |
| for _ in range(cfg.n_worlds) | |
| ] | |
| def _robot_bodies(model: mujoco.MjModel, root: int) -> np.ndarray: | |
| inside = np.zeros(model.nbody, bool) | |
| for b in range(1, model.nbody): | |
| a = b | |
| while a > 0 and a != root: | |
| a = model.body_parentid[a] | |
| inside[b] = a == root and root > 0 | |
| return inside | |
| def _free_objects(model: mujoco.MjModel, robot_bodies: np.ndarray) -> np.ndarray: | |
| """Bodies (and their children) hanging off a free joint that are not part of the robot.""" | |
| free = np.zeros(model.nbody, bool) | |
| for j in range(model.njnt): | |
| b = model.jnt_bodyid[j] | |
| if model.jnt_type[j] == mujoco.mjtJoint.mjJNT_FREE and not robot_bodies[b]: | |
| free[b] = True | |
| for b in range(1, model.nbody): # children of free objects | |
| if free[model.body_parentid[b]]: | |
| free[b] = True | |
| return free | |
| def _joint_noise_mask(model: mujoco.MjModel, robot_bodies: np.ndarray) -> np.ndarray: | |
| """qpos entries of the robot's hinge/slide joints (the free joint is left alone).""" | |
| mask = np.zeros(model.nq, bool) | |
| for j in range(model.njnt): | |
| if robot_bodies[model.jnt_bodyid[j]] and model.jnt_type[j] in ( | |
| mujoco.mjtJoint.mjJNT_HINGE, | |
| mujoco.mjtJoint.mjJNT_SLIDE, | |
| ): | |
| mask[model.jnt_qposadr[j]] = True | |
| return mask | |
| def _setup(scene: Path, robot: str | None, cfg: SweepConfig): | |
| model = mujoco.MjModel.from_xml_path(str(scene)) | |
| data = mujoco.MjData(model) | |
| run_cfg = config_for(robot, cfg.duration_s, cfg.controller, render=False) | |
| reset_to_robot_keyframe(model, data, run_cfg.root_body, run_cfg.home_qpos) | |
| root = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, run_cfg.root_body) if run_cfg.root_body else -1 | |
| robot_bodies = _robot_bodies(model, root) | |
| return model, data, run_cfg, root, robot_bodies | |
| def _world_params(model, worlds, robot_bodies, cfg: SweepConfig, qpos0: np.ndarray): | |
| """Per-world arrays: friction, masses, inertias, starting qpos, push force.""" | |
| n = len(worlds) | |
| rng = np.random.default_rng(cfg.seed + 1) | |
| objects = _free_objects(model, robot_bodies) | |
| noise_mask = _joint_noise_mask(model, robot_bodies) | |
| friction = np.repeat(model.geom_friction[None], n, 0).copy() | |
| mass = np.repeat(model.body_mass[None], n, 0).copy() | |
| inertia = np.repeat(model.body_inertia[None], n, 0).copy() | |
| qpos = np.repeat(qpos0[None], n, 0).copy() | |
| force = np.zeros((n, 3)) | |
| for i, w in enumerate(worlds): | |
| friction[i, :, 0] *= w.friction # sliding friction of every geom | |
| mass[i, objects] *= w.object_mass | |
| inertia[i, objects] *= w.object_mass | |
| qpos[i, noise_mask] += rng.uniform(-cfg.joint_noise, cfg.joint_noise, noise_mask.sum()) | |
| a = math.radians(w.push_dir_deg) | |
| force[i] = (w.push_n * math.cos(a), w.push_n * math.sin(a), 0.0) | |
| return friction, mass, inertia, qpos, force | |
| def _finish(worlds, heights, fell_at, final, fall_height) -> None: | |
| for i, w in enumerate(worlds): | |
| w.min_height = round(float(heights[i]), 3) | |
| w.final_height = round(float(final[i]), 3) | |
| if fall_height is not None: | |
| w.fell = fell_at[i] is not None | |
| w.fell_at = fell_at[i] | |
| # --------------------------------------------------------------------------- engines | |
| def run_cpu(scene: Path, robot: str | None, cfg: SweepConfig, worlds: list[World]) -> None: | |
| model, data, run_cfg, root, robot_bodies = _setup(scene, robot, cfg) | |
| friction, mass, inertia, qpos, force = _world_params(model, worlds, robot_bodies, cfg, data.qpos.copy()) | |
| base = (model.geom_friction.copy(), model.body_mass.copy(), model.body_inertia.copy()) | |
| steps = int(math.ceil(cfg.duration_s / model.opt.timestep)) | |
| push = (int(cfg.push_at / model.opt.timestep), int((cfg.push_at + cfg.push_s) / model.opt.timestep)) | |
| n = len(worlds) | |
| heights, final, fell_at = np.full(n, np.inf), np.zeros(n), [None] * n | |
| for i in range(n): | |
| model.geom_friction[:], model.body_mass[:], model.body_inertia[:] = friction[i], mass[i], inertia[i] | |
| mujoco.mj_resetData(model, data) | |
| data.qpos[:] = qpos[i] | |
| mujoco.mj_forward(model, data) | |
| hold = _hold_controller(model, data, run_cfg) if cfg.controller == "hold" else None | |
| for s in range(steps): | |
| if hold is not None: | |
| hold.apply(data) | |
| if root > 0: | |
| data.xfrc_applied[root, :3] = force[i] if push[0] <= s < push[1] else 0.0 | |
| mujoco.mj_step(model, data) | |
| if root > 0: | |
| z = float(data.xpos[root, 2]) | |
| heights[i] = min(heights[i], z) | |
| if fell_at[i] is None and run_cfg.fall_height is not None and z < run_cfg.fall_height: | |
| fell_at[i] = round(float(data.time), 2) | |
| final[i] = float(data.xpos[root, 2]) if root > 0 else 0.0 | |
| model.geom_friction[:], model.body_mass[:], model.body_inertia[:] = base | |
| _finish(worlds, heights, fell_at, final, run_cfg.fall_height) | |
| def run_warp( | |
| scene: Path, robot: str | None, cfg: SweepConfig, worlds: list[World], device: str | None = None | |
| ) -> str | None: | |
| """Returns a note when the model had to be adapted for MuJoCo Warp.""" | |
| import mujoco_warp as mjw | |
| import warp as wp | |
| wp.config.quiet = True | |
| model, data, run_cfg, root, robot_bodies = _setup(scene, robot, cfg) | |
| note = None | |
| friction, mass, inertia, qpos, force = _world_params(model, worlds, robot_bodies, cfg, data.qpos.copy()) | |
| n = len(worlds) | |
| hold = _hold_controller(model, data, run_cfg) if cfg.controller == "hold" else None | |
| with wp.ScopedDevice(device): | |
| batch = {"geom_friction": n, "body_mass": n, "body_inertia": n} | |
| try: | |
| m = mjw.put_model(model, batch_sizes=batch) | |
| except NotImplementedError as e: | |
| if "margin" not in str(e): | |
| raise | |
| # e.g. Menagerie's Go2 gives its feet 1 mm contact margins, which MuJoCo Warp's CCD does not support. | |
| note = f"contact margins up to {model.geom_margin.max() * 1000:g} mm set to 0 for MuJoCo Warp" | |
| model.geom_margin[:] = 0.0 | |
| m = mjw.put_model(model, batch_sizes=batch) | |
| m.geom_friction.assign(friction.astype(np.float32)) | |
| m.body_mass.assign(mass.astype(np.float32)) | |
| m.body_inertia.assign(inertia.astype(np.float32)) | |
| d = mjw.put_data(model, data, nworld=n, nconmax=max(64, 8 * model.ngeom), njmax=max(256, 40 * model.nv)) | |
| d.qpos.assign(qpos.astype(np.float32)) | |
| mjw.forward(m, d) | |
| ctrl_base = np.repeat((hold.ctrl if hold else np.zeros(model.nu))[None], n, 0).astype(np.float32) | |
| pd = hold is not None and len(hold.motors) > 0 | |
| if pd: | |
| motors = wp.array(hold.motors.astype(np.int32), dtype=int) | |
| qadr = wp.array(hold.qadr.astype(np.int32), dtype=int) | |
| vadr = wp.array(hold.vadr.astype(np.int32), dtype=int) | |
| target, kp, kd, lo, hi = ( | |
| wp.array(a.astype(np.float32), dtype=float) for a in (hold.target, hold.kp, hold.kd, hold.lo, hold.hi) | |
| ) | |
| d.ctrl.assign(ctrl_base) | |
| xfrc = np.zeros((n, model.nbody, 6), np.float32) | |
| steps = int(math.ceil(cfg.duration_s / model.opt.timestep)) | |
| push = (int(cfg.push_at / model.opt.timestep), int((cfg.push_at + cfg.push_s) / model.opt.timestep)) | |
| sample_every = max(1, int(0.02 / model.opt.timestep)) | |
| heights, fell_at = np.full(n, np.inf), [None] * n | |
| for s in range(steps): | |
| if root > 0 and s in push: | |
| xfrc[:, root, :3] = force if s == push[0] else 0.0 | |
| d.xfrc_applied.assign(xfrc) | |
| if pd: | |
| wp.launch( | |
| _pd_kernel(), | |
| dim=(n, len(hold.motors)), | |
| inputs=[d.qpos, d.qvel, motors, qadr, vadr, target, kp, kd, lo, hi], | |
| outputs=[d.ctrl], | |
| ) | |
| mjw.step(m, d) | |
| if root > 0 and (s % sample_every == 0 or s == steps - 1): | |
| z = d.xpos.numpy()[:, root, 2] | |
| heights = np.minimum(heights, z) | |
| if run_cfg.fall_height is not None: | |
| t = round((s + 1) * model.opt.timestep, 2) | |
| for i in np.nonzero(z < run_cfg.fall_height)[0]: | |
| if fell_at[i] is None: | |
| fell_at[i] = t | |
| final = d.xpos.numpy()[:, root, 2] if root > 0 else np.zeros(n) | |
| _finish(worlds, heights, fell_at, final, run_cfg.fall_height) | |
| return note | |
| def _make_pd_kernel(): | |
| import warp as wp | |
| def pd( | |
| qpos: wp.array2d(dtype=float), | |
| qvel: wp.array2d(dtype=float), | |
| motors: wp.array(dtype=int), | |
| qadr: wp.array(dtype=int), | |
| vadr: wp.array(dtype=int), | |
| target: wp.array(dtype=float), | |
| kp: wp.array(dtype=float), | |
| kd: wp.array(dtype=float), | |
| lo: wp.array(dtype=float), | |
| hi: wp.array(dtype=float), | |
| ctrl: wp.array2d(dtype=float), | |
| ): | |
| w, k = wp.tid() | |
| tau = kp[k] * (target[k] - qpos[w, qadr[k]]) - kd[k] * qvel[w, vadr[k]] | |
| ctrl[w, motors[k]] = wp.clamp(tau, lo[k], hi[k]) | |
| return pd | |
| _PD_KERNEL = None | |
| def _pd_kernel(): | |
| """Compiled on first use, so importing this module never needs Warp.""" | |
| global _PD_KERNEL | |
| if _PD_KERNEL is None: | |
| _PD_KERNEL = _make_pd_kernel() | |
| return _PD_KERNEL | |
| def cuda_available() -> bool: | |
| try: | |
| import warp as wp | |
| wp.config.quiet = True | |
| wp.init() | |
| return wp.is_cuda_available() | |
| except Exception: | |
| return False | |
| # --------------------------------------------------------------------------- job | |
| def sweep(scene: Path, robot: str | None, cfg: SweepConfig, engine: str = "auto") -> dict: | |
| if not 1 <= cfg.n_worlds <= MAX_WORLDS: | |
| raise ValueError(f"n_worlds must be 1..{MAX_WORLDS}") | |
| cfg.duration_s = min(cfg.duration_s, MAX_DURATION_S) | |
| if engine == "auto": | |
| engine = "warp" if cuda_available() else "mujoco" | |
| worlds = sample_worlds(cfg) | |
| started = time.monotonic() | |
| note = None | |
| if engine == "warp": | |
| try: | |
| note = run_warp(scene, robot, cfg, worlds, device="cuda:0" if cuda_available() else "cpu") | |
| except NotImplementedError as e: # a model feature MuJoCo Warp lacks: same sweep on the CPU | |
| engine, note = "mujoco", f"MuJoCo Warp cannot run this scene ({e}); ran on the CPU instead" | |
| worlds = sample_worlds(cfg) | |
| run_cpu(scene, robot, cfg, worlds) | |
| else: | |
| run_cpu(scene, robot, cfg, worlds) | |
| label = "MuJoCo (CPU)" | |
| if engine == "warp": | |
| label = "MuJoCo Warp (GPU)" if cuda_available() else "MuJoCo Warp (CPU device)" | |
| result = summarize(worlds, cfg, robot, label, time.monotonic() - started) | |
| if note: | |
| result["summary"]["engine_note"] = note | |
| return result | |
| def summarize(worlds: list[World], cfg: SweepConfig, robot: str | None, engine: str, wall: float) -> dict: | |
| spec = ROBOTS.get(robot) if robot else None | |
| detect = spec is not None and spec.fall_height is not None | |
| fell = [w for w in worlds if w.fell] | |
| n = len(worlds) | |
| summary = { | |
| "worlds": n, | |
| "engine": engine, | |
| "simulated_seconds_each": cfg.duration_s, | |
| "wall_seconds": round(wall, 1), | |
| "controller": cfg.controller, | |
| "ranges": { | |
| "friction_scale": list(cfg.friction), | |
| "object_mass_scale": list(cfg.object_mass), | |
| "push_newtons": list(cfg.push_n), | |
| "push_at_s": cfg.push_at, | |
| "push_duration_s": cfg.push_s, | |
| "joint_noise_rad": cfg.joint_noise, | |
| }, | |
| "seed": cfg.seed, | |
| } | |
| if detect: | |
| summary["fell"] = len(fell) | |
| summary["fall_rate"] = round(len(fell) / n, 3) | |
| if fell: | |
| summary["mean_push_n_fell"] = round(float(np.mean([w.push_n for w in fell])), 1) | |
| summary["min_push_n_fell"] = min(w.push_n for w in fell) | |
| stood = [w for w in worlds if not w.fell] | |
| if stood: | |
| summary["max_push_n_survived"] = max(w.push_n for w in stood) | |
| summary["headline"] = f"{len(fell)}/{n} fell ({100 * len(fell) / n:.0f}%)" | |
| if fell: | |
| summary["headline"] += f" · lightest push that toppled it: {summary['min_push_n_fell']:g} N" | |
| else: | |
| summary["headline"] = f"{n} worlds simulated (no fall detection for this robot)" | |
| return {"summary": summary, "worlds": [asdict(w) for w in worlds]} | |
| def run_sweep_job(job_input: dict, menagerie_dir: Path | None = None, engine: str = "auto") -> dict: | |
| """Worker entry point. Returns the summary and the per-world table as an inline artifact.""" | |
| import os | |
| import shutil | |
| from api.mjcf_check import SCENE_FILE, prepare_scene_dir, static_check | |
| mjcf, robot = job_input["mjcf"], job_input.get("robot") | |
| failed = static_check(mjcf, robot) | |
| if failed: | |
| raise ValueError(f"scene rejected: {failed.error}") | |
| fields = {k: job_input[k] for k in ("n_worlds", "duration_s", "controller", "seed") if k in job_input} | |
| for k in ("friction", "object_mass", "push_n"): | |
| if k in job_input: | |
| fields[k] = tuple(float(x) for x in job_input[k]) | |
| cfg = SweepConfig(**fields) | |
| menagerie = Path(menagerie_dir or os.environ.get("MENAGERIE_DIR", "/opt/mujoco_menagerie")) | |
| scene_dir = prepare_scene_dir(mjcf, robot, menagerie) | |
| try: | |
| result = sweep(scene_dir / SCENE_FILE, robot, cfg, engine) | |
| finally: | |
| shutil.rmtree(scene_dir, ignore_errors=True) | |
| worker_seconds = result["summary"]["wall_seconds"] | |
| return { | |
| "summary": result["summary"], | |
| "worker_seconds": worker_seconds, | |
| "artifacts": {"sweep.json": json.dumps(result, indent=1).encode()}, | |
| } | |