170 lines
6 KiB
Python
170 lines
6 KiB
Python
"""Re-evaluate saved checkpoints from a completed training run using MJX.
|
|
|
|
This script scans the checkpoint directory of a training run (the `checkpoints/`
|
|
folder inside a Hydra output directory), loads each `.flax` checkpoint, runs
|
|
one deterministic evaluation episode with `build_eval_rollout_fn`, and appends
|
|
the result to the run's `metrics/checkpoint_evaluation.csv`.
|
|
|
|
It is intended for post-training analysis when per-checkpoint evaluation was not
|
|
enabled during training (`evaluate_checkpoints: false`).
|
|
|
|
Usage:
|
|
python scripts/evaluate_checkpoints.py \
|
|
simulation.model_path=runs/2024-01-01/12-00-00/final_model.flax \
|
|
evaluation.eval_max_steps=5000 \
|
|
evaluation.eval_seed=0
|
|
|
|
The script resolves the run directory from `simulation.model_path`, discovers
|
|
all `*.flax` checkpoints under `checkpoints/`, and evaluates them in order.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import re
|
|
from pathlib import Path
|
|
|
|
import hydra
|
|
import jax
|
|
import numpy as np
|
|
from omegaconf import DictConfig, OmegaConf
|
|
|
|
from brittle_star_project.configs.main_config import BrittleStarConfig
|
|
from brittle_star_project.configs.register_configs import register_configs
|
|
from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
|
|
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,
|
|
load_params,
|
|
metadata_to_configs,
|
|
)
|
|
from brittle_star_project.evaluation.evaluate_mjx import (
|
|
append_checkpoint_eval_row,
|
|
build_eval_rollout_fn,
|
|
evaluate_checkpoint_mjx,
|
|
)
|
|
from brittle_star_project.trainers.PPOTrainer import reward_fn
|
|
|
|
|
|
def _parse_iteration(checkpoint_path: Path) -> int:
|
|
"""Parse the iteration number from a checkpoint filename like `checkpoint_0042.flax`."""
|
|
match = re.search(r"(\d+)", checkpoint_path.stem)
|
|
return int(match.group(1)) if match else -1
|
|
|
|
|
|
@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)
|
|
)
|
|
sim_cfg = cfg.simulation
|
|
eval_cfg = cfg.evaluation
|
|
|
|
# --- Resolve the model path to find the run directory ---
|
|
model_path_str = sim_cfg.model_path
|
|
if model_path_str is None:
|
|
raise ValueError(
|
|
"simulation.model_path must point to the final_model.flax of a training run."
|
|
)
|
|
|
|
model_path = Path(hydra.utils.to_absolute_path(model_path_str))
|
|
run_dir = model_path.parent
|
|
|
|
checkpoints_dir = run_dir / "checkpoints"
|
|
if not checkpoints_dir.exists():
|
|
raise FileNotFoundError(
|
|
f"No checkpoints/ directory found in run directory: {run_dir}\n"
|
|
"Make sure simulation.model_path points to a completed training run."
|
|
)
|
|
|
|
checkpoints = sorted(checkpoints_dir.glob("*.flax"), key=_parse_iteration)
|
|
if not checkpoints:
|
|
raise FileNotFoundError(f"No .flax checkpoints found in {checkpoints_dir}")
|
|
|
|
logger.info(f"Found {len(checkpoints)} checkpoint(s) in {checkpoints_dir}")
|
|
|
|
# --- Load sidecar metadata + reconstruct training config ---
|
|
metadata_override = (
|
|
Path(hydra.utils.to_absolute_path(sim_cfg.metadata_path))
|
|
if sim_cfg.metadata_path is not None
|
|
else None
|
|
)
|
|
metadata = load_metadata(model_path, metadata_override)
|
|
training = metadata_to_configs(metadata)
|
|
|
|
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,
|
|
)
|
|
|
|
env = BrittleStarJaxEnvWrapper(
|
|
morphology=training.morphology,
|
|
arena=training.arena,
|
|
env_config=training.environment,
|
|
num_envs=1,
|
|
)
|
|
|
|
action_low = np.asarray(env.single_action_space.low, dtype=np.float32)
|
|
action_high = np.asarray(env.single_action_space.high, dtype=np.float32)
|
|
|
|
from brittle_star_project.MLPs.mlps import Actor, GenericDenseLayersWithActivation
|
|
|
|
sensor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300])
|
|
actor = Actor(action_dim=env.single_action_space.shape[0])
|
|
sensor.apply = jax.jit(sensor.apply)
|
|
actor.apply = jax.jit(actor.apply)
|
|
|
|
eval_fn = build_eval_rollout_fn(
|
|
env=env,
|
|
obs_processor=obs_processor,
|
|
sensor_apply=sensor.apply,
|
|
actor_apply=actor.apply,
|
|
action_low=action_low,
|
|
action_high=action_high,
|
|
reward_fn=reward_fn,
|
|
)
|
|
|
|
seed = int(eval_cfg.eval_seed)
|
|
max_steps = int(eval_cfg.eval_max_steps)
|
|
|
|
logger.info(f"Evaluating each checkpoint (seed={seed}, max_steps={max_steps}).")
|
|
|
|
for checkpoint_path in checkpoints:
|
|
iteration = _parse_iteration(checkpoint_path)
|
|
try:
|
|
params = load_params(checkpoint_path)
|
|
except Exception as e:
|
|
logger.warning(f"Could not load {checkpoint_path.name}: {e}")
|
|
continue
|
|
|
|
result = evaluate_checkpoint_mjx(eval_fn, params, seed=seed, max_steps=max_steps)
|
|
csv_path = append_checkpoint_eval_row(
|
|
run_dir,
|
|
iteration=iteration,
|
|
trained_timesteps=0, # unknown without training logs
|
|
result=result,
|
|
)
|
|
|
|
logger.debug(
|
|
f"checkpoint={iteration:5d} | "
|
|
f"reached={str(result.reached_target):<5} | "
|
|
f"return={result.eval_return:+8.3f} | "
|
|
f"steps={result.steps:4d} | "
|
|
f"final_dist={result.final_xy_dist:.3f}"
|
|
)
|
|
|
|
logger.info(f"Done. CSV at: {csv_path}")
|
|
env.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
register_configs()
|
|
main()
|