diff --git a/.gitignore b/.gitignore index 83cfc3f..dd76b50 100644 --- a/.gitignore +++ b/.gitignore @@ -5,6 +5,7 @@ wandb/ outputs/ multirun/ metrics/ +adjacency_debug.txt # Python-generated files __pycache__/ diff --git a/configs/evaluation/poster.yaml b/configs/evaluation/poster.yaml index a09907f..0777cd1 100644 --- a/configs/evaluation/poster.yaml +++ b/configs/evaluation/poster.yaml @@ -15,3 +15,9 @@ comparison_output_csv: "runs/evaluation/comparison.csv" # Paths to the .cleanrl_model files to be compared (relative to workspace root). comparison_models: - "runs/input-space-2-arms/2026-05-02/08-14-58/final_model.flax" + +# Path to the morphologies to evaluate against. +comparison_morphologies: + - "configs/morphology/5_arms_full.yaml" + - "configs/morphology/3_arms.yaml" + - "configs/morphology/2_arms.yaml" diff --git a/scripts/compare_models.py b/scripts/compare_models.py index ceebb89..b3e8338 100644 --- a/scripts/compare_models.py +++ b/scripts/compare_models.py @@ -25,20 +25,23 @@ 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 import build_eval_env 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", + "architecture", + "arm_0", + "arm_1", + "arm_2", + "arm_3", + "arm_4", + "num_active_arms", "seed", "reached_target", "episode_length", @@ -100,7 +103,6 @@ def main(dict_cfg: DictConfig) -> None: 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: @@ -109,80 +111,52 @@ def main(dict_cfg: DictConfig) -> None: 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, - ) + # Determine morphologies to evaluate + # If comparison_morphologies is empty, use the model's training morphology + morphologies = [None] + if eval_cfg.comparison_morphologies: + morphologies = [ + Path(hydra.utils.to_absolute_path(m)) for m in eval_cfg.comparison_morphologies + ] - # --- 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, - ) + for morph_path in morphologies: + morph_label = morph_path.name if morph_path else "training" + logger.info(f" Morphology: {morph_label}") - 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)" + bundle = build_eval_env( + model_path=model_path, + training=training, + metadata=metadata, + morphology_override_path=morph_path, ) - writer.writerow( - { + for seed in seeds: + t0 = time.time() + result = rollout_headless( + env=bundle.env, + policy=bundle.policy, + seed=seed, + max_steps=max_steps, + action_low=bundle.action_low, + action_high=bundle.action_high, + action_mask=bundle.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"({elapsed:.1f}s)" + ) + + row = { "model_path": model_path_str, + "architecture": bundle.architecture, + "num_active_arms": bundle.num_active_arms, "seed": seed, "reached_target": result.reached_target, "episode_length": result.length, @@ -191,10 +165,14 @@ def main(dict_cfg: DictConfig) -> None: "final_xy_dist": result.final_xy_dist, "approx_max_velocity": velocity, } - ) - csv_file.flush() + # Add per-arm segments + for i, segs in enumerate(bundle.segments_per_arm): + row[f"arm_{i}"] = segs - env.close() + writer.writerow(row) + csv_file.flush() + + bundle.env.close() logger.info(f"Done. Results saved to {output_path}") diff --git a/src/brittle_star_project/configs/config_evaluation.py b/src/brittle_star_project/configs/config_evaluation.py index 69c6609..c9c248b 100644 --- a/src/brittle_star_project/configs/config_evaluation.py +++ b/src/brittle_star_project/configs/config_evaluation.py @@ -25,6 +25,10 @@ class EvaluationConfig: comparison_models: list[str] = field(default_factory=list) # Path where the comparison results CSV will be saved (relative to workspace root). comparison_output_csv: str = "metrics/model_comparison.csv" + # Morphology override YAML paths for cross-morphology comparison. + # Each path points to a file in configs/morphology/ (e.g., "configs/morphology/3_arms.yaml"). + # When empty, each model is evaluated only on its training morphology. + comparison_morphologies: list[str] = field(default_factory=list) def __post_init__(self) -> None: if self.evaluate_checkpoints and self.eval_max_steps <= 0: