1
Fork 0

feat: abstracted checkpoint eval configs

This commit is contained in:
Jona Reynaert 2026-05-07 13:45:03 +02:00
parent 67c459ca94
commit b91908bf7a
8 changed files with 55 additions and 30 deletions

View file

@ -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

View file

@ -10,7 +10,4 @@ save_checkpoints: true
checkpoint_frequency: 100
upload_final_model: false
upload_checkpoints: false
hf_entity: ""
evaluate_checkpoints: false
eval_max_steps: 5000
eval_seed: 0
hf_entity: ""

View file

@ -6,6 +6,7 @@ defaults:
- brittle_star_config
- experiment: base
- logging: default
- evaluation: default
- ppo: default
- architecture: centralized
- morphology: 5_arms_full

View file

@ -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."
)

View file

@ -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.

View file

@ -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

View file

@ -263,6 +263,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
@ -710,11 +711,22 @@ class PPOTrainer:
return csv_path
def _evaluate_checkpoint(self, iteration: int, *, trained_timesteps: int) -> None:
if not self.logging_cfg.evaluate_checkpoints:
if not self.evaluation_cfg.evaluate_checkpoints:
return
max_steps = int(self.logging_cfg.eval_max_steps)
seed = int(self.logging_cfg.eval_seed)
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
# Run evaluation best-effort; never fail training because evaluation failed.
try:

View file

@ -18,13 +18,6 @@ class LoggingConfig:
upload_final_model: bool = False
upload_checkpoints: bool = False
# Checkpoint evaluation (synchronous, in-process)
# 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
hf_entity: str = ""
def __post_init__(self):
@ -39,19 +32,5 @@ class LoggingConfig:
"both 'track' and 'save_checkpoints' to also be True."
)
if self.evaluate_checkpoints:
if not self.save_checkpoints:
raise ValueError(
"Configuration Error: 'evaluate_checkpoints' is True, but it requires "
"'save_checkpoints' to also be True."
)
if self.checkpoint_frequency <= 0:
raise ValueError(
"Configuration Error: 'evaluate_checkpoints' is True, but it requires "
"'checkpoint_frequency' to be > 0."
)
if self.eval_max_steps <= 0:
raise ValueError(
"Configuration Error: 'eval_max_steps' must be > 0 when "
"'evaluate_checkpoints' is enabled."
)
# NOTE: Checkpoint evaluation settings live under the project's
# `evaluation` config group (see brittle_star_project.configs).