scripts/run_physics_identification.py

#!/usr/bin/env python
"""S9 driver: physics-parameter identification, one episode or a whole batch, pooled.

Thin CLI over the same ``fpgm.physics`` modules :meth:`fpgm.datagen.pipeline.
EpisodePipeline._run_s9` wires into the per-episode pipeline -- but this script's
job is the one thing a per-episode pipeline stage structurally cannot do:
**pool evidence across episodes**. Every :class:`~fpgm.physics.types.
EpisodeLogLik` from every episode of the same physical object (grouped by
``(scene_id, label)``, matching how :mod:`fpgm.physics.priors`' VLM cache is
already keyed and why -- see that module's own docstring) is summed into one
posterior via ``fpgm.physics.inference.accumulate``/``pool``.

**Non-negotiable design constraint: no episode is ever filtered.** Every
episode with an S6 pose track (``poses.npz``) is scored and its per-particle
log-likelihood is folded into its object's pool -- there is no ``--min-motion``
flag, no skip-if-static branch, and none should ever be added. The reason this
is safe, not merely convenient, is arithmetic, not policy: an object that did
not move produces (up to floating point) the *same* predicted pose for every
particle a physically-faithful simulator would draw (see
``fpgm.physics.types``'s own module docstring and
``scripts/_mujoco_settle_worker.py``'s proof that a rigid body's fall-and-settle
trajectory is mass-independent), so ``EpisodeLogLik.loglik`` for that episode is
a near-constant vector. Adding a constant to every particle's log-weight before
a softmax leaves the softmax exactly unchanged -- so an uninformative episode
contributes nothing, automatically, with no threshold anywhere in this file
deciding which episodes get to vote. ``tests/test_run_physics_identification.py``
asserts this end to end: a pool that includes a fully static object's episode
produces a bit-identical posterior to the same pool without it.

That does not mean every episode is *interesting* -- it means every episode is
*included*. ``spread_nats`` (``EpisodeLogLik.spread_nats``, printed per episode
below) is the diagnostic for "was this one interesting": near zero means the
episode said nothing about theta. It is printed as a report for a human to
read, never consulted by this script to decide what to score.

Usage::

    # One episode:
    PYTHONPATH=src python scripts/run_physics_identification.py \\
        --uuid AUTOLab+0d4edc83+2023-10-21-19h-07m-04s --camera ext1

    # Every episode with S6 output under a batch's output_root, pooled:
    PYTHONPATH=src python scripts/run_physics_identification.py \\
        --from-batch outputs/datagen --particles 1024 --seed 0 --out outputs/physics
"""

from __future__ import annotations

import argparse
import json
import sys
from collections.abc import Sequence
from contextlib import nullcontext
from pathlib import Path
from typing import Any

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "src"))

from fpgm.datagen.episode_spec import EpisodeMetadata  # noqa: E402
from fpgm.utils.io import ensure_dir  # noqa: E402
from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402

logger = get_logger("run_physics_identification")

_DEFAULT_CONFIG = REPO_ROOT / "configs" / "datagen_droid.yaml"
_DEFAULT_PARTICLES = 512
_DEFAULT_SEED = 0
_DEFAULT_OUT_ROOT = REPO_ROOT / "outputs" / "physics"

__all__ = ["EpisodeRef", "discover_episodes_from_batch", "run_physics_identification"]


# --------------------------------------------------------------------------- #
# timer=None no-op, matching fpgm.physics.likelihood/inference's own helper
# --------------------------------------------------------------------------- #


def _step(timer: Any, label: str, **fields: Any):
    """``timer.step(label, **fields)`` if ``timer`` is given, else a clean no-op.

    Same pattern as ``fpgm.physics.likelihood._step``/``inference._step`` --
    duplicated rather than imported so this script has no import-time
    dependency on those modules (they are imported lazily inside
    :func:`run_physics_identification`, see that function's docstring).
    """
    if timer is None:
        return nullcontext()
    return timer.step(label, **fields)


# --------------------------------------------------------------------------- #
# Episode discovery
# --------------------------------------------------------------------------- #


class EpisodeRef:
    """One episode this run will score: ``(uuid, camera_serial, master_dir)``."""

    __slots__ = ("uuid", "camera_serial", "master_dir")

    def __init__(self, uuid: str, camera_serial: str, master_dir: Path) -> None:
        self.uuid = uuid
        self.camera_serial = camera_serial
        self.master_dir = Path(master_dir)

    def __repr__(self) -> str:  # pragma: no cover - debug convenience only
        return f"EpisodeRef({self.uuid!r}, {self.camera_serial!r}, {self.master_dir})"


