simrig 0.2.2__py3-none-any.whl

simrig/rendering.py

@@ -0,0 +1,127 @@
+"""MuJoCo offscreen rendering helpers for browser viewers."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import os
+import sys
+from typing import Any
+
+
+@dataclass
+class CameraState:
+    azimuth: float = 135.0
+    elevation: float = -20.0
+    distance: float = 2.4
+    interactive: bool = True
+
+    def apply(self, camera: Any) -> Any:
+        if not self.interactive or not hasattr(camera, "azimuth"):
+            return camera
+        camera.azimuth = float(self.azimuth)
+        camera.elevation = float(self.elevation)
+        camera.distance = float(self.distance)
+        return camera
+
+    def update_from_query(self, query: dict[str, list[str]]) -> None:
+        for key in ("azimuth", "elevation", "distance"):
+            if key in query:
+                setattr(self, key, float(query[key][0]))
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "azimuth": self.azimuth,
+            "elevation": self.elevation,
+            "distance": self.distance,
+            "interactive": self.interactive,
+        }
+
+
+def preferred_gl_backends() -> list[str | None]:
+    """Return MuJoCo GL backends to try, in order."""
+
+    configured = os.environ.get("MUJOCO_GL")
+    if configured:
+        return [configured]
+    if sys.platform == "darwin":
+        return ["glfw", None]
+    if os.environ.get("DISPLAY"):
+        return ["egl", "glfw", None]
+    return ["osmesa", "egl", None]
+
+
+def ensure_offscreen_framebuffer(model: Any, *, height: int, width: int) -> None:
+    """Grow MuJoCo's offscreen buffer so browser frames can exceed 640x480."""
+
+    vis_global = model.vis.global_
+    vis_global.offwidth = max(int(vis_global.offwidth), int(width))
+    vis_global.offheight = max(int(vis_global.offheight), int(height))
+
+
+def create_mujoco_renderer(
+    mujoco: Any,
+    model: Any,
+    *,
+    height: int,
+    width: int,
+) -> Any:
+    """Create a MuJoCo offscreen renderer, trying platform-appropriate GL backends."""
+
+    ensure_offscreen_framebuffer(model, height=height, width=width)
+    errors: list[str] = []
+    for backend in preferred_gl_backends():
+        if backend:
+            os.environ["MUJOCO_GL"] = backend
+        try:
+            return mujoco.Renderer(model, height=height, width=width)
+        except Exception as exc:
+            errors.append(f"{backend or 'default'}: {exc}")
+    raise RuntimeError(
+        "MuJoCo offscreen rendering failed. Tried "
+        + "; ".join(errors)
+        + ". On macOS, run from a desktop session and keep MUJOCO_GL=glfw."
+    )
+
+
+def make_tracking_camera(
+    mujoco: Any,
+    model: Any,
+    data: Any,
+    camera: str | int | None,
+) -> tuple[Any, CameraState]:
+    """Build a MuJoCo camera object and mutable browser camera state."""
+
+    if isinstance(camera, int):
+        return camera, CameraState(interactive=False)
+    if isinstance(camera, str) and camera.isdigit():
+        return int(camera), CameraState(interactive=False)
+    if camera not in (None, "track", "tracking"):
+        return camera, CameraState(interactive=False)
+
+    cam = mujoco.MjvCamera()
+    cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
+    cam.trackbodyid = tracking_body_id(mujoco, model, data)
+    cam.distance = 2.4
+    cam.azimuth = 135
+    cam.elevation = -20
+    state = CameraState(
+        azimuth=float(cam.azimuth),
+        elevation=float(cam.elevation),
+        distance=float(cam.distance),
+        interactive=True,
+    )
+    return cam, state
+
+
+def tracking_body_id(mujoco: Any, model: Any, data: Any | None = None) -> int:
+    """Pick a stable body for camera tracking."""
+
+    del data
+    for joint_id in range(model.njnt):
+        if model.jnt_type[joint_id] == mujoco.mjtJoint.mjJNT_FREE:
+            return int(model.jnt_bodyid[joint_id])
+    for name in ("trunk", "base", "torso", "torso_link", "pelvis", "body", "root"):
+        body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, name)
+        if body_id >= 0:
+            return int(body_id)
+    return 1 if model.nbody > 1 else 0

