feat(simulate): metadata config loading

src/brittle_star_project/configs/config_simulation.py

@@ -13,7 +13,9 @@ 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

src/brittle_star_project/environment/__init__.py

@@ -1,7 +1,9 @@
 from .env_config import ArenaConfig, EnvConfig, MorphologyConfig
 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,6 +12,7 @@ __all__ = [
     "Backend",
     "Task",
     "BrittleStarEnv",
-    "StepResult",
     "BrittleStarEnvFactory",
+    "create_obs_processor",
+    "compute_padding_masks",
 ]

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