Merge branch 'dev' into feat/training-evaluation

scripts/simulate.py

@@ -17,12 +17,14 @@ 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 MorphologyConfig
+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
@@ -32,6 +34,7 @@ from brittle_star_project.evaluation.video import (
     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")
@@ -78,9 +81,34 @@ def main(dict_cfg: DictConfig) -> None:
         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
@@ -104,11 +132,24 @@ def main(dict_cfg: DictConfig) -> None:
     state0 = env.reset(seed=seed)
 
     # Calculate the action dimension the model was trained with
-    trained_action_dim = sum(training.morphology.segments_per_arm) * 2
+    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
+        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

scripts/simulate.sh

@@ -0,0 +1,9 @@
+#!/usr/bin/env bash
+
+path=$1
+
+uv run simulate.py \
+    simulation.model_path="$path"/final_model.flax \
+    simulation.record_video=True \
+    simulation.video_output_path=../vids/simulation.mp4 \
+    simulation.max_steps=10000