def discover_episodes_from_batch(batch_root: Path) -> list[EpisodeRef]:
    """Every ``(uuid, camera_serial)`` under ``batch_root`` with an S6 pose track.

    "Has ``poses.npz``" is an *existence* check (did S6 run at all for this
    episode/camera), not an informativeness one -- see the module docstring's
    design constraint. An episode/camera pair with no ``poses.npz`` was never
    run through S6 in the first place; there is nothing here to score, which
    is a different statement from "this episode carried no information" (that
    one IS scored, and nets out to a no-op on its own -- see the module
    docstring).
    """
    batch_root = Path(batch_root)
    if not batch_root.is_dir():
        raise FileNotFoundError(f"--from-batch root does not exist or is not a dir: {batch_root}")
    refs: list[EpisodeRef] = []
    for uuid_dir in sorted(p for p in batch_root.iterdir() if p.is_dir()):
        for cam_dir in sorted(p for p in uuid_dir.iterdir() if p.is_dir()):
            master_dir = cam_dir / "master"
            if (master_dir / "poses.npz").exists():
                refs.append(EpisodeRef(uuid_dir.name, cam_dir.name, master_dir))
    return refs


def _scene_id_for(ref: EpisodeRef) -> str:
    """Read ``scene_id`` from the episode's own ``episode_spec.json``.

    Written unconditionally by every :meth:`~fpgm.datagen.pipeline.
    EpisodePipeline.run` call (see that class's own ``run`` method), so any
    episode with S6 output also has this file -- except a partial/corrupted
    ``master/`` directory, which is a real data problem. Rather than crash
    the whole batch pool over one episode's missing spec file, this episode
    becomes its own single-member "scene" (grouped under its own uuid) --
    still scored, still pooled (of one), just not merged with sibling
    episodes of what may be the same physical object. Recorded via a log
    line, not silently absorbed.
    """
    spec_path = ref.master_dir.parent / "episode_spec.json"
    if not spec_path.exists():
        # master_dir is .../<uuid>/<camera_serial>/master -- master_dir.parent
        # is .../<uuid>/<camera_serial>, and episode_spec.json is written
        # at .../<uuid>/<camera_serial>/episode_spec.json by EpisodePipeline.run.
        logger.warning(
            "%s/%s: no episode_spec.json at %s -- cannot group by scene_id; treating this "
            "episode as its own scene (pooled alone)",
            ref.uuid, ref.camera_serial, spec_path,
        )
        return f"uuid:{ref.uuid}"
    try:
        return str(json.loads(spec_path.read_text())["scene_id"])
    except (json.JSONDecodeError, KeyError, OSError) as exc:
        logger.warning(
            "%s/%s: could not read scene_id from %s (%s); treating this episode as its own "
            "scene (pooled alone)",
            ref.uuid, ref.camera_serial, spec_path, exc,
        )
        return f"uuid:{ref.uuid}"


# --------------------------------------------------------------------------- #
# Core, testable logic
# --------------------------------------------------------------------------- #


