src/fpgm/physics/simulate.py

"""S9's ``fpgm``-side entry point into the isolated MuJoCo worker: JSON/npz out, npz back.

Mirrors ``scripts/settle_after_release.py``'s own subprocess plumbing (read that file's
"Why MuJoCo lives in its own conda env, not fpgm" section first -- the isolation rationale
is identical and not repeated here) with one difference: that script drives exactly one
MJCF per subprocess call, while :class:`MujocoSimulator` drives a whole *particle batch*
through ``fpgm.physics.worker_mujoco`` (see that module's docstring for why batching one
compiled model over many particles, rather than one subprocess per particle, is the entire
point -- compiling an ``MjModel`` costs milliseconds and spawning a Python subprocess costs
tens of milliseconds more on top, so one-subprocess-per-particle would make subprocess
overhead dominate wall-clock at exactly the particle counts (thousands) S9's importance
sampling needs).

**Why raising on a non-zero exit code matters here specifically.** A worker crash and "every
particle diverged" look identical if a caller only checks ``result.ok``: both are an all-False
array. But they are completely different physical statements -- one says "the simulator could
not even run" (a bug, a bad mesh path, a malformed spec), the other says "every hypothesis in
this batch is inconsistent with basic physics" (a real, if surprising, finding). Swallowing a
crash into an empty/all-``False`` :class:`~fpgm.physics.types.SimResult` would silently turn
the first into the second, which is exactly backwards for a stage whose whole point is trusting
what its own likelihood evaluation says. :func:`MujocoSimulator.run` therefore never
constructs a result from a failed subprocess; it raises :class:`~fpgm.physics.types.PhysicsError`
with the worker's own stderr tail attached, so the failure is legible without re-running
anything by hand.
"""

from __future__ import annotations

import contextlib
import json
import os
import subprocess
import tempfile
import time
from pathlib import Path

import numpy as np

from fpgm.physics.types import ParamSpace, PhysicsError, SimResult, SimSpec
from fpgm.utils.logging import get_logger
from fpgm.utils.timing import StepTimer

logger = get_logger(__name__)

__all__ = ["MujocoSimulator", "DEFAULT_MUJOCO_PYTHON"]

REPO_ROOT = Path(__file__).resolve().parents[3]
_WORKER_MODULE = "fpgm.physics.worker_mujoco"

#: Overridable via the ``FPGM_MUJOCO_PYTHON`` env var (see module docstring's task
#: requirement) -- the same override pattern ``scripts/settle_after_release.py`` uses
#: for its own ``--mujoco-python`` CLI flag, just as an env var here since
#: :class:`MujocoSimulator` is a library entry point, not itself a CLI.
DEFAULT_MUJOCO_PYTHON = Path(
    os.environ.get("FPGM_MUJOCO_PYTHON", "/home/quang/miniconda3/envs/mujoco/bin/python")
)

#: Particles per worker subprocess call in :meth:`MujocoSimulator.simulate_batched`.
#: Not tuned against a measured wall-clock cliff -- chosen so a single `particles.npz`
#: plus the returned `(N, T, 7)` poses array stays comfortably in the tens-of-MB range
#: even for a long episode (hundreds of frames) with a generous particle count, and so
#: one crashed/diverged chunk does not force re-running an entire multi-thousand-particle
#: batch. Override via `MujocoSimulator(max_particles_per_chunk=...)` if a specific
#: episode's frame count needs a different balance.
DEFAULT_MAX_PARTICLES_PER_CHUNK = 512

#: How much of the worker's stderr to keep in a PhysicsError -- enough for a full
#: Python traceback, not so much that a log-spamming failure buries the actual cause.
_STDERR_TAIL_CHARS = 4000


