diff --git a/configs/evaluation/poster.yaml b/configs/evaluation/poster.yaml new file mode 100644 index 0000000..048a785 --- /dev/null +++ b/configs/evaluation/poster.yaml @@ -0,0 +1,20 @@ +# @package evaluation +# Configuration for the models used in the poster comparison. + +# Standard evaluation settings +evaluate_checkpoints: false +eval_max_steps: 5000 +eval_seed: 0 + +# Cross-model comparison settings +# We use 10 episodes to get a more robust average for the final poster results. +comparison_base_seed: 0 +comparison_num_episodes: 10 +comparison_output_csv: "metrics/poster_comparison.csv" + +# Paths to the .cleanrl_model files to be compared (relative to workspace root). +# These are placeholders; replace with actual trained model paths for the poster. +comparison_models: + - "experiments/poster/centralized.cleanrl_model" + - "experiments/poster/decentralized.cleanrl_model" + - "experiments/poster/decentralized_amputated.cleanrl_model" diff --git a/src/brittle_star_project/configs/config_evaluation.py b/src/brittle_star_project/configs/config_evaluation.py index 4eeadbc..69c6609 100644 --- a/src/brittle_star_project/configs/config_evaluation.py +++ b/src/brittle_star_project/configs/config_evaluation.py @@ -1,6 +1,6 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field @dataclass @@ -16,6 +16,16 @@ class EvaluationConfig: eval_max_steps: int = 5000 eval_seed: int = 0 + # Cross-model comparison settings. + # comparison_base_seed is the starting seed for generating episode seeds. + comparison_base_seed: int = 0 + # comparison_num_episodes controls how many target positions to evaluate for each model. + comparison_num_episodes: int = 5 + # comparison_models lists the paths (relative to workspace root) to the .cleanrl_model files. + 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" + def __post_init__(self) -> None: if self.evaluate_checkpoints and self.eval_max_steps <= 0: raise ValueError( diff --git a/src/brittle_star_project/evaluation/__init__.py b/src/brittle_star_project/evaluation/__init__.py index 2922d04..f9ed595 100644 --- a/src/brittle_star_project/evaluation/__init__.py +++ b/src/brittle_star_project/evaluation/__init__.py @@ -7,6 +7,7 @@ from .evaluate_mjx import ( build_eval_rollout_fn, evaluate_checkpoint_mjx, ) +from .evaluate import evaluate_policy from .policy import PolicyAgent, ControlPolicy from .rollout import rollout_headless, rollout_viewer, EpisodeResult from .video import record_episode, create_evaluation_dir, save_evaluation_metadata @@ -22,6 +23,8 @@ __all__ = [ "append_checkpoint_eval_row", "build_eval_rollout_fn", "evaluate_checkpoint_mjx", + # CPU evaluation + "evaluate_policy", # policy "PolicyAgent", "ControlPolicy", diff --git a/src/brittle_star_project/evaluation/evaluate.py b/src/brittle_star_project/evaluation/evaluate.py new file mode 100644 index 0000000..12cc6fc --- /dev/null +++ b/src/brittle_star_project/evaluation/evaluate.py @@ -0,0 +1,58 @@ +"""MJC-based (CPU) checkpoint evaluation. + +This module provides the CPU-bound evaluation path using the standard MJC backend. +It is primarily used by the `evaluate_checkpoints` CLI to compute metrics and +render videos. +""" + +from pathlib import Path + +import numpy as np + +from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper +from brittle_star_project.environment.obs_processing import create_obs_processor +from brittle_star_project.evaluation.policy import PolicyAgent +from brittle_star_project.evaluation.rollout import EpisodeResult, rollout_headless + + +def evaluate_policy( + env: BrittleStarJaxEnvWrapper, + policy_path: str | Path, + seed: int, + max_steps: int, +) -> EpisodeResult: + """Evaluate a trained policy in a CPU-bound environment. + + Args: + env: Initialised CPU environment (MJC backend). + policy_path: Path to the ``.cleanrl_model`` weights file. + seed: Random seed for environment reset. + max_steps: Maximum number of control steps. + + Returns: + Structured result containing return, length, and distance metrics. + """ + obs_processor = create_obs_processor( + bounds_dict=env.cfg.obs_bounds.to_bounds_dict(), + padding_masks=env.padding_masks, + ) + + action_dim = env.single_action_space.shape[0] + + policy = PolicyAgent.from_checkpoint( + model_path=Path(policy_path), + action_dim=action_dim, + obs_processor=obs_processor, + ) + + action_low = np.asarray(env.single_action_space.low, dtype=np.float32) + action_high = np.asarray(env.single_action_space.high, dtype=np.float32) + + return rollout_headless( + env=env, + policy=policy, + seed=seed, + max_steps=max_steps, + action_low=action_low, + action_high=action_high, + ) diff --git a/src/brittle_star_project/evaluation/rollout.py b/src/brittle_star_project/evaluation/rollout.py index 79c86ed..a8eb8db 100644 --- a/src/brittle_star_project/evaluation/rollout.py +++ b/src/brittle_star_project/evaluation/rollout.py @@ -17,6 +17,7 @@ class EpisodeResult: length: int reached_target: bool final_xy_dist: float | None + initial_target_distance: float | None def _get_observations(state: Any) -> dict[str, Any] | None: @@ -61,6 +62,7 @@ def rollout_headless( ep_return = 0.0 observations = _get_observations(state) prev_dist = _get_xy_distance_to_target(observations) if observations else None + initial_target_distance = prev_dist reached_target = _target_reached(state=state) steps = 0 @@ -91,6 +93,7 @@ def rollout_headless( length=steps, reached_target=reached_target, final_xy_dist=final_dist, + initial_target_distance=initial_target_distance, )