def run_physics_identification(
    episodes: Sequence[EpisodeRef],
    *,
    n_particles: int = _DEFAULT_PARTICLES,
    seed: int = _DEFAULT_SEED,
    out_root: Path,
    timer: Any = None,
) -> dict[str, Any]:
    """Score every episode in ``episodes``, pool by ``(scene_id, label)``, write reports.

    Imports every ``fpgm.physics.*`` module lazily (not at this file's module
    scope), for the same reason
    :meth:`~fpgm.datagen.pipeline.EpisodePipeline._run_s9` does -- see that
    method's own docstring: these modules were written by other agents in
    parallel with this script, and a module-scope import would make the
    entire CLI fail to even parse ``--help`` while any one of them was
    mid-edit.

    Returns:
        A JSON-able summary: ``{"episodes": [...], "groups": {...}}`` --
        also the exact payload written to ``physics_posterior.json``'s
        sibling summary file, and what a caller (or a test) inspects instead
        of re-reading every JSON file this function wrote.
    """
    import numpy as np

    from fpgm.physics.types import RIGID_PARAMS, EpisodeLogLik, ParamSpace, PhysicsError

    scene = _physics_import("scene")
    materials = _physics_import("materials")
    priors_mod = _physics_import("priors")
    simulate_mod = _physics_import("simulate")
    likelihood_mod = _physics_import("likelihood")
    inference_mod = _physics_import("inference")
    report_mod = _physics_import("report")

    out_root = ensure_dir(Path(out_root))
    rng = np.random.default_rng(seed)

    # One entry per (scene_id, label): the prior/particles drawn ONCE (see the
    # module docstring's "draw particles once" property, inherited from
    # fpgm.physics.types) and reused for every episode of that group; per-episode
    # PosteriorResults accumulate()-ed against that SAME particle array so
    # inference.pool()'s bit-identity check on `particles` never trips.
    groups: dict[tuple[str, str], dict[str, Any]] = {}
    episode_records: list[dict[str, Any]] = []

    for ref in episodes:
        scene_id = _scene_id_for(ref)
        with _step(timer, "load_poses", uuid=ref.uuid):
            tracks, unusable = scene.load_observed_tracks(
                ref.uuid, ref.camera_serial, timer=timer,
            )
        for bad_label, reason in unusable.items():
            print(f"[{ref.uuid}/{bad_label}] no usable measured track: {reason}")
        if not tracks:
            logger.info("%s/%s: no ObservedTrack -- nothing to score", ref.uuid, ref.camera_serial)
            continue

        for label, track in tracks.items():
            with _step(timer, "build_spec", uuid=ref.uuid, obj_label=label):
                sim_spec, spec_diag = scene.build_sim_spec(
                    ref.uuid, ref.camera_serial, label,
                    scratch_dir=ref.master_dir / "physics" / "scratch", timer=timer,
                )
            for note in spec_diag.get("limitations", []):
                print(f"[{ref.uuid}/{label}] {note}")
            body = sim_spec.bodies[0]
            key = (scene_id, label)
            group = groups.get(key)
            if group is None:
                space = ParamSpace(RIGID_PARAMS)
                if body.kind == "prismatic":
                    from fpgm.physics.types import PRISMATIC_PARAMS
                    space = space.extended(PRISMATIC_PARAMS)
                try:
                    with _step(timer, "vlm_prior", uuid=ref.uuid, obj_label=label):
                        crop_path = scene.object_crop(ref.uuid, ref.camera_serial, label)
                        verdict = priors_mod.VlmPriorProposer(scene_id=scene_id).propose(
                            {label: crop_path}, timer=timer
                        )[label]
                except PhysicsError as exc:
                    verdict = priors_mod.VlmPriorProposer.fallback_verdict(label)
                    print(
                        f"[{ref.uuid}/{label}] VLM material prior unavailable ({exc}); "
                        "using VlmPriorProposer.fallback_verdict"
                    )
                prior = materials.material_prior(verdict, space)
                particles = prior.sample(rng, n_particles)
                group = groups[key] = {
                    "space": space, "prior": prior, "particles": particles,
                    "verdict": verdict, "results": [], "episode_refs": [],
                }

            with _step(timer, "simulate", uuid=ref.uuid, obj_label=label, n=n_particles):
                simulator = simulate_mod.MujocoSimulator(
                    scratch_dir=ref.master_dir / "physics" / "scratch"
                )
                sim_result = simulator.simulate_batched(
                    sim_spec, group["particles"], group["space"], timer=timer
                )
            with _step(timer, "score", uuid=ref.uuid, obj_label=label, n=n_particles):
                loglik = likelihood_mod.log_likelihood(
                    track, sim_result.poses, sim_result.ok, timer=timer,
                )
            episode_ll = EpisodeLogLik(
                uuid=ref.uuid, camera_serial=ref.camera_serial, label=label,
                loglik=loglik, n_obs_frames=track.n_valid,
                n_diverged=int(np.sum(~np.asarray(sim_result.ok, dtype=bool))),
            )
            print(
                f"{ref.uuid}/{ref.camera_serial} {label}: spread_nats={episode_ll.spread_nats:.4f} "
                f"n_obs_frames={episode_ll.n_obs_frames} n_diverged={episode_ll.n_diverged}"
            )

            with _step(timer, "accumulate", uuid=ref.uuid, obj_label=label, n=n_particles):
                per_episode_posterior = inference_mod.accumulate(
                    group["prior"], group["particles"], [episode_ll], timer=timer
                )
            episode_out = ensure_dir(out_root / ref.uuid / ref.camera_serial / label)
            with _step(timer, "write_report", uuid=ref.uuid, obj_label=label):
                episode_report_path = report_mod.write_object_report(
                    per_episode_posterior, episode_out,
                    extra={"uuid": ref.uuid, "camera_serial": ref.camera_serial, "label": label,
                           "scene_id": scene_id},
                )
            group["results"].append(per_episode_posterior)
            group["episode_refs"].append((ref.uuid, ref.camera_serial))
            episode_records.append({
                "uuid": ref.uuid, "camera_serial": ref.camera_serial, "scene_id": scene_id,
                "label": label, "spread_nats": episode_ll.spread_nats,
                "n_obs_frames": episode_ll.n_obs_frames, "n_diverged": episode_ll.n_diverged,
                "physics_json": str(episode_report_path),
            })

    group_summaries: dict[str, Any] = {}
    for (scene_id, label), group in groups.items():
        with _step(timer, "accumulate", n=len(group["results"])):
            pooled = inference_mod.pool(group["results"], timer=timer)
        pooled_out = ensure_dir(out_root / f"{scene_id}__{label}")
        with _step(timer, "write_report"):
            pooled_path = report_mod.write_pooled_report(
                pooled, pooled_out,
                extra={"scene_id": scene_id, "label": label,
                       "n_episodes": len(group["results"]),
                       "episode_ids": [f"{u}/{c}" for u, c in group["episode_refs"]]},
            )
        group_summaries[f"{scene_id}/{label}"] = {
            "n_episodes": len(group["results"]),
            "ess": pooled.ess,
            "info_nats": pooled.info_nats,
            "physics_posterior_json": str(pooled_path),
        }

    summary = {"episodes": episode_records, "groups": group_summaries}
    (out_root / "physics_posterior.json").write_text(json.dumps(summary, indent=2))
    return summary


