feat(evaluate): evaluation scripts

configs/evaluation/poster.yaml

@@ -9,12 +9,9 @@ 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"
+comparison_num_episodes: 2
+comparison_output_csv: "runs/evaluation/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"
+  - "runs/input-space-2-arms/2026-05-02/08-14-58/final_model.flax"

experiments/debug-experiment-10042026/used_variables.md

@@ -1,107 +0,0 @@
-## Default envconfig
-task: Task = Task.DIRECTED_LOCOMOTION
-simulation_time: float = 500.0
-num_physics_steps_per_control_step: int = 10
-time_scale: int = 2
-camera_ids: list[int] = field(default_factory=lambda: [0, 1])
-render_size: tuple[int, int] = (480, 640)
-joint_randomization_noise_scale: float = 0.0
-target_distance: float = 3.0
-light_perlin_noise_scale: int = 0
-
-
-## Default ppoargs
-seed: int = 1
-torch_deterministic: bool = True
-cuda: bool = True
-track: bool = False
-checkpoint_frequency: int = 100
-learning_rate: float = 2.5e-4
-anneal_lr: bool = True
-gamma: float = 0.99
-gae_lambda: float = 0.95
-update_epochs: int = 4
-norm_adv: bool = True
-clip_vloss: bool = True
-max_grad_norm: float = 0.5
-target_kl: float | None = None
-batch_size: int = 0
-minibatch_size: int = 0
-num_iterations: int = 0
-
-## Used config file: (hpc/debug.yaml)
-exp_name: "debug-experiment"
-seed: 42
-track: true
-wandb_project_name: "Let's-find-that-bug"
-wandb_entity: "SEL3-2026-Groep-4"
-run_dir: "/data/gent/465/vsc46589"
-num_envs: 32
-num_steps: 32
-num_minibatches: 32
-total_timesteps: 409600 
-num_arms: 2
-cuda: true
-
-ent_coef: 0.005
-vf_coef: 1.0
-clip_coef: 0.2
-
-anneal_lr: true
-learning_rate: 0.0003
-
-## Arena config:
-size: tuple[float, float] = (10.0, 5.0)
-sand_ground_color: bool = True
-attach_target: bool = True
-wall_height: float = 1.5
-wall_thickness: float = 0.1
-
-## Morphology:
-num_segments_per_arm: int = 4
-use_p_control: bool = True
-use_torque_control: bool = False
-
-## MLPs:
-### Sensor & Feature_extractor:
-Both with 3 layers of 300 neurons per layer.
-
-class GenericDenseLayersWithActivation(nn.Module):
-    layer_sizes: Sequence[int] = field(default_factory=lambda: [64, 64])
-    activation: Callable = nn.tanh
-
-    @nn.compact
-    def __call__(self, x):
-        for size in self.layer_sizes:
-            x = nn.Dense(size, kernel_init=orthogonal(jnp.sqrt(2)))(x)
-            x = self.activation(x)
-        return x
-
-### Actor:
-class Actor(nn.Module):
-    action_dim: int
-    @nn.compact
-    def __call__(self, x):
-        mean = nn.Dense(self.action_dim, kernel_init=orthogonal(0.01), bias_init=constant(0.0))(x)
-        log_std = self.param("log_std", nn.initializers.zeros, (self.action_dim,))
-        return mean, log_std
-
-### Critic:
-class OneDenseLayerMLP(nn.Module):
-    @nn.compact
-    def __call__(self, x):
-        return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x)
-
-### Observations:
-_ALLOWED_OBS_KEYS = {
-    "joint_position",
-    "joint_velocity",
-    "joint_actuator_force",
-    "actuator_force",
-    "disk_position",
-    "disk_rotation",
-    "disk_linear_velocity",
-    "disk_angular_velocity",
-    "unit_xy_direction_to_target",
-    "xy_distance_to_target",
-}

scripts/compare_models.py

