From 14d719fe7443173613d0f94e3f46203494b7a8b3 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Thu, 7 May 2026 22:48:23 +0200 Subject: [PATCH] refactor: extract mjx evaluation --- .../evaluation/__init__.py | 15 ++ .../evaluation/evaluate_mjx.py | 255 ++++++++++++++++++ 2 files changed, 270 insertions(+) create mode 100644 src/brittle_star_project/evaluation/evaluate_mjx.py diff --git a/src/brittle_star_project/evaluation/__init__.py b/src/brittle_star_project/evaluation/__init__.py index 38a1c8b..2922d04 100644 --- a/src/brittle_star_project/evaluation/__init__.py +++ b/src/brittle_star_project/evaluation/__init__.py @@ -1,20 +1,35 @@ from __future__ import annotations from .checkpoint import load_metadata, load_params, metadata_to_configs, TrainingConfig +from .evaluate_mjx import ( + CheckpointEvalResult, + append_checkpoint_eval_row, + build_eval_rollout_fn, + evaluate_checkpoint_mjx, +) from .policy import PolicyAgent, ControlPolicy from .rollout import rollout_headless, rollout_viewer, EpisodeResult from .video import record_episode, create_evaluation_dir, save_evaluation_metadata __all__ = [ + # checkpoint loading "load_metadata", "load_params", "metadata_to_configs", "TrainingConfig", + # MJX evaluation + "CheckpointEvalResult", + "append_checkpoint_eval_row", + "build_eval_rollout_fn", + "evaluate_checkpoint_mjx", + # policy "PolicyAgent", "ControlPolicy", + # rollout "rollout_headless", "rollout_viewer", "EpisodeResult", + # video "record_episode", "create_evaluation_dir", "save_evaluation_metadata", diff --git a/src/brittle_star_project/evaluation/evaluate_mjx.py b/src/brittle_star_project/evaluation/evaluate_mjx.py new file mode 100644 index 0000000..f281535 --- /dev/null +++ b/src/brittle_star_project/evaluation/evaluate_mjx.py @@ -0,0 +1,255 @@ +"""MJX-based headless checkpoint evaluation. + +This module provides a fast, JIT-compiled evaluation path using the MJX +(JAX-accelerated MuJoCo) backend. It is intended for evaluating checkpoints +*during* or *after* a training run, where the environment and policy are +already fully initialised. + +The key functions are: + +- `build_eval_rollout_fn` — builds and JIT-compiles a single-episode rollout function from the + training environment and policy components. +- `evaluate_checkpoint_mjx` — runs that function for a given set of parameters and returns a typed + `CheckpointEvalResult`. +- `append_checkpoint_eval_row` — persists the result to the run's + ``metrics/checkpoint_evaluation.csv``, migrating old schemas automatically. +""" + +from __future__ import annotations + +import csv +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable + +import jax +import jax.numpy as jnp + + +@dataclass +class CheckpointEvalResult: + """Structured result from a single MJX checkpoint evaluation episode.""" + + steps: int + """Number of control steps taken (≤ max_steps).""" + + reached_target: bool + """Whether the robot reached the target (terminated) before max_steps.""" + + eval_return: float + """Accumulated shaped reward over the episode.""" + + final_xy_dist: float + """XY distance to target at episode end. 0.0 when ``reached_target`` is True.""" + + initial_xy_dist: float + """XY distance to target at episode start.""" + + +def build_eval_rollout_fn( + *, + env: Any, + obs_processor: Callable, + sensor_apply: Callable, + actor_apply: Callable, + action_low: jnp.ndarray, + action_high: jnp.ndarray, + reward_fn: Callable, +) -> Callable: + """Build and JIT-compile a single-episode MJX evaluation rollout. + + The returned function has the signature:: + + eval_fn(params: dict, seed: int, max_steps: int) + -> (steps, reached_target, eval_return, final_xy_dist, initial_xy_dist) + + All outputs are JAX arrays. Convert to Python scalars before logging. + + Args: + env: The training environment wrapper. Must expose ``env.raw`` with + ``reset`` and ``step`` methods compatible with ``jax.vmap``. + obs_processor: Observation normalisation / padding callable, as + returned by ``create_obs_processor``. + sensor_apply: The sensor network's ``apply`` method (JIT-compiled). + actor_apply: The actor network's ``apply`` method (JIT-compiled). + action_low: Per-joint action lower bound (JAX array, shape ``(action_dim,)``). + action_high: Per-joint action upper bound (JAX array, shape ``(action_dim,)``). + reward_fn: Shaped reward function with signature + ``reward_fn(env_state, next_env_state) -> jnp.ndarray``. + Typically the module-level ``reward_fn`` from ``PPOTrainer``. + + Returns: + A JIT-compiled callable that runs one deterministic evaluation episode. + """ + # vmap over a batch of 1 so the MJX API is satisfied without any + # extra bookkeeping in the caller. + reset_1 = jax.vmap(env.raw.reset) + step_1 = jax.vmap(env.raw.step) + + def _eval_rollout(params: dict, seed: int, max_steps: int): + rng = jax.random.PRNGKey(seed) + rngs = jnp.asarray(jax.random.split(rng, 1)) + state = reset_1(rng=rngs) + + initial_xy_dist = jnp.squeeze(state.observations["xy_distance_to_target"]) + + t0 = jnp.asarray(0, dtype=jnp.int32) + done0 = jnp.squeeze(state.terminated | state.truncated) + return0 = jnp.asarray(0.0, dtype=jnp.float32) + + def cond(carry): + t, _state, done, _return_ = carry + return jnp.logical_and(t < max_steps, jnp.logical_not(done)) + + def body(carry): + t, state, _done, return_ = carry + + obs = obs_processor(state.observations) + hidden = sensor_apply(params["sensor_params"], obs) + mean, _log_std = actor_apply(params["actor_params"], hidden) + + # Deterministic action: use the actor mean, no exploration noise. + action = jnp.clip(mean, action_low, action_high) + next_state = step_1(state=state, action=action) + + shaped_reward = reward_fn(state, next_state) + return_ = return_ + jnp.squeeze(shaped_reward) + + done_next = jnp.squeeze(next_state.terminated | next_state.truncated) + return (t + 1, next_state, done_next, return_) + + t, final_state, _done, return_ = jax.lax.while_loop(cond, body, (t0, state, done0, return0)) + + reached_target = jnp.squeeze(final_state.terminated) + final_xy_dist_raw = jnp.squeeze(final_state.observations["xy_distance_to_target"]) + # Clamp to 0 when the target was reached so downstream consumers + # don't have to special-case "terminated" themselves. + final_xy_dist = jnp.where(reached_target, 0.0, final_xy_dist_raw) + + return t, reached_target, return_, final_xy_dist, initial_xy_dist + + return jax.jit(_eval_rollout) + + +def evaluate_checkpoint_mjx( + eval_fn: Callable, + params: dict, + *, + seed: int, + max_steps: int, +) -> CheckpointEvalResult: + """Run one deterministic evaluation episode and return typed metrics. + + Args: + eval_fn: A JIT-compiled function as returned by :func:`build_eval_rollout_fn`. + params: Agent parameter dict (e.g. ``agent_state.params``). + seed: Random seed for environment reset (controls target placement). + max_steps: Maximum number of control steps before the episode is cut off. + + Returns: + A :class:`CheckpointEvalResult` with all JAX arrays converted to + plain Python scalars. + """ + steps, reached, eval_return, final_xy_dist, initial_xy_dist = eval_fn(params, seed, max_steps) + return CheckpointEvalResult( + steps=int(steps), + reached_target=bool(reached), + eval_return=float(eval_return), + final_xy_dist=float(final_xy_dist), + initial_xy_dist=float(initial_xy_dist), + ) + + +_FIELDNAMES = [ + "checkpoint", + "trained_timesteps", + "eval_steps", + "eval_return", + "final_xy_dist", + "initial_xy_dist", + "reached_target", +] + + +def _migrate_csv_if_needed(csv_path: Path) -> None: + """Rewrite the CSV with the canonical field names if the schema changed. + + Best-effort: any exception is silently swallowed so that a schema mismatch + never causes a training crash. + """ + try: + with open(csv_path, "r", newline="") as f: + header = next(csv.reader(f), None) + + if header is None or list(header) == _FIELDNAMES: + return # Nothing to migrate. + + migrated_rows: list[dict[str, Any]] = [] + with open(csv_path, "r", newline="") as f: + for row in csv.DictReader(f): + migrated_rows.append( + { + "checkpoint": row.get("checkpoint", row.get("iteration")), + "trained_timesteps": row.get("trained_timesteps"), + "eval_steps": row.get("eval_steps", row.get("steps_to_target")), + "eval_return": row.get("eval_return"), + "final_xy_dist": row.get("final_xy_dist"), + "initial_xy_dist": row.get("initial_xy_dist"), + "reached_target": row.get("reached_target"), + } + ) + + with open(csv_path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=_FIELDNAMES) + writer.writeheader() + writer.writerows(migrated_rows) + except Exception: + pass # Never crash training on a migration issue. + + +def append_checkpoint_eval_row( + run_dir: str | Path, + *, + iteration: int, + trained_timesteps: int, + result: CheckpointEvalResult, +) -> Path: + """Append one evaluation row to ``/metrics/checkpoint_evaluation.csv``. + + Creates the file (including the ``metrics/`` directory) if it does not yet + exist. Migrates the file to the current schema if the header has changed. + + Args: + run_dir: Root directory of the training run (Hydra's output dir). + iteration: Training iteration number, used as the checkpoint identifier. + trained_timesteps: Total environment steps taken at this checkpoint. + result: Evaluation result as returned by :func:`evaluate_checkpoint_mjx`. + + Returns: + Absolute path to the CSV file (useful for W&B sync). + """ + metrics_dir = Path(run_dir) / "metrics" + metrics_dir.mkdir(parents=True, exist_ok=True) + csv_path = metrics_dir / "checkpoint_evaluation.csv" + + if csv_path.exists(): + _migrate_csv_if_needed(csv_path) + + file_exists = csv_path.exists() + with open(csv_path, "a", newline="") as f: + writer = csv.DictWriter(f, fieldnames=_FIELDNAMES) + if not file_exists: + writer.writeheader() + writer.writerow( + { + "checkpoint": int(iteration), + "trained_timesteps": int(trained_timesteps), + "eval_steps": result.steps, + "eval_return": result.eval_return, + "final_xy_dist": result.final_xy_dist, + "initial_xy_dist": result.initial_xy_dist, + "reached_target": result.reached_target, + } + ) + + return csv_path