simrig 0.2.2__py3-none-any.whl

simrig/cli.py

@@ -0,0 +1,412 @@
+"""Command-line interface for SimRig."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import Any
+
+from simrig._version import __version__
+from simrig.core import report_markdown, to_dict
+from simrig.huggingface import resolve_policy_checkpoint
+from simrig.io import save_json, save_report_pair, slugify
+from simrig.mujoco_backend import inspect_model, list_models
+from simrig.paths import ensure_project_dirs
+from simrig.playground_backend import (
+    demo_policy,
+    eval_policy,
+    inspect_env,
+    list_envs,
+    smoke_env,
+    train_ppo,
+)
+from simrig.scaffold import new_env
+from simrig.validate_env import validate_env
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    try:
+        result = args.func(args)
+    except Exception as exc:
+        print(f"simrig: error: {exc}", file=sys.stderr)
+        return 1
+    return 0 if result is None else int(result)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="simrig",
+        description="Physical AI simulation training starter workflows.",
+    )
+    parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    init = sub.add_parser("init", help="Create local SimRig output folders.")
+    init.add_argument("--root", type=Path, default=Path("."))
+    init.set_defaults(func=_cmd_init)
+
+    list_models_parser = sub.add_parser("list-models", help="List MuJoCo Menagerie models.")
+    list_models_parser.add_argument("--menagerie", type=Path)
+    list_models_parser.add_argument("--json", action="store_true")
+    list_models_parser.set_defaults(func=_cmd_list_models)
+
+    inspect_model_parser = sub.add_parser("inspect-model", help="Inspect a MuJoCo XML/model.")
+    inspect_model_parser.add_argument("model_or_xml")
+    inspect_model_parser.add_argument("--menagerie", type=Path)
+    inspect_model_parser.add_argument("--steps", type=int, default=25)
+    inspect_model_parser.add_argument("--json", action="store_true")
+    inspect_model_parser.add_argument("--save-report", action="store_true")
+    inspect_model_parser.set_defaults(func=_cmd_inspect_model)
+
+    list_envs_parser = sub.add_parser("list-envs", help="List trainable backend envs.")
+    list_envs_parser.add_argument("--backend", default="mujoco-playground")
+    list_envs_parser.add_argument("--json", action="store_true")
+    list_envs_parser.set_defaults(func=_cmd_list_envs)
+
+    inspect_env_parser = sub.add_parser("inspect-env", help="Inspect a Playground or custom env.")
+    inspect_env_parser.add_argument(
+        "env_name",
+        help="Playground env name or path to a custom *.py env module.",
+    )
+    inspect_env_parser.add_argument("--backend", default="mujoco-playground")
+    inspect_env_parser.add_argument("--json", action="store_true")
+    inspect_env_parser.add_argument("--save-report", action="store_true")
+    inspect_env_parser.set_defaults(func=_cmd_inspect_env)
+
+    smoke_parser = sub.add_parser("smoke", help="Run a short env reset/step smoke test.")
+    smoke_parser.add_argument(
+        "env_name",
+        help="Playground env name or path to a custom *.py env module.",
+    )
+    smoke_parser.add_argument("--backend", default="mujoco-playground")
+    smoke_parser.add_argument("--steps", type=int, default=10)
+    smoke_parser.add_argument("--json", action="store_true")
+    smoke_parser.set_defaults(func=_cmd_smoke)
+
+    train_parser = sub.add_parser(
+        "train",
+        help="Train a Playground env or custom *.py module with Brax PPO.",
+    )
+    train_parser.add_argument(
+        "env_name",
+        help="Playground env name or path to a custom *.py env module.",
+    )
+    train_parser.add_argument("--backend", default="mujoco-playground")
+    train_parser.add_argument("--preset", choices=("smoke", "local", "cloud"), default="smoke")
+    train_parser.add_argument("--output", type=Path)
+    train_parser.add_argument("--timesteps", type=int)
+    train_parser.add_argument("--num-envs", type=int)
+    train_parser.add_argument("--batch-size", type=int)
+    train_parser.set_defaults(func=_cmd_train)
+
+    eval_parser = sub.add_parser("eval", help="Headless policy eval.")
+    eval_parser.add_argument("checkpoint")
+    eval_parser.add_argument(
+        "--env",
+        dest="env_name",
+        required=True,
+        help="Playground env name or path to a custom *.py env module.",
+    )
+    eval_parser.add_argument("--backend", default="mujoco-playground")
+    eval_parser.add_argument("--steps", type=int, default=500)
+    eval_parser.add_argument(
+        "--seed",
+        type=int,
+        default=0,
+        help="Deterministic environment and policy rollout seed.",
+    )
+    eval_parser.add_argument(
+        "--command",
+        type=float,
+        nargs="+",
+        help="Fix command-like environment state, for example X Y YAW.",
+    )
+    eval_parser.add_argument("--small-network", action=argparse.BooleanOptionalAction, default=None)
+    eval_parser.add_argument("--hf-revision", help="Revision for hf:// policy checkpoints.")
+    eval_parser.add_argument("--hf-token", help="Hugging Face token for private policy repos.")
+    eval_parser.add_argument("--json", action="store_true")
+    eval_parser.set_defaults(func=_cmd_eval)
+
+    demo_parser = sub.add_parser("demo", help="Run a trained policy in a desktop MuJoCo viewer.")
+    demo_parser.add_argument("checkpoint")
+    demo_parser.add_argument("--env", dest="env_name", required=True)
+    demo_parser.add_argument("--backend", default="mujoco-playground")
+    demo_parser.add_argument("--steps", type=int, default=5000)
+    demo_parser.add_argument("--small-network", action=argparse.BooleanOptionalAction, default=None)
+    demo_parser.add_argument("--hf-revision", help="Revision for hf:// policy checkpoints.")
+    demo_parser.add_argument("--hf-token", help="Hugging Face token for private policy repos.")
+    demo_parser.add_argument("--command", type=float, nargs="+")
+    demo_parser.add_argument("--speed", type=float, default=1.0)
+    demo_parser.add_argument("--camera-distance", type=float)
+    demo_parser.add_argument("--json", action="store_true")
+    demo_parser.set_defaults(func=_cmd_demo)
+
+    preview_parser = sub.add_parser("preview", help="Serve a trained policy preview in the browser.")
+    preview_parser.add_argument("checkpoint")
+    preview_parser.add_argument("--env", dest="env_name", required=True)
+    preview_parser.add_argument("--backend", default="mujoco-playground")
+    preview_parser.add_argument("--host", default="127.0.0.1")
+    preview_parser.add_argument("--port", type=int, default=8765)
+    preview_parser.add_argument("--width", type=int, default=960)
+    preview_parser.add_argument("--height", type=int, default=540)
+    preview_parser.add_argument("--frame-skip", type=int, default=1)
+    preview_parser.add_argument("--fps", type=int, default=24, help="Browser render loop target FPS.")
+    preview_parser.add_argument("--small-network", action=argparse.BooleanOptionalAction, default=None)
+    preview_parser.add_argument("--hf-revision", help="Revision for hf:// policy checkpoints.")
+    preview_parser.add_argument("--hf-token", help="Hugging Face token for private policy repos.")
+    preview_parser.add_argument("--command", type=float, nargs="+")
+    preview_parser.add_argument("--camera")
+    preview_parser.add_argument("--paused", action="store_true", help="Start the browser preview paused.")
+    preview_parser.add_argument(
+        "--render-mode",
+        choices=("threejs", "mujoco", "topdown"),
+        default="threejs",
+        help=(
+            "Browser render mode. threejs renders rollout geometry interactively in WebGL; "
+            "mujoco streams offscreen frames; topdown is a schematic debug view."
+        ),
+    )
+    preview_parser.set_defaults(func=_cmd_preview)
+
+    view_model_parser = sub.add_parser(
+        "view-model",
+        help="Serve a MuJoCo model in the browser with per-joint controls.",
+    )
+    view_model_parser.add_argument("model_or_xml")
+    view_model_parser.add_argument("--menagerie", type=Path)
+    view_model_parser.add_argument("--host", default="127.0.0.1")
+    view_model_parser.add_argument("--port", type=int, default=8766)
+    view_model_parser.add_argument("--width", type=int, default=960)
+    view_model_parser.add_argument("--height", type=int, default=540)
+    view_model_parser.add_argument("--fps", type=int, default=24, help="Browser render loop target FPS.")
+    view_model_parser.add_argument("--camera")
+    view_model_parser.add_argument(
+        "--render-mode",
+        choices=("threejs", "mujoco", "topdown"),
+        default="threejs",
+        help=(
+            "Browser render mode. threejs renders geometry interactively in WebGL; "
+            "mujoco streams offscreen frames; topdown is a schematic debug view."
+        ),
+    )
+    view_model_parser.set_defaults(func=_cmd_view_model)
+
+    new_env_parser = sub.add_parser("new-env", help="Create an editable custom env starter.")
+    new_env_parser.add_argument("name")
+    new_env_parser.add_argument("--model", required=True)
+    new_env_parser.add_argument("--template", default="mjx")
+    new_env_parser.add_argument("--root", type=Path, default=Path("envs"))
+    new_env_parser.set_defaults(func=_cmd_new_env)
+
+    validate_env_parser = sub.add_parser(
+        "validate-env",
+        help="Validate a custom env module (static checklist; optional runtime).",
+    )
+    validate_env_parser.add_argument("path", type=Path)
+    validate_env_parser.add_argument(
+        "--runtime",
+        action="store_true",
+        help="Import the module and run construct/reset/step checks when possible.",
+    )
+    validate_env_parser.add_argument("--json", action="store_true")
+    validate_env_parser.set_defaults(func=_cmd_validate_env)
+
+    return parser
+
+
+def _cmd_init(args: argparse.Namespace) -> None:
+    paths = ensure_project_dirs(args.root)
+    for path in paths:
+        print(path)
+
+
+def _cmd_list_models(args: argparse.Namespace) -> None:
+    entries = list_models(args.menagerie)
+    _print(entries, as_json=args.json)
+
+
+def _cmd_inspect_model(args: argparse.Namespace) -> None:
+    report = inspect_model(args.model_or_xml, menagerie=args.menagerie, steps=args.steps)
+    if args.save_report:
+        md_path, json_path = save_report_pair(
+            report,
+            name=report.name,
+            title=f"Model Inspection: {report.name}",
+        )
+        print(f"saved {md_path}")
+        print(f"saved {json_path}")
+    _print_report(f"Model Inspection: {report.name}", report, as_json=args.json)
+
+
+def _cmd_list_envs(args: argparse.Namespace) -> None:
+    _print(list_envs(args.backend), as_json=args.json)
+
+
+def _cmd_inspect_env(args: argparse.Namespace) -> None:
+    report = inspect_env(args.env_name, backend=args.backend)
+    if args.save_report:
+        md_path, json_path = save_report_pair(
+            report,
+            name=report.name,
+            title=f"Env Inspection: {report.name}",
+        )
+        print(f"saved {md_path}")
+        print(f"saved {json_path}")
+    _print_report(f"Env Inspection: {report.name}", report, as_json=args.json)
+
+
+def _cmd_smoke(args: argparse.Namespace) -> int:
+    result = smoke_env(args.env_name, backend=args.backend, steps=args.steps)
+    _print(result, as_json=args.json)
+    return 0 if result.passed else 1
+
+
+def _cmd_train(args: argparse.Namespace) -> None:
+    overrides = _training_overrides(args)
+    run_config = train_ppo(
+        args.env_name,
+        backend=args.backend,
+        preset_name=args.preset,
+        output=args.output,
+        overrides=overrides,
+    )
+    print(f"saved run: {run_config.output_dir}")
+
+
+def _cmd_eval(args: argparse.Namespace) -> None:
+    command = tuple(args.command) if args.command is not None else None
+    checkpoint = resolve_policy_checkpoint(
+        args.checkpoint,
+        hf_revision=args.hf_revision,
+        hf_token=args.hf_token,
+    )
+    result = eval_policy(
+        checkpoint,
+        env_name=args.env_name,
+        backend=args.backend,
+        steps=args.steps,
+        small_network=args.small_network,
+        seed=args.seed,
+        command=command,
+    )
+    save_json(Path("reports") / f"{slugify(args.env_name)}_eval.json", result)
+    _print(result, as_json=args.json)
+
+
+def _cmd_demo(args: argparse.Namespace) -> None:
+    command = tuple(args.command) if args.command is not None else None
+    checkpoint = resolve_policy_checkpoint(
+        args.checkpoint,
+        hf_revision=args.hf_revision,
+        hf_token=args.hf_token,
+    )
+    result = demo_policy(
+        checkpoint,
+        env_name=args.env_name,
+        backend=args.backend,
+        steps=args.steps,
+        small_network=args.small_network,
+        command=command,
+        speed=args.speed,
+        camera_distance=args.camera_distance,
+    )
+    _print(result, as_json=args.json)
+
+
+def _cmd_view_model(args: argparse.Namespace) -> None:
+    from simrig.model_view import serve_model_view
+
+    serve_model_view(
+        args.model_or_xml,
+        menagerie=args.menagerie,
+        host=args.host,
+        port=args.port,
+        width=args.width,
+        height=args.height,
+        render_mode=args.render_mode,
+        camera=args.camera,
+        fps=args.fps,
+    )
+
+
+def _cmd_preview(args: argparse.Namespace) -> None:
+    from simrig.preview import serve_policy_preview
+
+    command = tuple(args.command) if args.command is not None else None
+    checkpoint = resolve_policy_checkpoint(
+        args.checkpoint,
+        hf_revision=args.hf_revision,
+        hf_token=args.hf_token,
+    )
+    serve_policy_preview(
+        checkpoint,
+        env_name=args.env_name,
+        backend=args.backend,
+        host=args.host,
+        port=args.port,
+        width=args.width,
+        height=args.height,
+        frame_skip=args.frame_skip,
+        small_network=args.small_network,
+        command=command,
+        camera=args.camera,
+        render_mode=args.render_mode,
+        paused=args.paused,
+        fps=args.fps,
+    )
+
+
+def _cmd_new_env(args: argparse.Namespace) -> None:
+    path = new_env(args.name, args.model, template=args.template, root=args.root)
+    print(path)
+
+
+def _cmd_validate_env(args: argparse.Namespace) -> int:
+    result = validate_env(args.path, runtime=bool(args.runtime))
+    _print(result, as_json=args.json)
+    if not args.json:
+        status = "passed" if result.passed else "failed"
+        print(f"validate-env: {status} (trainable={result.trainable})")
+    return 0 if result.passed else 1
+
+
+def _training_overrides(args: argparse.Namespace) -> dict[str, Any]:
+    overrides: dict[str, Any] = {}
+    for cli_name, key in (
+        ("timesteps", "timesteps"),
+        ("num_envs", "num_envs"),
+        ("batch_size", "batch_size"),
+    ):
+        value = getattr(args, cli_name)
+        if value is not None:
+            overrides[key] = value
+    return overrides
+
+
+def _print_report(title: str, report: Any, *, as_json: bool) -> None:
+    if as_json:
+        _print(report, as_json=True)
+    else:
+        print(report_markdown(title, report), end="")
+
+
+def _print(value: Any, *, as_json: bool) -> None:
+    if as_json:
+        print(json.dumps(to_dict(value), indent=2, sort_keys=True))
+        return
+    if isinstance(value, list):
+        for item in value:
+            if isinstance(item, dict) and "name" in item:
+                print(item["name"])
+            else:
+                print(item)
+        return
+    print(json.dumps(to_dict(value), indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

simrig/core.py

@@ -0,0 +1,144 @@
+"""Backend-neutral SimRig data contracts."""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field, is_dataclass
+from enum import Enum
+from pathlib import Path
+from typing import Any
+
+
+class TrainabilityStatus(str, Enum):
+    """Coarse status for what SimRig can safely do with an asset."""
+
+    UNKNOWN = "unknown"
+    INSPECTABLE = "inspectable"
+    SIMULATABLE = "simulatable"
+    TRAINABLE_EXISTING_ENV = "trainable_existing_env"
+    NEEDS_CUSTOM_ENV = "needs_custom_env"
+    FAILED = "failed"
+
+
+@dataclass(frozen=True)
+class BackendInfo:
+    """Information about a simulation/training backend."""
+
+    name: str
+    available: bool
+    version: str | None = None
+    detail: str | None = None
+
+
+@dataclass(frozen=True)
+class ModelInspectionReport:
+    """Summary of a MuJoCo model inspection."""
+
+    name: str
+    path: str
+    backend: str
+    status: TrainabilityStatus
+    compiled: bool
+    stepped: bool
+    bodies: int = 0
+    joints: int = 0
+    dofs: int = 0
+    actuators: int = 0
+    sensors: int = 0
+    keyframes: int = 0
+    has_freejoint: bool = False
+    has_mjx_hint: bool = False
+    warnings: list[str] = field(default_factory=list)
+    errors: list[str] = field(default_factory=list)
+    notes: list[str] = field(default_factory=list)
+
+
+@dataclass(frozen=True)
+class EnvInspectionReport:
+    """Summary of a training environment inspection."""
+
+    name: str
+    backend: str
+    status: TrainabilityStatus
+    available: bool
+    loaded: bool
+    observation_size: Any = None
+    action_size: int | None = None
+    xml_path: str | None = None
+    model_bodies: int | None = None
+    model_actuators: int | None = None
+    has_domain_randomizer: bool | None = None
+    warnings: list[str] = field(default_factory=list)
+    errors: list[str] = field(default_factory=list)
+    notes: list[str] = field(default_factory=list)
+
+
+@dataclass(frozen=True)
+class RunConfig:
+    """Resolved training or evaluation run metadata."""
+
+    env_name: str
+    backend: str
+    preset: str
+    output_dir: str
+    config: dict[str, Any] = field(default_factory=dict)
+    command: list[str] = field(default_factory=list)
+
+
+@dataclass(frozen=True)
+class SmokeResult:
+    """Result of a short environment smoke test."""
+
+    env_name: str
+    backend: str
+    steps_requested: int
+    steps_completed: int
+    passed: bool
+    action_size: int | None = None
+    observation_size: Any = None
+    final_reward: float | None = None
+    final_done: bool | None = None
+    errors: list[str] = field(default_factory=list)
+
+
+def to_dict(value: Any) -> Any:
+    """Convert dataclasses, enums, paths, and containers into JSONable values."""
+    if isinstance(value, Enum):
+        return value.value
+    if isinstance(value, Path):
+        return str(value)
+    if is_dataclass(value):
+        return {key: to_dict(item) for key, item in asdict(value).items()}
+    if isinstance(value, dict):
+        return {str(key): to_dict(item) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [to_dict(item) for item in value]
+    if isinstance(value, (str, int, float, bool)) or value is None:
+        return value
+
+    # JAX/NumPy scalars and arrays.
+    item = getattr(value, "item", None)
+    if callable(item):
+        shape = getattr(value, "shape", None)
+        if shape == ():
+            return to_dict(item())
+    tolist = getattr(value, "tolist", None)
+    if callable(tolist):
+        return to_dict(tolist())
+
+    return value
+
+
+def report_markdown(title: str, report: Any) -> str:
+    """Render a compact Markdown report for humans and agents."""
+    data = to_dict(report)
+    lines = [f"# {title}", ""]
+    for key, value in data.items():
+        label = key.replace("_", " ").title()
+        if isinstance(value, list):
+            rendered = ", ".join(str(item) for item in value) if value else "none"
+        else:
+            rendered = str(value)
+        lines.append(f"- **{label}:** {rendered}")
+    lines.append("")
+    return "\n".join(lines)
+

simrig/custom_env.py

@@ -0,0 +1,109 @@
+"""Load user-authored custom environment modules."""
+
+from __future__ import annotations
+
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+from typing import Any
+
+
+def is_env_module_path(env_ref: str | Path) -> bool:
+    """Return True when env_ref points at a custom *.py env module."""
+    text = str(env_ref).strip()
+    if not text.endswith(".py"):
+        return False
+    return True
+
+
+def resolve_env_label(env_ref: str | Path) -> str:
+    """Stable label for reports/run dirs (basename without .py for modules)."""
+    if is_env_module_path(env_ref):
+        return Path(str(env_ref)).expanduser().stem
+    return str(env_ref)
+
+
+def import_env_module(path: Path | str) -> ModuleType:
+    """Import a custom env module from a filesystem path."""
+    env_path = Path(path).expanduser().resolve()
+    if not env_path.is_file():
+        raise FileNotFoundError(f"Custom env module not found: {env_path}")
+    if env_path.suffix != ".py":
+        raise ValueError(f"Custom env module must be a .py file: {env_path}")
+
+    module_name = f"simrig_custom_env_{env_path.stem}_{abs(hash(str(env_path)))}"
+    spec = importlib.util.spec_from_file_location(module_name, env_path)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Could not load custom env module: {env_path}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def load_custom_env(
+    path: Path | str,
+    *,
+    class_name: str = "CustomEnv",
+    config_overrides: dict[str, Any] | None = None,
+) -> Any:
+    """Instantiate a custom env from a module path.
+
+    Supports:
+    - make_env(config_overrides=...) factory if present
+    - CustomEnv(config=...) / CustomEnv(config=..., config_overrides=...)
+    """
+    module = import_env_module(path)
+    overrides = dict(config_overrides or {})
+
+    make_env = getattr(module, "make_env", None)
+    if callable(make_env):
+        try:
+            return make_env(config_overrides=overrides or None)
+        except TypeError:
+            return make_env()
+
+    cls = getattr(module, class_name, None)
+    if cls is None:
+        raise AttributeError(
+            f"Custom env module {path} must define make_env() or class {class_name}."
+        )
+
+    config = None
+    default_config = getattr(module, "default_config", None)
+    if callable(default_config):
+        config = default_config()
+
+    if isinstance(config, dict):
+        merged = dict(config)
+        merged.update(overrides)
+        return _construct_env(cls, config=merged, config_overrides=None)
+
+    if config is not None and overrides:
+        return _construct_env(cls, config=config, config_overrides=overrides)
+    if config is not None:
+        return _construct_env(cls, config=config, config_overrides=None)
+    if overrides:
+        return _construct_env(cls, config=overrides, config_overrides=None)
+    return _construct_env(cls, config=None, config_overrides=None)
+
+
+def _construct_env(cls: type, *, config: Any, config_overrides: dict[str, Any] | None) -> Any:
+    if config is None and config_overrides is None:
+        return cls()
+    if config_overrides is None:
+        try:
+            return cls(config=config)
+        except TypeError:
+            return cls(config)
+    try:
+        return cls(config=config, config_overrides=config_overrides)
+    except TypeError:
+        try:
+            return cls(config, config_overrides)
+        except TypeError as exc:
+            raise TypeError(
+                f"Could not construct {cls.__name__} with config/config_overrides. "
+                "Prefer make_env(config_overrides=...) in the module."
+            ) from exc