creature-lab 0.3.0__py3-none-any.whl

creature_lab/autopsy.py

@@ -0,0 +1,376 @@
+"""Failure-first experiment autopsy, independent of any physics backend."""
+
+from __future__ import annotations
+
+import html
+import json
+from collections.abc import Callable
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+
+from creature_lab.diagnosis import diagnose
+from creature_lab.diagnostics import summarize_episode
+from creature_lab.qualification import (
+    BUILTIN_PROFILES,
+    QualificationProfile,
+    task_setup_check,
+    task_success,
+)
+from creature_lab.robustness import RobustnessResult, run_trials
+from creature_lab.schema import CreatureSpec, EpisodeTrace, TaskSpec
+
+SimulateFn = Callable[[CreatureSpec, TaskSpec], EpisodeTrace]
+
+
+@dataclass(frozen=True)
+class AutopsyEvidence:
+    name: str
+    passed: bool | None
+    detail: str
+
+
+@dataclass(frozen=True)
+class TrialRecording:
+    creature: CreatureSpec
+    task: TaskSpec
+    trace: EpisodeTrace
+
+
+@dataclass(frozen=True)
+class AutopsyResult:
+    creature: str
+    task: str
+    profile: str
+    primary_cause: str
+    confidence: str
+    summary: str
+    evidence: list[AutopsyEvidence]
+    recommendations: list[str]
+    baseline_trace: EpisodeTrace = field(repr=False)
+    reference_trace: EpisodeTrace | None = field(repr=False)
+    other_backend_trace: EpisodeTrace | None = field(default=None, repr=False)
+    robustness: RobustnessResult | None = field(default=None, repr=False)
+    trial_recordings: list[TrialRecording] = field(default_factory=list, repr=False)
+
+    def to_json_dict(self) -> dict[str, object]:
+        robustness = None
+        if self.robustness is not None:
+            robustness = {
+                "trials": len(self.robustness.trials),
+                "mean_score": self.robustness.mean_score,
+                "std_score": self.robustness.std_score,
+                "fail_rate": self.robustness.fail_rate,
+                "trial_details": [asdict(trial) for trial in self.robustness.trials],
+            }
+        return {
+            "creature": self.creature,
+            "task": self.task,
+            "profile": self.profile,
+            "primary_cause": self.primary_cause,
+            "confidence": self.confidence,
+            "confidence_method": "heuristic; not a calibrated probability",
+            "summary": self.summary,
+            "evidence": [evidence.__dict__ for evidence in self.evidence],
+            "recommendations": self.recommendations,
+            "baseline": {
+                "run_id": self.baseline_trace.run_id,
+                "backend": self.baseline_trace.backend,
+                "score": self.baseline_trace.score,
+            },
+            "reference": (
+                {
+                    "run_id": self.reference_trace.run_id,
+                    "backend": self.reference_trace.backend,
+                    "score": self.reference_trace.score,
+                }
+                if self.reference_trace is not None
+                else None
+            ),
+            "other_backend": (
+                {
+                    "run_id": self.other_backend_trace.run_id,
+                    "backend": self.other_backend_trace.backend,
+                    "score": self.other_backend_trace.score,
+                }
+                if self.other_backend_trace is not None
+                else None
+            ),
+            "robustness": robustness,
+            "evidence_packs": {
+                "reference": "evidence/reference" if self.reference_trace is not None else None,
+                "other_backend": (
+                    "evidence/other_backend" if self.other_backend_trace is not None else None
+                ),
+                "trials": [
+                    {"path": f"evidence/trial_{index + 1:03d}", "run_id": item.trace.run_id}
+                    for index, item in enumerate(self.trial_recordings)
+                ],
+            },
+        }
+
+
+def infer_profile(task: TaskSpec) -> QualificationProfile:
+    if task.target is not None:
+        return BUILTIN_PROFILES["target-reach"]
+    if task.damage_event is not None or task.impulse_event is not None:
+        return BUILTIN_PROFILES["push-recovery"]
+    return BUILTIN_PROFILES["basic-locomotion"]
+
+
+def distinct_curated_reference(creature: CreatureSpec, controller: str) -> bool:
+    """Only claim a counterfactual when exact controller snapshots differ."""
+    from creature_lab.controllers.factory import curated_controller
+    from creature_lab.exporting import file_hash, resolve_controller_bundle
+
+    reference = curated_controller(creature)
+    selected = reference if controller == "curated" else controller
+    try:
+        a = resolve_controller_bundle(creature, selected)
+        b = resolve_controller_bundle(creature, reference)
+        if not a.exact or not b.exact:
+            return False
+
+        def behavior(bundle):
+            data = bundle.spec.model_dump(mode="json", exclude_none=True)
+            for key in ("name", "creature_hash", "task_hash", "runtime_versions"):
+                data.pop(key, None)
+            if bundle.policy_source is not None:
+                data["policy_file"] = file_hash(bundle.policy_source)
+            return data
+
+        return behavior(a) != behavior(b)
+    except (OSError, ValueError):
+        return False
+
+
+def run_autopsy(
+    creature: CreatureSpec,
+    task: TaskSpec,
+    *,
+    simulate: SimulateFn,
+    simulate_reference: SimulateFn | None,
+    profile: QualificationProfile | None = None,
+    simulate_other_backend: SimulateFn | None = None,
+    robustness_trials: int = 5,
+) -> AutopsyResult:
+    """Run baseline/counterfactual/perturbation/backend evidence and attribute failure."""
+    profile = profile or infer_profile(task)
+    setup = task_setup_check(profile, task)
+    if setup is not None and not setup.passed:
+        raise ValueError(f"Task/profile mismatch: {setup.detail}")
+    if profile.max_sim2sim_gap is not None and simulate_other_backend is None:
+        raise ValueError("This profile requires a second physics backend.")
+    baseline = simulate(creature, task)
+    baseline_summary = summarize_episode(baseline, task, creature)
+    baseline_ok = task_success(profile, baseline, baseline_summary)
+
+    reference = simulate_reference(creature, task) if simulate_reference is not None else None
+    reference_summary = summarize_episode(reference, task, creature) if reference else None
+    reference_ok = task_success(profile, reference, reference_summary) if reference else None
+    trial_recordings: list[TrialRecording] = []
+
+    def evaluate(c: CreatureSpec, t: TaskSpec) -> tuple[float, bool, bool]:
+        trace = simulate(c, t)
+        trial_recordings.append(TrialRecording(c, t, trace))
+        summary = summarize_episode(trace, t, c)
+        return trace.score, bool(summary.fell), task_success(profile, trace, summary)
+
+    robustness = run_trials(
+        creature,
+        task,
+        evaluate,
+        trials=robustness_trials,
+        mass_jitter=profile.robustness_mass_jitter,
+        friction_jitter=profile.robustness_friction_jitter,
+    )
+    robust_ok = robustness.fail_rate <= profile.max_fail_rate
+
+    other = simulate_other_backend(creature, task) if simulate_other_backend is not None else None
+    portability_ok = True
+    normalized_gap = None
+    if other is not None:
+        normalized_gap = abs(baseline.score - other.score) / max(
+            1.0, abs(baseline.score), abs(other.score)
+        )
+        threshold = profile.max_sim2sim_gap if profile.max_sim2sim_gap is not None else 0.2
+        portability_ok = normalized_gap <= threshold
+
+    diagnosis = diagnose(baseline, creature, task)
+    evidence = [
+        AutopsyEvidence(
+            "Selected controller task success",
+            baseline_ok,
+            f"score={baseline.score:.4f}, fell={bool(baseline_summary.fell)}",
+        ),
+        AutopsyEvidence(
+            "Curated-controller counterfactual",
+            reference_ok,
+            (
+                f"score={reference.score:.4f}, fell={bool(reference_summary.fell)}"
+                if reference is not None and reference_summary is not None
+                else "Not run: selected and curated controllers cannot be distinguished."
+            ),
+        ),
+        AutopsyEvidence(
+            "Perturbation robustness",
+            robust_ok,
+            f"task failure rate={robustness.fail_rate:.0%} over {len(robustness.trials)} trials",
+        ),
+    ]
+    if other is not None:
+        evidence.append(
+            AutopsyEvidence(
+                "Backend agreement",
+                portability_ok,
+                f"normalized score gap={normalized_gap:.3f} "
+                f"({baseline.backend} vs {other.backend})",
+            )
+        )
+
+    if not baseline_ok and reference_ok is None:
+        cause = "undetermined"
+        confidence = "low"
+        summary = "No distinct controller counterfactual was available; the cause is unknown."
+        recommendations = ["Select a different controller and rerun the analysis."]
+    elif not baseline_ok and reference_ok:
+        cause = "controller"
+        confidence = "high"
+        summary = "The same body and task succeed with the curated controller."
+        recommendations = [
+            "Compare joint targets and phase relationships against the curated controller.",
+            "Reduce gait amplitude/frequency or migrate the curated gait into controller.json.",
+        ]
+    elif not baseline_ok and not reference_ok:
+        cause = "undetermined"
+        confidence = "low"
+        summary = (
+            "Both tested controllers failed. The evidence does not distinguish a controller "
+            "problem from a body or task problem."
+        )
+        recommendations = [
+            "Try another controller or change one body or task variable at a time.",
+            "Check task scale, target/event placement, support polygon, and center of mass.",
+        ]
+    elif not robust_ok:
+        cause = "fragility"
+        confidence = "high"
+        summary = "The baseline works, but small mass/friction changes frequently break it."
+        recommendations = [
+            "Widen the stance, lower the center of mass, or slow the controller.",
+            "Tune against a distribution of perturbations instead of one nominal setup.",
+        ]
+    elif not portability_ok:
+        cause = "backend_sensitivity"
+        confidence = "medium"
+        summary = "Nominal behavior depends strongly on one backend's contact dynamics."
+        recommendations = [
+            "Inspect trajectory divergence and reduce contact-model-sensitive gait tuning.",
+            "Treat the artifact as structurally portable, not behaviorally portable yet.",
+        ]
+    else:
+        cause = "no_failure_detected"
+        confidence = "high"
+        summary = "The selected controller passed nominal, perturbation, and requested checks."
+        recommendations = [
+            "Increase perturbation strength or add a more demanding task-specific threshold."
+        ]
+
+    if diagnosis.patterns:
+        recommendations.extend(diagnosis.suggestions[:2])
+    return AutopsyResult(
+        creature=creature.name,
+        task=task.name,
+        profile=profile.name,
+        primary_cause=cause,
+        confidence=confidence,
+        summary=summary,
+        evidence=evidence,
+        recommendations=list(dict.fromkeys(recommendations)),
+        baseline_trace=baseline,
+        reference_trace=reference,
+        other_backend_trace=other,
+        robustness=robustness,
+        trial_recordings=trial_recordings,
+    )
+
+
+def write_autopsy_evidence(
+    result: AutopsyResult,
+    creature: CreatureSpec,
+    task: TaskSpec,
+    out: Path,
+    *,
+    overwrite: bool = False,
+) -> None:
+    """Save hash-verifiable packs for each recorded counterfactual and trial."""
+    import shutil
+
+    from creature_lab.exporting import export_design_pack, verify_design_pack
+
+    evidence_dir = out / "evidence"
+    if evidence_dir.exists():
+        if not overwrite:
+            raise FileExistsError(f"evidence directory already exists: {evidence_dir}")
+        if evidence_dir.is_symlink() or evidence_dir.resolve().parent != out.resolve():
+            raise ValueError("refusing to replace an evidence directory outside the report")
+        shutil.rmtree(evidence_dir)
+    recordings: list[tuple[str, CreatureSpec, TaskSpec, EpisodeTrace]] = []
+    if result.reference_trace is not None:
+        recordings.append(("reference", creature, task, result.reference_trace))
+    if result.other_backend_trace is not None:
+        recordings.append(("other_backend", creature, task, result.other_backend_trace))
+    for index, trial in enumerate(result.trial_recordings):
+        recordings.append((f"trial_{index + 1:03d}", trial.creature, trial.task, trial.trace))
+    for name, body, scenario, trace in recordings:
+        destination = evidence_dir / name
+        export_design_pack(body, scenario, trace, out_dir=destination)
+        verification = verify_design_pack(destination)
+        if not verification.valid:
+            raise ValueError(
+                f"{name} evidence verification failed: {'; '.join(verification.errors)}"
+            )
+
+
+def autopsy_to_markdown(result: AutopsyResult) -> str:
+    lines = [
+        f"# Creature Lab Experiment Autopsy: {result.creature}",
+        "",
+        f"**Candidate explanation:** `{result.primary_cause}` (heuristic)",
+        "",
+        "This rule-based assessment is not a calibrated probability or proof of causation.",
+        "",
+        result.summary,
+        "",
+        "## Evidence",
+        "",
+    ]
+    for item in result.evidence:
+        status = "SKIPPED" if item.passed is None else "PASS" if item.passed else "FAIL"
+        lines.append(f"- {status} — **{item.name}:** {item.detail}")
+    lines.extend(["", "## Recommended next experiments", ""])
+    lines.extend(f"- {item}" for item in result.recommendations)
+    lines.append("")
+    return "\n".join(lines)
+
+
+def autopsy_to_html(result: AutopsyResult) -> str:
+    payload = result.to_json_dict()
+    evidence = "".join(
+        f"<li class='{'skipped' if item.passed is None else 'pass' if item.passed else 'fail'}'>"
+        "<strong>"
+        f"{html.escape(item.name)}</strong>: {html.escape(item.detail)}</li>"
+        for item in result.evidence
+    )
+    recommendations = "".join(f"<li>{html.escape(item)}</li>" for item in result.recommendations)
+    return f"""<!doctype html>
+<html><head><meta charset="utf-8"><title>Creature Lab Autopsy</title>
+<style>body{{font:16px system-ui;max-width:900px;margin:3rem auto;padding:0 1rem;color:#172033}}
+.card{{border:1px solid #ccd3df;border-radius:12px;padding:1.2rem;margin:1rem 0}}
+.pass{{color:#176b3a}}.fail{{color:#a52b2b}}code{{background:#eef1f6;padding:.15rem .35rem}}
+</style></head><body><h1>Experiment Autopsy: {html.escape(result.creature)}</h1>
+<div class="card"><h2>Candidate explanation: <code>{html.escape(result.primary_cause)}</code></h2>
+<p>{html.escape(result.summary)}</p><p>This rule-based assessment is not a calibrated
+probability or proof of causation.</p></div>
+<h2>Evidence</h2><ul>{evidence}</ul><h2>Recommended next experiments</h2>
+<ul>{recommendations}</ul><details><summary>Machine-readable result</summary>
+<pre>{html.escape(json.dumps(payload, indent=2))}</pre></details></body></html>"""

