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