simrig/scaffold.py

@@ -0,0 +1,127 @@
+"""Editable starter files for custom environments."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from simrig.io import slugify
+
+# Markers used by validate-env static checklist (keep in sync with template).
+REQUIRED_SECTION_MARKERS = (
+    "SECTION: model loading",
+    "SECTION: reset",
+    "SECTION: action mapping",
+    "SECTION: observations",
+    "SECTION: rewards",
+    "SECTION: termination",
+)
+
+REQUIRED_CLASS_METHODS = (
+    "__init__",
+    "reset",
+    "step",
+)
+
+
+def new_env(name: str, model: str | Path, *, template: str = "mjx", root: Path | str = "envs") -> Path:
+    """Create an editable starter env module."""
+    if template != "mjx":
+        raise ValueError("SimRig v0 supports only the 'mjx' env template.")
+    module_name = slugify(name).replace("-", "_").replace(".", "_")
+    path = Path(root) / f"{module_name}.py"
+    if path.exists():
+        raise FileExistsError(f"Refusing to overwrite existing env template: {path}")
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(_mjx_template(name=name, model=str(model)), encoding="utf-8")
+    return path
+
+
+def _mjx_template(*, name: str, model: str) -> str:
+    return f'''"""Editable SimRig MJX environment starter for {name}.
+
+NOT TRAINABLE YET.
+Fill the SECTION blocks below, remove this banner when reset/step work, then:
+
+    simrig validate-env PATH --runtime
+    simrig smoke PATH --steps 10
+    simrig train PATH --preset smoke
+
+Prefer subclassing mujoco_playground._src.mjx_env.MjxEnv when Playground is
+installed. Return obs as dict keys `state` and `privileged_state` for SimRig PPO.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+
+MODEL_PATH = Path({model!r}).expanduser()
+ENV_NAME = {name!r}
+
+
+def default_config() -> dict[str, Any]:
+    return {{
+        "episode_length": 1000,
+        "action_scale": 1.0,
+        "ctrl_dt": 0.02,
+        "sim_dt": 0.002,
+        "impl": "jax",
+    }}
+
+
+def make_env(config_overrides: dict[str, Any] | None = None) -> "CustomEnv":
+    """Factory used by `simrig smoke/train/eval` for this module."""
+    config = default_config()
+    if config_overrides:
+        config.update(config_overrides)
+    return CustomEnv(config=config)
+
+
+class CustomEnv:
+    """Replace this starter with a real mjx_env.MjxEnv implementation.
+
+    Required for SimRig smoke/train/eval:
+    - reset(rng) -> state with .obs, .reward, .done, .data
+    - step(state, action) -> state
+    - observation_size (dict with `state` and `privileged_state` preferred)
+    - action_size
+    - mj_model / mjx_model (for demos and previews)
+    """
+
+    def __init__(self, config: dict[str, Any] | None = None) -> None:
+        self.config = config or default_config()
+        # SECTION: model loading
+        # Load MODEL_PATH into mj_model / mjx_model and cache actuator limits.
+        raise NotImplementedError(
+            "Define model loading, reset, step, rewards, observations, and termination."
+        )
+
+    def reset(self, rng: Any) -> Any:
+        # SECTION: reset
+        # Sample initial qpos/qvel (and any task state). Return a Brax/MJX state.
+        raise NotImplementedError("Implement reset randomization.")
+
+    def step(self, state: Any, action: Any) -> Any:
+        # SECTION: action mapping
+        # Scale/clip `action` into actuator controls.
+
+        # SECTION: observations
+        # Build policy obs (`state`) and privileged obs (`privileged_state`).
+
+        # SECTION: rewards
+        # Compute dense/sparse reward terms; do not invent them from the model name.
+
+        # SECTION: termination
+        # Set done / truncations from falls, limits, success, or time.
+
+        raise NotImplementedError("Implement step, including action mapping, obs, reward, termination.")
+
+    @property
+    def observation_size(self) -> Any:
+        raise NotImplementedError("Return dict observation sizes for state/privileged_state.")
+
+    @property
+    def action_size(self) -> int:
+        raise NotImplementedError("Return the actuator/action dimension.")
+'''

simrig/three_scene.py

@@ -0,0 +1,107 @@
+"""Serialize compiled MuJoCo geometry for browser-side Three.js rendering."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+
+def scene_payload(
+    mujoco: Any,
+    model: Any,
+    data: Any,
+    *,
+    model_name: str,
+) -> dict[str, Any]:
+    """Return visual geometry and current world transforms for a model."""
+
+    visible_geom_ids = [
+        geom_id
+        for geom_id in range(model.ngeom)
+        if int(model.geom_group[geom_id]) <= 2
+    ]
+    used_mesh_ids = sorted(
+        {
+            int(model.geom_dataid[geom_id])
+            for geom_id in visible_geom_ids
+            if int(model.geom_type[geom_id]) == int(mujoco.mjtGeom.mjGEOM_MESH)
+            and int(model.geom_dataid[geom_id]) >= 0
+        }
+    )
+    return {
+        "model_name": model_name,
+        "coordinate_system": "z-up",
+        "meshes": [_mesh_payload(mujoco, model, mesh_id) for mesh_id in used_mesh_ids],
+        "geoms": [_geom_payload(mujoco, model, geom_id) for geom_id in visible_geom_ids],
+        "transforms": geom_transforms(model, data),
+    }
+
+
+def geom_transforms(model: Any, data: Any) -> list[dict[str, Any]]:
+    """Return each geom's current MuJoCo world position and rotation matrix."""
+
+    return [
+        {
+            "id": geom_id,
+            "position": np.asarray(data.geom_xpos[geom_id], dtype=float).tolist(),
+            "matrix": np.asarray(data.geom_xmat[geom_id], dtype=float).reshape(-1).tolist(),
+        }
+        for geom_id in range(model.ngeom)
+    ]
+
+
+def _mesh_payload(mujoco: Any, model: Any, mesh_id: int) -> dict[str, Any]:
+    vert_adr = int(model.mesh_vertadr[mesh_id])
+    vert_num = int(model.mesh_vertnum[mesh_id])
+    face_adr = int(model.mesh_faceadr[mesh_id])
+    face_num = int(model.mesh_facenum[mesh_id])
+    name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_MESH, mesh_id)
+    return {
+        "id": mesh_id,
+        "name": name or f"mesh_{mesh_id}",
+        "vertices": np.asarray(
+            model.mesh_vert[vert_adr : vert_adr + vert_num],
+            dtype=np.float32,
+        )
+        .reshape(-1)
+        .tolist(),
+        "indices": np.asarray(
+            model.mesh_face[face_adr : face_adr + face_num],
+            dtype=np.uint32,
+        )
+        .reshape(-1)
+        .tolist(),
+    }
+
+
+def _geom_payload(mujoco: Any, model: Any, geom_id: int) -> dict[str, Any]:
+    material_id = int(model.geom_matid[geom_id])
+    if material_id >= 0:
+        rgba = np.asarray(model.mat_rgba[material_id], dtype=float)
+        specular = float(model.mat_specular[material_id])
+        shininess = float(model.mat_shininess[material_id])
+        reflectance = float(model.mat_reflectance[material_id])
+        emission = float(model.mat_emission[material_id])
+    else:
+        rgba = np.asarray(model.geom_rgba[geom_id], dtype=float)
+        specular = 0.25
+        shininess = 0.25
+        reflectance = 0.0
+        emission = 0.0
+    name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_GEOM, geom_id)
+    return {
+        "id": geom_id,
+        "name": name or f"geom_{geom_id}",
+        "type": int(model.geom_type[geom_id]),
+        "mesh_id": int(model.geom_dataid[geom_id]),
+        "group": int(model.geom_group[geom_id]),
+        "size": np.asarray(model.geom_size[geom_id], dtype=float).tolist(),
+        "rgba": rgba.tolist(),
+        "material": {
+            "specular": specular,
+            "shininess": shininess,
+            "reflectance": reflectance,
+            "emission": emission,
+        },
+    }

simrig/validate_env.py

@@ -0,0 +1,211 @@
+"""Validation for custom env starter modules (static + optional runtime)."""
+
+from __future__ import annotations
+
+import ast
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from simrig.scaffold import REQUIRED_CLASS_METHODS, REQUIRED_SECTION_MARKERS
+
+
+@dataclass(frozen=True)
+class EnvValidationResult:
+    """Result of custom-env validation.
+
+    ``trainable`` is True only when ``--runtime`` reset/step checks succeed.
+    Static-only passes never claim trainability.
+    """
+
+    path: str
+    passed: bool
+    trainable: bool = False
+    missing: list[str] = field(default_factory=list)
+    warnings: list[str] = field(default_factory=list)
+    notes: list[str] = field(default_factory=list)
+
+
+def validate_env(path: Path | str, *, runtime: bool = False) -> EnvValidationResult:
+    """Validate a custom env module.
+
+    Static mode checks scaffold structure. Runtime mode also imports the module,
+    constructs the env, and runs a short reset/step smoke when JAX is available.
+    """
+    env_path = Path(path).expanduser()
+    missing: list[str] = []
+    warnings: list[str] = []
+    notes = [
+        "Static checklist checks structure only.",
+        "Use --runtime before proposing simrig smoke/train on a custom module.",
+    ]
+
+    if not env_path.is_file():
+        return EnvValidationResult(
+            path=str(env_path),
+            passed=False,
+            missing=[f"file not found: {env_path}"],
+            notes=notes,
+        )
+
+    source = env_path.read_text(encoding="utf-8")
+    for marker in REQUIRED_SECTION_MARKERS:
+        if marker not in source:
+            missing.append(f"section marker missing: {marker}")
+
+    try:
+        tree = ast.parse(source, filename=str(env_path))
+    except SyntaxError as exc:
+        return EnvValidationResult(
+            path=str(env_path),
+            passed=False,
+            missing=[f"syntax error: {exc}"],
+            notes=notes,
+        )
+
+    if "MODEL_PATH" not in source:
+        missing.append("symbol missing: MODEL_PATH")
+    if "ENV_NAME" not in source:
+        missing.append("symbol missing: ENV_NAME")
+    if "def default_config" not in source and "def make_env" not in source:
+        missing.append("function missing: default_config or make_env")
+
+    custom_env = _find_class(tree, "CustomEnv")
+    has_make_env = any(
+        isinstance(node, ast.FunctionDef) and node.name == "make_env" for node in tree.body
+    )
+    if custom_env is None and not has_make_env:
+        missing.append("class missing: CustomEnv (or define make_env)")
+    elif custom_env is not None:
+        method_names = {
+            node.name
+            for node in custom_env.body
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
+        }
+        for name in REQUIRED_CLASS_METHODS:
+            if name not in method_names:
+                missing.append(f"CustomEnv method missing: {name}")
+        for prop in ("observation_size", "action_size"):
+            if prop not in method_names:
+                missing.append(f"CustomEnv property missing: {prop}")
+
+    if "NotImplementedError" in source:
+        warnings.append("NotImplementedError still present; implementation is incomplete.")
+    if "NOT TRAINABLE YET" in source:
+        warnings.append("File still marked NOT TRAINABLE YET.")
+
+    passed = not missing
+    trainable = False
+
+    if runtime and passed:
+        runtime_missing, runtime_warnings, runtime_notes, trainable = _runtime_checks(env_path)
+        missing.extend(runtime_missing)
+        warnings.extend(runtime_warnings)
+        notes.extend(runtime_notes)
+        passed = not missing
+    elif passed:
+        notes.append(
+            "Checklist structure looks complete. Fill SECTION bodies, then "
+            "`simrig validate-env PATH --runtime` and `simrig smoke PATH`."
+        )
+
+    return EnvValidationResult(
+        path=str(env_path),
+        passed=passed,
+        trainable=trainable,
+        missing=missing,
+        warnings=warnings,
+        notes=notes,
+    )
+
+
+def _runtime_checks(env_path: Path) -> tuple[list[str], list[str], list[str], bool]:
+    missing: list[str] = []
+    warnings: list[str] = []
+    notes: list[str] = []
+    trainable = False
+
+    try:
+        from simrig.custom_env import load_custom_env
+
+        env = load_custom_env(env_path)
+    except Exception as exc:
+        missing.append(f"runtime load failed: {exc}")
+        return missing, warnings, notes, False
+
+    notes.append("Custom env module constructed successfully.")
+    action_size = getattr(env, "action_size", None)
+    observation_size = getattr(env, "observation_size", None)
+
+    if action_size is None:
+        missing.append("runtime: action_size missing")
+    obs_warnings = _obs_size_warnings(observation_size)
+    warnings.extend(obs_warnings)
+
+    for attr in ("mj_model", "mjx_model"):
+        if getattr(env, attr, None) is None:
+            warnings.append(f"runtime: {attr} missing (needed for Playground/Brax demos).")
+
+    try:
+        import jax  # type: ignore
+        import jax.numpy as jp  # type: ignore
+    except ImportError:
+        warnings.append("runtime: JAX not installed; skipped reset/step smoke.")
+        notes.append("Install the playground extra to run reset/step validation.")
+        return missing, warnings, notes, False
+
+    try:
+        state = env.reset(jax.random.PRNGKey(0))
+        obs = getattr(state, "obs", None)
+        if isinstance(obs, dict):
+            if "state" not in obs:
+                missing.append("runtime: state.obs missing key `state`")
+            if "privileged_state" not in obs:
+                warnings.append("runtime: state.obs missing key `privileged_state`")
+        elif obs is None:
+            missing.append("runtime: state.obs missing")
+        else:
+            warnings.append(
+                "runtime: state.obs is flat; SimRig PPO defaults expect dict keys "
+                "`state` and `privileged_state`."
+            )
+
+        if action_size is None:
+            return missing, warnings, notes, False
+        state = env.step(state, jp.zeros(int(action_size)))
+        _ = float(state.reward), bool(state.done)
+    except Exception as exc:
+        missing.append(f"runtime reset/step failed: {exc}")
+        return missing, warnings, notes, False
+
+    if missing:
+        return missing, warnings, notes, False
+
+    trainable = True
+    notes.append("Runtime reset/step succeeded. Safe to try `simrig smoke` then `simrig train`.")
+    return missing, warnings, notes, trainable
+
+
+def _obs_size_warnings(observation_size: Any) -> list[str]:
+    warnings: list[str] = []
+    if observation_size is None:
+        warnings.append("runtime: observation_size missing")
+        return warnings
+    if isinstance(observation_size, dict):
+        if "state" not in observation_size:
+            warnings.append("runtime: observation_size missing key `state`")
+        if "privileged_state" not in observation_size:
+            warnings.append("runtime: observation_size missing key `privileged_state`")
+    else:
+        warnings.append(
+            "runtime: observation_size is flat; SimRig PPO defaults expect "
+            "`state` / `privileged_state`."
+        )
+    return warnings
+
+
+def _find_class(tree: ast.AST, name: str) -> ast.ClassDef | None:
+    for node in ast.walk(tree):
+        if isinstance(node, ast.ClassDef) and node.name == name:
+            return node
+    return None