From 90c2c3d1f0e6d3d3633757744a2392f39b99c521 Mon Sep 17 00:00:00 2001 From: Jona Reynaert Date: Sun, 3 May 2026 15:53:26 +0200 Subject: [PATCH 1/8] feat: added flags for checkpoint evaluation --- configs/logging/default.yaml | 3 +++ src/experiment_logger/config_logger.py | 24 ++++++++++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/configs/logging/default.yaml b/configs/logging/default.yaml index 1ecf7fd..af9bc6d 100644 --- a/configs/logging/default.yaml +++ b/configs/logging/default.yaml @@ -11,3 +11,6 @@ checkpoint_frequency: 100 upload_final_model: false upload_checkpoints: false hf_entity: "" +evaluate_checkpoints: false +eval_max_steps: 5000 +eval_seed: 0 \ No newline at end of file diff --git a/src/experiment_logger/config_logger.py b/src/experiment_logger/config_logger.py index c34c162..d158e84 100644 --- a/src/experiment_logger/config_logger.py +++ b/src/experiment_logger/config_logger.py @@ -18,6 +18,13 @@ class LoggingConfig: upload_final_model: bool = False upload_checkpoints: bool = False + # Checkpoint evaluation (synchronous, in-process) + # When enabled, each saved checkpoint is evaluated headlessly and the results + # are appended to a CSV in the run's metrics/ folder. + evaluate_checkpoints: bool = False + eval_max_steps: int = 5000 + eval_seed: int = 0 + hf_entity: str = "" def __post_init__(self): @@ -31,3 +38,20 @@ class LoggingConfig: "Configuration Error: 'upload_checkpoints' is True, but it requires " "both 'track' and 'save_checkpoints' to also be True." ) + + if self.evaluate_checkpoints: + if not self.save_checkpoints: + raise ValueError( + "Configuration Error: 'evaluate_checkpoints' is True, but it requires " + "'save_checkpoints' to also be True." + ) + if self.checkpoint_frequency <= 0: + raise ValueError( + "Configuration Error: 'evaluate_checkpoints' is True, but it requires " + "'checkpoint_frequency' to be > 0." + ) + if self.eval_max_steps <= 0: + raise ValueError( + "Configuration Error: 'eval_max_steps' must be > 0 when " + "'evaluate_checkpoints' is enabled." + ) From a3d6b32eaa78442f3f39a1f1c2330122a0dff58e Mon Sep 17 00:00:00 2001 From: Jona Reynaert Date: Sun, 3 May 2026 15:53:47 +0200 Subject: [PATCH 2/8] feat: implemented checkpoint evaluation --- .../trainers/PPOTrainer.py | 145 ++++++++++++++++++ 1 file changed, 145 insertions(+) diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index a3e4eac..725f29b 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -1,8 +1,10 @@ import datetime import random import time +import csv from dataclasses import asdict, dataclass from functools import partial +from pathlib import Path from typing import Any import jax @@ -26,6 +28,10 @@ from brittle_star_project.MLPs.mlps import ( ) from brittle_star_project.ppo import PPO +from brittle_star_project import Backend, BrittleStarEnv, BrittleStarEnvFactory +from brittle_star_project.evaluation.policy import PolicyAgent +from brittle_star_project.evaluation.rollout import rollout_headless + # TODO: clip scaled reward? @@ -538,6 +544,144 @@ class PPOTrainer: params=self.agent_state.params, step=iteration, metadata=asdict(self.cfg) ) + def _maybe_sync_csv_to_wandb(self, csv_path: Path) -> None: + if not self.logging_cfg.track: + return + try: + import wandb + + if wandb.run is None: + return + + # "Simple sync" behavior: wandb will copy this file into the run. + wandb.save(str(csv_path), base_path=str(csv_path.parent)) + except Exception as e: + self.logger.warning(f"[EVAL]: Failed to sync CSV to wandb: {e}") + + def _append_checkpoint_eval_row( + self, + *, + iteration: int, + reached_target: bool, + steps_to_target: int, + max_steps: int, + seed: int, + ) -> Path: + metrics_dir = Path(self.run_dir) / "metrics" + metrics_dir.mkdir(parents=True, exist_ok=True) + csv_path = metrics_dir / "checkpoint_evaluation.csv" + + file_exists = csv_path.exists() + with open(csv_path, "a", newline="") as f: + writer = csv.DictWriter( + f, + fieldnames=[ + "iteration", + "steps_to_target", + "reached_target", + "max_steps", + "seed", + ], + ) + if not file_exists: + writer.writeheader() + writer.writerow( + { + "iteration": int(iteration), + "steps_to_target": int(steps_to_target), + "reached_target": bool(reached_target), + "max_steps": int(max_steps), + "seed": int(seed), + } + ) + + return csv_path + + def _evaluate_checkpoint(self, iteration: int) -> None: + if not self.logging_cfg.evaluate_checkpoints: + return + + checkpoint_path = Path(self.run_dir) / "checkpoints" / f"checkpoint_step_{iteration}.flax" + if not checkpoint_path.exists(): + self.logger.warning( + f"[EVAL]: Checkpoint not found at {checkpoint_path} (skipping evaluation)" + ) + return + + max_steps = int(self.logging_cfg.eval_max_steps) + seed = int(self.logging_cfg.eval_seed) + + # Run evaluation best-effort; never fail training because evaluation failed. + env = None + try: + backend = Backend.MJC + + factory = BrittleStarEnvFactory() + raw_env = factory.create_environment( + backend, + self.cfg.morphology, + self.cfg.arena, + self.cfg.environment, + ) + env = BrittleStarEnv( + raw_env, + backend=backend, + config=self.cfg.environment, + morphology_config=self.cfg.morphology, + ) + + action_space = getattr(raw_env, "action_space", None) + action_dim = ( + int(np.asarray(action_space.shape).reshape(-1)[0]) + if action_space is not None and hasattr(action_space, "shape") + else sum(self.cfg.morphology.segments_per_arm) * 2 + ) + + action_low = ( + None + if action_space is None + else np.asarray(action_space.low, dtype=np.float32).ravel() + ) + action_high = ( + None + if action_space is None + else np.asarray(action_space.high, dtype=np.float32).ravel() + ) + + policy = PolicyAgent.from_checkpoint( + checkpoint_path, + action_dim=action_dim, + obs_processor=self.obs_processor, + ) + + result = rollout_headless( + env=env, + policy=policy, + seed=seed, + max_steps=max_steps, + action_low=action_low, + action_high=action_high, + action_mask=None, + ) + + steps_to_target = int(result.length) if result.reached_target else int(max_steps) + csv_path = self._append_checkpoint_eval_row( + iteration=iteration, + reached_target=bool(result.reached_target), + steps_to_target=steps_to_target, + max_steps=max_steps, + seed=seed, + ) + self._maybe_sync_csv_to_wandb(csv_path) + except Exception as e: + self.logger.warning(f"[EVAL]: Checkpoint evaluation failed: {e}") + finally: + try: + if env is not None: + env.close() + except Exception: + pass + def train(self): """ Train the PPO agent for a specified number of iterations. @@ -591,6 +735,7 @@ class PPOTrainer: if self.logging_cfg.save_checkpoints and self.logging_cfg.checkpoint_frequency > 0: if iteration % self.logging_cfg.checkpoint_frequency == 0: self._save_checkpoint(iteration) + self._evaluate_checkpoint(iteration) if getattr(self.cfg.experiment, "debug_sanity", False): self.logger.info("\n[SANITY CHECK] Successfully completed 1 epoch") From 0b845dd7a9a2783565943564a76625f2355ada67 Mon Sep 17 00:00:00 2001 From: Jona Reynaert Date: Sun, 3 May 2026 18:26:24 +0200 Subject: [PATCH 3/8] fix: improved speed of checkpoint eval --- src/brittle_star_project/evaluation/policy.py | 22 +++ .../trainers/PPOTrainer.py | 179 ++++++++++-------- 2 files changed, 120 insertions(+), 81 deletions(-) diff --git a/src/brittle_star_project/evaluation/policy.py b/src/brittle_star_project/evaluation/policy.py index 5aa2d6e..75b6a22 100644 --- a/src/brittle_star_project/evaluation/policy.py +++ b/src/brittle_star_project/evaluation/policy.py @@ -61,6 +61,28 @@ class PolicyAgent: } self._obs_processor = obs_processor + @classmethod + def from_params( + cls, + *, + sensor_params: Any, + actor_params: Any, + action_dim: int, + obs_processor: Any, + ) -> "PolicyAgent": + """Construct a PolicyAgent directly from in-memory parameters.""" + return cls( + sensor_params=sensor_params, + actor_params=actor_params, + action_dim=action_dim, + obs_processor=obs_processor, + ) + + def set_params(self, *, sensor_params: Any, actor_params: Any) -> None: + """Update parameters for evaluation without rebuilding the model.""" + self._params["sensor_params"] = sensor_params + self._params["actor_params"] = actor_params + @classmethod def from_checkpoint( cls, diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 725f29b..54653a1 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -28,9 +28,7 @@ from brittle_star_project.MLPs.mlps import ( ) from brittle_star_project.ppo import PPO -from brittle_star_project import Backend, BrittleStarEnv, BrittleStarEnvFactory -from brittle_star_project.evaluation.policy import PolicyAgent -from brittle_star_project.evaluation.rollout import rollout_headless +from brittle_star_project.environment.env_types import Backend # TODO: clip scaled reward? @@ -290,6 +288,8 @@ class PPOTrainer: action_low = jnp.asarray(self.env.single_action_space.low, dtype=jnp.float32) action_high = jnp.asarray(self.env.single_action_space.high, dtype=jnp.float32) + self._action_low = action_low + self._action_high = action_high self._rollout_jit = jax.jit( partial( @@ -326,6 +326,60 @@ class PPOTrainer: self.episode_stats = self._init_episode_stats() self._init_random() + # Lazily created MJX/JAX evaluation rollout (compiled on first use) + self._eval_rollout_mjx_fn = None + + def _get_or_create_eval_rollout_mjx_fn(self): + if self._eval_rollout_mjx_fn is not None: + return self._eval_rollout_mjx_fn + + # Use the same backend as training (typically MJX). + if getattr(self.env, "backend", None) != Backend.MJX: + self.logger.warning( + f"[EVAL]: Training env backend is {self.env.backend}; " + "MJX evaluation may be unavailable/slow." + ) + + # We vmap over a single environment (batch size 1) for simplicity. + reset_1 = jax.vmap(self.env.raw.reset) + step_1 = jax.vmap(self.env.raw.step) + + action_low = self._action_low + action_high = self._action_high + obs_processor = self.obs_processor + sensor_apply = self.sensor.apply + actor_apply = self.actor.apply + + def _eval_rollout(params, seed: int, max_steps: int): + rng = jax.random.PRNGKey(seed) + rngs = jnp.asarray(jax.random.split(rng, 1)) + state = reset_1(rng=rngs) + + t0 = jnp.asarray(0, dtype=jnp.int32) + done0 = jnp.squeeze(state.terminated | state.truncated) + + def cond(carry): + t, _state, done = carry + return jnp.logical_and(t < max_steps, jnp.logical_not(done)) + + def body(carry): + t, state, _done = carry + + obs = obs_processor(state.observations) + hidden = sensor_apply(params["sensor_params"], obs) + mean, _log_std = actor_apply(params["actor_params"], hidden) + + action = jnp.clip(mean, action_low, action_high) + next_state = step_1(state=state, action=action) + + done_next = jnp.squeeze(next_state.terminated | next_state.truncated) + return (t + 1, next_state, done_next) + + t, _state, done = jax.lax.while_loop(cond, body, (t0, state, done0)) + return t, done + + self._eval_rollout_mjx_fn = jax.jit(_eval_rollout) + return self._eval_rollout_mjx_fn def _init_random(self): self.logger.info(f"[RANDOM]: Setting random seed to {self.experiment.seed}") @@ -564,24 +618,51 @@ class PPOTrainer: iteration: int, reached_target: bool, steps_to_target: int, - max_steps: int, - seed: int, ) -> Path: metrics_dir = Path(self.run_dir) / "metrics" metrics_dir.mkdir(parents=True, exist_ok=True) csv_path = metrics_dir / "checkpoint_evaluation.csv" + fieldnames = [ + "iteration", + "steps_to_target", + "reached_target", + ] + + # If a previous version created this CSV with a different header, migrate it. + if csv_path.exists(): + try: + with open(csv_path, "r", newline="") as f: + reader = csv.reader(f) + header = next(reader, None) + + if header is not None and list(header) != fieldnames: + migrated_rows: list[dict[str, Any]] = [] + with open(csv_path, "r", newline="") as f: + dict_reader = csv.DictReader(f) + for row in dict_reader: + migrated_rows.append( + { + "iteration": row.get("iteration"), + "steps_to_target": row.get("steps_to_target"), + "reached_target": row.get("reached_target"), + } + ) + + with open(csv_path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + for row in migrated_rows: + writer.writerow(row) + except Exception: + # Best-effort only; do not fail training on migration issues. + pass + file_exists = csv_path.exists() with open(csv_path, "a", newline="") as f: writer = csv.DictWriter( f, - fieldnames=[ - "iteration", - "steps_to_target", - "reached_target", - "max_steps", - "seed", - ], + fieldnames=fieldnames, ) if not file_exists: writer.writeheader() @@ -590,8 +671,6 @@ class PPOTrainer: "iteration": int(iteration), "steps_to_target": int(steps_to_target), "reached_target": bool(reached_target), - "max_steps": int(max_steps), - "seed": int(seed), } ) @@ -601,86 +680,24 @@ class PPOTrainer: if not self.logging_cfg.evaluate_checkpoints: return - checkpoint_path = Path(self.run_dir) / "checkpoints" / f"checkpoint_step_{iteration}.flax" - if not checkpoint_path.exists(): - self.logger.warning( - f"[EVAL]: Checkpoint not found at {checkpoint_path} (skipping evaluation)" - ) - return - max_steps = int(self.logging_cfg.eval_max_steps) seed = int(self.logging_cfg.eval_seed) # Run evaluation best-effort; never fail training because evaluation failed. - env = None try: - backend = Backend.MJC + eval_fn = self._get_or_create_eval_rollout_mjx_fn() + steps, reached = eval_fn(self.agent_state.params, seed, max_steps) + steps_to_target = int(steps) + reached_target = bool(reached) - factory = BrittleStarEnvFactory() - raw_env = factory.create_environment( - backend, - self.cfg.morphology, - self.cfg.arena, - self.cfg.environment, - ) - env = BrittleStarEnv( - raw_env, - backend=backend, - config=self.cfg.environment, - morphology_config=self.cfg.morphology, - ) - - action_space = getattr(raw_env, "action_space", None) - action_dim = ( - int(np.asarray(action_space.shape).reshape(-1)[0]) - if action_space is not None and hasattr(action_space, "shape") - else sum(self.cfg.morphology.segments_per_arm) * 2 - ) - - action_low = ( - None - if action_space is None - else np.asarray(action_space.low, dtype=np.float32).ravel() - ) - action_high = ( - None - if action_space is None - else np.asarray(action_space.high, dtype=np.float32).ravel() - ) - - policy = PolicyAgent.from_checkpoint( - checkpoint_path, - action_dim=action_dim, - obs_processor=self.obs_processor, - ) - - result = rollout_headless( - env=env, - policy=policy, - seed=seed, - max_steps=max_steps, - action_low=action_low, - action_high=action_high, - action_mask=None, - ) - - steps_to_target = int(result.length) if result.reached_target else int(max_steps) csv_path = self._append_checkpoint_eval_row( iteration=iteration, - reached_target=bool(result.reached_target), + reached_target=reached_target, steps_to_target=steps_to_target, - max_steps=max_steps, - seed=seed, ) self._maybe_sync_csv_to_wandb(csv_path) except Exception as e: self.logger.warning(f"[EVAL]: Checkpoint evaluation failed: {e}") - finally: - try: - if env is not None: - env.close() - except Exception: - pass def train(self): """ From 67c459ca94d230f2f9ce4fdfbe2dd76975cbff6f Mon Sep 17 00:00:00 2001 From: Jona Reynaert Date: Thu, 7 May 2026 12:03:53 +0200 Subject: [PATCH 4/8] fix: changed checkpoint evaluation csv fields --- .../trainers/PPOTrainer.py | 85 +++++++++++++++---- 1 file changed, 67 insertions(+), 18 deletions(-) diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 54653a1..9629e60 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -355,15 +355,18 @@ class PPOTrainer: 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 = carry + t, _state, done, _return_ = carry return jnp.logical_and(t < max_steps, jnp.logical_not(done)) def body(carry): - t, state, _done = carry + t, state, _done, return_ = carry obs = obs_processor(state.observations) hidden = sensor_apply(params["sensor_params"], obs) @@ -372,11 +375,22 @@ class PPOTrainer: action = jnp.clip(mean, action_low, action_high) next_state = step_1(state=state, action=action) - done_next = jnp.squeeze(next_state.terminated | next_state.truncated) - return (t + 1, next_state, done_next) + # Match training's shaped reward as closely as possible. + shaped_reward = _reward_fn(state, next_state) + return_ = return_ + jnp.squeeze(shaped_reward) - t, _state, done = jax.lax.while_loop(cond, body, (t0, state, done0)) - return t, done + 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"]) + final_xy_dist = jnp.where(reached_target, 0.0, final_xy_dist_raw) + + return t, reached_target, return_, final_xy_dist, initial_xy_dist self._eval_rollout_mjx_fn = jax.jit(_eval_rollout) return self._eval_rollout_mjx_fn @@ -616,16 +630,24 @@ class PPOTrainer: self, *, iteration: int, + trained_timesteps: int, + eval_steps: int, + eval_return: float, + final_xy_dist: float, + initial_xy_dist: float, reached_target: bool, - steps_to_target: int, ) -> Path: metrics_dir = Path(self.run_dir) / "metrics" metrics_dir.mkdir(parents=True, exist_ok=True) csv_path = metrics_dir / "checkpoint_evaluation.csv" fieldnames = [ - "iteration", - "steps_to_target", + "checkpoint", + "trained_timesteps", + "eval_steps", + "eval_return", + "final_xy_dist", + "initial_xy_dist", "reached_target", ] @@ -641,10 +663,17 @@ class PPOTrainer: with open(csv_path, "r", newline="") as f: dict_reader = csv.DictReader(f) for row in dict_reader: + # Support older schemas best-effort. + checkpoint = row.get("checkpoint", row.get("iteration")) + steps = row.get("eval_steps", row.get("steps_to_target")) migrated_rows.append( { - "iteration": row.get("iteration"), - "steps_to_target": row.get("steps_to_target"), + "checkpoint": checkpoint, + "trained_timesteps": row.get("trained_timesteps"), + "eval_steps": steps, + "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"), } ) @@ -668,15 +697,19 @@ class PPOTrainer: writer.writeheader() writer.writerow( { - "iteration": int(iteration), - "steps_to_target": int(steps_to_target), + "checkpoint": int(iteration), + "trained_timesteps": int(trained_timesteps), + "eval_steps": int(eval_steps), + "eval_return": float(eval_return), + "final_xy_dist": float(final_xy_dist), + "initial_xy_dist": float(initial_xy_dist), "reached_target": bool(reached_target), } ) return csv_path - def _evaluate_checkpoint(self, iteration: int) -> None: + def _evaluate_checkpoint(self, iteration: int, *, trained_timesteps: int) -> None: if not self.logging_cfg.evaluate_checkpoints: return @@ -686,14 +719,30 @@ class PPOTrainer: # Run evaluation best-effort; never fail training because evaluation failed. try: eval_fn = self._get_or_create_eval_rollout_mjx_fn() - steps, reached = eval_fn(self.agent_state.params, seed, max_steps) - steps_to_target = int(steps) + ( + steps, + reached, + eval_return, + final_xy_dist, + initial_xy_dist, + ) = eval_fn(self.agent_state.params, seed, max_steps) + + eval_steps = int(steps) reached_target = bool(reached) + # Keep numeric conversions explicit (JAX scalars -> Python scalars). + eval_return_f = float(eval_return) + final_xy_dist_f = float(final_xy_dist) + initial_xy_dist_f = float(initial_xy_dist) + csv_path = self._append_checkpoint_eval_row( iteration=iteration, + trained_timesteps=int(trained_timesteps), + eval_steps=eval_steps, + eval_return=eval_return_f, + final_xy_dist=final_xy_dist_f, + initial_xy_dist=initial_xy_dist_f, reached_target=reached_target, - steps_to_target=steps_to_target, ) self._maybe_sync_csv_to_wandb(csv_path) except Exception as e: @@ -752,7 +801,7 @@ class PPOTrainer: if self.logging_cfg.save_checkpoints and self.logging_cfg.checkpoint_frequency > 0: if iteration % self.logging_cfg.checkpoint_frequency == 0: self._save_checkpoint(iteration) - self._evaluate_checkpoint(iteration) + self._evaluate_checkpoint(iteration, trained_timesteps=global_step) if getattr(self.cfg.experiment, "debug_sanity", False): self.logger.info("\n[SANITY CHECK] Successfully completed 1 epoch") From b91908bf7a9fd73a94c31a7e03f345356fa8bce1 Mon Sep 17 00:00:00 2001 From: Jona Reynaert Date: Thu, 7 May 2026 13:45:03 +0200 Subject: [PATCH 5/8] feat: abstracted checkpoint eval configs --- configs/evaluation/default.yaml | 8 ++++++ configs/logging/default.yaml | 5 +--- configs/main_config.yaml | 1 + .../configs/config_evaluation.py | 24 ++++++++++++++++++ .../configs/main_config.py | 2 ++ .../configs/register_configs.py | 2 ++ .../trainers/PPOTrainer.py | 18 ++++++++++--- src/experiment_logger/config_logger.py | 25 ++----------------- 8 files changed, 55 insertions(+), 30 deletions(-) create mode 100644 configs/evaluation/default.yaml create mode 100644 src/brittle_star_project/configs/config_evaluation.py diff --git a/configs/evaluation/default.yaml b/configs/evaluation/default.yaml new file mode 100644 index 0000000..ac49e46 --- /dev/null +++ b/configs/evaluation/default.yaml @@ -0,0 +1,8 @@ +# Default Evaluation Configuration +# Settings used for checkpoint evaluation during training. + +evaluate_checkpoints: false +# Max number of control steps during evaluation rollout. +eval_max_steps: 5000 +# Seed for deterministic evaluation reset. +eval_seed: 0 diff --git a/configs/logging/default.yaml b/configs/logging/default.yaml index af9bc6d..2f7de8c 100644 --- a/configs/logging/default.yaml +++ b/configs/logging/default.yaml @@ -10,7 +10,4 @@ save_checkpoints: true checkpoint_frequency: 100 upload_final_model: false upload_checkpoints: false -hf_entity: "" -evaluate_checkpoints: false -eval_max_steps: 5000 -eval_seed: 0 \ No newline at end of file +hf_entity: "" \ No newline at end of file diff --git a/configs/main_config.yaml b/configs/main_config.yaml index b7cb620..44451f4 100644 --- a/configs/main_config.yaml +++ b/configs/main_config.yaml @@ -6,6 +6,7 @@ defaults: - brittle_star_config - experiment: base - logging: default + - evaluation: default - ppo: default - architecture: centralized - morphology: 5_arms_full diff --git a/src/brittle_star_project/configs/config_evaluation.py b/src/brittle_star_project/configs/config_evaluation.py new file mode 100644 index 0000000..4eeadbc --- /dev/null +++ b/src/brittle_star_project/configs/config_evaluation.py @@ -0,0 +1,24 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class EvaluationConfig: + """Evaluation settings. + + Currently used for synchronous checkpoint evaluation during training. + """ + + # When enabled, each saved checkpoint is evaluated headlessly and the results + # are appended to a CSV in the run's metrics/ folder. + evaluate_checkpoints: bool = False + eval_max_steps: int = 5000 + eval_seed: int = 0 + + def __post_init__(self) -> None: + if self.evaluate_checkpoints and self.eval_max_steps <= 0: + raise ValueError( + "Configuration Error: 'eval_max_steps' must be > 0 when " + "'evaluate_checkpoints' is enabled." + ) diff --git a/src/brittle_star_project/configs/main_config.py b/src/brittle_star_project/configs/main_config.py index 10fd22e..7caa1aa 100644 --- a/src/brittle_star_project/configs/main_config.py +++ b/src/brittle_star_project/configs/main_config.py @@ -2,6 +2,7 @@ from dataclasses import dataclass, field from experiment_logger.config_logger import LoggingConfig from brittle_star_project.configs.config_experiment import ExperimentConfig +from brittle_star_project.configs.config_evaluation import EvaluationConfig from brittle_star_project.configs.config_ppo import PPOConfig from brittle_star_project.configs.config_architecture import ArchitectureConfig from brittle_star_project.configs.config_simulation import SimulationSettings @@ -23,6 +24,7 @@ class BrittleStarConfig: experiment: ExperimentConfig = field(default_factory=ExperimentConfig) logging: LoggingConfig = field(default_factory=LoggingConfig) + evaluation: EvaluationConfig = field(default_factory=EvaluationConfig) ppo: PPOConfig = field(default_factory=PPOConfig) # This field is polymorphic; defaults to the base class to allow subclasses # (CentralizedConfig, DecentralizedConfig) to be merged in via Hydra. diff --git a/src/brittle_star_project/configs/register_configs.py b/src/brittle_star_project/configs/register_configs.py index 147da10..9522212 100644 --- a/src/brittle_star_project/configs/register_configs.py +++ b/src/brittle_star_project/configs/register_configs.py @@ -2,6 +2,7 @@ from hydra.core.config_store import ConfigStore from experiment_logger.config_logger import LoggingConfig from brittle_star_project.configs.config_experiment import ExperimentConfig +from brittle_star_project.configs.config_evaluation import EvaluationConfig from brittle_star_project.configs.config_ppo import PPOConfig from brittle_star_project.configs.config_architecture import ( CentralizedConfig, @@ -32,6 +33,7 @@ def register_configs() -> None: # Sub-config groups — each group corresponds to a configs/ subdirectory. cs.store(group="experiment", name="base_experiment", node=ExperimentConfig) cs.store(group="logging", name="base_logging", node=LoggingConfig) + cs.store(group="evaluation", name="base_evaluation", node=EvaluationConfig) cs.store(group="ppo", name="base_ppo", node=PPOConfig) # Architecture variants — swap via CLI: architecture=decentralized diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 9629e60..5ae20ec 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -263,6 +263,7 @@ class PPOTrainer: self.ppo = cfg.ppo self.experiment = cfg.experiment self.logging_cfg = cfg.logging + self.evaluation_cfg = cfg.evaluation self.env = env self.run_dir = run_dir self.run_name = run_name @@ -710,11 +711,22 @@ class PPOTrainer: return csv_path def _evaluate_checkpoint(self, iteration: int, *, trained_timesteps: int) -> None: - if not self.logging_cfg.evaluate_checkpoints: + if not self.evaluation_cfg.evaluate_checkpoints: return - max_steps = int(self.logging_cfg.eval_max_steps) - seed = int(self.logging_cfg.eval_seed) + max_steps = int(self.evaluation_cfg.eval_max_steps) + seed = int(self.evaluation_cfg.eval_seed) + + if max_steps <= 0: + self.logger.warning("[EVAL]: eval_max_steps must be > 0; skipping evaluation") + return + + if not self.logging_cfg.save_checkpoints or self.logging_cfg.checkpoint_frequency <= 0: + self.logger.warning( + "[EVAL]: evaluate_checkpoints is enabled but checkpoint saving is disabled; " + "skipping evaluation" + ) + return # Run evaluation best-effort; never fail training because evaluation failed. try: diff --git a/src/experiment_logger/config_logger.py b/src/experiment_logger/config_logger.py index d158e84..fd77a28 100644 --- a/src/experiment_logger/config_logger.py +++ b/src/experiment_logger/config_logger.py @@ -18,13 +18,6 @@ class LoggingConfig: upload_final_model: bool = False upload_checkpoints: bool = False - # Checkpoint evaluation (synchronous, in-process) - # When enabled, each saved checkpoint is evaluated headlessly and the results - # are appended to a CSV in the run's metrics/ folder. - evaluate_checkpoints: bool = False - eval_max_steps: int = 5000 - eval_seed: int = 0 - hf_entity: str = "" def __post_init__(self): @@ -39,19 +32,5 @@ class LoggingConfig: "both 'track' and 'save_checkpoints' to also be True." ) - if self.evaluate_checkpoints: - if not self.save_checkpoints: - raise ValueError( - "Configuration Error: 'evaluate_checkpoints' is True, but it requires " - "'save_checkpoints' to also be True." - ) - if self.checkpoint_frequency <= 0: - raise ValueError( - "Configuration Error: 'evaluate_checkpoints' is True, but it requires " - "'checkpoint_frequency' to be > 0." - ) - if self.eval_max_steps <= 0: - raise ValueError( - "Configuration Error: 'eval_max_steps' must be > 0 when " - "'evaluate_checkpoints' is enabled." - ) + # NOTE: Checkpoint evaluation settings live under the project's + # `evaluation` config group (see brittle_star_project.configs). From 14d719fe7443173613d0f94e3f46203494b7a8b3 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Thu, 7 May 2026 22:48:23 +0200 Subject: [PATCH 6/8] 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 From a7db4c77bbec06649653a9314289c39516933d40 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Thu, 7 May 2026 22:50:38 +0200 Subject: [PATCH 7/8] refactor(log): extract w&b file sync --- src/experiment_logger/simple_logger.py | 4 ++++ src/experiment_logger/unified_logger.py | 15 +++++++++++++++ 2 files changed, 19 insertions(+) diff --git a/src/experiment_logger/simple_logger.py b/src/experiment_logger/simple_logger.py index 7e4a816..71844dc 100644 --- a/src/experiment_logger/simple_logger.py +++ b/src/experiment_logger/simple_logger.py @@ -71,6 +71,10 @@ class SimpleLogger: def save_final_model(self, params: Any, metadata: Optional[Dict[str, Any]] = None): print("[SAVE] Final model would be saved (SimpleLogger: No-Op)") + def sync_file(self, path: Any): + """No-op for SimpleLogger.""" + pass + def finish(self): print(f"[FINISH] SimpleLogger finished for run: {self.run_name}") diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py index f37136b..e308d12 100644 --- a/src/experiment_logger/unified_logger.py +++ b/src/experiment_logger/unified_logger.py @@ -432,6 +432,21 @@ class UnifiedLogger: except Exception as e: self.error(f"Error saving final model: {e}") + def sync_file(self, path: Path) -> None: + """Upload a file to W&B if tracking is enabled. + + Best-effort: logs a warning on failure, never raises. + """ + if self.wandb_run is None: + return + try: + import wandb + + # "Simple sync" behavior: wandb will copy this file into the run. + wandb.save(str(path), base_path=str(path.parent)) + except Exception as e: + self.warning(f"Failed to sync file to W&B: {e}") + def finish(self): """Finalize logging and cleanup.""" # Flush remaining metrics From 1f14ccfd4ac9d4b0ff5666f0bbe53986d0198693 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Thu, 7 May 2026 23:01:56 +0200 Subject: [PATCH 8/8] refactor: use mjx eval in PPOTrainer --- .../evaluation/evaluate_mjx.py | 11 +- .../trainers/PPOTrainer.py | 241 ++++-------------- 2 files changed, 49 insertions(+), 203 deletions(-) diff --git a/src/brittle_star_project/evaluation/evaluate_mjx.py b/src/brittle_star_project/evaluation/evaluate_mjx.py index f281535..0193d99 100644 --- a/src/brittle_star_project/evaluation/evaluate_mjx.py +++ b/src/brittle_star_project/evaluation/evaluate_mjx.py @@ -58,11 +58,6 @@ def build_eval_rollout_fn( ) -> 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: @@ -141,13 +136,13 @@ def evaluate_checkpoint_mjx( """Run one deterministic evaluation episode and return typed metrics. Args: - eval_fn: A JIT-compiled function as returned by :func:`build_eval_rollout_fn`. + eval_fn: A JIT-compiled function as returned by `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 + A `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) @@ -223,7 +218,7 @@ def append_checkpoint_eval_row( 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`. + result: Evaluation result as returned by `evaluate_checkpoint_mjx`. Returns: Absolute path to the CSV file (useful for W&B sync). diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index 5ae20ec..8934c83 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -1,17 +1,15 @@ import datetime import random import time -import csv from dataclasses import asdict, dataclass from functools import partial -from pathlib import Path -from typing import Any import jax import jax.numpy as jnp import numpy as np import optax from flax.training.train_state import TrainState +from typing import Any from experiment_logger import get_logger @@ -19,6 +17,11 @@ from brittle_star_project.configs.main_config import BrittleStarConfig from brittle_star_project.dataclasses import EpisodeStatistics from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper from brittle_star_project.environment.obs_processing import create_obs_processor +from brittle_star_project.evaluation.evaluate_mjx import ( + append_checkpoint_eval_row, + build_eval_rollout_fn, + evaluate_checkpoint_mjx, +) from brittle_star_project.MLPs.mlps import ( Actor, AgentParams, @@ -122,8 +125,13 @@ def _step_once( return (agent_state, episode_stats, next_obs, next_done, key, env_state), storage -def _reward_fn(env_state, next_env_state): - # if delta distance positive ==> brittle star walking away from target +def reward_fn(env_state, next_env_state): + """Shaped reward used during training and checkpoint evaluation. + + Public so that ``evaluation.evaluate_mjx`` can import it and produce + metrics that are directly comparable to training-time returns. + """ + # Positive delta_distance means the brittle star is moving *away* from target. delta_distance = ( next_env_state.observations["xy_distance_to_target"] - env_state.observations["xy_distance_to_target"] @@ -142,7 +150,7 @@ def _reward_fn(env_state, next_env_state): def _step_env_wrapped(episode_stats, env_state, action, env_step_fn, obs_processor): next_env_state = env_step_fn(env_state, action) - reward = _reward_fn(env_state, next_env_state) + reward = reward_fn(env_state, next_env_state) terminated = next_env_state.terminated truncated = next_env_state.truncated done = terminated | truncated @@ -327,74 +335,8 @@ class PPOTrainer: self.episode_stats = self._init_episode_stats() self._init_random() - # Lazily created MJX/JAX evaluation rollout (compiled on first use) - self._eval_rollout_mjx_fn = None - - def _get_or_create_eval_rollout_mjx_fn(self): - if self._eval_rollout_mjx_fn is not None: - return self._eval_rollout_mjx_fn - - # Use the same backend as training (typically MJX). - if getattr(self.env, "backend", None) != Backend.MJX: - self.logger.warning( - f"[EVAL]: Training env backend is {self.env.backend}; " - "MJX evaluation may be unavailable/slow." - ) - - # We vmap over a single environment (batch size 1) for simplicity. - reset_1 = jax.vmap(self.env.raw.reset) - step_1 = jax.vmap(self.env.raw.step) - - action_low = self._action_low - action_high = self._action_high - obs_processor = self.obs_processor - sensor_apply = self.sensor.apply - actor_apply = self.actor.apply - - def _eval_rollout(params, 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) - - action = jnp.clip(mean, action_low, action_high) - next_state = step_1(state=state, action=action) - - # Match training's shaped reward as closely as possible. - 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"]) - final_xy_dist = jnp.where(reached_target, 0.0, final_xy_dist_raw) - - return t, reached_target, return_, final_xy_dist, initial_xy_dist - - self._eval_rollout_mjx_fn = jax.jit(_eval_rollout) - return self._eval_rollout_mjx_fn + # Lazily-built JIT-compiled MJX eval rollout, created on first evaluation. + self._eval_fn = None def _init_random(self): self.logger.info(f"[RANDOM]: Setting random seed to {self.experiment.seed}") @@ -613,104 +555,12 @@ class PPOTrainer: params=self.agent_state.params, step=iteration, metadata=asdict(self.cfg) ) - def _maybe_sync_csv_to_wandb(self, csv_path: Path) -> None: - if not self.logging_cfg.track: - return - try: - import wandb - - if wandb.run is None: - return - - # "Simple sync" behavior: wandb will copy this file into the run. - wandb.save(str(csv_path), base_path=str(csv_path.parent)) - except Exception as e: - self.logger.warning(f"[EVAL]: Failed to sync CSV to wandb: {e}") - - def _append_checkpoint_eval_row( - self, - *, - iteration: int, - trained_timesteps: int, - eval_steps: int, - eval_return: float, - final_xy_dist: float, - initial_xy_dist: float, - reached_target: bool, - ) -> Path: - metrics_dir = Path(self.run_dir) / "metrics" - metrics_dir.mkdir(parents=True, exist_ok=True) - csv_path = metrics_dir / "checkpoint_evaluation.csv" - - fieldnames = [ - "checkpoint", - "trained_timesteps", - "eval_steps", - "eval_return", - "final_xy_dist", - "initial_xy_dist", - "reached_target", - ] - - # If a previous version created this CSV with a different header, migrate it. - if csv_path.exists(): - try: - with open(csv_path, "r", newline="") as f: - reader = csv.reader(f) - header = next(reader, None) - - if header is not None and list(header) != fieldnames: - migrated_rows: list[dict[str, Any]] = [] - with open(csv_path, "r", newline="") as f: - dict_reader = csv.DictReader(f) - for row in dict_reader: - # Support older schemas best-effort. - checkpoint = row.get("checkpoint", row.get("iteration")) - steps = row.get("eval_steps", row.get("steps_to_target")) - migrated_rows.append( - { - "checkpoint": checkpoint, - "trained_timesteps": row.get("trained_timesteps"), - "eval_steps": steps, - "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() - for row in migrated_rows: - writer.writerow(row) - except Exception: - # Best-effort only; do not fail training on migration issues. - pass - - 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": int(eval_steps), - "eval_return": float(eval_return), - "final_xy_dist": float(final_xy_dist), - "initial_xy_dist": float(initial_xy_dist), - "reached_target": bool(reached_target), - } - ) - - return csv_path - def _evaluate_checkpoint(self, iteration: int, *, trained_timesteps: int) -> None: + """Evaluate the current checkpoint and persist metrics to CSV. + + Delegates all evaluation logic to `evaluation.evaluate_mjx`. + Best-effort: a failure here must never abort training. + """ if not self.evaluation_cfg.evaluate_checkpoints: return @@ -728,35 +578,36 @@ class PPOTrainer: ) return - # Run evaluation best-effort; never fail training because evaluation failed. try: - eval_fn = self._get_or_create_eval_rollout_mjx_fn() - ( - steps, - reached, - eval_return, - final_xy_dist, - initial_xy_dist, - ) = eval_fn(self.agent_state.params, seed, max_steps) + if self._eval_fn is None: + if getattr(self.env, "backend", None) != Backend.MJX: + self.logger.warning( + f"[EVAL]: Training env backend is {self.env.backend}; " + "MJX evaluation may be unavailable/slow." + ) + self._eval_fn = build_eval_rollout_fn( + env=self.env, + obs_processor=self.obs_processor, + sensor_apply=self.sensor.apply, + actor_apply=self.actor.apply, + action_low=self._action_low, + action_high=self._action_high, + reward_fn=reward_fn, + ) - eval_steps = int(steps) - reached_target = bool(reached) - - # Keep numeric conversions explicit (JAX scalars -> Python scalars). - eval_return_f = float(eval_return) - final_xy_dist_f = float(final_xy_dist) - initial_xy_dist_f = float(initial_xy_dist) - - csv_path = self._append_checkpoint_eval_row( + result = evaluate_checkpoint_mjx( + self._eval_fn, + self.agent_state.params, + seed=seed, + max_steps=max_steps, + ) + csv_path = append_checkpoint_eval_row( + self.run_dir, iteration=iteration, trained_timesteps=int(trained_timesteps), - eval_steps=eval_steps, - eval_return=eval_return_f, - final_xy_dist=final_xy_dist_f, - initial_xy_dist=initial_xy_dist_f, - reached_target=reached_target, + result=result, ) - self._maybe_sync_csv_to_wandb(csv_path) + self.logger.sync_file(csv_path) except Exception as e: self.logger.warning(f"[EVAL]: Checkpoint evaluation failed: {e}")