@@ -0,0 +1,204 @@
+"""Compare multiple trained policies across shared evaluation conditions.
+
+For each model listed in evaluation.comparison_models, this script runs
+`comparison_num_episodes` headless rollouts (seeded sequentially from
+`comparison_base_seed`) and writes a results CSV to `comparison_output_csv`.
+
+Results include two metrics per episode:
+- `eval_return`    — shaped reward (same function used during training)
+- `max_velocity`  — approximated as initial_xy_dist / steps taken
+
+Usage:
+    # With the default evaluation config
+    python scripts/compare_models.py evaluation=poster
+
+    # Override the output path on the fly
+    python scripts/compare_models.py evaluation=poster \\
+        evaluation.comparison_output_csv=metrics/quick_comparison.csv
+"""
+
+from __future__ import annotations
+
+import csv
+import logging
+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.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",
+    "seed",
+    "reached_target",
+    "episode_length",
+    "eval_return",
+    "initial_target_distance",
+    "final_xy_dist",
+    "approx_max_velocity",
+]
+
+
+def _approx_max_velocity(result) -> float | None:
+    """Approximate max velocity as distance covered per step.
+
+    This is a rough upper bound: (initial_dist - final_dist) / steps.
+    """
+    if result.initial_target_distance is None or result.final_xy_dist is None or result.length <= 0:
+        return None
+    dist_covered = result.initial_target_distance - result.final_xy_dist
+    return dist_covered / result.length
+
+
+@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)
+    )
+    eval_cfg = cfg.evaluation
+
+    model_paths = [str(p) for p in eval_cfg.comparison_models]
+    if not model_paths:
+        raise ValueError(
+            "evaluation.comparison_models is empty. "
+            "Add at least one model path in your evaluation config."
+        )
+
+    base_seed = int(eval_cfg.comparison_base_seed)
+    num_episodes = int(eval_cfg.comparison_num_episodes)
+    max_steps = int(eval_cfg.eval_max_steps)
+
+    seeds = list(range(base_seed, base_seed + num_episodes))
+
+    output_path = Path(hydra.utils.to_absolute_path(eval_cfg.comparison_output_csv))
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    logger.info(
+        f"Comparing {len(model_paths)} models over {num_episodes} episodes "
+        f"(seeds {seeds[0]}–{seeds[-1]})."
+    )
+    logger.info(f"Results will be written to: {output_path}")
+
+    with open(output_path, "w", newline="") as csv_file:
+        writer = csv.DictWriter(csv_file, fieldnames=_FIELDNAMES)
+        writer.writeheader()
+
+        for model_path_str in model_paths:
+            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:
+                logger.warning(f"Skipping model — {e}")
+                continue
+
+            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,
+            )
+
+            # --- 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,
+            )
+
+            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)"
+                )
+
+                writer.writerow(
+                    {
+                        "model_path": model_path_str,
+                        "seed": seed,
+                        "reached_target": result.reached_target,
+                        "episode_length": result.length,
+                        "eval_return": result.return_,
+                        "initial_target_distance": result.initial_target_distance,
+                        "final_xy_dist": result.final_xy_dist,
+                        "approx_max_velocity": velocity,
+                    }
+                )
+                csv_file.flush()
+
+            env.close()
+
+    logger.info(f"Done. Results saved to {output_path}")
+
+
+if __name__ == "__main__":
+    register_configs()
+    main()

scripts/evaluate_checkpoints.py

@@ -0,0 +1,170 @@
+"""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()

src/brittle_star_project/evaluation/evaluate.py

@@ -25,7 +25,7 @@ def evaluate_policy(
 
     Args:
         env: Initialised CPU environment (MJC backend).
-        policy_path: Path to the ``.cleanrl_model`` weights file.
+        policy_path: Path to the `.cleanrl_model` weights file.
         seed: Random seed for environment reset.
         max_steps: Maximum number of control steps.
 

src/brittle_star_project/evaluation/evaluate_mjx.py

@@ -12,7 +12,7 @@ The key functions are:
 - `evaluate_checkpoint_mjx` — runs that function for a given set of parameters and returns a typed
   `CheckpointEvalResult`.
 - `append_checkpoint_eval_row` — persists the result to the run's
-  ``metrics/checkpoint_evaluation.csv``, migrating old schemas automatically.
+  `metrics/checkpoint_evaluation.csv`, migrating old schemas automatically.
 """
 
 from __future__ import annotations
@@ -61,17 +61,17 @@ def build_eval_rollout_fn(
     All outputs are JAX arrays. Convert to Python scalars before logging.
 
     Args:
-        env: The training environment wrapper. Must expose ``env.raw`` with
-            ``reset`` and ``step`` methods compatible with ``jax.vmap``.
+        env: The training environment wrapper. Must expose `env.raw` with
+            `reset` and `step` methods compatible with `jax.vmap`.
         obs_processor: Observation normalisation / padding callable, as
-            returned by ``create_obs_processor``.
-        sensor_apply: The sensor network's ``apply`` method (JIT-compiled).
-        actor_apply: The actor network's ``apply`` method (JIT-compiled).
-        action_low: Per-joint action lower bound (JAX array, shape ``(action_dim,)``).
-        action_high: Per-joint action upper bound (JAX array, shape ``(action_dim,)``).
+            returned by `create_obs_processor`.
+        sensor_apply: The sensor network's `apply` method (JIT-compiled).
+        actor_apply: The actor network's `apply` method (JIT-compiled).
+        action_low: Per-joint action lower bound (JAX array, shape `(action_dim,)`).
+        action_high: Per-joint action upper bound (JAX array, shape `(action_dim,)`).
         reward_fn: Shaped reward function with signature
-            ``reward_fn(env_state, next_env_state) -> jnp.ndarray``.
-            Typically the module-level ``reward_fn`` from ``PPOTrainer``.
+            `reward_fn(env_state, next_env_state) -> jnp.ndarray`.
+            Typically the module-level `reward_fn` from `PPOTrainer`.
 
     Returns:
         A JIT-compiled callable that runs one deterministic evaluation episode.
@@ -209,9 +209,9 @@ def append_checkpoint_eval_row(
     trained_timesteps: int,
     result: CheckpointEvalResult,
 ) -> Path:
-    """Append one evaluation row to ``<run_dir>/metrics/checkpoint_evaluation.csv``.
+    """Append one evaluation row to `<run_dir>/metrics/checkpoint_evaluation.csv`.
 
-    Creates the file (including the ``metrics/`` directory) if it does not yet
+    Creates the file (including the `metrics/` directory) if it does not yet
     exist. Migrates the file to the current schema if the header has changed.
 
     Args: