diff --git a/configs/evaluation/poster.yaml b/configs/evaluation/poster.yaml index 048a785..a09907f 100644 --- a/configs/evaluation/poster.yaml +++ b/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" diff --git a/experiments/debug-experiment-10042026/used_variables.md b/experiments/debug-experiment-10042026/used_variables.md deleted file mode 100644 index 92f23a1..0000000 --- a/experiments/debug-experiment-10042026/used_variables.md +++ /dev/null @@ -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", -} \ No newline at end of file diff --git a/scripts/compare_models.py b/scripts/compare_models.py new file mode 100644 index 0000000..ceebb89 --- /dev/null +++ b/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() diff --git a/scripts/evaluate_checkpoints.py b/scripts/evaluate_checkpoints.py new file mode 100644 index 0000000..82475a1 --- /dev/null +++ b/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() diff --git a/src/brittle_star_project/evaluation/evaluate.py b/src/brittle_star_project/evaluation/evaluate.py index 12cc6fc..9d62880 100644 --- a/src/brittle_star_project/evaluation/evaluate.py +++ b/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. diff --git a/src/brittle_star_project/evaluation/evaluate_mjx.py b/src/brittle_star_project/evaluation/evaluate_mjx.py index 0193d99..5655744 100644 --- a/src/brittle_star_project/evaluation/evaluate_mjx.py +++ b/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 ``/metrics/checkpoint_evaluation.csv``. + """Append one evaluation row to `/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: