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 1ecf7fd..2f7de8c 100644 --- a/configs/logging/default.yaml +++ b/configs/logging/default.yaml @@ -10,4 +10,4 @@ save_checkpoints: true checkpoint_frequency: 100 upload_final_model: false upload_checkpoints: false -hf_entity: "" +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/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..0193d99 --- /dev/null +++ b/src/brittle_star_project/evaluation/evaluate_mjx.py @@ -0,0 +1,250 @@ +"""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. + + 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 `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 `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 `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 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 9f22ae3..7e7f7b3 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -3,13 +3,13 @@ import random import time from dataclasses import asdict, dataclass from functools import partial -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 @@ -17,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, @@ -26,6 +31,8 @@ from brittle_star_project.MLPs.mlps import ( ) from brittle_star_project.ppo import PPO +from brittle_star_project.environment.env_types import Backend + # TODO: clip scaled reward? @@ -137,8 +144,13 @@ def _step_once( ), 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"] @@ -165,7 +177,7 @@ def _step_env_wrapped( ): next_env_state_pre_reset = env_step_fn(env_state, action) - reward = _reward_fn(env_state, next_env_state_pre_reset) + reward = reward_fn(env_state, next_env_state_pre_reset) terminated = next_env_state_pre_reset.terminated truncated = next_env_state_pre_reset.truncated done = terminated | truncated @@ -351,6 +363,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 @@ -376,6 +389,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( @@ -414,6 +429,8 @@ class PPOTrainer: self.episode_stats = self._init_episode_stats() self._init_random() + # 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}") @@ -634,6 +651,62 @@ class PPOTrainer: params=self.agent_state.params, step=iteration, metadata=asdict(self.cfg) ) + 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 + + 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 + + try: + 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, + ) + + 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), + result=result, + ) + self.logger.sync_file(csv_path) + except Exception as e: + self.logger.warning(f"[EVAL]: Checkpoint evaluation failed: {e}") + def train(self): """ Train the PPO agent for a specified number of iterations. @@ -687,6 +760,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, trained_timesteps=global_step) if getattr(self.cfg.experiment, "debug_sanity", False): self.logger.info("\n[SANITY CHECK] Successfully completed 1 epoch") diff --git a/src/experiment_logger/config_logger.py b/src/experiment_logger/config_logger.py index c34c162..fd77a28 100644 --- a/src/experiment_logger/config_logger.py +++ b/src/experiment_logger/config_logger.py @@ -31,3 +31,6 @@ class LoggingConfig: "Configuration Error: 'upload_checkpoints' is True, but it requires " "both 'track' and 'save_checkpoints' to also be True." ) + + # NOTE: Checkpoint evaluation settings live under the project's + # `evaluation` config group (see brittle_star_project.configs). 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