"""Compare multiple trained policies across shared evaluation conditions. For each model listed in evaluation.comparison_models, this script runs `comparison_num_episodes` headless rollouts (seeded sequentially from `comparison_base_seed`) and writes a results CSV to `comparison_output_csv`. Results include two metrics per episode: - `eval_return` — shaped reward (same function used during training) - `max_velocity` — approximated as initial_xy_dist / steps taken Usage: # With the default evaluation config python scripts/compare_models.py evaluation=poster # Override the output path on the fly python scripts/compare_models.py evaluation=poster \\ evaluation.comparison_output_csv=metrics/quick_comparison.csv """ from __future__ import annotations import csv import logging import time from pathlib import Path import hydra import numpy as np from omegaconf import DictConfig, OmegaConf from brittle_star_project import Backend, BrittleStarEnv, BrittleStarEnvFactory from brittle_star_project.configs.main_config import BrittleStarConfig from brittle_star_project.configs.register_configs import register_configs from brittle_star_project.environment.obs_processing import create_obs_processor from brittle_star_project.environment.padded_obs_wrapper import compute_padding_masks from brittle_star_project.evaluation.checkpoint import load_metadata, metadata_to_configs from brittle_star_project.evaluation.policy import PolicyAgent from brittle_star_project.evaluation.rollout import rollout_headless _FIELDNAMES = [ "model_path", "seed", "reached_target", "episode_length", "eval_return", "initial_target_distance", "final_xy_dist", "approx_max_velocity", ] def _approx_max_velocity(result) -> float | None: """Approximate max velocity as distance covered per step. This is a rough upper bound: (initial_dist - final_dist) / steps. """ if result.initial_target_distance is None or result.final_xy_dist is None or result.length <= 0: return None dist_covered = result.initial_target_distance - result.final_xy_dist return dist_covered / result.length @hydra.main(config_path="../configs", config_name="main_config", version_base="1.3") def main(dict_cfg: DictConfig) -> None: logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") logger = logging.getLogger(__name__) cfg: BrittleStarConfig = OmegaConf.to_object( OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg) ) eval_cfg = cfg.evaluation model_paths = [str(p) for p in eval_cfg.comparison_models] if not model_paths: raise ValueError( "evaluation.comparison_models is empty. " "Add at least one model path in your evaluation config." ) base_seed = int(eval_cfg.comparison_base_seed) num_episodes = int(eval_cfg.comparison_num_episodes) max_steps = int(eval_cfg.eval_max_steps) seeds = list(range(base_seed, base_seed + num_episodes)) output_path = Path(hydra.utils.to_absolute_path(eval_cfg.comparison_output_csv)) output_path.parent.mkdir(parents=True, exist_ok=True) logger.info( f"Comparing {len(model_paths)} models over {num_episodes} episodes " f"(seeds {seeds[0]}–{seeds[-1]})." ) logger.info(f"Results will be written to: {output_path}") with open(output_path, "w", newline="") as csv_file: writer = csv.DictWriter(csv_file, fieldnames=_FIELDNAMES) writer.writeheader() for model_path_str in model_paths: model_path = Path(hydra.utils.to_absolute_path(model_path_str)) logger.info(f"Evaluating model: {model_path.name}") # --- Load sidecar metadata + reconstruct configs --- try: metadata = load_metadata(model_path) except FileNotFoundError as e: logger.warning(f"Skipping model — {e}") continue training = metadata_to_configs(metadata) # --- Build padding masks and obs_processor --- padding_masks = compute_padding_masks( segments_per_arm=training.morphology.segments_per_arm, reference_segments_per_arm=training.morphology.segments_per_arm, ) obs_processor = create_obs_processor( bounds_dict=training.obs_bounds.to_bounds_dict(), padding_masks=padding_masks, ) # --- Build the CPU environment from training config --- factory = BrittleStarEnvFactory() raw_env = factory.create_environment( Backend.MJC, training.morphology, training.arena, training.environment, ) env = BrittleStarEnv( raw_env, backend=Backend.MJC, config=training.environment, morphology_config=training.morphology, ) trained_action_dim = sum(training.morphology.segments_per_arm) * 2 action_mask = np.asarray(padding_masks["mask_2x"]) action_space = getattr(raw_env, "action_space", None) 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( model_path, action_dim=trained_action_dim, obs_processor=obs_processor, ) # --- Run episodes --- for seed in seeds: t0 = time.time() result = rollout_headless( env=env, policy=policy, seed=seed, max_steps=max_steps, action_low=action_low, action_high=action_high, action_mask=action_mask, ) elapsed = time.time() - t0 velocity = _approx_max_velocity(result) logger.debug( f"seed={seed:3d} | " f"reached={str(result.reached_target):<5} | " f"return={result.return_:+8.3f} | " f"steps={result.length:4d} | " f"final_dist=" f"{'n/a' if result.final_xy_dist is None else f'{result.final_xy_dist:.3f}'} | " f"({elapsed:.1f}s)" ) writer.writerow( { "model_path": model_path_str, "seed": seed, "reached_target": result.reached_target, "episode_length": result.length, "eval_return": result.return_, "initial_target_distance": result.initial_target_distance, "final_xy_dist": result.final_xy_dist, "approx_max_velocity": velocity, } ) csv_file.flush() env.close() logger.info(f"Done. Results saved to {output_path}") if __name__ == "__main__": register_configs() main()