Merge branch 'dev' into feat/message_passing

configs/evaluation/default.yaml

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

configs/logging/default.yaml

@@ -10,4 +10,4 @@ save_checkpoints: true
 checkpoint_frequency: 100
 upload_final_model: false
 upload_checkpoints: false
-hf_entity: ""
+hf_entity: ""

configs/main_config.yaml

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

src/brittle_star_project/configs/config_evaluation.py

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

src/brittle_star_project/configs/main_config.py

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

src/brittle_star_project/configs/register_configs.py

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

src/brittle_star_project/evaluation/__init__.py

@@ -1,20 +1,35 @@
 from __future__ import annotations
 
 from .checkpoint import load_metadata, load_params, metadata_to_configs, TrainingConfig
+from .evaluate_mjx import (
+    CheckpointEvalResult,
+    append_checkpoint_eval_row,
+    build_eval_rollout_fn,
+    evaluate_checkpoint_mjx,
+)
 from .policy import PolicyAgent, ControlPolicy
 from .rollout import rollout_headless, rollout_viewer, EpisodeResult
 from .video import record_episode, create_evaluation_dir, save_evaluation_metadata
 
 __all__ = [
+    # checkpoint loading
     "load_metadata",
     "load_params",
     "metadata_to_configs",
     "TrainingConfig",
+    # MJX evaluation
+    "CheckpointEvalResult",
+    "append_checkpoint_eval_row",
+    "build_eval_rollout_fn",
+    "evaluate_checkpoint_mjx",
+    # policy
     "PolicyAgent",
     "ControlPolicy",
+    # rollout
     "rollout_headless",
     "rollout_viewer",
     "EpisodeResult",
+    # video
     "record_episode",
     "create_evaluation_dir",
     "save_evaluation_metadata",

src/brittle_star_project/evaluation/evaluate_mjx.py

@@ -0,0 +1,250 @@
+"""MJX-based headless checkpoint evaluation.
+
+This module provides a fast, JIT-compiled evaluation path using the MJX
+(JAX-accelerated MuJoCo) backend. It is intended for evaluating checkpoints
+*during* or *after* a training run, where the environment and policy are
+already fully initialised.
+
+The key functions are:
+
+- `build_eval_rollout_fn` — builds and JIT-compiles a single-episode rollout function from the
+  training environment and policy components.
+- `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.
+"""
+
+from __future__ import annotations
+
+import csv
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable
+
+import jax
+import jax.numpy as jnp
+
+
+@dataclass
+class CheckpointEvalResult:
+    """Structured result from a single MJX checkpoint evaluation episode."""
+
+    steps: int
+    """Number of control steps taken (≤ max_steps)."""
+
+    reached_target: bool
+    """Whether the robot reached the target (terminated) before max_steps."""
+
+    eval_return: float
+    """Accumulated shaped reward over the episode."""
+
+    final_xy_dist: float
+    """XY distance to target at episode end. 0.0 when ``reached_target`` is True."""
+
+    initial_xy_dist: float
+    """XY distance to target at episode start."""
+
+
+def build_eval_rollout_fn(
+    *,
+    env: Any,
+    obs_processor: Callable,
+    sensor_apply: Callable,
+    actor_apply: Callable,
+    action_low: jnp.ndarray,
+    action_high: jnp.ndarray,
+    reward_fn: Callable,
+) -> Callable:
+    """Build and JIT-compile a single-episode MJX evaluation rollout.
+
+    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``.
+        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,)``).
+        reward_fn: Shaped reward function with signature
+            ``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.
+    """
+    # vmap over a batch of 1 so the MJX API is satisfied without any
+    # extra bookkeeping in the caller.
+    reset_1 = jax.vmap(env.raw.reset)
+    step_1 = jax.vmap(env.raw.step)
+
+    def _eval_rollout(params: dict, seed: int, max_steps: int):
+        rng = jax.random.PRNGKey(seed)
+        rngs = jnp.asarray(jax.random.split(rng, 1))
+        state = reset_1(rng=rngs)
+
+        initial_xy_dist = jnp.squeeze(state.observations["xy_distance_to_target"])
+
+        t0 = jnp.asarray(0, dtype=jnp.int32)
+        done0 = jnp.squeeze(state.terminated | state.truncated)
+        return0 = jnp.asarray(0.0, dtype=jnp.float32)
+
+        def cond(carry):
+            t, _state, done, _return_ = carry
+            return jnp.logical_and(t < max_steps, jnp.logical_not(done))
+
+        def body(carry):
+            t, state, _done, return_ = carry
+
+            obs = obs_processor(state.observations)
+            hidden = sensor_apply(params["sensor_params"], obs)
+            mean, _log_std = actor_apply(params["actor_params"], hidden)
+
+            # Deterministic action: use the actor mean, no exploration noise.
+            action = jnp.clip(mean, action_low, action_high)
+            next_state = step_1(state=state, action=action)
+
+            shaped_reward = reward_fn(state, next_state)
+            return_ = return_ + jnp.squeeze(shaped_reward)
+
+            done_next = jnp.squeeze(next_state.terminated | next_state.truncated)
+            return (t + 1, next_state, done_next, return_)
+
+        t, final_state, _done, return_ = jax.lax.while_loop(cond, body, (t0, state, done0, return0))
+
+        reached_target = jnp.squeeze(final_state.terminated)
+        final_xy_dist_raw = jnp.squeeze(final_state.observations["xy_distance_to_target"])
+        # Clamp to 0 when the target was reached so downstream consumers
+        # don't have to special-case "terminated" themselves.
+        final_xy_dist = jnp.where(reached_target, 0.0, final_xy_dist_raw)
+
+        return t, reached_target, return_, final_xy_dist, initial_xy_dist
+
+    return jax.jit(_eval_rollout)
+
+
+def evaluate_checkpoint_mjx(
+    eval_fn: Callable,
+    params: dict,
+    *,
+    seed: int,
+    max_steps: int,
+) -> CheckpointEvalResult:
+    """Run one deterministic evaluation episode and return typed metrics.
+
+    Args:
+        eval_fn: A JIT-compiled function as returned by `build_eval_rollout_fn`.
+        params: Agent parameter dict (e.g. ``agent_state.params``).
+        seed: Random seed for environment reset (controls target placement).
+        max_steps: Maximum number of control steps before the episode is cut off.
+
+    Returns:
+        A `CheckpointEvalResult` with all JAX arrays converted to
+        plain Python scalars.
+    """
+    steps, reached, eval_return, final_xy_dist, initial_xy_dist = eval_fn(params, seed, max_steps)
+    return CheckpointEvalResult(
+        steps=int(steps),
+        reached_target=bool(reached),
+        eval_return=float(eval_return),
+        final_xy_dist=float(final_xy_dist),
+        initial_xy_dist=float(initial_xy_dist),
+    )
+
+
+_FIELDNAMES = [
+    "checkpoint",
+    "trained_timesteps",
+    "eval_steps",
+    "eval_return",
+    "final_xy_dist",
+    "initial_xy_dist",
+    "reached_target",
+]
+
+
+def _migrate_csv_if_needed(csv_path: Path) -> None:
+    """Rewrite the CSV with the canonical field names if the schema changed.
+
+    Best-effort: any exception is silently swallowed so that a schema mismatch
+    never causes a training crash.
+    """
+    try:
+        with open(csv_path, "r", newline="") as f:
+            header = next(csv.reader(f), None)
+
+        if header is None or list(header) == _FIELDNAMES:
+            return  # Nothing to migrate.
+
+        migrated_rows: list[dict[str, Any]] = []
+        with open(csv_path, "r", newline="") as f:
+            for row in csv.DictReader(f):
+                migrated_rows.append(
+                    {
+                        "checkpoint": row.get("checkpoint", row.get("iteration")),
+                        "trained_timesteps": row.get("trained_timesteps"),
+                        "eval_steps": row.get("eval_steps", row.get("steps_to_target")),
+                        "eval_return": row.get("eval_return"),
+                        "final_xy_dist": row.get("final_xy_dist"),
+                        "initial_xy_dist": row.get("initial_xy_dist"),
+                        "reached_target": row.get("reached_target"),
+                    }
+                )
+
+        with open(csv_path, "w", newline="") as f:
+            writer = csv.DictWriter(f, fieldnames=_FIELDNAMES)
+            writer.writeheader()
+            writer.writerows(migrated_rows)
+    except Exception:
+        pass  # Never crash training on a migration issue.
+
+
+def append_checkpoint_eval_row(
+    run_dir: str | Path,
+    *,
+    iteration: int,
+    trained_timesteps: int,
+    result: CheckpointEvalResult,
+) -> Path:
+    """Append one evaluation row to ``<run_dir>/metrics/checkpoint_evaluation.csv``.
+
+    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:
+        run_dir: Root directory of the training run (Hydra's output dir).
+        iteration: Training iteration number, used as the checkpoint identifier.
+        trained_timesteps: Total environment steps taken at this checkpoint.
+        result: Evaluation result as returned by `evaluate_checkpoint_mjx`.
+
+    Returns:
+        Absolute path to the CSV file (useful for W&B sync).
+    """
+    metrics_dir = Path(run_dir) / "metrics"
+    metrics_dir.mkdir(parents=True, exist_ok=True)
+    csv_path = metrics_dir / "checkpoint_evaluation.csv"
+
+    if csv_path.exists():
+        _migrate_csv_if_needed(csv_path)
+
+    file_exists = csv_path.exists()
+    with open(csv_path, "a", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=_FIELDNAMES)
+        if not file_exists:
+            writer.writeheader()
+        writer.writerow(
+            {
+                "checkpoint": int(iteration),
+                "trained_timesteps": int(trained_timesteps),
+                "eval_steps": result.steps,
+                "eval_return": result.eval_return,
+                "final_xy_dist": result.final_xy_dist,
+                "initial_xy_dist": result.initial_xy_dist,
+                "reached_target": result.reached_target,
+            }
+        )
+
+    return csv_path

src/brittle_star_project/evaluation/policy.py

@@ -62,6 +62,28 @@ class PolicyAgent:
         }
         self._obs_processor = obs_processor
 
