feat(evaluate): evaluate_policy base
This commit is contained in:
parent
cf6751ad15
commit
b88a7ac660
5 changed files with 95 additions and 1 deletions
20
configs/evaluation/poster.yaml
Normal file
20
configs/evaluation/poster.yaml
Normal file
|
|
@ -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"
|
||||||
|
|
@ -1,6 +1,6 @@
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass, field
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
|
|
@ -16,6 +16,16 @@ class EvaluationConfig:
|
||||||
eval_max_steps: int = 5000
|
eval_max_steps: int = 5000
|
||||||
eval_seed: int = 0
|
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:
|
def __post_init__(self) -> None:
|
||||||
if self.evaluate_checkpoints and self.eval_max_steps <= 0:
|
if self.evaluate_checkpoints and self.eval_max_steps <= 0:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
|
|
|
||||||
|
|
@ -7,6 +7,7 @@ from .evaluate_mjx import (
|
||||||
build_eval_rollout_fn,
|
build_eval_rollout_fn,
|
||||||
evaluate_checkpoint_mjx,
|
evaluate_checkpoint_mjx,
|
||||||
)
|
)
|
||||||
|
from .evaluate import evaluate_policy
|
||||||
from .policy import PolicyAgent, ControlPolicy
|
from .policy import PolicyAgent, ControlPolicy
|
||||||
from .rollout import rollout_headless, rollout_viewer, EpisodeResult
|
from .rollout import rollout_headless, rollout_viewer, EpisodeResult
|
||||||
from .video import record_episode, create_evaluation_dir, save_evaluation_metadata
|
from .video import record_episode, create_evaluation_dir, save_evaluation_metadata
|
||||||
|
|
@ -22,6 +23,8 @@ __all__ = [
|
||||||
"append_checkpoint_eval_row",
|
"append_checkpoint_eval_row",
|
||||||
"build_eval_rollout_fn",
|
"build_eval_rollout_fn",
|
||||||
"evaluate_checkpoint_mjx",
|
"evaluate_checkpoint_mjx",
|
||||||
|
# CPU evaluation
|
||||||
|
"evaluate_policy",
|
||||||
# policy
|
# policy
|
||||||
"PolicyAgent",
|
"PolicyAgent",
|
||||||
"ControlPolicy",
|
"ControlPolicy",
|
||||||
|
|
|
||||||
58
src/brittle_star_project/evaluation/evaluate.py
Normal file
58
src/brittle_star_project/evaluation/evaluate.py
Normal file
|
|
@ -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,
|
||||||
|
)
|
||||||
|
|
@ -17,6 +17,7 @@ class EpisodeResult:
|
||||||
length: int
|
length: int
|
||||||
reached_target: bool
|
reached_target: bool
|
||||||
final_xy_dist: float | None
|
final_xy_dist: float | None
|
||||||
|
initial_target_distance: float | None
|
||||||
|
|
||||||
|
|
||||||
def _get_observations(state: Any) -> dict[str, Any] | None:
|
def _get_observations(state: Any) -> dict[str, Any] | None:
|
||||||
|
|
@ -61,6 +62,7 @@ def rollout_headless(
|
||||||
ep_return = 0.0
|
ep_return = 0.0
|
||||||
observations = _get_observations(state)
|
observations = _get_observations(state)
|
||||||
prev_dist = _get_xy_distance_to_target(observations) if observations else None
|
prev_dist = _get_xy_distance_to_target(observations) if observations else None
|
||||||
|
initial_target_distance = prev_dist
|
||||||
reached_target = _target_reached(state=state)
|
reached_target = _target_reached(state=state)
|
||||||
|
|
||||||
steps = 0
|
steps = 0
|
||||||
|
|
@ -91,6 +93,7 @@ def rollout_headless(
|
||||||
length=steps,
|
length=steps,
|
||||||
reached_target=reached_target,
|
reached_target=reached_target,
|
||||||
final_xy_dist=final_dist,
|
final_xy_dist=final_dist,
|
||||||
|
initial_target_distance=initial_target_distance,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
Reference in a new issue