requirements_manager ¶

def step(
    self,
    state: MjState,
    *,
    signal_manager: SignalManager | None,
    mojo_model: MojoModel | None,
) -> None:
    """Evaluates live requirements for the current step, posts telemetry, and raises `SimulationStopped` if a terminating check fails."""
    if not self._live_requirements or mojo_model is None:
        return

    self._step_count += 1
    terminate_msgs: list[str] = []
    satisfied_msgs: list[str] = []

    for spec in self._live_requirements:
        assert spec.every is not None

        if self._step_count % spec.every == 0:
            already_latched = self._latched.get(spec.name)
            if already_latched is None:
                req_msg = ""
                try:
                    passed, req_msg = spec.fn(mojo_model, state, None)
                except Exception as e:
                    logger.warning(
                        f"Requirement {spec.name} raised an exception during its evaluation: {e}"
                    )
                    passed = False
                logger.debug(
                    f"requirement '{spec.name}' evaluated at t={state.data.time:.6f}: "
                    f"{'undetermined' if passed is None else 'passed' if passed else 'failed'} "
                    f"({req_msg or 'no message'})"
                )

                # cache every verdict, including undetermined (None), so
                # last_passed() can tell "evaluated, no verdict yet" apart
                # from "not evaluated at this sim time". None never latches
                # a failure, terminates the trial, or posts telemetry,
                # though. Once latched (see below), this stops: neither
                # run_end_of_trial_evaluation() nor last_passed() consult
                # the cache for a latched requirement anymore (both read
                # straight from _latched), so caching a fresh entry on
                # every subsequent tick would just grow the dict forever
                # with nothing left to read it.
                self._live_cache[(spec.name, state.data.time)] = passed
                if passed is not None:
                    self._last_live_result[spec.name] = passed
                    if not passed and spec.name not in self._first_live_failure:
                        self._first_live_failure[spec.name] = (
                            state.data.time,
                            req_msg or "no message attached",
                        )
                    if not passed and spec.terminate_on_fail:
                        terminate_msgs.append(
                            f"requirement '{spec.name}' triggered early termination ({req_msg or 'no message attached'})"
                        )
                    if passed and spec.terminate_on_pass:
                        satisfied_msgs.append(
                            f"requirement '{spec.name}' was satisfied and ended the trial early ({req_msg or 'no message attached'})"
                        )
                    if (passed and spec.latch_on_pass) or (
                        not passed and spec.latch_on_fail
                    ):
                        self._latched[spec.name] = (
                            passed,
                            req_msg or "no message attached",
                            state.data.time,
                        )
                        logger.debug(
                            f"requirement '{spec.name}' latched {'pass' if passed else 'fail'} "
                            f"at t={state.data.time:.6f}; further evaluations will be skipped"
                        )
            # else: verdict already locked in. Replay it silently instead
            # of paying to call (possibly expensive) fn again. This was
            # already logged once, when the latch was first set above, not
            # on every replay tick. Nothing else to do: _last_live_result
            # already holds the latched value (telemetry below reads it
            # unchanged), and _live_cache/_latched don't need updating.

        # post last known result on every step so the telemetry signal is continuous;
        # skip entirely if this requirement has never been evaluated yet.
        if (
            signal_manager
            and spec.column is not None
            and spec.name in self._last_live_result
        ):
            signal_manager.post(
                value=1.0 if self._last_live_result[spec.name] else 0.0,
                column=spec.column,
            )

    # a failing terminator wins over a satisfied one on the same step
    if terminate_msgs:
        combined = "; ".join(terminate_msgs)
        self._termination_reason = combined
        logger.info(combined)
        raise RequirementTerminated(combined)
    if satisfied_msgs:
        combined = "; ".join(satisfied_msgs)
        self._termination_reason = combined
        logger.info(combined)
        raise RequirementSatisfied(combined)

def run_end_of_trial_evaluation(
    self,
    last_state: MjState,
    *,
    signal_manager: SignalManager,
    mojo_model: MojoModel,
) -> None:
    """Runs all requirements against the final simulation state and telemetry parquet, stores results, and writes `requirements.json` to the trial directory."""
    from mujoco_mojo.utils.dataframe import _MojoFrame

    parquet_path = signal_manager.export_path
    df = _MojoFrame.read_parquet(parquet_path) if parquet_path.exists() else None

    results: list[RequirementResult] = []
    for spec in self._requirements.values():
        # the requirement's static configuration, shown alongside the
        # per-trial verdict so callers can tell e.g. "this failed" apart
        # from "this failed and would have kept re-checking every 10 steps"
        config = {
            "every": spec.every,
            "terminate_on_fail": spec.terminate_on_fail,
            "terminate_on_pass": spec.terminate_on_pass,
            "latch_on_fail": spec.latch_on_fail,
            "latch_on_pass": spec.latch_on_pass,
        }

        if (latched := self._latched.get(spec.name)) is not None:
            # verdict already locked in by latch_on_pass/latch_on_fail:
            # never call fn again, just report the locked-in verdict
            passed, latch_msg, latch_time = latched
            message = (
                f"latched {'passed' if passed else 'failed'} at t={latch_time:.6f} "
                f"({latch_msg}); evaluation skipped afterward to save compute"
            )
            results.append(
                RequirementResult(
                    name=spec.name,
                    passed=passed,
                    message=message,
                    decided_at=latch_time,
                    **config,
                )
            )
            continue

        try:
            passed, message = spec.fn(mojo_model, last_state, df)
        except Exception as e:
            passed = False
            message = f"requirement raised {type(e).__name__}: {e}"

        # end of trial is the last chance to decide: a check that is still
        # undetermined here never passed, so it fails
        if passed is None:
            passed = False
            message = f"undetermined at end of trial ({message})"

        # a plain end-of-trial verdict is decided at the final sim time;
        # overridden below if a live failure decided it earlier
        decided_at = last_state.data.time

        # a live requirement that failed at any point during the run is a
        # failure for the trial, even if its end-of-trial evaluation passes
        if (failure := self._first_live_failure.get(spec.name)) is not None:
            fail_time, fail_msg = failure
            # only count determinate (True/False) evaluations; undetermined
            # (None) verdicts are cached but never counted as checks or failures
            n_evals = sum(
                1
                for (n, _), ok in self._live_cache.items()
                if n == spec.name and ok is not None
            )
            n_fails = sum(
                1
                for (n, _), ok in self._live_cache.items()
                if n == spec.name and ok is False
            )
            message = (
                f"failed {n_fails}/{n_evals} live checks, first at t={fail_time:.6f} ({fail_msg}); "
                f"end-of-trial evaluation {'passed' if passed else 'failed'} ({message})"
            )
            passed = False
            decided_at = fail_time

        results.append(
            RequirementResult(
                name=spec.name,
                passed=passed,
                message=message,
                decided_at=decided_at,
                **config,
            )
        )

    self.results = results
    trial_dir = parquet_path.parent
    (trial_dir / REQUIREMENTS_FNAME).write_text(
        json.dumps([r.model_dump() for r in results], indent=2),
        encoding="utf-8",
    )