class MujocoSimulator:
    """Runs :mod:`fpgm.physics.worker_mujoco` in the ``mujoco`` conda env as a subprocess.

    Args:
        mujoco_python: Path to the ``mujoco`` env's interpreter. Defaults to
            :data:`DEFAULT_MUJOCO_PYTHON` (itself overridable via ``FPGM_MUJOCO_PYTHON``).
        max_particles_per_chunk: See :data:`DEFAULT_MAX_PARTICLES_PER_CHUNK`; only
            consulted by :meth:`simulate_batched`, not :meth:`run`.
        scratch_dir: Where per-call ``sim_spec.json``/``particles.npz``/``result.npz``
            are written. Defaults to a fresh :class:`tempfile.TemporaryDirectory` per
            call (cleaned up immediately after) -- pass an explicit, persistent
            directory to keep the wire files around for debugging a failure by hand.
    """

    def __init__(
        self,
        *,
        mujoco_python: Path | None = None,
        max_particles_per_chunk: int = DEFAULT_MAX_PARTICLES_PER_CHUNK,
        scratch_dir: Path | None = None,
    ) -> None:
        self.mujoco_python = (
            Path(mujoco_python) if mujoco_python is not None else DEFAULT_MUJOCO_PYTHON
        )
        self.max_particles_per_chunk = int(max_particles_per_chunk)
        self._scratch_dir = scratch_dir

    # -- single subprocess call -------------------------------------------------- #

    def run(
        self,
        spec: SimSpec,
        particles: np.ndarray,
        space: ParamSpace,
        *,
        timer: StepTimer | None = None,
    ) -> SimResult:
        """One worker subprocess call over the whole ``particles`` array.

        Args:
            spec: The scene to simulate -- see :func:`fpgm.physics.scene.build_sim_spec`.
            particles: ``(N, D)`` particle array, ``D == space.dim``, columns ordered
                per ``space.names`` -- this ordering is what the worker's
                ``param_names``/``theta`` npz keys carry across the env boundary (see
                ``fpgm.physics.worker_mujoco``'s module docstring on why that contract
                is duplicated rather than imported).
            space: Names ``particles``'s columns are indexed by.
            timer: See module-level ``Timing`` note; ``None`` is a no-op.

        Returns:
            A :class:`~fpgm.physics.types.SimResult` for ``spec.bodies[0]`` (the
            worker's documented "body of interest" convention).

        Raises:
            PhysicsError: If ``particles``'s shape disagrees with ``space``, or the
                worker subprocess exits non-zero (its stderr tail is included).
        """
        particles = np.asarray(particles, dtype=np.float64)
        if particles.ndim != 2 or particles.shape[1] != space.dim:
            raise PhysicsError(
                f"particles shape {particles.shape} disagrees with ParamSpace{space.names} "
                f"(dim={space.dim})"
            )
        n_particles = particles.shape[0]

        step_cm = (
            timer.step("simulate", n=n_particles, uuid=spec.uuid, body_label=spec.bodies[0].label)
            if timer is not None
            else contextlib.nullcontext()
        )
        with step_cm:
            result = self._run_once(spec, particles, space)

        if timer is not None:
            self._mark_worker_timings(timer, result, n_particles)
        return result

    def _run_once(self, spec: SimSpec, particles: np.ndarray, space: ParamSpace) -> SimResult:
        with self._scratch(spec.uuid) as scratch:
            spec_path = scratch / "sim_spec.json"
            particles_path = scratch / "particles.npz"
            out_path = scratch / "result.npz"

            spec_path.write_text(json.dumps(spec.to_json_dict()))
            np.savez(particles_path, theta=particles, param_names=np.array(space.names))

            cmd = [
                str(self.mujoco_python), "-m", _WORKER_MODULE,
                "--spec", str(spec_path), "--particles", str(particles_path),
                "--out", str(out_path),
            ]
            env = os.environ.copy()
            src_dir = str(REPO_ROOT / "src")
            existing_pp = env.get("PYTHONPATH")
            env["PYTHONPATH"] = (
                src_dir if not existing_pp else f"{src_dir}{os.pathsep}{existing_pp}"
            )

            logger.info(
                "invoking mujoco worker (%d particles): %s", particles.shape[0], " ".join(cmd)
            )
            t0 = time.perf_counter()
            try:
                proc = subprocess.run(
                    cmd, capture_output=True, text=True, cwd=str(REPO_ROOT), env=env
                )
            except OSError as exc:
                raise PhysicsError(
                    f"could not launch mujoco worker with interpreter {self.mujoco_python} "
                    f"(uuid={spec.uuid!r}, label={spec.bodies[0].label!r}): {exc}. Check "
                    f"FPGM_MUJOCO_PYTHON / MujocoSimulator(mujoco_python=...)."
                ) from exc
            wall_seconds = time.perf_counter() - t0
            logger.info(
                "mujoco worker finished in %.2fs wall, returncode=%d", wall_seconds, proc.returncode
            )
            if proc.stdout:
                logger.info("[mujoco worker stdout] %s", proc.stdout.strip())

            if proc.returncode != 0:
                tail = (proc.stderr or "")[-_STDERR_TAIL_CHARS:]
                raise PhysicsError(
                    f"mujoco worker failed (uuid={spec.uuid!r}, label={spec.bodies[0].label!r}, "
                    f"{particles.shape[0]} particles, returncode={proc.returncode}). "
                    f"An empty/all-diverged result here would be a silent lie -- see this "
                    f"module's docstring. Worker stderr tail:\n{tail}"
                )
            if not out_path.exists():
                raise PhysicsError(
                    f"mujoco worker exited 0 but wrote no output at {out_path} "
                    f"(uuid={spec.uuid!r}, label={spec.bodies[0].label!r})"
                )

            with np.load(out_path) as npz:
                poses = np.asarray(npz["poses"])
                ok = np.asarray(npz["ok"])
                sim_seconds = float(npz["sim_seconds"])
                n_substeps = int(npz["n_substeps"])
                label = str(npz["label"])
                build_seconds = float(npz["build_seconds"])
                io_seconds = float(npz["io_seconds"])

        result = SimResult(
            label=label, poses=poses, ok=ok, sim_seconds=sim_seconds, n_substeps=n_substeps
        )
        # Stash the extra child-side timing breakdown for _mark_worker_timings; not
        # part of SimResult's own wire contract (physics/types.py is out of scope for
        # this agent), so it rides along as a plain attribute rather than a field.
        result._worker_build_seconds = build_seconds  # type: ignore[attr-defined]
        result._worker_io_seconds = io_seconds  # type: ignore[attr-defined]
        result._worker_wall_seconds = wall_seconds  # type: ignore[attr-defined]
        return result

    @staticmethod
    def _mark_worker_timings(timer: StepTimer, result: SimResult, n_particles: int) -> None:
        """Fold the worker's self-reported breakdown into the parent's own table.

        Both views end up in one `timer.summary()`: `simulate` (this class's own
        `timer.step`, wall-clock including subprocess spawn + JSON/npz I/O) and these
        `mark`ed rows (the worker's own perf_counter measurements of just its build/
        step/io phases) -- exactly the "child stepping time and parent wall-clock" pair
        the Timing requirement asks for.
        """
        timer.mark("worker_build", result._worker_build_seconds, n=1)  # type: ignore[attr-defined]
        timer.mark("worker_step", result.sim_seconds, n=n_particles)
        timer.mark("worker_io", result._worker_io_seconds, n=1)  # type: ignore[attr-defined]

    # -- chunked batches ----------------------------------------------------------- #

    def simulate_batched(
        self,
        spec: SimSpec,
        particles: np.ndarray,
        space: ParamSpace,
        *,
        timer: StepTimer | None = None,
    ) -> SimResult:
        """Like :meth:`run`, but splits very large ``particles`` into several worker calls.

        See :data:`DEFAULT_MAX_PARTICLES_PER_CHUNK` for why a cap exists at all. Chunks
        are simulated sequentially (not in parallel -- each already saturates the single
        MuJoCo worker process; running several at once would just contend for CPU) and
        their results concatenated in order, so the returned particle axis matches the
        input's row order exactly.
        """
        particles = np.asarray(particles, dtype=np.float64)
        n_particles = particles.shape[0]
        chunk = self.max_particles_per_chunk
        if n_particles <= chunk:
            return self.run(spec, particles, space, timer=timer)

        n_chunks = -(-n_particles // chunk)  # ceil div
        logger.info(
            "simulate_batched: %d particles > max_particles_per_chunk=%d -> %d chunks",
            n_particles, chunk, n_chunks,
        )
        poses_parts: list[np.ndarray] = []
        ok_parts: list[np.ndarray] = []
        sim_seconds_total = 0.0
        label = None
        n_substeps = None
        for c in range(n_chunks):
            lo, hi = c * chunk, min((c + 1) * chunk, n_particles)
            logger.info(
                "simulate_batched: chunk %d/%d (particles [%d, %d))", c + 1, n_chunks, lo, hi
            )
            part = self.run(spec, particles[lo:hi], space, timer=timer)
            poses_parts.append(part.poses)
            ok_parts.append(part.ok)
            sim_seconds_total += part.sim_seconds
            label = part.label
            n_substeps = part.n_substeps

        return SimResult(
            label=label,
            poses=np.concatenate(poses_parts, axis=0),
            ok=np.concatenate(ok_parts, axis=0),
            sim_seconds=sim_seconds_total,
            n_substeps=n_substeps,
        )

    # -- scratch dir --------------------------------------------------------------- #

    @contextlib.contextmanager
    def _scratch(self, uuid: str):
        if self._scratch_dir is not None:
            out = self._scratch_dir / uuid
            out.mkdir(parents=True, exist_ok=True)
            yield out
        else:
            with tempfile.TemporaryDirectory(prefix=f"fpgm_s9_{uuid}_") as td:
                yield Path(td)