refactor: lift eval env builder

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

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",
 ]

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,
+    )