creature_lab/backends/__init__.py

@@ -0,0 +1 @@
+"""Simulator adapters. Only backend modules may import a physics engine."""

creature_lab/backends/base.py

@@ -0,0 +1,25 @@
+"""Backend-neutral simulator protocol.
+
+Every physics engine adapter (PyBullet first, others later) implements this
+protocol. Code outside `backends/` must only depend on this interface, never
+on a specific engine.
+"""
+
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from creature_lab.schema import CreatureSpec, FrameState, TaskSpec
+
+
+@runtime_checkable
+class SimBackend(Protocol):
+    def build(self, creature: CreatureSpec, task: TaskSpec, *, seed: int | None = None) -> None: ...
+    def reset(self) -> None: ...
+    def step(self, dt: float) -> FrameState: ...
+    def observe(self) -> FrameState: ...
+    def apply_motor_targets(self, targets: dict[str, float]) -> None: ...
+    def apply_joint_control(self, targets: dict[str, float], mode: str = "position") -> None: ...
+    def damage_part(self, part_id: str) -> None: ...
+    def score_summary(self) -> dict[str, float]: ...
+    def close(self) -> None: ...

creature_lab/backends/mujoco_backend.py

@@ -0,0 +1,253 @@
+"""MuJoCo implementation of the SimBackend protocol.
+
+Only this module imports ``mujoco``. The model is built at runtime from the MJCF
+exporter (``export_mjcf``), so the creature spec stays the single source of truth.
+Scoring reuses the backend-neutral ``scoring`` module, exactly like the PyBullet
+backend, so the two engines score episodes the same way (even though their exact
+dynamics differ — see docs/archive/MVP_PLAN.md on portability).
+"""
+
+from __future__ import annotations
+
+import math
+
+import mujoco
+import numpy as np
+
+from creature_lab.export.mjcf import export_mjcf
+from creature_lab.schema import CreatureSpec, PartPose, TaskSpec
+from creature_lab.schema.creature import JointType
+from creature_lab.schema.trace import ContactSpec, FrameState
+from creature_lab.scoring import score_components
+from creature_lab.terrain import is_flat, normalized_heightfield_data
+
+
+def backend_version() -> str:
+    """Human-readable MuJoCo backend version."""
+    return f"mujoco {mujoco.__version__}"
+
+
+class MuJoCoBackend:
+    """Builds a creature as a MuJoCo model and steps it."""
+
+    def __init__(self, gui: bool = False) -> None:
+        # gui is accepted for protocol parity; this backend renders via traces.
+        self._creature: CreatureSpec | None = None
+        self._task: TaskSpec | None = None
+        self._model: mujoco.MjModel | None = None
+        self._data: mujoco.MjData | None = None
+        self._root_body = ""
+        self._ground_geom = -1
+        self._geom_part: dict[int, str] = {}
+        self._damaged: set[str] = set()
+        self._t = 0.0
+        self._energy = 0.0
+        self._initial_root_x = 0.0
+        self._initial_target_distance = 0.0
+        self._last_score_components: dict[str, float] = {}
+        self._damage_fired = False
+        self._impulse_fired = False
+        self._seed: int | None = None
+
+    def build(self, creature: CreatureSpec, task: TaskSpec, *, seed: int | None = None) -> None:
+        self._creature = creature
+        self._task = task
+        self._seed = seed
+        xml = export_mjcf(
+            creature, friction=task.terrain.friction, timestep=task.timestep, terrain=task.terrain
+        )
+        self._model = mujoco.MjModel.from_xml_string(xml)
+        if not is_flat(task.terrain):
+            self._model.hfield_data[:] = normalized_heightfield_data(task.terrain)
+        self._data = mujoco.MjData(self._model)
+
+        child_ids = {joint.child for joint in creature.joints}
+        self._root_body = next(p.id for p in creature.parts if p.id not in child_ids)
+        self._ground_geom = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_GEOM, "ground")
+        # geom index -> part id (each part has exactly one geom, added in part order).
+        self._geom_part = {}
+        for part in creature.parts:
+            body_id = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_BODY, part.id)
+            for geom in range(
+                self._model.body_geomadr[body_id],
+                self._model.body_geomadr[body_id] + self._model.body_geomnum[body_id],
+            ):
+                self._geom_part[geom] = part.id
+
+        self._t = 0.0
+        self._energy = 0.0
+        self._damaged.clear()
+        self._damage_fired = False
+        self._impulse_fired = False
+        mujoco.mj_forward(self._model, self._data)
+        self._initial_root_x = float(self._root_pos()[0])
+        self._initial_target_distance = (
+            math.dist(self._root_pos(), task.target.position) if task.target is not None else 0.0
+        )
+
+    def reset(self) -> None:
+        if self._creature is None or self._task is None:
+            raise RuntimeError("build() must be called before reset()")
+        self.build(self._creature, self._task, seed=self._seed)
+
+    def step(self, dt: float) -> FrameState:
+        if self._model is None or self._data is None or self._task is None:
+            raise RuntimeError("build() must be called before step()")
+        self._t += dt
+        events: list[str] = []
+
+        damage = self._task.damage_event
+        if damage is not None and not self._damage_fired and self._t >= damage.time:
+            self.damage_part(damage.part_id)
+            self._damage_fired = True
+            events.append(f"damage:{damage.part_id}")
+
+        impulse = self._task.impulse_event
+        if impulse is not None and not self._impulse_fired and self._t >= impulse.time:
+            body_id = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_BODY, impulse.part_id)
+            self._data.xfrc_applied[body_id, :3] = impulse.force
+            self._impulse_fired = True
+            events.append(f"impulse:{impulse.part_id}")
+
+        # ``step(dt)`` is part of the backend contract. Keep MuJoCo's actual
+        # integrator step aligned with the requested duration instead of advancing
+        # physics by a stale model timestep while only bookkeeping uses ``dt``.
+        self._model.opt.timestep = dt
+        mujoco.mj_step(self._model, self._data)
+        if impulse is not None and self._impulse_fired:
+            body_id = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_BODY, impulse.part_id)
+            self._data.xfrc_applied[body_id, :3] = 0.0  # one-step push only
+        self._accumulate_energy(dt)
+
+        frame = self._read_frame()
+        if events:
+            frame = frame.model_copy(update={"events": events})
+        return frame
+
+    def apply_motor_targets(self, targets: dict[str, float]) -> None:
+        if self._model is None or self._data is None:
+            raise RuntimeError("build() must be called before apply_motor_targets()")
+        for joint_id, value in targets.items():
+            if joint_id in self._damaged:
+                continue
+            actuator_id = mujoco.mj_name2id(
+                self._model, mujoco.mjtObj.mjOBJ_ACTUATOR, f"act_{joint_id}"
+            )
+            if actuator_id >= 0:
+                self._data.ctrl[actuator_id] = value
+
+    def apply_joint_control(self, targets: dict[str, float], mode: str = "position") -> None:
+        """Position control of motored joints (used by CreatureEnv).
+
+        The MJCF exports ``<position>`` servos, so only position mode is supported;
+        velocity/torque control would need a different actuator model.
+        """
+        if mode != "position":
+            raise NotImplementedError(
+                f"MuJoCo backend supports position control only, not {mode!r}"
+            )
+        self.apply_motor_targets(targets)
+
+    def damage_part(self, part_id: str) -> None:
+        if self._creature is None:
+            raise RuntimeError("build() must be called before damage_part()")
+        # Disable the actuator on the joint that drives this part (best-effort parity
+        # with PyBullet: the limb goes limp).
+        for joint in self._creature.joints:
+            if joint.child == part_id:
+                self._damaged.add(joint.id)
+                actuator_id = mujoco.mj_name2id(
+                    self._model, mujoco.mjtObj.mjOBJ_ACTUATOR, f"act_{joint.id}"
+                )
+                if actuator_id >= 0:
+                    self._data.ctrl[actuator_id] = 0.0
+
+    def close(self) -> None:
+        self._model = None
+        self._data = None
+
+    def score_summary(self) -> dict[str, float]:
+        return dict(self._last_score_components)
+
+    def observe(self) -> FrameState:
+        return self._read_frame()
+
+    # -- internals --------------------------------------------------------------
+
+    def _root_pos(self) -> tuple[float, float, float]:
+        body_id = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_BODY, self._root_body)
+        return tuple(float(v) for v in self._data.xpos[body_id])
+
+    def _accumulate_energy(self, dt: float) -> None:
+        for joint in self._creature.joints:
+            if joint.type != JointType.HINGE:
+                continue
+            jid = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_JOINT, joint.id)
+            if jid < 0:
+                continue
+            velocity = float(self._data.qvel[self._model.jnt_dofadr[jid]])
+            self._energy += velocity * velocity * dt
+
+    def _read_frame(self) -> FrameState:
+        parts: dict[str, PartPose] = {}
+        for part in self._creature.parts:
+            body_id = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_BODY, part.id)
+            position = tuple(float(v) for v in self._data.xpos[body_id])
+            orientation = tuple(float(v) for v in self._data.xquat[body_id])  # MuJoCo: w,x,y,z
+            parts[part.id] = PartPose(position=position, orientation=orientation)
+
+        joint_angles: dict[str, float] = {}
+        for joint in self._creature.joints:
+            if joint.type != JointType.HINGE:
+                continue
+            jid = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_JOINT, joint.id)
+            if jid >= 0:
+                joint_angles[joint.id] = float(self._data.qpos[self._model.jnt_qposadr[jid]])
+
+        root_position = self._root_pos()
+        forward_distance = root_position[0] - self._initial_root_x
+        target_progress = 0.0
+        if self._task.target is not None:
+            target_progress = self._initial_target_distance - math.dist(
+                root_position, self._task.target.position
+            )
+        root_body = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_BODY, self._root_body)
+        up_z = float(self._data.xmat[root_body].reshape(3, 3)[2, 2])
+
+        self._last_score_components = score_components(
+            self._task.reward,
+            forward_distance=forward_distance,
+            target_progress=target_progress,
+            energy=self._energy,
+            fallen=up_z < 0.5,
+        )
+        return FrameState(
+            t=self._t,
+            parts=parts,
+            joint_angles=joint_angles,
+            contacts=self._read_contacts(),
+            score=self._last_score_components["total"],
+        )
+
+    def _read_contacts(self) -> list[ContactSpec]:
+        contacts: list[ContactSpec] = []
+        force = np.zeros(6)
+        for i in range(self._data.ncon):
+            contact = self._data.contact[i]
+            part_geom = None
+            if contact.geom1 == self._ground_geom:
+                part_geom = contact.geom2
+            elif contact.geom2 == self._ground_geom:
+                part_geom = contact.geom1
+            if part_geom is None or part_geom not in self._geom_part:
+                continue
+            mujoco.mj_contactForce(self._model, self._data, i, force)
+            contacts.append(
+                ContactSpec(
+                    part_id=self._geom_part[part_geom],
+                    position=tuple(float(v) for v in contact.pos),
+                    normal=tuple(float(v) for v in contact.frame[:3]),
+                    force=max(0.0, float(force[0])),  # normal component
+                )
+            )
+        return contacts