Merge branch 'dev' into feat/message_passing

src/brittle_star_project/__init__.py

@@ -2,7 +2,14 @@ from .environment.env_types import Backend, Task
 from .environment.env_config import ArenaConfig, EnvConfig, MorphologyConfig
 from .environment.factory import BrittleStarEnvFactory
 from .environment.env_wrapper import BrittleStarEnv
-from .render import simulate_policy, SimulationConfig, ControlPolicy
+from .evaluation import (
+    PolicyAgent,
+    ControlPolicy,
+    load_metadata,
+    rollout_headless,
+    rollout_viewer,
+    EpisodeResult,
+)
 
 __all__ = [
     "ArenaConfig",
@@ -12,7 +19,10 @@ __all__ = [
     "EnvConfig",
     "MorphologyConfig",
     "Task",
-    "simulate_policy",
-    "SimulationConfig",
+    "PolicyAgent",
     "ControlPolicy",
+    "load_metadata",
+    "rollout_headless",
+    "rollout_viewer",
+    "EpisodeResult",
 ]

src/brittle_star_project/configs/config_simulation.py

@@ -13,7 +13,20 @@ class SimulationSettings:
     # If None, viewer mode runs until window closed or target reached.
     max_steps: Optional[int] = None
 
-    # Optional: point to a Hydra config.yaml from a training run (e.g. runs/.../.hydra/config.yaml).
-    # When set, the simulation script can override
-    # morphology/arena/environment/architecture to match.
-    trained_config_path: Optional[str] = None
+    # Override morphology for amputation experiments.
+    # When set, the environment uses this morphology instead of the trained one.
+    # Points to a morphology config YAML file (e.g. configs/morphology/3_arms.yaml).
+    # Observations are padded from the override morphology UP TO the training
+    # morphology's shape via compute_padding_masks(override, reference=training).
+    morphology_override: Optional[str] = None
+
+    # Video recording (requires [evaluation] extra)
+    record_video: bool = False
+    # When None, video is saved in a per-model evaluation folder alongside the model.
+    video_output_path: Optional[str] = None
+    # Camera ID to use for video recording (1 is usually the close-up camera)
+    camera_id: int = 1
+
+    # Optional override for the sidecar metadata YAML file.
+    # If None, it defaults to the model_path with a `_metadata.yaml` suffix.
+    metadata_path: Optional[str] = None

src/brittle_star_project/configs/main_config.py

@@ -5,7 +5,12 @@ from brittle_star_project.configs.config_experiment import ExperimentConfig
 from brittle_star_project.configs.config_ppo import PPOConfig
 from brittle_star_project.configs.config_architecture import ArchitectureConfig
 from brittle_star_project.configs.config_simulation import SimulationSettings
-from brittle_star_project.environment.env_config import MorphologyConfig, ArenaConfig, EnvConfig
+from brittle_star_project.environment.env_config import (
+    MorphologyConfig,
+    ArenaConfig,
+    EnvConfig,
+    ObservationBoundsConfig,
+)
 
 
 @dataclass
@@ -25,4 +30,5 @@ class BrittleStarConfig:
     morphology: MorphologyConfig = field(default_factory=MorphologyConfig)
     arena: ArenaConfig = field(default_factory=ArenaConfig)
     environment: EnvConfig = field(default_factory=EnvConfig)
+    obs_bounds: ObservationBoundsConfig = field(default_factory=ObservationBoundsConfig)
     simulation: SimulationSettings = field(default_factory=SimulationSettings)

src/brittle_star_project/configs/register_configs.py

@@ -8,7 +8,12 @@ from brittle_star_project.configs.config_architecture import (
     DecentralizedConfig,
 )
 from brittle_star_project.configs.config_simulation import SimulationSettings
-from brittle_star_project.environment.env_config import MorphologyConfig, ArenaConfig, EnvConfig
+from brittle_star_project.environment.env_config import (
+    MorphologyConfig,
+    ArenaConfig,
+    EnvConfig,
+    ObservationBoundsConfig,
+)
 from brittle_star_project.configs.main_config import BrittleStarConfig
 
 
@@ -37,4 +42,5 @@ def register_configs() -> None:
     cs.store(group="morphology", name="base_morphology", node=MorphologyConfig)
     cs.store(group="arena", name="base_arena", node=ArenaConfig)
     cs.store(group="environment", name="base_environment", node=EnvConfig)
+    cs.store(group="obs_bounds", name="base_obs_bounds", node=ObservationBoundsConfig)
     cs.store(group="simulation", name="base_simulation", node=SimulationSettings)

src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py

@@ -5,7 +5,7 @@ from experiment_logger import get_logger
 from .env_config import EnvConfig, MorphologyConfig, ArenaConfig
 from .env_types import Backend
 from .factory import BrittleStarEnvFactory
-from .padded_obs_wrapper import compute_padding_masks, pad_observations_batched
+from .padded_obs_wrapper import compute_padding_masks
 
 
 class BrittleStarJaxEnvWrapper:
@@ -48,6 +48,15 @@ class BrittleStarJaxEnvWrapper:
     def raw(self):
         return self._env
 
+    @property
+    def padding_masks(self) -> dict:
+        """Pre-computed boolean masks for amputated limb padding.
+
+        Pass to create_obs_processor so the processor handles padding
+        after normalization in the correct pipeline order.
+        """
+        return self._padding_masks
+
     @property
     def single_action_space(self):
         return self._env.action_space
@@ -61,10 +70,6 @@ class BrittleStarJaxEnvWrapper:
         self._action_rng, env_rng = jax.random.split(jax.random.PRNGKey(seed), 2)
         env_rngs = jnp.array(jax.random.split(env_rng, self._num_envs))
         state = self._vectorized_reset(rng=env_rngs)
-
-        state = state.replace(
-            observations=pad_observations_batched(state.observations, self._padding_masks)
-        )
         return state
 
     def sample_actions(self):
@@ -75,12 +80,7 @@ class BrittleStarJaxEnvWrapper:
         return self._vectorized_action_sample(rng=jnp.array(sub_rngs))
 
     def step(self, state, action):
-        next_state = self._vectorized_step(state=state, action=action)
-
-        next_state = next_state.replace(
-            observations=pad_observations_batched(next_state.observations, self._padding_masks)
-        )
-        return next_state
+        return self._vectorized_step(state=state, action=action)
 
     def close(self):
         self._env.close()

src/brittle_star_project/environment/__init__.py

@@ -1,7 +1,9 @@
 from .env_config import ArenaConfig, EnvConfig, MorphologyConfig, MorphMode
 from .env_types import Backend, Task
-from .env_wrapper import BrittleStarEnv, StepResult
+from .env_wrapper import BrittleStarEnv
 from .factory import BrittleStarEnvFactory
+from .obs_processing import create_obs_processor
+from .padded_obs_wrapper import compute_padding_masks
 
 __all__ = [
     "ArenaConfig",
@@ -10,7 +12,8 @@ __all__ = [
     "Backend",
     "Task",
     "BrittleStarEnv",
-    "StepResult",
     "BrittleStarEnvFactory",
     "MorphMode",
+    "create_obs_processor",
+    "compute_padding_masks",
 ]

src/brittle_star_project/environment/env_config.py

@@ -70,4 +70,30 @@ class EnvConfig:
     # Per docs in upstream env config: integer factors of 200.
     light_perlin_noise_scale: int = 0
 
-    morph_mode: MorphMode = MorphMode.CENTRALIZED
+      
+@dataclass
+class ObservationBoundsConfig:
+    """Physical observation bounds for deterministic min-max normalization."""
+
+    # Empirical testing based on the extract_observation_bounds.py script run for 1.000.000 steps
+
+    # Based on max. ctrlrange (0.78539816339744828) in XML, but empirical testing went slightly over
+    joint_position: list[float] = field(default_factory=lambda: [-0.8, 0.8])
+    # Empirical testing showed max. 3.22, adding buffer to be safe. Consider higher values "fast".
+    joint_velocity: list[float] = field(default_factory=lambda: [-5.0, 5.0])
+    # Based on max. forceRange in XML, verified with empirical testing
+    joint_actuator_force: list[float] = field(default_factory=lambda: [-3.75, 3.75])
+    # Based on intuition and reasoning
+    segment_contact: list[float] = field(default_factory=lambda: [0.0, 1.0])
+    robot_direction_to_target: list[float] = field(default_factory=lambda: [-1.0, 1.0])
+    disk_z_tilt: list[float] = field(default_factory=lambda: [0.0, 3.141592653589793])
+
+    def to_bounds_dict(self) -> dict[str, tuple[float, float]]:
+        return {
+            "disk_z_tilt": tuple(self.disk_z_tilt),
+            "joint_actuator_force": tuple(self.joint_actuator_force),
+            "joint_position": tuple(self.joint_position),
+            "joint_velocity": tuple(self.joint_velocity),
+            "robot_direction_to_target": tuple(self.robot_direction_to_target),
+            "segment_contact": tuple(self.segment_contact),
+        }

src/brittle_star_project/environment/obs_processing.py

@@ -0,0 +1,91 @@
+import jax
+import jax.numpy as jnp
+from typing import Dict, Tuple, Optional
+
+_JOINT_SCALED_KEYS = frozenset(
+    {
+        "joint_position",
+        "joint_velocity",
+        "joint_actuator_force",
+        "actuator_force",
+    }
+)
+
+_SEGMENT_SCALED_KEYS = frozenset(
+    {
+        "segment_contact",
+    }
+)
+
+
+def create_obs_processor(
+    bounds_dict: Dict[str, Tuple[float, float]], padding_masks: Optional[Dict] = None
+):
+    def _add_derived_features(obs: dict) -> dict:
+        new_obs = dict(obs)
+        if "disk_rotation" in new_obs:
+            rot = new_obs["disk_rotation"]
+            new_obs["disk_z_tilt"] = jnp.sqrt(jnp.pow(rot[0], 2) + jnp.pow(rot[1], 2))
+
+            if "unit_xy_direction_to_target" in new_obs:
+                yaw = rot[2]
+                unit_x, unit_y = new_obs["unit_xy_direction_to_target"]
+                cos_yaw, sin_yaw = jnp.cos(yaw), jnp.sin(yaw)
+                new_x = unit_x * cos_yaw + unit_y * sin_yaw
+                new_y = -unit_x * sin_yaw + unit_y * cos_yaw
+                new_obs["robot_direction_to_target"] = jnp.stack([new_x, new_y])
+
+        return new_obs
+
+    def _normalize_features(obs: dict) -> dict:
+        normalized = {}
+        for key, arr in obs.items():
+            if key in bounds_dict:
+                low, high = bounds_dict[key]
+                if low == -1.0 and high == 1.0:
+                    normalized[key] = jnp.clip(arr, -1.0, 1.0)
+                else:
+                    arr_clipped = jnp.clip(arr, low, high)
+                    normalized[key] = 2.0 * (arr_clipped - low) / (high - low) - 1.0
+            else:
+                normalized[key] = arr
+        return normalized
+
+    def _pad_features(obs: dict) -> dict:
+        padded = {}
+        for key, arr in obs.items():
+            if key in _JOINT_SCALED_KEYS:
+                padded_arr = jnp.zeros(padding_masks["target_size_2x"], dtype=arr.dtype)
+                padded[key] = padded_arr.at[padding_masks["mask_2x"]].set(arr)
+            elif key in _SEGMENT_SCALED_KEYS:
+                padded_arr = jnp.zeros(padding_masks["target_size_1x"], dtype=arr.dtype)
+                padded[key] = padded_arr.at[padding_masks["mask_1x"]].set(arr)
+            else:
+                padded[key] = arr
+        return padded
+
+    def _flatten_features(obs: dict) -> jnp.ndarray:
+        ordered_keys = [
+            "disk_z_tilt",
+            "joint_actuator_force",
+            "joint_position",
+            "joint_velocity",
+            "robot_direction_to_target",
+            "segment_contact",
+        ]
+        values = []
+        for key in ordered_keys:
+            if key in obs:
+                arr = jnp.asarray(obs[key]).flatten()
+                if arr.size > 0:
+                    values.append(arr)
+        return jnp.concatenate(values)
+
+    def _process_single(obs_dict: dict) -> jnp.ndarray:
+        processed = _add_derived_features(obs_dict)
+        processed = _normalize_features(processed)
+        if padding_masks is not None:
+            processed = _pad_features(processed)
+        return _flatten_features(processed)
+
+    return jax.jit(jax.vmap(_process_single))

src/brittle_star_project/environment/padded_obs_wrapper.py

@@ -1,35 +1,11 @@
-"""Observation padding wrapper for amputated brittle star morphologies.
-
-When using a centralized controller, the global observation vector must remain
-a constant size regardless of how many segments are amputated. This wrapper pads
-the observation dictionary values with zeros using spatial insertion so that the
-flattened observation maintains the correct physical mapping to the neural network.
-"""
+"""Observation padding masks for amputated brittle star morphologies."""
 
 from __future__ import annotations
 
 from typing import Any, Sequence
 
-import jax
 import jax.numpy as jnp
 
-# Observation keys whose size scales with the number of joints (2 per segment).
-_JOINT_SCALED_KEYS = frozenset(
-    {
-        "joint_position",
-        "joint_velocity",
-        "joint_actuator_force",
-        "actuator_force",
-    }
-)
-
-# Observation keys whose size scales with the number of segments (1 per segment).
-_SEGMENT_SCALED_KEYS = frozenset(
-    {
-        "segment_contact",
-    }
-)
-
 
 def compute_padding_masks(
     segments_per_arm: Sequence[int],
@@ -60,7 +36,6 @@ def compute_padding_masks(
                 f"actual segments ({actual}) must be between 0 and reference ({ref})."
             )
         # 1x scaling (e.g., contacts: 1 value per segment)
-        # 1x scaling (e.g., contacts: 1 value per segment)
         mask_1x.extend([True] * actual + [False] * (ref - actual))
         # 2x scaling (e.g., joints: 2 values per segment)
         mask_2x.extend([True] * (actual * 2) + [False] * ((ref - actual) * 2))
@@ -71,60 +46,3 @@ def compute_padding_masks(
         "target_size_1x": sum(reference_segments_per_arm),
         "target_size_2x": sum(reference_segments_per_arm) * 2,
     }
-
-
-def pad_observation(
-    obs: dict[str, Any],
-    masks: dict[str, Any],
-) -> dict[str, Any]:
-    """Pad an observation dict using spatial insertion."""
-    padded = {}
-    for key, value in obs.items():
-        padded_dtype = _padding_dtype(value)
-        if key in _JOINT_SCALED_KEYS:
-            out = jnp.zeros(masks["target_size_2x"], dtype=padded_dtype)
-            padded[key] = out.at[masks["mask_2x"]].set(value)
-        elif key in _SEGMENT_SCALED_KEYS:
-            out = jnp.zeros(masks["target_size_1x"], dtype=padded_dtype)
-            padded[key] = out.at[masks["mask_1x"]].set(value)
-        else:
-            padded[key] = value
-    return padded
-
-
-def pad_observations_batched(
-    obs: dict[str, Any],
-    masks: dict[str, Any],
-) -> dict[str, Any]:
-    """Pad a batched observation dict (leading batch dimension) using spatial insertion."""
-    padded = {}
-    for key, value in obs.items():
-        batch_size = value.shape[0]
-        padded_dtype = _padding_dtype(value)
-        if key in _JOINT_SCALED_KEYS:
-            out = jnp.zeros((batch_size, masks["target_size_2x"]), dtype=padded_dtype)
-            padded[key] = out.at[:, masks["mask_2x"]].set(value)
-        elif key in _SEGMENT_SCALED_KEYS:
-            out = jnp.zeros((batch_size, masks["target_size_1x"]), dtype=padded_dtype)
-            padded[key] = out.at[:, masks["mask_1x"]].set(value)
-        else:
-            padded[key] = value
-    return padded
-
-
-def _padding_dtype(value: Any) -> jnp.dtype:
-    """Choose a JAX-safe dtype for padding arrays.
-
-    When JAX x64 is disabled, allocating float64 zeros emits a warning. We
-    preserve the original dtype whenever it is supported, and otherwise fall
-    back to float32 for padding buffers.
-    """
-
-    dtype = getattr(value, "dtype", None)
-    if dtype is None:
-        dtype = jnp.asarray(value).dtype
-    else:
-        dtype = jnp.dtype(dtype)
-    if dtype == jnp.float64 and not jax.config.read("jax_enable_x64"):
-        return jnp.float32
-    return dtype

src/brittle_star_project/evaluation/__init__.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from .checkpoint import load_metadata, load_params, metadata_to_configs, TrainingConfig
+from .policy import PolicyAgent, ControlPolicy
+from .rollout import rollout_headless, rollout_viewer, EpisodeResult
+from .video import record_episode, create_evaluation_dir, save_evaluation_metadata
+
+__all__ = [
+    "load_metadata",
+    "load_params",
+    "metadata_to_configs",
+    "TrainingConfig",
+    "PolicyAgent",
+    "ControlPolicy",
+    "rollout_headless",
+    "rollout_viewer",
+    "EpisodeResult",
+    "record_episode",
+    "create_evaluation_dir",
+    "save_evaluation_metadata",
+]

src/brittle_star_project/evaluation/checkpoint.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+import yaml
+from dataclasses import dataclass
+from pathlib import Path
+
+import flax
+from omegaconf import OmegaConf
+
+from brittle_star_project.environment.env_config import (
+    MorphologyConfig,
+    ArenaConfig,
+    EnvConfig,
+    ObservationBoundsConfig,
+)
+
+
+@dataclass
+class TrainingConfig:
+    """Holds typed configurations extracted from a training run's metadata."""
+
+    morphology: MorphologyConfig
+    arena: ArenaConfig
+    environment: EnvConfig
+    obs_bounds: ObservationBoundsConfig
+
+
+def load_params(path: Path) -> dict:
+    """Load model parameters from a .flax checkpoint file."""
+    payload = path.read_bytes()
+    restored = flax.serialization.msgpack_restore(payload)
+
+    sensor_params = None
+    actor_params = None
+
+    # Extract params from restored checkpoint
+    if isinstance(restored, dict):
+        params_sub = restored.get("params", {})
+        sensor_params = restored.get("sensor_params") or params_sub.get("sensor_params")
+        actor_params = restored.get("actor_params") or params_sub.get("actor_params")
+    elif isinstance(restored, (list, tuple)) and len(restored) >= 2:
+        params_part = restored[1]
+        if isinstance(params_part, dict):
+            sensor_params = params_part.get("0", params_part.get(0))
+            actor_params = params_part.get("1", params_part.get(1))
+        elif isinstance(params_part, (list, tuple)) and len(params_part) >= 2:
+            sensor_params = params_part[0]
+            actor_params = params_part[1]
+
+    if sensor_params is None or actor_params is None:
+        raise ValueError(f"Could not extract sensor and actor params from checkpoint: {path}")
+
+    return {
+        "sensor_params": sensor_params,
+        "actor_params": actor_params,
+    }
+
+
+def load_metadata(model_path: Path, metadata_override_path: Path | None = None) -> dict:
+    """Discover and load the sidecar metadata YAML file."""
+    if metadata_override_path is not None:
+        metadata_path = metadata_override_path
+    else:
+        metadata_path = model_path.with_name(model_path.stem + "_metadata.yaml")
+
+    if not metadata_path.exists():
+        raise FileNotFoundError(f"Could not find metadata YAML at {metadata_path}")
+    with open(metadata_path, "r") as f:
+        return yaml.safe_load(f)
+
+
+def metadata_to_configs(metadata: dict) -> TrainingConfig:
+    """Reconstruct typed configuration objects from a metadata dictionary."""
+    trained_morphology = OmegaConf.to_object(
+        OmegaConf.merge(OmegaConf.structured(MorphologyConfig), metadata.get("morphology", {}))
+    )
+    trained_arena = OmegaConf.to_object(
+        OmegaConf.merge(OmegaConf.structured(ArenaConfig), metadata.get("arena", {}))
+    )
+
+    env_dict = metadata.get("environment", {})
+    if isinstance(env_dict.get("task"), str):
+        from brittle_star_project.environment.env_types import Task
+
+        try:
+            env_dict["task"] = Task[env_dict["task"]].name
+        except Exception:
+            try:
+                env_dict["task"] = Task(env_dict["task"]).name
+            except Exception:
+                pass
+
+    trained_environment = OmegaConf.to_object(
+        OmegaConf.merge(OmegaConf.structured(EnvConfig), env_dict)
+    )
+    trained_obs_bounds = OmegaConf.to_object(
+        OmegaConf.merge(
+            OmegaConf.structured(ObservationBoundsConfig), metadata.get("obs_bounds", {})
+        )
+    )
+
+    return TrainingConfig(
+        morphology=trained_morphology,
+        arena=trained_arena,
+        environment=trained_environment,
+        obs_bounds=trained_obs_bounds,
+    )

src/brittle_star_project/evaluation/policy.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Protocol
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from brittle_star_project.evaluation.checkpoint import load_params
+
+
+class ControlPolicy(Protocol):
+    """Protocol for any policy that can produce actions from observations."""
+
+    def act(self, *, observations: dict[str, Any]) -> np.ndarray: ...
+
+
+class PolicyAgent:
+    """Wraps a trained Flax actor for deterministic inference."""
+
+    def __init__(
+        self,
+        *,
+        sensor_params: Any,
+        actor_params: Any,
+        action_dim: int,
+        obs_processor: Any,
+    ) -> None:
+        from brittle_star_project.MLPs.mlps import Actor, GenericDenseLayersWithActivation
+
+        # Infer layer sizes from params
+        try:
+            dense_params = (
+                sensor_params.get("params", {})
+                if isinstance(sensor_params, dict)
+                else sensor_params["params"]
+            )
+        except Exception:
+            dense_params = sensor_params
+
+        layer_sizes = []
+        idx = 0
+        while True:
+            key = f"Dense_{idx}"
+            if key not in dense_params:
+                break
+            layer_sizes.append(int(np.asarray(dense_params[key]["kernel"]).shape[1]))
+            idx += 1
+
+        if not layer_sizes:
+            raise ValueError("Could not infer Dense_* layers from sensor params")
+
+        self._sensor = GenericDenseLayersWithActivation(layer_sizes=layer_sizes)
+        self._actor = Actor(action_dim=action_dim)
+        self._sensor_apply = jax.jit(self._sensor.apply)
+        self._actor_apply = jax.jit(self._actor.apply)
+        self._params = {
+            "sensor_params": sensor_params,
+            "actor_params": actor_params,
+        }
+        self._obs_processor = obs_processor
+
+    @classmethod
+    def from_checkpoint(
+        cls,
+        model_path: Path,
+        *,
+        action_dim: int,
+        obs_processor: Any,
+    ) -> "PolicyAgent":
+        """Load params from .flax and construct the agent."""
+        params = load_params(model_path)
+
+        return cls(
+            sensor_params=params["sensor_params"],
+            actor_params=params["actor_params"],
+            action_dim=action_dim,
+            obs_processor=obs_processor,
+        )
+
+    def act(self, *, observations: dict[str, Any]) -> np.ndarray:
+        """Return deterministic action (actor mean, no exploration noise)."""
+        batched_obs = jax.tree.map(lambda x: jnp.asarray(x)[None, ...], observations)
+        obs = self._obs_processor(batched_obs)[0]
+        hidden = self._sensor_apply(self._params["sensor_params"], obs)
+        mean, _log_std = self._actor_apply(self._params["actor_params"], hidden)
+
+        return np.asarray(mean, dtype=np.float32).ravel()

src/brittle_star_project/evaluation/rollout.py

@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+import itertools
+import time
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+from brittle_star_project import BrittleStarEnv
+from brittle_star_project.evaluation.policy import ControlPolicy
+
+
+@dataclass
+class EpisodeResult:
+    return_: float
+    length: int
+    reached_target: bool
+    final_xy_dist: float | None
+
+
+def _get_observations(state: Any) -> dict[str, Any] | None:
+    return getattr(state, "observations", None)
+
+
+def _get_xy_distance_to_target(observations: dict[str, Any]) -> float | None:
+    return float(np.asarray(observations["xy_distance_to_target"]).reshape(-1)[0])
+
+
+def _target_reached(*, state: Any) -> bool:
+    return bool(getattr(state, "terminated", False) or getattr(state, "truncated", False))
+
+
+def _maybe_clip_action(
+    action: np.ndarray,
+    low: np.ndarray | None,
+    high: np.ndarray | None,
+) -> np.ndarray:
+    if low is None or high is None:
+        return action
+    low = np.asarray(low, dtype=np.float32).ravel()
+    high = np.asarray(high, dtype=np.float32).ravel()
+    if low.shape != action.shape or high.shape != action.shape:
+        return action
+    return np.clip(action, low, high)
+
+
+def rollout_headless(
+    *,
+    env: BrittleStarEnv,
+    policy: ControlPolicy,
+    seed: int,
+    max_steps: int,
+    action_low: np.ndarray | None,
+    action_high: np.ndarray | None,
+    action_mask: np.ndarray | None = None,
+) -> EpisodeResult:
+    """Run an episode headlessly and return the result."""
+    state = env.reset(seed=seed)
+
+    ep_return = 0.0
+    observations = _get_observations(state)
+    prev_dist = _get_xy_distance_to_target(observations) if observations else None
+    reached_target = _target_reached(state=state)
+
+    steps = 0
+    for _ in range(int(max_steps)):
+        obs_dict = observations or {}
+
+        action = policy.act(observations=obs_dict)
+        if action_mask is not None:
+            action = action[action_mask]
+        action = _maybe_clip_action(action, action_low, action_high)
+
+        state = env.step(state=state, action=action)
+        steps += 1
+
+        observations = _get_observations(state)
+        cur_dist = _get_xy_distance_to_target(observations) if observations else None
+        if prev_dist is not None and cur_dist is not None:
+            ep_return += prev_dist - cur_dist
+        prev_dist = cur_dist
+
+        reached_target = _target_reached(state=state)
+        if reached_target:
+            break
+
+    final_dist = _get_xy_distance_to_target(observations) if observations else None
+    return EpisodeResult(
+        return_=ep_return,
+        length=steps,
+        reached_target=reached_target,
+        final_xy_dist=final_dist,
+    )
+
+
+def rollout_viewer(
+    *,
+    env: BrittleStarEnv,
+    policy: ControlPolicy,
+    seed: int,
+    state: Any,
+    control_dt: float,
+    max_steps: int | None,
+    action_low: np.ndarray | None,
+    action_high: np.ndarray | None,
+    action_mask: np.ndarray | None = None,
+) -> None:
+    """Run an episode using the interactive MuJoCo viewer."""
+    import mujoco.viewer
+
+    model = state.mj_model
+    data = state.mj_data
+
+    episode_return = 0.0
+    observations = _get_observations(state)
+    prev_dist = _get_xy_distance_to_target(observations) if observations else None
+    reached_target = _target_reached(state=state)
+
+    steps = 0
+    with mujoco.viewer.launch_passive(model, data) as viewer:
+        step_iter = range(int(max_steps)) if max_steps is not None else itertools.count()
+        for _step_idx in step_iter:
+            if not viewer.is_running():
+                break
+            step_start = time.time()
+
+            obs_dict = observations or {}
+            action = policy.act(observations=obs_dict)
+            if action_mask is not None:
+                action = action[action_mask]
+            action = _maybe_clip_action(action, action_low, action_high)
+
+            with viewer.lock():
+                state = env.step(state=state, action=action)
+
+            if not viewer.is_running():
+                break
+            viewer.sync()
+
+            steps += 1
+
+            observations = _get_observations(state)
+            cur_dist = _get_xy_distance_to_target(observations) if observations else None
+            if prev_dist is not None and cur_dist is not None:
+                episode_return += prev_dist - cur_dist
+            prev_dist = cur_dist
+
+            reached_target = _target_reached(state=state)
+            if reached_target:
+                break
+
+            remaining = control_dt - (time.time() - step_start)
+            if remaining > 0:
+                time.sleep(remaining)
+
+    dist = _get_xy_distance_to_target(observations) if observations else None
+    dist_str = "n/a" if dist is None else f"{dist:.3f}"
+    print(
+        "episode done: "
+        f"return={episode_return:.6f}, len={steps}, "
+        f"target_reached={reached_target}, final_xy_dist={dist_str}"
+    )

src/brittle_star_project/evaluation/video.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import datetime
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+from brittle_star_project import BrittleStarEnv
+from brittle_star_project.evaluation.policy import ControlPolicy
+from brittle_star_project.evaluation.rollout import (
+    EpisodeResult,
+    _get_observations,
+    _get_xy_distance_to_target,
+    _target_reached,
+    _maybe_clip_action,
+)
+
+
+def create_evaluation_dir(model_path: Path) -> Path:
+    """Create a unique timestamped directory for saving evaluation results."""
+    timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    eval_dir = model_path.parent / f"{model_path.stem}_evaluations" / f"eval_{timestamp}"
+    eval_dir.mkdir(parents=True, exist_ok=True)
+    return eval_dir
+
+
+def save_evaluation_metadata(
+    eval_dir: Path,
+    *,
+    morphology_override_path: str | None,
+    seed: int,
+    max_steps: int | None,
+    result: EpisodeResult,
+) -> None:
+    """Save metadata about the evaluation run."""
+    metadata = {
+        "timestamp": datetime.datetime.now().isoformat(),
+        "morphology_override": morphology_override_path,
+        "seed": seed,
+        "max_steps": max_steps,
+        "result": {
+            "return": float(result.return_),
+            "length": int(result.length),
+            "reached_target": bool(result.reached_target),
+            "final_xy_dist": float(result.final_xy_dist)
+            if result.final_xy_dist is not None
+            else None,
+        },
+    }
+    with open(eval_dir / "evaluation_metadata.yaml", "w") as f:
+        yaml.safe_dump(metadata, f, sort_keys=False)
+
+
+def record_episode(
+    *,
+    env: BrittleStarEnv,
+    policy: ControlPolicy,
+    seed: int,
+    max_steps: int,
+    action_low: np.ndarray | None,
+    action_high: np.ndarray | None,
+    action_mask: np.ndarray | None = None,
+    output_path: Path,
+    camera_id: int = 1,
+    fps: int = 60,
+    width: int = 640,
+    height: int = 480,
+) -> EpisodeResult:
+    """Run an episode headlessly and record a video using MuJoCo's Renderer and imageio.
+
+    Args:
+        env: The environment.
+        policy: The policy agent.
+        seed: Random seed.
+        max_steps: Maximum number of steps.
+        action_low: Minimum action values.
+        action_high: Maximum action values.
+        action_mask: Boolean mask for the actions.
+        output_path: Where to save the .mp4 file.
+        camera_id: Camera index to use for rendering (1 is usually close-up).
+        fps: Frames per second for the video.
+        width: Video width.
+        height: Video height.
+    """
+    try:
+        import imageio
+        import mujoco
+    except ImportError as e:
+        raise ImportError(
+            "Video recording requires 'imageio' and 'mujoco'. "
+            "Please install the evaluation dependencies: `uv pip install .[evaluation]`"
+        ) from e
+
+    state = env.reset(seed=seed)
+    model = state.mj_model
+    data = state.mj_data
+
+    renderer = mujoco.Renderer(model, width=width, height=height)
+
+    ep_return = 0.0
+    observations = _get_observations(state)
+    prev_dist = _get_xy_distance_to_target(observations) if observations else None
+    reached_target = _target_reached(state=state)
+
+    frames = []
+    steps = 0
+
+    for _ in range(int(max_steps)):
+        # Capture frame
+        renderer.update_scene(data, camera=camera_id)
+        frames.append(renderer.render())
+
+        # Step environment
+        obs_dict = observations or {}
+        action = policy.act(observations=obs_dict)
+        if action_mask is not None:
+            action = action[action_mask]
+        action = _maybe_clip_action(action, action_low, action_high)
+
+        state = env.step(state=state, action=action)
+        steps += 1
+
+        observations = _get_observations(state)
+        cur_dist = _get_xy_distance_to_target(observations) if observations else None
+        if prev_dist is not None and cur_dist is not None:
+            ep_return += prev_dist - cur_dist
+        prev_dist = cur_dist
+
+        reached_target = _target_reached(state=state)
+        if reached_target:
+            break
+
+    # Capture final frame
+    renderer.update_scene(data, camera=camera_id)
+    frames.append(renderer.render())
+    renderer.close()
+
+    # Save video
+    imageio.mimsave(str(output_path), frames, fps=fps)
+
+    final_dist = _get_xy_distance_to_target(observations) if observations else None
+    return EpisodeResult(
+        return_=ep_return,
+        length=steps,
+        reached_target=reached_target,
+        final_xy_dist=final_dist,
+    )

src/brittle_star_project/render/__init__.py

@@ -1,3 +0,0 @@
-from .renderer import simulate_policy, SimulationConfig, ControlPolicy
-
-__all__ = ["simulate_policy", "SimulationConfig", "ControlPolicy"]

src/brittle_star_project/render/renderer.py

@@ -1,78 +0,0 @@
-from __future__ import annotations
-
-import time
-from dataclasses import dataclass
-from typing import Any, Protocol
-
-import numpy as np
-
-
-@dataclass
-class SimulationConfig:
-    realtime: bool = True
-    seed: int = 0
-
-
-class ControlPolicy(Protocol):
-    def act(self, *, obs: np.ndarray | None = None, t: float = 0.0) -> np.ndarray: ...
-
-
-def _default_observations(data: Any) -> np.ndarray:
-    qpos = np.asarray(data.qpos, dtype=np.float32).ravel()
-    qvel = np.asarray(data.qvel, dtype=np.float32).ravel()
-    return np.concatenate([qpos, qvel], axis=0)
-
-
-def simulate_policy(
-    policy: ControlPolicy,
-    config: SimulationConfig,
-    state: Any | None = None,
-) -> None:
-    """Open MuJoCo's native viewer and step using actions from a policy.
-
-    This path drives MuJoCo physics directly (mj_step) and uses the policy output
-    as `data.ctrl`.
-    """
-
-    import mujoco.viewer
-
-    if state is None:
-        raise ValueError("A valid environment state must be provided.")
-
-    model = state.mj_model
-    data = state.mj_data
-
-    start = time.time()
-    with mujoco.viewer.launch_passive(model, data) as viewer:
-        while viewer.is_running():
-            step_start = time.time()
-
-            t = time.time() - start
-
-            # Input vector for the policy
-            # TODO: custom input
-            obs = _default_observations(data)
-
-            # Policy action
-            ctrl = policy.act(obs=obs, t=t)
-
-            # Check if the policy output vector give an input for each actuator (nu)
-            # TODO: what if model trained on full morphology but we want to test on a damaged one?
-            # (nu mismatch)
-            if model.nu > 0:
-                ctrl = np.asarray(ctrl, dtype=np.float32).ravel()
-                if ctrl.shape != (model.nu,):
-                    raise ValueError(
-                        f"Policy returned ctrl shape {ctrl.shape}, expected ({model.nu},)"
-                    )
-                data.ctrl[:] = ctrl
-
-            # Step the simulation and update the viewer
-            mujoco.mj_step(model, data)
-            viewer.sync()
-
-            # If we're running in realtime mode, sleep to maintain real-time pacing.
-            if config.realtime:
-                remaining = model.opt.timestep - (time.time() - step_start)
-                if remaining > 0:
-                    time.sleep(remaining)

src/brittle_star_project/trainers/PPOTrainer.py

@@ -17,6 +17,7 @@ from experiment_logger import get_logger
 from brittle_star_project.configs.main_config import BrittleStarConfig
 from brittle_star_project.dataclasses import EpisodeStatistics
 from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
+from brittle_star_project.environment.obs_processing import create_obs_processor
 from brittle_star_project.MLPs.mlps import (
     Actor,
     AgentParams,
@@ -132,7 +133,7 @@ def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, lear
 def _normalize_obs(obs, mean, var, eps=1e-8):
     return jnp.clip((obs - mean) / jnp.sqrt(var + eps), -10.0, 10.0)
 
-
+# TODO: update to work with new obs_processor
 def _convert_obs_dict_to_array_morphology(obs_dict, morph_mode, num_segments: int, num_arms: int):
     @logged_jit
     def _filter_and_flatten(o) -> jnp.ndarray:
@@ -203,7 +204,6 @@ _SEGMENT_SCALED_KEYS = frozenset(
 )
 
 
-# TODO: update to work with extra dimension + message passing
 def _get_action_and_value_noise(
     sensor: nn.Module,
     feature_extractor: nn.Module,
@@ -329,7 +329,7 @@ def _reward_fn(env_state, next_env_state):
 
 
 def _step_env_wrapped(
-    episode_stats, env_state, action, env_step_fn, morph_mode, num_segments: int, num_arms: int
+    episode_stats, env_state, action, env_step_fn, morph_mode, num_segments: int, num_arms: int, # TODO: obs_processor
 ):
     next_env_state = env_step_fn(env_state, action)
 
@@ -361,6 +361,7 @@ def _step_env_wrapped(
             reward,
             done,
         ),
+        # TODO (obs_processor(next_env_state.observations), reward, done),
     )
 
 
@@ -501,8 +502,8 @@ class PPOTrainer:
         self.key = jax.random.PRNGKey(self.experiment.seed)
 
         self.morph_mode = self.cfg.morphology.morph_mode
+        
         self.segments_per_arm = jnp.asarray(self.cfg.morphology.segments_per_arm, dtype=jnp.int32)
-
         self.num_segments = self.segments_per_arm.sum().item()
         self.num_arms = jnp.where(self.segments_per_arm > 0, 1, 0).sum().item()
 
@@ -523,6 +524,12 @@ class PPOTrainer:
         self.actor.apply = logged_jit(self.actor.apply)
         self.critic.apply = logged_jit(self.critic.apply)
 
+        # Build the centralized observation processor: derive -> normalize -> pad -> flatten.
+        self.obs_processor = create_obs_processor(
+            bounds_dict=self.cfg.obs_bounds.to_bounds_dict(),
+            padding_masks=self.env.padding_masks,
+        )
+
         action_low = jnp.asarray(self.env.single_action_space.low, dtype=jnp.float32)
         action_high = jnp.asarray(self.env.single_action_space.high, dtype=jnp.float32)
 
@@ -536,6 +543,7 @@ class PPOTrainer:
                     morph_mode=self.morph_mode,
                     num_segments=self.num_segments,
                     num_arms=self.num_arms,
+                    # TODO: obs_processor=self.obs_processor,
                 ),
                 sensor=self.sensor,
                 feature_extractor=self.feature_extractor,
@@ -570,9 +578,6 @@ class PPOTrainer:
         def apply_feature(p, x):
             return apply_shared(self.feature_extractor, p, x)
 
-        def apply_message_passer(p, x):
-            return apply_shared(self.message_passer, p, x)
-
         self._ppo = PPO(self.ppo, apply_sensor, apply_actor, apply_critic, apply_feature)
 
         self.agent_state = self._init_agent_state()
@@ -623,6 +628,7 @@ class PPOTrainer:
         )
 
         dummy_reset = self.env.reset(seed=0)
+        
         for k, v in dummy_reset.observations.items():
             self.logger.debug(k, v.shape)
         sample_obs = _convert_obs_dict_to_array_morphology(
@@ -697,6 +703,7 @@ class PPOTrainer:
         critic_params = self.critic.init(critic_key, critic_input)
         self.logger.debug(
             f"[_init_agent_state] critic_params: {jax.tree.map(lambda x: x.shape, critic_params)}"
+        # TODO: sample_obs = self.obs_processor(dummy_reset.observations)[0]  # take first env
         )
 
         return TrainState.create(
@@ -924,6 +931,7 @@ class PPOTrainer:
         self.logger.log_non_interactive(f"Initial reset started: {time.ctime()}")
 
         env_state = self.env.reset(seed=self.experiment.seed)
+        
         next_obs = _convert_obs_dict_to_array_morphology(
             env_state.observations,
             self.morph_mode,
@@ -931,6 +939,7 @@ class PPOTrainer:
             self.num_arms,
         )
         self.logger.info(f"[train] next_obs: {next_obs.shape}")
+        # TODO: next_obs = self.obs_processor(env_state.observations)
         next_done = jnp.zeros(self.ppo.num_envs, dtype=jnp.bool_)
 
         self.logger.log_non_interactive(f"Initial reset completed: {time.ctime()}")
@@ -945,7 +954,7 @@ class PPOTrainer:
             env_state, next_obs, next_done, training_measurements, storage = self._step(
                 env_state, next_obs, next_done, iteration=iteration
             )
-            self.logger.info(f"[train] next_obs (post-step): {next_obs.shape}")
+            self.logger.debug(f"[train] next_obs (post-step): {next_obs.shape}")
             self._update_obs_stats(next_obs)
             next_obs = _normalize_obs(next_obs, self.obs_mean, self.obs_var)