From 7ce9dc1216ff662935c2c54d34867b38812cd49e Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 12 May 2026 13:33:19 +0200 Subject: [PATCH] refactor: lift eval env builder --- scripts/simulate.py | 117 ++------------ .../evaluation/__init__.py | 4 + .../evaluation/eval_env_builder.py | 149 ++++++++++++++++++ 3 files changed, 166 insertions(+), 104 deletions(-) create mode 100644 src/brittle_star_project/evaluation/eval_env_builder.py diff --git a/scripts/simulate.py b/scripts/simulate.py index a56e970..01a1af9 100644 --- a/scripts/simulate.py +++ b/scripts/simulate.py @@ -13,28 +13,20 @@ from __future__ import annotations from pathlib import Path import hydra -import numpy as np from omegaconf import DictConfig, OmegaConf -import yaml -import jax.numpy as jnp -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.padded_obs_wrapper import compute_padding_masks -from brittle_star_project.environment.obs_processing import create_obs_processor -from brittle_star_project.environment.env_config import MorphMode, MorphologyConfig 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.eval_env_builder import build_eval_env from brittle_star_project.evaluation.rollout import rollout_headless, rollout_viewer from brittle_star_project.evaluation.video import ( record_episode, create_evaluation_dir, save_evaluation_metadata, ) -from brittle_star_project.MLPs.adjancency_builder import build_adjacency @hydra.main(config_path="../configs", config_name="main_config", version_base="1.3") @@ -63,107 +55,24 @@ def main(dict_cfg: DictConfig) -> None: # 3. Reconstruct typed configs from metadata training = metadata_to_configs(metadata) - # 4. Determine environment morphology - if sim_cfg.morphology_override is not None: - override_path = Path(hydra.utils.to_absolute_path(sim_cfg.morphology_override)) - if not override_path.exists(): - raise FileNotFoundError(f"Could not find morphology override YAML at {override_path}") - with open(override_path, "r") as f: - override_dict = yaml.safe_load(f) - env_morphology = OmegaConf.to_object( - OmegaConf.merge(OmegaConf.structured(MorphologyConfig), override_dict) - ) - else: - env_morphology = training.morphology - - # 5. Build obs_processor with TRAINING morphology padding masks always - padding_masks = compute_padding_masks( - segments_per_arm=env_morphology.segments_per_arm, - reference_segments_per_arm=training.morphology.segments_per_arm, - ) - - segs_per_arm = jnp.array(env_morphology.segments_per_arm) - - needed_copies = 0 - agent_indices = [0, 1, 2, 3, 4] - match env_morphology.morph_mode: - case MorphMode.CENTRALIZED: - needed_copies = 1 - case MorphMode.FULLY_CONNECTED | MorphMode.RING: - agent_mask = segs_per_arm > 0 - agent_indices = jnp.where(agent_mask)[0] - needed_copies = jnp.where(segs_per_arm > 0, 1, 0).sum().item() - case MorphMode.SEGMENT: - agent_mask = segs_per_arm > 0 - agent_indices = jnp.where(agent_mask)[0] - needed_copies = jnp.where(segs_per_arm > 0, 1, 0).sum().item() - needed_copies = (segs_per_arm.sum() + jnp.where(segs_per_arm > 0, 1, 0).sum()).item() - - num_arms = jnp.where(segs_per_arm > 0, 1, 0).sum().item() - - obs_processor = create_obs_processor( - bounds_dict=training.obs_bounds.to_bounds_dict(), - padding_masks=padding_masks, - needed_copies=needed_copies, - num_arms=num_arms, - morph_mode=env_morphology.morph_mode, - segments_per_arm=env_morphology.segments_per_arm, - agent_indices=agent_indices, - ) - - # 6. Build environment - backend = Backend.MJC seed = int(cfg.experiment.seed) - factory = BrittleStarEnvFactory() - raw_env = factory.create_environment( - backend, - env_morphology, - training.arena, - training.environment, - ) - env = BrittleStarEnv( - raw_env, - backend=backend, - config=training.environment, - morphology_config=env_morphology, + # 4-7. Build evaluation environment and policy + bundle = build_eval_env( + model_path=model_path, + training=training, + metadata=metadata, + morphology_override_path=sim_cfg.morphology_override, ) + env = bundle.env + policy = bundle.policy + action_low = bundle.action_low + action_high = bundle.action_high + action_mask = bundle.action_mask + state0 = env.reset(seed=seed) - # Calculate the action dimension the model was trained with - trained_action_dim = raw_env.action_space.shape[0] // needed_copies - - # 7. Load policy - message_passing_steps = (metadata.get("architecture", {}) or {}).get("message_passing_steps") - if message_passing_steps is None: - message_passing_steps = 4 - message_passing_steps = int(message_passing_steps) - - adj_matrix = None - if env_morphology.morph_mode != MorphMode.CENTRALIZED: - adj_matrix = build_adjacency(env_morphology.segments_per_arm, env_morphology.morph_mode) - - policy = PolicyAgent.from_checkpoint( - model_path, - action_dim=trained_action_dim, - obs_processor=obs_processor, - message_passing_steps=message_passing_steps, - adj_matrix=adj_matrix, - ) - - # Convert the JAX boolean mask to a numpy array for easy indexing - action_mask = np.asarray(padding_masks["mask_2x"]) - - # Match training's action clipping behavior. - 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() - ) - # 8. Run simulation headless = bool(sim_cfg.headless) max_steps = sim_cfg.max_steps diff --git a/src/brittle_star_project/evaluation/__init__.py b/src/brittle_star_project/evaluation/__init__.py index f9ed595..e529115 100644 --- a/src/brittle_star_project/evaluation/__init__.py +++ b/src/brittle_star_project/evaluation/__init__.py @@ -11,6 +11,7 @@ from .evaluate import evaluate_policy from .policy import PolicyAgent, ControlPolicy from .rollout import rollout_headless, rollout_viewer, EpisodeResult from .video import record_episode, create_evaluation_dir, save_evaluation_metadata +from .eval_env_builder import EvalEnvBundle, build_eval_env __all__ = [ # checkpoint loading @@ -36,4 +37,7 @@ __all__ = [ "record_episode", "create_evaluation_dir", "save_evaluation_metadata", + # env builder + "EvalEnvBundle", + "build_eval_env", ] diff --git a/src/brittle_star_project/evaluation/eval_env_builder.py b/src/brittle_star_project/evaluation/eval_env_builder.py new file mode 100644 index 0000000..80c148e --- /dev/null +++ b/src/brittle_star_project/evaluation/eval_env_builder.py @@ -0,0 +1,149 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import jax.numpy as jnp +import numpy as np +import yaml +from omegaconf import OmegaConf + +from brittle_star_project import Backend, BrittleStarEnv, BrittleStarEnvFactory +from brittle_star_project.environment.env_config import MorphMode, MorphologyConfig +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 TrainingConfig +from brittle_star_project.evaluation.policy import PolicyAgent +from brittle_star_project.MLPs.adjancency_builder import build_adjacency + + +@dataclass +class EvalEnvBundle: + """Everything needed to run a headless evaluation episode.""" + + env: BrittleStarEnv + policy: PolicyAgent + action_low: np.ndarray | None + action_high: np.ndarray | None + action_mask: np.ndarray | None + segments_per_arm: list[int] + num_active_arms: int + architecture: str + + +def build_eval_env( + *, + model_path: Path, + training: TrainingConfig, + metadata: dict, + morphology_override_path: Path | str | None = None, +) -> EvalEnvBundle: + """Build environment + policy for evaluation, optionally with a morphology override.""" + + # 1. Determine environment morphology + if morphology_override_path is not None: + override_path = Path(morphology_override_path) + if not override_path.exists(): + raise FileNotFoundError(f"Could not find morphology override YAML at {override_path}") + with open(override_path, "r") as f: + override_dict = yaml.safe_load(f) + env_morphology = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(MorphologyConfig), override_dict) + ) + # Force morph_mode to be inherited from training since it's baked into weights + env_morphology.morph_mode = training.morphology.morph_mode + else: + env_morphology = training.morphology + + # 2. Build obs_processor with TRAINING morphology padding masks always + padding_masks = compute_padding_masks( + segments_per_arm=env_morphology.segments_per_arm, + reference_segments_per_arm=training.morphology.segments_per_arm, + ) + + segs_per_arm = jnp.array(env_morphology.segments_per_arm) + + needed_copies = 0 + agent_indices = [0, 1, 2, 3, 4] + match env_morphology.morph_mode: + case MorphMode.CENTRALIZED: + needed_copies = 1 + case MorphMode.FULLY_CONNECTED | MorphMode.RING: + agent_mask = segs_per_arm > 0 + agent_indices = jnp.where(agent_mask)[0] + needed_copies = jnp.where(segs_per_arm > 0, 1, 0).sum().item() + case MorphMode.SEGMENT: + agent_mask = segs_per_arm > 0 + agent_indices = jnp.where(agent_mask)[0] + needed_copies = (segs_per_arm.sum() + jnp.where(segs_per_arm > 0, 1, 0).sum()).item() + + num_arms = jnp.where(segs_per_arm > 0, 1, 0).sum().item() + + obs_processor = create_obs_processor( + bounds_dict=training.obs_bounds.to_bounds_dict(), + padding_masks=padding_masks, + needed_copies=needed_copies, + num_arms=num_arms, + morph_mode=env_morphology.morph_mode, + segments_per_arm=env_morphology.segments_per_arm, + agent_indices=agent_indices, + ) + + # 3. Build environment + backend = Backend.MJC + factory = BrittleStarEnvFactory() + raw_env = factory.create_environment( + backend, + env_morphology, + training.arena, + training.environment, + ) + env = BrittleStarEnv( + raw_env, + backend=backend, + config=training.environment, + morphology_config=env_morphology, + ) + + # Calculate the action dimension the model was trained with + trained_action_dim = raw_env.action_space.shape[0] // needed_copies + + # 4. Load policy + message_passing_steps = (metadata.get("architecture", {}) or {}).get("message_passing_steps") + if message_passing_steps is None: + message_passing_steps = 4 + message_passing_steps = int(message_passing_steps) + + adj_matrix = None + if env_morphology.morph_mode != MorphMode.CENTRALIZED: + adj_matrix = build_adjacency(env_morphology.segments_per_arm, env_morphology.morph_mode) + + policy = PolicyAgent.from_checkpoint( + model_path, + action_dim=trained_action_dim, + obs_processor=obs_processor, + message_passing_steps=message_passing_steps, + adj_matrix=adj_matrix, + ) + + # 5. Build action clipping and masks + 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() + ) + + return EvalEnvBundle( + env=env, + policy=policy, + action_low=action_low, + action_high=action_high, + action_mask=action_mask, + segments_per_arm=env_morphology.segments_per_arm, + num_active_arms=num_arms, + architecture=env_morphology.morph_mode.name, + )