+    @classmethod
+    def from_params(
+        cls,
+        *,
+        sensor_params: Any,
+        actor_params: Any,
+        action_dim: int,
+        obs_processor: Any,
+    ) -> "PolicyAgent":
+        """Construct a PolicyAgent directly from in-memory parameters."""
+        return cls(
+            sensor_params=sensor_params,
+            actor_params=actor_params,
+            action_dim=action_dim,
+            obs_processor=obs_processor,
+        )
+
+    def set_params(self, *, sensor_params: Any, actor_params: Any) -> None:
+        """Update parameters for evaluation without rebuilding the model."""
+        self._params["sensor_params"] = sensor_params
+        self._params["actor_params"] = actor_params
+
     @classmethod
     def from_checkpoint(
         cls,

src/brittle_star_project/trainers/PPOTrainer.py

@@ -18,7 +18,12 @@ from brittle_star_project.configs.main_config import BrittleStarConfig
 from brittle_star_project.dataclasses import EpisodeStatistics
 from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
 from brittle_star_project.environment.obs_processing import create_obs_processor
-from brittle_star_project.MLPs import (
+from brittle_star_project.evaluation.evaluate_mjx import (
+    append_checkpoint_eval_row,
+    build_eval_rollout_fn,
+    evaluate_checkpoint_mjx,
+)
+from brittle_star_project.MLPs.mlps import (
     Actor,
     AgentParams,
     GenericDenseLayersWithActivation,
@@ -31,6 +36,10 @@ from brittle_star_project.ppo import PPO
 from brittle_star_project.environment import MorphMode
 from brittle_star_project.utils import logged_jit
 
+from brittle_star_project.environment.env_types import Backend
+
+# TODO: clip scaled reward?
+
 
 @logged_jit
 def _clip_action(action: jnp.ndarray, low: jnp.ndarray, high: jnp.ndarray) -> jnp.ndarray:
@@ -176,8 +185,13 @@ def _step_once(
     ), storage
 
 
-def _reward_fn(env_state, next_env_state):
-    # if delta distance positive ==> brittle star walking away from target
+def reward_fn(env_state, next_env_state):
+    """Shaped reward used during training and checkpoint evaluation.
+
+    Public so that ``evaluation.evaluate_mjx`` can import it and produce
+    metrics that are directly comparable to training-time returns.
+    """
+    # Positive delta_distance means the brittle star is moving *away* from target.
     delta_distance = (
         next_env_state.observations["xy_distance_to_target"]
         - env_state.observations["xy_distance_to_target"]
@@ -204,7 +218,7 @@ def _step_env_wrapped(
 ):
     next_env_state_pre_reset = env_step_fn(env_state, action)
 
-    reward = _reward_fn(env_state, next_env_state_pre_reset)
+    reward = reward_fn(env_state, next_env_state_pre_reset)
     terminated = next_env_state_pre_reset.terminated
     truncated = next_env_state_pre_reset.truncated
     done = terminated | truncated
@@ -416,6 +430,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
@@ -464,6 +479,8 @@ class PPOTrainer:
 
         action_low = jnp.asarray(self.env.single_action_space.low, dtype=jnp.float32)
         action_high = jnp.asarray(self.env.single_action_space.high, dtype=jnp.float32)
+        self._action_low = action_low
+        self._action_high = action_high
 
         self._rollout_jit = logged_jit(
             partial(
@@ -526,6 +543,8 @@ class PPOTrainer:
         self.episode_stats = self._init_episode_stats()
 
         self._init_random()
+        # Lazily-built JIT-compiled MJX eval rollout, created on first evaluation.
+        self._eval_fn = None
 
     def _init_random(self):
         self.logger.info(f"[RANDOM]: Setting random seed to {self.experiment.seed}")
@@ -849,6 +868,62 @@ class PPOTrainer:
             params=self.agent_state.params, step=iteration, metadata=asdict(self.cfg)
         )
 
+    def _evaluate_checkpoint(self, iteration: int, *, trained_timesteps: int) -> None:
+        """Evaluate the current checkpoint and persist metrics to CSV.
+
+        Delegates all evaluation logic to `evaluation.evaluate_mjx`.
+        Best-effort: a failure here must never abort training.
+        """
+        if not self.evaluation_cfg.evaluate_checkpoints:
+            return
+
+        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
+
+        try:
+            if self._eval_fn is None:
+                if getattr(self.env, "backend", None) != Backend.MJX:
+                    self.logger.warning(
+                        f"[EVAL]: Training env backend is {self.env.backend}; "
+                        "MJX evaluation may be unavailable/slow."
+                    )
+                self._eval_fn = build_eval_rollout_fn(
+                    env=self.env,
+                    obs_processor=self.obs_processor,
+                    sensor_apply=self.sensor.apply,
+                    actor_apply=self.actor.apply,
+                    action_low=self._action_low,
+                    action_high=self._action_high,
+                    reward_fn=reward_fn,
+                )
+
+            result = evaluate_checkpoint_mjx(
+                self._eval_fn,
+                self.agent_state.params,
+                seed=seed,
+                max_steps=max_steps,
+            )
+            csv_path = append_checkpoint_eval_row(
+                self.run_dir,
+                iteration=iteration,
+                trained_timesteps=int(trained_timesteps),
+                result=result,
+            )
+            self.logger.sync_file(csv_path)
+        except Exception as e:
+            self.logger.warning(f"[EVAL]: Checkpoint evaluation failed: {e}")
+
     def train(self):
         """
         Train the PPO agent for a specified number of iterations.
@@ -905,6 +980,7 @@ class PPOTrainer:
             if self.logging_cfg.save_checkpoints and self.logging_cfg.checkpoint_frequency > 0:
                 if iteration % self.logging_cfg.checkpoint_frequency == 0:
                     self._save_checkpoint(iteration)
+                    self._evaluate_checkpoint(iteration, trained_timesteps=global_step)
 
             if getattr(self.cfg.experiment, "debug_sanity", False):
                 self.logger.info("\n[SANITY CHECK] Successfully completed 1 epoch")

src/experiment_logger/config_logger.py

@@ -31,3 +31,6 @@ class LoggingConfig:
                 "Configuration Error: 'upload_checkpoints' is True, but it requires "
                 "both 'track' and 'save_checkpoints' to also be True."
             )
+
+        # NOTE: Checkpoint evaluation settings live under the project's
+        # `evaluation` config group (see brittle_star_project.configs).

src/experiment_logger/simple_logger.py

@@ -71,6 +71,10 @@ class SimpleLogger:
     def save_final_model(self, params: Any, metadata: Optional[Dict[str, Any]] = None):
         print("[SAVE] Final model would be saved (SimpleLogger: No-Op)")
 
+    def sync_file(self, path: Any):
+        """No-op for SimpleLogger."""
+        pass
+
     def finish(self):
         print(f"[FINISH] SimpleLogger finished for run: {self.run_name}")
 

src/experiment_logger/unified_logger.py

@@ -432,6 +432,21 @@ class UnifiedLogger:
         except Exception as e:
             self.error(f"Error saving final model: {e}")
 
+    def sync_file(self, path: Path) -> None:
+        """Upload a file to W&B if tracking is enabled.
+
+        Best-effort: logs a warning on failure, never raises.
+        """
+        if self.wandb_run is None:
+            return
+        try:
+            import wandb
+
+            # "Simple sync" behavior: wandb will copy this file into the run.
+            wandb.save(str(path), base_path=str(path.parent))
+        except Exception as e:
+            self.warning(f"Failed to sync file to W&B: {e}")
+
     def finish(self):
         """Finalize logging and cleanup."""
         # Flush remaining metrics