def _physics_import(module_name: str):
    """Import ``fpgm.physics.<module_name>``, or raise a clear, actionable error.

    Identical rationale and behaviour to
    :meth:`fpgm.datagen.pipeline.EpisodePipeline._physics_import` -- kept as
    its own copy here rather than imported from ``fpgm.datagen.pipeline``
    because this script has no other reason to depend on the datagen
    pipeline package at all, and one three-line helper is cheaper to
    duplicate than to couple two otherwise-independent entry points.
    """
    import importlib

    from fpgm.physics.types import PhysicsError

    full_name = f"fpgm.physics.{module_name}"
    try:
        return importlib.import_module(full_name)
    except ImportError as exc:
        raise PhysicsError(
            f"run_physics_identification requires {full_name} "
            f"(src/fpgm/physics/{module_name}.py), which is not importable: {exc}."
        ) from exc


# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    p.add_argument("--uuid", default=None, help="DROID episode uuid (single-episode mode)")
    p.add_argument(
        "--camera", default=None,
        help="camera serial or role (ext1/ext2); required with --uuid",
    )
    p.add_argument(
        "--from-batch", type=Path, default=None,
        help="sweep every episode with S6 output under this batch output_root "
             "(e.g. outputs/datagen); mutually exclusive with --uuid/--camera",
    )
    p.add_argument(
        "--particles", type=int, default=_DEFAULT_PARTICLES,
        help=f"particles drawn per (scene_id, label) group (default {_DEFAULT_PARTICLES})",
    )
    p.add_argument("--seed", type=int, default=_DEFAULT_SEED, help="particle-draw RNG seed")
    p.add_argument("--out", type=Path, default=_DEFAULT_OUT_ROOT, help="output root")
    p.add_argument("--config", type=Path, default=_DEFAULT_CONFIG, help="DatagenProfile YAML")
    p.add_argument("--log-level", default="INFO")
    args = p.parse_args()
    if args.from_batch is None and not (args.uuid and args.camera):
        p.error("either --from-batch, or both --uuid and --camera, must be given")
    if args.from_batch is not None and (args.uuid or args.camera):
        p.error("--from-batch is mutually exclusive with --uuid/--camera")
    return args


def _single_episode_ref(args: argparse.Namespace) -> EpisodeRef:
    from fpgm.config_datagen import DatagenProfile

    profile = DatagenProfile.from_yaml(args.config)
    metadata = EpisodeMetadata.from_flows_h5(profile.paths.flows_h5(args.uuid))
    camera_serial = metadata.camera_serial(args.camera)
    master_dir = profile.paths.master_dir(args.uuid, camera_serial)
    return EpisodeRef(args.uuid, camera_serial, master_dir)


def main() -> int:
    args = parse_args()
    setup_logging(args.log_level)

    from fpgm.utils.timing import StepTimer

    timer = StepTimer("run_physics_identification")

    if args.from_batch is not None:
        episodes = discover_episodes_from_batch(args.from_batch)
        logger.info("found %d episode(s) with S6 output under %s", len(episodes), args.from_batch)
    else:
        episodes = [_single_episode_ref(args)]

    if not episodes:
        print("no episodes to score (0 found) -- nothing written")
        print(timer.report())
        return 0

    summary = run_physics_identification(
        episodes, n_particles=args.particles, seed=args.seed, out_root=args.out, timer=timer,
    )
    out_path = args.out / "physics_posterior.json"
    print(
        f"\nscored {len(summary['episodes'])} episode-object pair(s) across "
        f"{len(summary['groups'])} (scene_id, label) group(s); wrote {out_path}"
    )
    for key, g in sorted(summary["groups"].items()):
        print(
            f"  {key}: n_episodes={g['n_episodes']} ess={g['ess']:.1f} "
            f"info_nats={g['info_nats']:.3f}"
        )
    print()
    print(timer.report())
    return 0


if __name__ == "__main__":
    raise SystemExit(main())