Deployed e4869e0 with MkDocs version: 1.6.1

src/brittle_star_project/MLPs/__init__.py

@@ -0,0 +1,19 @@
+from .mlps import (
+    GenericDenseLayersWithActivation,
+    OneDenseLayerMLP,
+    Actor,
+    MessagePasser,
+    AgentParams,
+    Storage,
+)
+from .adjancency_builder import build_adjacency
+
+__all__ = [
+    "GenericDenseLayersWithActivation",
+    "OneDenseLayerMLP",
+    "Actor",
+    "MessagePasser",
+    "AgentParams",
+    "Storage",
+    "build_adjacency",
+]

src/brittle_star_project/MLPs/adjancency_builder.py

@@ -0,0 +1,67 @@
+from brittle_star_project.environment.env_config import MorphMode
+import jax.numpy as jnp
+
+
+def build_adjacency(segments_per_arm, mode: MorphMode):
+    num_arms = sum(1 for s in segments_per_arm if s > 0)
+    num_segments = sum(segments_per_arm)
+
+    # FOR NOW SEMI HARDCODE:
+    # CENTRALIZED: 1 agent, no stress, adja = 1,1 = [[1]]
+    # FULLY CONNECTED: 5 agents: adj = alle 1
+    # CENTRAL DISK:#arms= 5 agents, only neighbor as adjacent so diagonal kinda..
+    # ARM = #segments agents: diago kinda, but extra, center ring too, put center mlps first or..
+
+    if mode == MorphMode.CENTRALIZED:
+        return jnp.ones((1, 1))
+
+    if mode == MorphMode.FULLY_CONNECTED:
+        adj = jnp.ones((num_arms, num_arms))  # everybody adjacent everybody
+        return adj
+
+    if mode == MorphMode.RING:  # ring
+        adj = jnp.zeros((num_arms, num_arms))
+        for i in range(num_arms):
+            adj = adj.at[i, i].set(1)  # self
+            adj = adj.at[i, (i - 1) % num_arms].set(1)
+            adj = adj.at[i, (i + 1) % num_arms].set(1)  # left and right..
+        return adj
+
+    if mode == MorphMode.SEGMENT:
+        num_nodes = num_arms + num_segments
+        adj = jnp.zeros((num_nodes, num_nodes))
+
+        # first ring
+        for i in range(num_arms):
+            # self
+            adj = adj.at[i, i].set(1)
+
+            # ring neighbors
+            adj = adj.at[i, (i - 1) % num_arms].set(1)
+            adj = adj.at[i, (i + 1) % num_arms].set(1)
+
+        # then segment chains
+        idx = 0
+        for arm_idx, seg_count in enumerate(segments_per_arm):
+            for i in range(seg_count):
+                seg_node = num_arms + idx + i
+
+                adj = adj.at[seg_node, seg_node].set(1)
+                if i > 0:
+                    adj = adj.at[seg_node, seg_node - 1].set(1)
+                if i < seg_count - 1:
+                    adj = adj.at[seg_node, seg_node + 1].set(1)
+
+            idx += seg_count
+
+        idx = 0
+        for arm_idx, seg_count in enumerate(segments_per_arm):
+            first_seg = num_arms + idx  # first segment of this arm
+
+            # connect ring node first segment
+            adj = adj.at[arm_idx, first_seg].set(1)
+            adj = adj.at[first_seg, arm_idx].set(1)
+
+            idx += seg_count
+
+        return adj

src/brittle_star_project/MLPs/mlps.py

@@ -0,0 +1,93 @@
+from dataclasses import dataclass, fields, field
+
+import flax.linen as nn
+import jax.numpy as jnp
+import jax.tree_util
+from typing import Sequence, Callable
+from flax.linen.initializers import constant, orthogonal
+from flax.core import FrozenDict
+
+
+# semi generic so we can easily make a config for it in experiments
+class GenericDenseLayersWithActivation(nn.Module):
+    layer_sizes: Sequence[int] = field(default_factory=lambda: [64, 64])
+    activation: Callable = nn.tanh
+
+    @nn.compact
+    def __call__(self, x):
+        for size in self.layer_sizes:
+            x = nn.Dense(size, kernel_init=orthogonal(jnp.sqrt(2)))(x)
+            x = self.activation(x)
+        return x
+
+
+class OneDenseLayerMLP(nn.Module):
+    @nn.compact
+    def __call__(self, x):
+        return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x)
+
+
+class Actor(nn.Module):
+    action_dim: int
+
+    @nn.compact
+    def __call__(self, x):
+        mean = nn.Dense(self.action_dim, kernel_init=orthogonal(0.01), bias_init=constant(0.0))(x)
+        log_std = self.param("log_std", nn.initializers.zeros, (self.action_dim,))
+        return mean, log_std
+
+
+class MessagePasser(nn.Module):
+    hidden_dim: int
+    num_propagation_steps: int
+    adj_matrix: jnp.ndarray
+
+    @nn.compact
+    def __call__(self, x: jnp.ndarray):
+        for _ in range(self.num_propagation_steps):
+            # (n_nodes, feat)
+            messages = nn.Dense(self.hidden_dim)(x)
+            messages = nn.tanh(messages)
+
+            # note: if mean is wanted: adj_matrix / (adj.sum(axis=-1, keepdims=True) + 1e-8)
+            agg = self.adj_matrix
+            aggregated = agg @ messages
+
+            x_concat = jnp.concatenate([x, aggregated], axis=-1)
+
+            gate = nn.sigmoid(nn.Dense(self.hidden_dim)(x_concat))
+            candidate = nn.tanh(nn.Dense(self.hidden_dim)(x_concat))
+            x = gate * x + (1 - gate) * candidate
+
+        return x
+
+
+@jax.tree_util.register_dataclass
+@dataclass
+class AgentParams:
+    sensor_params: FrozenDict | dict
+    actor_params: FrozenDict | dict
+    critic_params: FrozenDict | dict
+    feature_extractor_params: FrozenDict | dict
+    message_passer_params: FrozenDict | dict
+
+
+@jax.tree_util.register_dataclass
+@dataclass
+class Storage:
+    obs: jnp.ndarray
+    actions: jnp.ndarray
+    logprobs: jnp.ndarray
+    dones: jnp.ndarray
+    values: jnp.ndarray
+    advantages: jnp.ndarray
+    returns: jnp.ndarray
+    rewards: jnp.ndarray
+
+    raw_actions: jnp.ndarray | None = None  # before clipping
+    means: jnp.ndarray | None = None  # policy mean
+    stds: jnp.ndarray | None = None  # policy std
+
+    def replace(self, **kwargs) -> "Storage":
+        fs = fields(self)
+        return Storage(**{f.name: kwargs.get(f.name, getattr(self, f.name)) for f in fs})

src/brittle_star_project/MLPs/routing.py

@@ -0,0 +1,22 @@
+"""Shared JAX routing utilities for decentralized multi-agent models."""
+
+import jax
+
+
+def apply_per_node(apply_fn, params, x):
+    """Apply a Flax module independently to each node.
+
+    Args:
+        apply_fn: The module's ``apply`` method (e.g. ``sensor.apply``).
+        params: Per-node parameters with shape ``(num_nodes, ...)``.
+        x: Input tensor with shape ``(batch, num_nodes, features)``.
+
+    Returns:
+        Output tensor with shape ``(batch, num_nodes, out_features)``.
+    """
+
+    def apply_single_node(p, x_node):
+        # x_node: (batch, feat) — one node's input across the batch
+        return jax.vmap(lambda xi: apply_fn(p, xi))(x_node)
+
+    return jax.vmap(apply_single_node, in_axes=(0, 1), out_axes=1)(params, x)

src/brittle_star_project/__init__.py

@@ -0,0 +1,28 @@
+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 .evaluation import (
+    PolicyAgent,
+    ControlPolicy,
+    load_metadata,
+    rollout_headless,
+    rollout_viewer,
+    EpisodeResult,
+)
+
+__all__ = [
+    "ArenaConfig",
+    "Backend",
+    "BrittleStarEnv",
+    "BrittleStarEnvFactory",
+    "EnvConfig",
+    "MorphologyConfig",
+    "Task",
+    "PolicyAgent",
+    "ControlPolicy",
+    "load_metadata",
+    "rollout_headless",
+    "rollout_viewer",
+    "EpisodeResult",
+]

src/brittle_star_project/configs/config_architecture.py

@@ -0,0 +1,66 @@
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+
+@dataclass
+class LayerConfig:
+    hidden_dims: List[int] = field(default_factory=lambda: [64, 64])
+    activation: str = "tanh"
+
+
+@dataclass
+class ArchitectureConfig:
+    """Base class for actor-critic network configurations.
+
+    Both centralized and decentralized architectures share a centralized critic
+    composed of a feature extractor followed by a shallow output layer.
+
+    See docs/design/actor-critic.md for the full design rationale.
+    """
+
+    name: str = "base"
+
+    # Actor pipeline
+    sensor: Optional[LayerConfig] = None
+    propagator: Optional[LayerConfig] = None
+    motor: Optional[LayerConfig] = None
+
+    # Critic pipeline
+    feature_extractor: Optional[LayerConfig] = None
+    critic: Optional[LayerConfig] = None
+
+    # Decentralized
+    message_passing_steps: Optional[int] = None
+    topology_type: Optional[str] = None  # Supported values: "ring", "fully_connected"
+
+
+@dataclass
+class CentralizedConfig(ArchitectureConfig):
+    """Centralized actor-critic architecture (baseline).
+
+    The actor is a single global policy composed of a sensor (input network)
+    and a motor (output network). The sensor receives the full concatenated
+    global observation; the motor projects the hidden state to all joint actions.
+
+    See docs/design/actor-critic.md for the full design rationale.
+    """
+
+    name: str = "centralized"
+
+
+@dataclass
+class DecentralizedConfig(ArchitectureConfig):
+    """Decentralized actor architecture (NerveNet-MLP variant).
+
+    Each node runs a local sensor, exchanges messages with neighbours via a
+    propagator for a fixed number of steps, and then a local motor produces
+    the joint offset for that node only.
+
+    The critic remains centralized (shared with the base class): it receives the
+    full concatenated global observation and outputs a single scalar.
+
+    See docs/design/actor-critic.md and docs/design/communication.md for the
+    full design rationale.
+    """
+
+    name: str = "decentralized"

src/brittle_star_project/configs/config_evaluation.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class EvaluationConfig:
+    """Evaluation settings.
+
+    Currently used for synchronous checkpoint evaluation during training.
+    """
+
+    # When enabled, each saved checkpoint is evaluated headlessly and the results
+    # are appended to a CSV in the run's metrics/ folder.
+    evaluate_checkpoints: bool = False
+    eval_max_steps: int = 5000
+    eval_seed: int = 0
+
+    # Cross-model comparison settings.
+    # comparison_base_seed is the starting seed for generating episode seeds.
+    comparison_base_seed: int = 0
+    # comparison_num_episodes controls how many target positions to evaluate for each model.
+    comparison_num_episodes: int = 5
+    # comparison_models lists the paths (relative to workspace root) to the .cleanrl_model files.
+    comparison_models: list[str] = field(default_factory=list)
+    # Path where the comparison results CSV will be saved (relative to workspace root).
+    comparison_output_csv: str = "metrics/model_comparison.csv"
+    # Morphology override YAML paths for cross-morphology comparison.
+    # Each path points to a file in configs/morphology/ (e.g., "configs/morphology/3_arms.yaml").
+    # When empty, each model is evaluated only on its training morphology.
+    comparison_morphologies: list[str] = field(default_factory=list)
+
+    def __post_init__(self) -> None:
+        if self.evaluate_checkpoints and self.eval_max_steps <= 0:
+            raise ValueError(
+                "Configuration Error: 'eval_max_steps' must be > 0 when "
+                "'evaluate_checkpoints' is enabled."
+            )

src/brittle_star_project/configs/config_experiment.py

@@ -0,0 +1,11 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class ExperimentConfig:
+    exp_name: str = "brittle_star_ppo"
+    seed: int = 1
+    torch_deterministic: bool = True
+    cuda: bool = True
+    debug_sanity: bool = False
+    base_run_dir: str = "runs"

src/brittle_star_project/configs/config_ppo.py

@@ -0,0 +1,22 @@
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class PPOConfig:
+    learning_rate: float = 2.5e-4
+    total_timesteps: int = 10000000
+    num_envs: int = 100
+    num_steps: int = 128
+    anneal_lr: bool = True
+    gamma: float = 0.99
+    gae_lambda: float = 0.95
+    num_minibatches: int = 4
+    update_epochs: int = 4
+    norm_adv: bool = True
+    clip_coef: float = 0.1
+    clip_vloss: bool = True
+    ent_coef: float = 0.01
+    vf_coef: float = 0.5
+    max_grad_norm: float = 0.5
+    target_kl: Optional[float] = None

src/brittle_star_project/configs/config_simulation.py

@@ -0,0 +1,32 @@
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class SimulationSettings:
+    """Settings for the simulation script."""
+
+    model_path: Optional[str] = None
+
+    # Script behavior
+    headless: bool = False
+    # If None, viewer mode runs until window closed or target reached.
+    max_steps: Optional[int] = 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

@@ -0,0 +1,36 @@
+from dataclasses import dataclass, field
+
+from experiment_logger.config_logger import LoggingConfig
+from brittle_star_project.configs.config_experiment import ExperimentConfig
+from brittle_star_project.configs.config_evaluation import EvaluationConfig
+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,
+    ObservationBoundsConfig,
+)
+
+
+@dataclass
+class BrittleStarConfig:
+    """Root configuration for a brittle star training run.
+
+    Composed of strictly separated sub-configs. Each sub-config can be swapped
+    independently via CLI or a different YAML file. See configs/README.md.
+    """
+
+    experiment: ExperimentConfig = field(default_factory=ExperimentConfig)
+    logging: LoggingConfig = field(default_factory=LoggingConfig)
+    evaluation: EvaluationConfig = field(default_factory=EvaluationConfig)
+    ppo: PPOConfig = field(default_factory=PPOConfig)
+    # This field is polymorphic; defaults to the base class to allow subclasses
+    # (CentralizedConfig, DecentralizedConfig) to be merged in via Hydra.
+    architecture: ArchitectureConfig = field(default_factory=ArchitectureConfig)
+    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

@@ -0,0 +1,48 @@
+from hydra.core.config_store import ConfigStore
+
+from experiment_logger.config_logger import LoggingConfig
+from brittle_star_project.configs.config_experiment import ExperimentConfig
+from brittle_star_project.configs.config_evaluation import EvaluationConfig
+from brittle_star_project.configs.config_ppo import PPOConfig
+from brittle_star_project.configs.config_architecture import (
+    CentralizedConfig,
+    DecentralizedConfig,
+)
+from brittle_star_project.configs.config_simulation import SimulationSettings
+from brittle_star_project.environment.env_config import (
+    MorphologyConfig,
+    ArenaConfig,
+    EnvConfig,
+    ObservationBoundsConfig,
+)
+from brittle_star_project.configs.main_config import BrittleStarConfig
+
+
+def register_configs() -> None:
+    """Register all dataclasses with Hydra's ConfigStore.
+
+    This must be called before hydra.main() processes the config, ensuring
+    every structured config is validated against its Python schema. Typos in
+    YAML keys will raise ConfigAttributeError at startup.
+    """
+    cs = ConfigStore.instance()
+
+    # Root schema
+    cs.store(name="brittle_star_config", node=BrittleStarConfig)
+
+    # Sub-config groups — each group corresponds to a configs/ subdirectory.
+    cs.store(group="experiment", name="base_experiment", node=ExperimentConfig)
+    cs.store(group="logging", name="base_logging", node=LoggingConfig)
+    cs.store(group="evaluation", name="base_evaluation", node=EvaluationConfig)
+    cs.store(group="ppo", name="base_ppo", node=PPOConfig)
+
+    # Architecture variants — swap via CLI: architecture=decentralized
+    cs.store(group="architecture", name="centralized_schema", node=CentralizedConfig)
+    cs.store(group="architecture", name="decentralized_schema", node=DecentralizedConfig)
+
+    # Environment configs
+    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/dataclasses/EpisodeStatistics.py

@@ -0,0 +1,10 @@
+import flax.struct
+import jax.numpy as jnp
+
+
+@flax.struct.dataclass
+class EpisodeStatistics:
+    episode_returns: jnp.ndarray
+    episode_lengths: jnp.ndarray
+    returned_episode_returns: jnp.ndarray
+    returned_episode_lengths: jnp.ndarray

src/brittle_star_project/dataclasses/__init__.py

@@ -0,0 +1,6 @@
+from .EpisodeStatistics import EpisodeStatistics
+
+
+__all__ = [
+    "EpisodeStatistics",
+]

src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py

@@ -0,0 +1,104 @@
+import jax
+import jax.numpy as jnp
+
+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
+
+
+class BrittleStarJaxEnvWrapper:
+    def __init__(
+        self,
+        morphology: MorphologyConfig,
+        arena: ArenaConfig,
+        env_config: EnvConfig,
+        num_envs: int,
+        backend: Backend = Backend.MJX,
+    ):
+        self._morphology = morphology
+        self._arena = arena
+        self._env_config = env_config
+        self._backend = backend
+        self._num_envs = num_envs
+        self._env = BrittleStarEnvFactory.create_environment(
+            self._backend, self._morphology, self._arena, self._env_config
+        )
+
+        # Pre-compute masks for observation padding
+        self._padding_masks = compute_padding_masks(self._morphology.segments_per_arm)
+
+        self._vectorized_reset = jax.jit(jax.vmap(self._env.reset))
+        self._vectorized_step = jax.jit(jax.vmap(self._env.step))
+        self._vectorized_action_sample = jax.jit(jax.vmap(self._env.action_space.sample))
+
+        self._action_rng = None
+
+        self.logger = get_logger()
+        self.logger.info(
+            f"Initialized BrittleStarJaxEnvWrapper with {num_envs} envs on {backend.value}"
+        )
+
+    @property
+    def backend(self):
+        return self._backend
+
+    @property
+    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
+
+    @property
+    def single_observation_space(self):
+        return self._env.observation_space
+
+    def reset(self, seed: int = 0):
+        self.logger.info(f"Resetting vectorized environment environments with seed {seed}")
+        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)
+        return state
+
+    def sample_actions(self):
+        assert self._action_rng is not None, "Call reset() before sample_actions()"
+        self._action_rng, *sub_rngs = jnp.array(
+            jax.random.split(self._action_rng, self._num_envs + 1)
+        )
+        return self._vectorized_action_sample(rng=jnp.array(sub_rngs))
+
+    def step(self, state, action):
+        return self._vectorized_step(state=state, action=action)
+
+    def close(self):
+        self._env.close()
+
+    @staticmethod
+    def default(num_envs: int, backend: Backend = Backend.MJX) -> "BrittleStarJaxEnvWrapper":
+        morphology = MorphologyConfig()
+        arena = ArenaConfig()
+        env_config = EnvConfig()
+        return BrittleStarJaxEnvWrapper(
+            morphology, arena, env_config, num_envs=num_envs, backend=backend
+        )
+
+    def __str__(self):
+        morphology_str = str(self._morphology)
+        arena_str = str(self._arena)
+        env_config_str = str(self._env_config)
+        return (
+            f"BrittleStarJaxEnvWrapper(backend={self._backend}, num_envs={self._num_envs}, "
+            + f"morphology={morphology_str}, arena={arena_str}, env_config={env_config_str})"
+        )

src/brittle_star_project/environment/__init__.py

@@ -0,0 +1,19 @@
+from .env_config import ArenaConfig, EnvConfig, MorphologyConfig, MorphMode
+from .env_types import Backend, Task
+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",
+    "EnvConfig",
+    "MorphologyConfig",
+    "Backend",
+    "Task",
+    "BrittleStarEnv",
+    "BrittleStarEnvFactory",
+    "MorphMode",
+    "create_obs_processor",
+    "compute_padding_masks",
+]

src/brittle_star_project/environment/env_config.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+
+from .env_types import Task
+
+
+class MorphMode(Enum):
+    CENTRALIZED = 0
+    FULLY_CONNECTED = 1
+    RING = 2
+    SEGMENT = 3
+
+
+@dataclass
+class MorphologyConfig:
+    """Brittle star morphology configuration.
+
+    segments_per_arm defines the number of segments for each arm. The length of
+    this list implicitly sets the number of arms. Use 0 segments to represent
+    a fully amputated arm (e.g., [4, 0, 4, 2, 4] for a 5-arm morphology with
+    arm 1 removed and arm 3 shortened).
+
+    The upstream biorobot library natively supports per-arm segment counts.
+    """
+
+    segments_per_arm: list[int] = field(default_factory=lambda: [4, 4, 4, 4, 4])
+    use_p_control: bool = True
+    use_torque_control: bool = False
+    morph_mode: MorphMode = MorphMode.CENTRALIZED
+
+    @property
+    def num_arms(self) -> int:
+        return len(self.segments_per_arm)
+
+
+@dataclass
+class ArenaConfig:
+    size: list[float] = field(default_factory=lambda: [10.0, 5.0])
+    sand_ground_color: bool = True
+    attach_target: bool = True
+    wall_height: float = 1.5
+    wall_thickness: float = 0.1
+
+
+@dataclass
+class EnvConfig:
+    """Shared environment settings.
+
+    Note: Some tasks have additional parameters (see fields below).
+    """
+
+    task: Task = Task.DIRECTED_LOCOMOTION
+
+    simulation_time: float = 10000.0
+    num_physics_steps_per_control_step: int = 10
+    time_scale: int = 2
+
+    camera_ids: list[int] = field(default_factory=lambda: [0, 1])
+    # (height, width)
+    render_size: list[int] = field(default_factory=lambda: [480, 640])
+
+    joint_randomization_noise_scale: float = 0.0
+
+    # Directed locomotion
+    target_distance: float = 3.0
+
+    # Light escape
+    # Per docs in upstream env config: integer factors of 200.
+    light_perlin_noise_scale: int = 0
+
+
+@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/env_types.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from enum import Enum
+
+
+class Backend(str, Enum):
+    """Physics backend.
+
+    - MJC: MuJoCo C engine
+    - MJX: MuJoCo XLA (JAX) engine
+    """
+
+    MJC = "MJC"
+    MJX = "MJX"
+
+
+class Task(str, Enum):
+    """Which brittle-star task/environment to instantiate."""
+
+    DIRECTED_LOCOMOTION = "directed_locomotion"
+    LIGHT_ESCAPE = "light_escape"

src/brittle_star_project/environment/env_wrapper.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+from .env_config import EnvConfig, MorphologyConfig
+from .env_types import Backend
+
+
+@dataclass(slots=True)
+class StepResult:
+    state: Any
+    reward: float | None = None
+    terminated: bool | None = None
+    truncated: bool | None = None
+    info: dict[str, Any] | None = None
+
+
+class BrittleStarEnv:
+    """Thin wrapper around the underlying DualMuJoCoEnvironment.
+
+    Goal: hide backend-specific RNG setup and provide a stable place to plug in RL.
+    """
+
+    def __init__(
+        self,
+        env: Any,
+        *,
+        backend: Backend,
+        config: EnvConfig,
+        morphology_config: MorphologyConfig | None = None,
+    ) -> None:
+        self._env = env
+        self._backend = backend
+        self._config = config
+        self._morphology_config = morphology_config
+
+    @property
+    def raw(self) -> Any:
+        return self._env
+
+    @property
+    def backend(self) -> Backend:
+        return self._backend
+
+    @property
+    def config(self) -> EnvConfig:
+        return self._config
+
+    @property
+    def morphology_config(self) -> MorphologyConfig | None:
+        return self._morphology_config
+
+    def make_rng(self, seed: int):
+        if self._backend == Backend.MJC:
+            return np.random.RandomState(seed)
+
+        import jax
+
+        return jax.random.PRNGKey(seed)
+
+    def reset(self, *, seed: int = 0):
+        rng = self.make_rng(seed)
+        state = self._env.reset(rng=rng)
+        return state
+
+    def render(self, *, state: Any):
+        return self._env.render(state=state)
+
+    def close(self) -> None:
+        self._env.close()
+
+    def step(self, *, state: Any, action: Any, rng: Any | None = None) -> StepResult:
+        """Best-effort step wrapper.
+
+        Different env libraries return different tuples; we normalize common cases.
+        """
+
+        if not hasattr(self._env, "step"):
+            raise AttributeError("Underlying env has no step() method")
+
+        step_fn = self._env.step
+        sig = inspect.signature(step_fn)
+        params = list(sig.parameters)
+
+        # Common patterns:
+        # - step(state, action)
+        # - step(state, action, rng)
+        # - step(state, action, key)
+        # We pass rng only if the callable accepts a 3rd arg.
+        if len(params) >= 3 and rng is not None:
+            out = step_fn(state, action, rng)
+        else:
+            out = step_fn(state, action)
+
+        return out

src/brittle_star_project/environment/factory.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+from dataclasses import asdict
+
+from moojoco.environment.dual import DualMuJoCoEnvironment
+
+from .env_config import ArenaConfig, EnvConfig, MorphologyConfig
+from .env_types import Backend, Task
+
+
+class BrittleStarEnvFactory:
+    """Creates brittle-star morphology, arena, and task environment instances."""
+
+    @staticmethod
+    def create_morphology(config: MorphologyConfig):
+        from biorobot.brittle_star.mjcf.morphology.morphology import (
+            MJCFBrittleStarMorphology,
+        )
+        from biorobot.brittle_star.mjcf.morphology.specification.default import (
+            default_brittle_star_morphology_specification,
+        )
+
+        spec = default_brittle_star_morphology_specification(
+            num_arms=config.num_arms,
+            num_segments_per_arm=list(config.segments_per_arm),
+            use_p_control=config.use_p_control,
+            use_torque_control=config.use_torque_control,
+        )
+        return MJCFBrittleStarMorphology(specification=spec)
+
+    @staticmethod
+    def create_arena(config: ArenaConfig):
+        from biorobot.brittle_star.mjcf.arena.aquarium import (
+            AquariumArenaConfiguration,
+            MJCFAquariumArena,
+        )
+
+        arena_config = AquariumArenaConfiguration(**asdict(config))
+        return MJCFAquariumArena(configuration=arena_config)
+
+    @staticmethod
+    def create_environment_configuration(config: EnvConfig):
+        # Import locally so the project can still be imported without these deps.
+        from biorobot.brittle_star.environment.directed_locomotion.shared import (
+            BrittleStarDirectedLocomotionEnvironmentConfiguration,
+        )
+        from biorobot.brittle_star.environment.light_escape.shared import (
+            BrittleStarLightEscapeEnvironmentConfiguration,
+        )
+
+        common = dict(
+            joint_randomization_noise_scale=config.joint_randomization_noise_scale,
+            render_mode="human",
+            simulation_time=config.simulation_time,
+            num_physics_steps_per_control_step=config.num_physics_steps_per_control_step,
+            time_scale=config.time_scale,
+            camera_ids=config.camera_ids,
+            render_size=config.render_size,
+        )
+
+        match config.task:
+            case Task.DIRECTED_LOCOMOTION:
+                return BrittleStarDirectedLocomotionEnvironmentConfiguration(
+                    target_distance=config.target_distance,
+                    **common,
+                )
+            case Task.LIGHT_ESCAPE:
+                return BrittleStarLightEscapeEnvironmentConfiguration(
+                    light_perlin_noise_scale=config.light_perlin_noise_scale,
+                    **common,
+                )
+            case _:
+                raise ValueError(f"Unsupported task: {config.task}")
+
+    @staticmethod
+    def create_environment(
+        backend: Backend,
+        morphology_config: MorphologyConfig,
+        arena_config: ArenaConfig,
+        env_config: EnvConfig,
+    ) -> DualMuJoCoEnvironment:
+        from biorobot.brittle_star.environment.directed_locomotion.dual import (
+            BrittleStarDirectedLocomotionEnvironment,
+        )
+        from biorobot.brittle_star.environment.light_escape.dual import (
+            BrittleStarLightEscapeEnvironment,
+        )
+
+        morphology = BrittleStarEnvFactory.create_morphology(morphology_config)
+        arena = BrittleStarEnvFactory.create_arena(arena_config)
+        env_configuration = BrittleStarEnvFactory.create_environment_configuration(env_config)
+
+        match env_config.task:
+            case Task.DIRECTED_LOCOMOTION:
+                env_class = BrittleStarDirectedLocomotionEnvironment
+            case Task.LIGHT_ESCAPE:
+                env_class = BrittleStarLightEscapeEnvironment
+            case _:
+                raise ValueError(f"Unsupported task: {env_config.task}")
+
+        env = env_class.from_morphology_and_arena(
+            morphology=morphology,
+            arena=arena,
+            configuration=env_configuration,
+            backend=backend.value,
+        )
+
+        from experiment_logger import get_logger
+
+        get_logger().info(f"Created {env_config.task.value} env on backend {backend.value}")
+
+        return env

src/brittle_star_project/environment/obs_processing.py

@@ -0,0 +1,192 @@
+import jax
+import jax.numpy as jnp
+from typing import Dict, Tuple, Optional
+
+from brittle_star_project.environment.env_config import MorphMode
+
+from experiment_logger import get_logger
+
+logger = get_logger()
+
+_JOINT_SCALED_KEYS = frozenset(
+    {
+        "joint_position",
+        "joint_velocity",
+        "joint_actuator_force",
+        "actuator_force",
+    }
+)
+
+_SEGMENT_SCALED_KEYS = frozenset(
+    {
+        "segment_contact",
+    }
+)
+
+
+def _build_joint_indices(segments_per_arm, indices_mlp):
+    indices = []
+    start = 0
+    for i, segs in enumerate(segments_per_arm):
+        # 2 joints per segment
+        if i in indices_mlp:
+            count = segs * 2
+            idx = jnp.arange(start, start + count)
+            indices.append(idx)
+            start += count
+    return indices
+
+
+def _build_segment_indices(segments_per_arm, indices_mlp):
+    indices = []
+    start = 0
+    for i, segs in enumerate(segments_per_arm):
+        if i in indices_mlp:
+            idx = jnp.arange(start, start + segs)
+            indices.append(idx)
+            start += segs
+    return indices
+
+
+def create_obs_processor(
+    bounds_dict: Dict[str, Tuple[float, float]],
+    num_arms: int,
+    needed_copies: int,
+    padding_masks: Optional[Dict] = None,
+    morph_mode: MorphMode = MorphMode.CENTRALIZED,
+    segments_per_arm=[4, 4, 4, 4, 4],
+    agent_indices=[0, 1, 2, 3, 4],
+):
+    # made a set to allow O(1) search
+    ordered_keys = frozenset(
+        [
+            "disk_z_tilt",
+            "joint_actuator_force",
+            "joint_position",
+            "joint_velocity",
+            "robot_direction_to_target",
+            "segment_contact",
+        ]
+    )
+    segment_indices = _build_segment_indices(segments_per_arm, agent_indices)
+    joint_indices = _build_joint_indices(segments_per_arm, agent_indices)
+
+    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 _split_to_agents(obs: dict, morph_mode) -> dict:
+        output = {}
+        num_agents = needed_copies  # IMPORTANT: number of MLPs
+
+        segs_per_arm = 4
+        joints_per_segment = 2
+        joints_per_arm = segs_per_arm * joints_per_segment
+        for key, arr in obs.items():
+            arr = jnp.asarray(arr)
+            if arr.size == 0:
+                continue
+
+            if arr.ndim == 0:
+                arr = arr.reshape(1)
+
+            if key in _SEGMENT_SCALED_KEYS:
+                per_agent = []
+                for i, _ in enumerate(agent_indices):
+                    idx = segment_indices[i]
+                    taken = jnp.take(arr, idx, axis=0)
+                    pad_len = segs_per_arm - taken.shape[0]
+                    padded = jnp.pad(taken, [(0, pad_len)] + [(0, 0)] * (taken.ndim - 1))
+
+                    per_agent.append(padded.reshape(-1))
+                arr = jnp.stack(per_agent)
+            elif key in _JOINT_SCALED_KEYS:
+                per_agent = []
+                for i, _ in enumerate(agent_indices):
+                    idx = joint_indices[i]
+                    taken = jnp.take(arr, idx, axis=0)
+                    pad_len = joints_per_arm - taken.shape[0]
+                    padded = jnp.pad(taken, [(0, pad_len)] + [(0, 0)] * (taken.ndim - 1))
+
+                    per_agent.append(padded.reshape(-1))
+                arr = jnp.stack(per_agent)
+            else:
+                arr = jnp.repeat(arr[None, :], num_agents, axis=0)
+
+            if morph_mode == MorphMode.CENTRALIZED:
+                output[key] = arr.reshape(1, -1)
+            elif key in _JOINT_SCALED_KEYS:
+                output[key] = arr.reshape(num_agents, -1)
+            elif key in _SEGMENT_SCALED_KEYS:
+                output[key] = arr[:, None]
+            else:
+                output[key] = arr
+
+        return output
+
+    def _flatten_features(obs: dict) -> jnp.ndarray:
+        """
+        Input:
+            key -> (num_arms, feat_per_key)
+
+        Output:
+            (num_arms, total_features)
+        """
+        values = []
+
+        for key in sorted(ordered_keys):
+            if key not in obs:
+                continue
+
+            arr = jnp.asarray(obs[key])  # (num_arms, feat)
+
+            if arr.size == 0:
+                continue
+
+            if arr.ndim == 1:
+                arr = arr[:, None]
+
+            arr = arr.reshape(arr.shape[0], -1)
+
+            values.append(arr)
+
+        return jnp.concatenate(values, axis=-1)  # (num_arms, total_feat)
+
+    def _process_single(obs_dict: dict) -> jnp.ndarray:
+        processed = _add_derived_features(obs_dict)
+        processed = _normalize_features(processed)
+        processed = _split_to_agents(processed, morph_mode)
+        flat = _flatten_features(processed)  # (num_arms, total_feat)
+
+        logger.debug(f"[FLATTENED FINAL] shape: {flat.shape}")
+        logger.debug(f"[PER AGENT] example row 0 shape: {flat[0].shape}")
+
+        return flat  # (agents, feat)
+
+    return jax.jit(jax.vmap(_process_single))

src/brittle_star_project/environment/padded_obs_wrapper.py

@@ -0,0 +1,54 @@
+"""Observation padding masks for amputated brittle star morphologies."""
+
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+import jax.numpy as jnp
+
+
+def compute_padding_masks(
+    segments_per_arm: Sequence[int],
+    reference_segments_per_arm: Sequence[int] = (4, 4, 4, 4, 4),
+) -> dict[str, Any]:
+    """Pre-compute boolean masks for spatial insertion of observations.
+
+    Args:
+        segments_per_arm: The current (possibly amputated) morphology.
+        reference_segments_per_arm: The full morphology that defines the expected size.
+
+    Returns:
+        A dict containing 1D boolean masks and target sizes.
+    """
+    if len(segments_per_arm) != len(reference_segments_per_arm):
+        raise ValueError(
+            f"Morphology mismatch: current has {len(segments_per_arm)} arms, "
+            f"but reference requires {len(reference_segments_per_arm)} arms."
+        )
+
+    mask_1x = []
+    mask_2x = []
+
+    for arm_idx, (actual, ref) in enumerate(zip(segments_per_arm, reference_segments_per_arm)):
+        if not isinstance(actual, int):
+            actual = actual.item()
+
+        if not isinstance(ref, int):
+            ref = ref.item()
+
+        if not (0 <= actual <= ref):
+            raise ValueError(
+                f"Invalid amputation at arm {arm_idx}: "
+                f"actual segments ({actual}) must be between 0 and reference ({ref})."
+            )
+        # 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))
+
+    return {
+        "mask_1x": jnp.array(mask_1x, dtype=bool),
+        "mask_2x": jnp.array(mask_2x, dtype=bool),
+        "target_size_1x": sum(reference_segments_per_arm),
+        "target_size_2x": sum(reference_segments_per_arm) * 2,
+    }

src/brittle_star_project/evaluation/__init__.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from .checkpoint import load_metadata, load_params, metadata_to_configs, TrainingConfig
+from .evaluate_mjx import (
+    CheckpointEvalResult,
+    append_checkpoint_eval_row,
+    build_eval_rollout_fn,
+    evaluate_checkpoint_mjx,
+)
+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
+    "load_metadata",
+    "load_params",
+    "metadata_to_configs",
+    "TrainingConfig",
+    # MJX evaluation
+    "CheckpointEvalResult",
+    "append_checkpoint_eval_row",
+    "build_eval_rollout_fn",
+    "evaluate_checkpoint_mjx",
+    # CPU evaluation
+    "evaluate_policy",
+    # policy
+    "PolicyAgent",
+    "ControlPolicy",
+    # rollout
+    "rollout_headless",
+    "rollout_viewer",
+    "EpisodeResult",
+    # video
+    "record_episode",
+    "create_evaluation_dir",
+    "save_evaluation_metadata",
+    # env builder
+    "EvalEnvBundle",
+    "build_eval_env",
+]

src/brittle_star_project/evaluation/checkpoint.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+import yaml
+from dataclasses import dataclass
+from pathlib import Path
+
+from collections.abc import Mapping
+
+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
+    message_passer_params = None
+
+    # Extract params from restored checkpoint
+    if isinstance(restored, Mapping):
+        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")
+        message_passer_params = restored.get("message_passer_params") or params_sub.get(
+            "message_passer_params"
+        )
+    elif isinstance(restored, (list, tuple)) and len(restored) >= 2:
+        params_part = restored[1]
+        if isinstance(params_part, Mapping):
+            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,
+        "message_passer_params": message_passer_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/eval_env_builder.py

@@ -0,0 +1,176 @@
+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,
+    )
+
+    training_segs_per_arm = jnp.array(training.morphology.segments_per_arm)
+
+    needed_copies = 0
+    agent_indices = [0, 1, 2, 3, 4]
+    match training.morphology.morph_mode:
+        case MorphMode.CENTRALIZED:
+            needed_copies = 1
+        case MorphMode.FULLY_CONNECTED | MorphMode.RING:
+            agent_mask = training_segs_per_arm > 0
+            agent_indices = jnp.where(agent_mask)[0].tolist()
+            needed_copies = jnp.where(training_segs_per_arm > 0, 1, 0).sum().item()
+        case MorphMode.SEGMENT:
+            agent_mask = training_segs_per_arm > 0
+            agent_indices = jnp.where(agent_mask)[0].tolist()
+            needed_copies = (
+                training_segs_per_arm.sum() + jnp.where(training_segs_per_arm > 0, 1, 0).sum()
+            ).item()
+
+    num_arms_training = jnp.where(training_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_training,
+        morph_mode=training.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
+    training_total_actions = sum(training.morphology.segments_per_arm) * 2
+    trained_action_dim = training_total_actions // 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 training.morphology.morph_mode != MorphMode.CENTRALIZED:
+        adj_matrix = build_adjacency(
+            training.morphology.segments_per_arm, training.morphology.morph_mode
+        )
+
+        override_segs = env_morphology.segments_per_arm
+        if training.morphology.morph_mode in (MorphMode.FULLY_CONNECTED, MorphMode.RING):
+            for i, segs in enumerate(override_segs):
+                if segs == 0 and i < adj_matrix.shape[0]:
+                    adj_matrix = adj_matrix.at[i, :].set(0)
+                    adj_matrix = adj_matrix.at[:, i].set(0)
+        elif training.morphology.morph_mode == MorphMode.SEGMENT:
+            for i, segs in enumerate(override_segs):
+                if segs == 0 and i < num_arms_training:
+                    adj_matrix = adj_matrix.at[i, :].set(0)
+                    adj_matrix = adj_matrix.at[:, i].set(0)
+
+            idx = 0
+            for arm_idx, seg_count in enumerate(training.morphology.segments_per_arm):
+                if override_segs[arm_idx] == 0:
+                    for i in range(seg_count):
+                        seg_node = num_arms_training + idx + i
+                        if seg_node < adj_matrix.shape[0]:
+                            adj_matrix = adj_matrix.at[seg_node, :].set(0)
+                            adj_matrix = adj_matrix.at[:, seg_node].set(0)
+                idx += seg_count
+
+    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=sum(1 for s in env_morphology.segments_per_arm if s > 0),
+        architecture=env_morphology.morph_mode.name,
+    )

src/brittle_star_project/evaluation/evaluate.py

@@ -0,0 +1,58 @@
+"""MJC-based (CPU) checkpoint evaluation.
+
+This module provides the CPU-bound evaluation path using the standard MJC backend.
+It is primarily used by the `evaluate_checkpoints` CLI to compute metrics and
+render videos.
+"""
+
+from pathlib import Path
+
+import numpy as np
+
+from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
+from brittle_star_project.environment.obs_processing import create_obs_processor
+from brittle_star_project.evaluation.policy import PolicyAgent
+from brittle_star_project.evaluation.rollout import EpisodeResult, rollout_headless
+
+
+def evaluate_policy(
+    env: BrittleStarJaxEnvWrapper,
+    policy_path: str | Path,
+    seed: int,
+    max_steps: int,
+) -> EpisodeResult:
+    """Evaluate a trained policy in a CPU-bound environment.
+
+    Args:
+        env: Initialised CPU environment (MJC backend).
+        policy_path: Path to the `.cleanrl_model` weights file.
+        seed: Random seed for environment reset.
+        max_steps: Maximum number of control steps.
+
+    Returns:
+        Structured result containing return, length, and distance metrics.
+    """
+    obs_processor = create_obs_processor(
+        bounds_dict=env.cfg.obs_bounds.to_bounds_dict(),
+        padding_masks=env.padding_masks,
+    )
+
+    action_dim = env.single_action_space.shape[0]
+
+    policy = PolicyAgent.from_checkpoint(
+        model_path=Path(policy_path),
+        action_dim=action_dim,
+        obs_processor=obs_processor,
+    )
+
+    action_low = np.asarray(env.single_action_space.low, dtype=np.float32)
+    action_high = np.asarray(env.single_action_space.high, dtype=np.float32)
+
+    return rollout_headless(
+        env=env,
+        policy=policy,
+        seed=seed,
+        max_steps=max_steps,
+        action_low=action_low,
+        action_high=action_high,
+    )

src/brittle_star_project/evaluation/evaluate_mjx.py

@@ -0,0 +1,258 @@
+"""MJX-based headless checkpoint evaluation.
+
+This module provides a fast, JIT-compiled evaluation path using the MJX
+(JAX-accelerated MuJoCo) backend. It is intended for evaluating checkpoints
+*during* or *after* a training run, where the environment and policy are
+already fully initialised.
+
+The key functions are:
+
+- `build_eval_rollout_fn` — builds and JIT-compiles a single-episode rollout function from the
+  training environment and policy components.
+- `evaluate_checkpoint_mjx` — runs that function for a given set of parameters and returns a typed
+  `CheckpointEvalResult`.
+- `append_checkpoint_eval_row` — persists the result to the run's
+  `metrics/checkpoint_evaluation.csv`, migrating old schemas automatically.
+"""
+
+from __future__ import annotations
+
+import csv
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable
+
+import jax
+import jax.numpy as jnp
+
+
+@dataclass
+class CheckpointEvalResult:
+    """Structured result from a single MJX checkpoint evaluation episode."""
+
+    steps: int
+    """Number of control steps taken (≤ max_steps)."""
+
+    reached_target: bool
+    """Whether the robot reached the target (terminated) before max_steps."""
+
+    eval_return: float
+    """Accumulated shaped reward over the episode."""
+
+    final_xy_dist: float
+    """XY distance to target at episode end. 0.0 when ``reached_target`` is True."""
+
+    initial_xy_dist: float
+    """XY distance to target at episode start."""
+
+
+def build_eval_rollout_fn(
+    *,
+    env: Any,
+    obs_processor: Callable,
+    sensor_apply: Callable,
+    actor_apply: Callable,
+    message_passer_apply: Callable | None = None,
+    action_low: jnp.ndarray,
+    action_high: jnp.ndarray,
+    reward_fn: Callable,
+) -> Callable:
+    """Build and JIT-compile a single-episode MJX evaluation rollout.
+
+    All outputs are JAX arrays. Convert to Python scalars before logging.
+
+    Args:
+        env: The training environment wrapper. Must expose `env.raw` with
+            `reset` and `step` methods compatible with `jax.vmap`.
+        obs_processor: Observation normalisation / padding callable, as
+            returned by `create_obs_processor`.
+        sensor_apply: The sensor network's `apply` method (JIT-compiled).
+        actor_apply: The actor network's `apply` method (JIT-compiled).
+        message_passer_apply: Optional message-passing module apply method.
+            When provided, it is applied between the sensor and actor, using
+            `params["message_passer_params"]`.
+        action_low: Per-joint action lower bound (JAX array, shape `(action_dim,)`).
+        action_high: Per-joint action upper bound (JAX array, shape `(action_dim,)`).
+        reward_fn: Shaped reward function with signature
+            `reward_fn(env_state, next_env_state) -> jnp.ndarray`.
+            Typically, the module-level `reward_fn` from `PPOTrainer`.
+
+    Returns:
+        A JIT-compiled callable that runs one deterministic evaluation episode.
+    """
+    # vmap over a batch of 1 so the MJX API is satisfied without any
+    # extra bookkeeping in the caller.
+    reset_1 = jax.vmap(env.raw.reset)
+    step_1 = jax.vmap(env.raw.step)
+
+    def _eval_rollout(params: dict, seed: int, max_steps: int):
+        rng = jax.random.PRNGKey(seed)
+        rngs = jnp.asarray(jax.random.split(rng, 1))
+        state = reset_1(rng=rngs)
+
+        initial_xy_dist = jnp.squeeze(state.observations["xy_distance_to_target"])
+
+        t0 = jnp.asarray(0, dtype=jnp.int32)
+        done0 = jnp.squeeze(state.terminated | state.truncated)
+        return0 = jnp.asarray(0.0, dtype=jnp.float32)
+
+        def cond(carry):
+            t, _state, done, _return_ = carry
+            return jnp.logical_and(t < max_steps, jnp.logical_not(done))
+
+        def body(carry):
+            t, state, _done, return_ = carry
+
+            obs = obs_processor(state.observations)
+            hidden = sensor_apply(params["sensor_params"], obs)
+            if message_passer_apply is not None:
+                mp_params = params["message_passer_params"]
+                hidden = jax.vmap(lambda x: message_passer_apply(mp_params, x))(hidden)
+            mean, _log_std = actor_apply(params["actor_params"], hidden)
+
+            # Deterministic action: use the actor mean, no exploration noise.
+            flat_mean = mean.reshape(mean.shape[0], -1)
+            action = jnp.clip(flat_mean, action_low, action_high)
+            next_state = step_1(state=state, action=action)
+
+            shaped_reward = reward_fn(state, next_state)
+            return_ = return_ + jnp.squeeze(shaped_reward)
+
+            done_next = jnp.squeeze(next_state.terminated | next_state.truncated)
+            return (t + 1, next_state, done_next, return_)
+
+        t, final_state, _done, return_ = jax.lax.while_loop(cond, body, (t0, state, done0, return0))
+
+        reached_target = jnp.squeeze(final_state.terminated)
+        final_xy_dist_raw = jnp.squeeze(final_state.observations["xy_distance_to_target"])
+        # Clamp to 0 when the target was reached so downstream consumers
+        # don't have to special-case "terminated" themselves.
+        final_xy_dist = jnp.where(reached_target, 0.0, final_xy_dist_raw)
+
+        return t, reached_target, return_, final_xy_dist, initial_xy_dist
+
+    return jax.jit(_eval_rollout)
+
+
+def evaluate_checkpoint_mjx(
+    eval_fn: Callable,
+    params: dict,
+    *,
+    seed: int,
+    max_steps: int,
+) -> CheckpointEvalResult:
+    """Run one deterministic evaluation episode and return typed metrics.
+
+    Args:
+        eval_fn: A JIT-compiled function as returned by `build_eval_rollout_fn`.
+        params: Agent parameter dict (e.g. ``agent_state.params``).
+        seed: Random seed for environment reset (controls target placement).
+        max_steps: Maximum number of control steps before the episode is cut off.
+
+    Returns:
+        A `CheckpointEvalResult` with all JAX arrays converted to
+        plain Python scalars.
+    """
+    steps, reached, eval_return, final_xy_dist, initial_xy_dist = eval_fn(params, seed, max_steps)
+    return CheckpointEvalResult(
+        steps=int(steps),
+        reached_target=bool(reached),
+        eval_return=float(eval_return),
+        final_xy_dist=float(final_xy_dist),
+        initial_xy_dist=float(initial_xy_dist),
+    )
+
+
+_FIELDNAMES = [
+    "checkpoint",
+    "trained_timesteps",
+    "eval_steps",
+    "eval_return",
+    "final_xy_dist",
+    "initial_xy_dist",
+    "reached_target",
+]
+
+
+def _migrate_csv_if_needed(csv_path: Path) -> None:
+    """Rewrite the CSV with the canonical field names if the schema changed.
+
+    Best-effort: any exception is silently swallowed so that a schema mismatch
+    never causes a training crash.
+    """
+    try:
+        with open(csv_path, "r", newline="") as f:
+            header = next(csv.reader(f), None)
+
+        if header is None or list(header) == _FIELDNAMES:
+            return  # Nothing to migrate.
+
+        migrated_rows: list[dict[str, Any]] = []
+        with open(csv_path, "r", newline="") as f:
+            for row in csv.DictReader(f):
+                migrated_rows.append(
+                    {
+                        "checkpoint": row.get("checkpoint", row.get("iteration")),
+                        "trained_timesteps": row.get("trained_timesteps"),
+                        "eval_steps": row.get("eval_steps", row.get("steps_to_target")),
+                        "eval_return": row.get("eval_return"),
+                        "final_xy_dist": row.get("final_xy_dist"),
+                        "initial_xy_dist": row.get("initial_xy_dist"),
+                        "reached_target": row.get("reached_target"),
+                    }
+                )
+
+        with open(csv_path, "w", newline="") as f:
+            writer = csv.DictWriter(f, fieldnames=_FIELDNAMES)
+            writer.writeheader()
+            writer.writerows(migrated_rows)
+    except Exception:
+        pass  # Never crash training on a migration issue.
+
+
+def append_checkpoint_eval_row(
+    run_dir: str | Path,
+    *,
+    iteration: int,
+    trained_timesteps: int,
+    result: CheckpointEvalResult,
+) -> Path:
+    """Append one evaluation row to `<run_dir>/metrics/checkpoint_evaluation.csv`.
+
+    Creates the file (including the `metrics/` directory) if it does not yet
+    exist. Migrates the file to the current schema if the header has changed.
+
+    Args:
+        run_dir: Root directory of the training run (Hydra's output dir).
+        iteration: Training iteration number, used as the checkpoint identifier.
+        trained_timesteps: Total environment steps taken at this checkpoint.
+        result: Evaluation result as returned by `evaluate_checkpoint_mjx`.
+
+    Returns:
+        Absolute path to the CSV file (useful for W&B sync).
+    """
+    metrics_dir = Path(run_dir) / "metrics"
+    metrics_dir.mkdir(parents=True, exist_ok=True)
+    csv_path = metrics_dir / "checkpoint_evaluation.csv"
+
+    if csv_path.exists():
+        _migrate_csv_if_needed(csv_path)
+
+    file_exists = csv_path.exists()
+    with open(csv_path, "a", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=_FIELDNAMES)
+        if not file_exists:
+            writer.writeheader()
+        writer.writerow(
+            {
+                "checkpoint": int(iteration),
+                "trained_timesteps": int(trained_timesteps),
+                "eval_steps": result.steps,
+                "eval_return": result.eval_return,
+                "final_xy_dist": result.final_xy_dist,
+                "initial_xy_dist": result.initial_xy_dist,
+                "reached_target": result.reached_target,
+            }
+        )
+
+    return csv_path

src/brittle_star_project/evaluation/policy.py

@@ -0,0 +1,168 @@
+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.MLPs.routing import apply_per_node
+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,
+        message_passer_params: Any | None = None,
+        message_passing_steps: int | None = None,
+        adj_matrix: Any | None = None,
+        action_dim: int,
+        obs_processor: Any,
+    ) -> None:
+        from brittle_star_project.MLPs.mlps import (
+            Actor,
+            GenericDenseLayersWithActivation,
+            MessagePasser,
+        )
+
+        # 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._message_passer = None
+        if message_passer_params is not None and not (
+            isinstance(message_passer_params, dict) and len(message_passer_params) == 0
+        ):
+            if message_passing_steps is None or adj_matrix is None:
+                raise ValueError(
+                    "Checkpoint contains message_passer_params but PolicyAgent was not given "
+                    "message_passing_steps and adj_matrix. Pass these when constructing the agent "
+                    "so decentralized evaluation matches training."
+                )
+
+            hidden_dim = int(layer_sizes[-1])
+            self._message_passer = MessagePasser(
+                hidden_dim=hidden_dim,
+                num_propagation_steps=int(message_passing_steps),
+                adj_matrix=jnp.asarray(adj_matrix),
+            )
+            self._message_passer.apply = jax.jit(self._message_passer.apply)
+        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,
+            "message_passer_params": message_passer_params,
+        }
+        self._obs_processor = obs_processor
+
+    @classmethod
+    def from_params(
+        cls,
+        *,
+        sensor_params: Any,
+        actor_params: Any,
+        message_passer_params: Any | None = None,
+        message_passing_steps: int | None = None,
+        adj_matrix: Any | None = None,
+        action_dim: int,
+        obs_processor: Any,
+    ) -> "PolicyAgent":
+        """Construct a PolicyAgent directly from in-memory parameters."""
+        return cls(
+            sensor_params=sensor_params,
+            actor_params=actor_params,
+            message_passer_params=message_passer_params,
+            message_passing_steps=message_passing_steps,
+            adj_matrix=adj_matrix,
+            action_dim=action_dim,
+            obs_processor=obs_processor,
+        )
+
+    def set_params(
+        self,
+        *,
+        sensor_params: Any,
+        actor_params: Any,
+        message_passer_params: Any | None = None,
+    ) -> None:
+        """Update parameters for evaluation without rebuilding the model."""
+        self._params["sensor_params"] = sensor_params
+        self._params["actor_params"] = actor_params
+        self._params["message_passer_params"] = message_passer_params
+
+    @classmethod
+    def from_checkpoint(
+        cls,
+        model_path: Path,
+        *,
+        action_dim: int,
+        obs_processor: Any,
+        message_passing_steps: int | None = None,
+        adj_matrix: Any | None = None,
+    ) -> "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"],
+            message_passer_params=params.get("message_passer_params"),
+            message_passing_steps=message_passing_steps,
+            adj_matrix=adj_matrix,
+            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)
+
+        hidden = apply_per_node(self._sensor.apply, self._params["sensor_params"], obs)
+
+        if self._message_passer is not None:
+            mp_params = self._params.get("message_passer_params")
+            if mp_params is None or (isinstance(mp_params, dict) and len(mp_params) == 0):
+                raise ValueError(
+                    "PolicyAgent has a message passer but message_passer_params are missing/empty."
+                )
+            hidden = jax.vmap(lambda x: self._message_passer.apply(mp_params, x))(hidden)
+
+        mean, _log_std = apply_per_node(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,168 @@
+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
+    initial_target_distance: 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
+    initial_target_distance = prev_dist
+    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,
+        initial_target_distance=initial_target_distance,
+    )
+
+
+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 _ 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,149 @@
+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
+    initial_dist = prev_dist
+    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,
+        initial_target_distance=initial_dist,
+    )

src/brittle_star_project/ppo.py

@@ -0,0 +1,201 @@
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+from jax import debug
+from flax.core import FrozenDict
+from experiment_logger import get_logger
+from brittle_star_project.utils import logged_jit
+
+logger = get_logger()
+
+
+# Chose to use a class as it seemed the easiest way to integrate the CleanRL code style
+# with our need to seperate concerns
+class PPO:
+    def __init__(
+        self,
+        args,
+        sensor_apply,
+        actor_apply,
+        critic_apply,
+        feature_extractor_apply,
+        message_passer=None,
+    ):
+        self.args = args
+
+        if not message_passer:
+            message_passer = identity
+
+        self.ppo_loss_grad_fn = jax.value_and_grad(
+            partial(
+                ppo_loss,
+                args=args,
+                sensor_apply=sensor_apply,
+                actor_apply=actor_apply,
+                critic_apply=critic_apply,
+                feature_extractor_apply=feature_extractor_apply,
+                message_passer=message_passer,
+            ),
+            has_aux=True,
+        )
+
+    # This PPO class should be initialized only once,
+    # or this function will need to recompile
+    @partial(logged_jit, static_argnums=0)
+    def update_ppo(self, agent_state, storage, key):
+        debug.callback(logger.debug, f"[PPO] storage.obs shape: {storage.obs.shape}")
+        debug.callback(logger.debug, f"[PPO] storage.actions shape: {storage.actions.shape}")
+        debug.callback(logger.debug, f"[PPO] storage.logprobs shape: {storage.logprobs.shape}")
+        debug.callback(logger.debug, f"[PPO] storage.advantages shape: {storage.advantages.shape}")
+        debug.callback(logger.debug, f"[PPO] storage.returns shape: {storage.returns.shape}")
+
+        args = self.args
+        ppo_loss_grad_fn = self.ppo_loss_grad_fn
+
+        def update_epoch(carry, _):
+            agent_state, key = carry
+            key, subkey = jax.random.split(key)
+
+            def flatten(x):
+                return x.reshape((-1,) + x.shape[2:])
+
+            def convert_data(x):
+                x = jax.random.permutation(subkey, x)
+                return jnp.reshape(x, (args.num_minibatches, -1) + x.shape[1:])
+
+            flatten_storage = jax.tree.map(flatten, storage)
+            shuffled_storage = jax.tree.map(convert_data, flatten_storage)
+
+            def update_minibatch(agent_state, minibatch):
+                debug.callback(logger.debug, f"[PPO] minibatch.obs: {minibatch.obs.shape}")
+                debug.callback(logger.debug, f"[PPO] minibatch.actions: {minibatch.actions.shape}")
+                debug.callback(
+                    logger.debug, f"[PPO] minibatch.logprobs: {minibatch.logprobs.shape}"
+                )
+                debug.callback(
+                    logger.debug, f"[PPO] minibatch.advantages: {minibatch.advantages.shape}"
+                )
+                debug.callback(logger.debug, f"[PPO] minibatch.returns: {minibatch.returns.shape}")
+
+                (loss, (pg_loss, v_loss, entropy_loss, approx_kl)), grads = ppo_loss_grad_fn(
+                    agent_state.params,
+                    minibatch.obs,
+                    minibatch.actions,
+                    minibatch.logprobs,
+                    minibatch.advantages,
+                    minibatch.returns,
+                )
+                agent_state = agent_state.apply_gradients(grads=grads)
+                return agent_state, (loss, pg_loss, v_loss, entropy_loss, approx_kl)
+
+            agent_state, metrics = jax.lax.scan(update_minibatch, agent_state, shuffled_storage)
+            return (agent_state, key), metrics
+
+        (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl) = jax.lax.scan(
+            update_epoch, (agent_state, key), (), length=args.update_epochs
+        )
+        return agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key
+
+
+"""
+Should be ok to use partial here, since the references to network,
+actor and critic should not change at runtime
+The cost of seperating concerns is to somehow pass these values
+that are now not in the same scope
+"""
+
+
+@partial(logged_jit, static_argnums=(0, 1, 2, 3, 4))
+def get_action_and_value(
+    sensor_apply,
+    actor_apply,
+    message_passer,
+    critic_apply,
+    feature_extractor_apply,
+    params: FrozenDict,
+    x: jnp.ndarray,
+    action: jnp.ndarray,
+):
+    hidden_sensor = sensor_apply(params["sensor_params"], x)
+    hidden_critic = feature_extractor_apply(params["feature_extractor_params"], x)
+
+    # only apply message passing in decentralized context
+    if message_passer is not None:
+        hidden_sensor = message_passer(params["message_passer_params"], hidden_sensor)
+
+    debug.callback(logger.debug, f"[SHAPE] hidden_sensor: {hidden_sensor.shape}")
+    debug.callback(logger.debug, f"[SHAPE] hidden_critic: {hidden_critic.shape}")
+
+    mean, log_std = actor_apply(params["actor_params"], hidden_sensor)
+
+    debug.callback(logger.debug, f"[SHAPE] mean: {mean.shape}")
+    debug.callback(logger.debug, f"[SHAPE] log_std: {log_std.shape}")
+    debug.callback(logger.debug, f"[SHAPE] action: {action.shape}")
+
+    log_std = jnp.clip(log_std, -5, 2)
+    std = jnp.exp(log_std)
+
+    logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi))
+    debug.callback(logger.debug, f"[SHAPE] logprob pre-sum: {logprob.shape}")
+
+    logprob = logprob.sum(axis=(-2, -1))
+    debug.callback(logger.debug, f"[SHAPE] logprob final: {logprob.shape}")
+
+    entropy = (0.5 + 0.5 * jnp.log(2 * jnp.pi) + log_std).sum(axis=(-2, -1))
+    value = critic_apply(params["critic_params"], hidden_critic).squeeze(-1)
+    debug.callback(logger.debug, f"[SHAPE] value: {value.shape}")
+
+    return logprob, entropy, value
+
+
+def ppo_loss(
+    params,
+    x,
+    a,
+    logp,
+    mb_advantages,
+    mb_returns,
+    args,
+    sensor_apply,
+    actor_apply,
+    message_passer,
+    critic_apply,
+    feature_extractor_apply,
+):
+    newlogprob, entropy, newvalue = get_action_and_value(
+        sensor_apply,
+        actor_apply,
+        message_passer,
+        critic_apply,
+        feature_extractor_apply,
+        params,
+        x,
+        a,
+    )
+    logratio = newlogprob - logp
+    ratio = jnp.exp(logratio)
+    approx_kl = ((ratio - 1) - logratio).mean()
+
+    if args.norm_adv:
+        mb_advantages = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8)
+
+    pg_loss1 = -mb_advantages * ratio
+    pg_loss2 = -mb_advantages * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef)
+    pg_loss = jnp.maximum(pg_loss1, pg_loss2).mean()
+
+    v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean()
+    entropy_loss = entropy.mean()
+    loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef
+    return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl))
+
+
+def identity(_, hidden):
+    """
+    Used for seamless jax integration,
+    avoids having branching inside jitted function,
+    used as message_passer in case it is not given,
+    (in case of centralized lvl)
+    """
+
+    return hidden

src/brittle_star_project/trainers/PPOTrainer.py

@@ -0,0 +1,988 @@
+import datetime
+import random
+import time
+from dataclasses import asdict, dataclass
+from functools import partial
+from typing import Any, Optional
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import optax
+import flax.linen as nn
+from flax.training.train_state import TrainState
+
+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.evaluation.evaluate_mjx import (
+    append_checkpoint_eval_row,
+    build_eval_rollout_fn,
+    evaluate_checkpoint_mjx,
+)
+from brittle_star_project.MLPs.routing import apply_per_node
+from brittle_star_project.MLPs.mlps import (
+    Actor,
+    AgentParams,
+    GenericDenseLayersWithActivation,
+    MessagePasser,
+    OneDenseLayerMLP,
+    Storage,
+)
+from brittle_star_project.MLPs.adjancency_builder import build_adjacency
+from brittle_star_project.ppo import PPO
+from brittle_star_project.environment import MorphMode
+from brittle_star_project.utils import logged_jit
+
+from brittle_star_project.environment.env_types import Backend
+
+# TODO: clip scaled reward?
+
+
+@logged_jit
+def _clip_action(action: jnp.ndarray, low: jnp.ndarray, high: jnp.ndarray) -> jnp.ndarray:
+    return jnp.clip(action, low, high)
+
+
+def _compute_explained_variance(values: jnp.ndarray, returns: jnp.ndarray) -> float:
+    var_returns = jnp.var(returns)
+    explained_var = 1.0 - jnp.var(returns - values) / (var_returns + 1e-8)
+    return float(explained_var)
+
+
+@logged_jit
+def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, learning_rate):
+    frac = 1.0 - (count // (minibatch_count * update_epochs)) / num_iterations
+    return learning_rate * frac
+
+
+def _get_action_and_value_noise(
+    sensor: nn.Module,
+    feature_extractor: nn.Module,
+    actor: nn.Module,
+    critic: nn.Module,
+    message_passer: Optional[nn.Module],
+    agent_state: TrainState,
+    next_obs: jnp.ndarray,
+    key,
+    action_low,
+    action_high,
+):
+    # (B, n_nodes, feat)
+    hidden = apply_per_node(sensor.apply, agent_state.params["sensor_params"], next_obs)
+
+    if message_passer is not None:
+        params = agent_state.params["message_passer_params"]
+        # (n_nodes, feat) --> let each node talk with its neighbours ==> vmap over B dimension
+        hidden = jax.vmap(lambda x: message_passer.apply(params, x))(hidden)
+
+    hidden_critic = apply_shared(
+        feature_extractor, agent_state.params["feature_extractor_params"], next_obs
+    )
+
+    mean, log_std = apply_per_node(actor.apply, agent_state.params["actor_params"], hidden)
+    log_std = jnp.clip(log_std, -5, 2)
+    key, subkey = jax.random.split(key)
+    noise = jax.random.normal(subkey, shape=mean.shape)
+    std = jnp.exp(log_std)
+
+    raw_action = mean + noise * std
+    flat_action = raw_action.reshape(
+        raw_action.shape[0], -1
+    )  # concat the per agent, keep the envs dim (batch, agent * action)
+    flat_clipped_action = _clip_action(flat_action, action_low, action_high)
+
+    logprob = -0.5 * (((raw_action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(
+        axis=(-2, -1)
+    )
+    value = apply_shared(critic, agent_state.params["critic_params"], hidden_critic)
+
+    return flat_clipped_action, raw_action, logprob, value.squeeze(-1), mean, std, key
+
+
+def _step_once(
+    carry,
+    _,
+    env_step_fn,
+    num_envs: int,
+    sensor: nn.Module,
+    feature_extractor: nn.Module,
+    actor: nn.Module,
+    critic: nn.Module,
+    message_passer: Optional[nn.Module],
+    action_low,
+    action_high,
+):
+    agent_state, episode_stats, obs, done, key, env_state, terminated_any, truncated_any = carry
+    flat_clipped_action, raw_action, logprob, value, mean, std, key = _get_action_and_value_noise(
+        sensor,
+        feature_extractor,
+        actor,
+        critic,
+        message_passer,
+        agent_state,
+        obs,
+        key,
+        action_low,
+        action_high,
+    )
+    logger = get_logger()
+
+    logger.debug(f"[_step_once] raw_action: {raw_action.shape}")
+    logger.debug(f"[_step_once] clipped_action: {flat_clipped_action.shape}")
+
+    # Supporting signals (often where mismatch originates)
+    logger.debug(f"[_step_once] logprob: {logprob.shape}")
+    logger.debug(f"[_step_once] value: {value.shape}")
+    logger.debug(f"[_step_once] mean: {mean.shape}")
+    logger.debug(f"[_step_once] std: {std.shape}")
+
+    key, reset_key = jax.random.split(key)
+    reset_rngs = jax.random.split(reset_key, num_envs)
+
+    # ---- ENV STEP ----
+    key, reset_key = jax.random.split(key)
+    reset_rngs = jax.random.split(reset_key, num_envs)
+
+    episode_stats, env_state, (next_obs, reward, next_done, terminated, truncated) = env_step_fn(
+        episode_stats,
+        env_state,
+        flat_clipped_action,
+        reset_rngs,
+    )
+
+    terminated_any = terminated_any | terminated
+    truncated_any = truncated_any | truncated
+
+    logger.debug(f"[_step_once] next_obs: {next_obs.shape}")
+    logger.debug(f"[_step_once] reward: {reward.shape}")
+    logger.debug(f"[_step_once] next_done: {next_done.shape}")
+
+    storage = Storage(
+        obs=obs,
+        actions=raw_action,
+        raw_actions=raw_action,
+        logprobs=logprob,
+        dones=done,
+        values=value,
+        rewards=reward,
+        means=mean,
+        stds=std,
+        returns=jnp.zeros_like(reward),
+        advantages=jnp.zeros_like(reward),
+    )
+    return (
+        agent_state,
+        episode_stats,
+        next_obs,
+        next_done,
+        key,
+        env_state,
+        terminated_any,
+        truncated_any,
+    ), storage
+
+
+def reward_fn(env_state, next_env_state):
+    """Shaped reward used during training and checkpoint evaluation.
+
+    Public so that ``evaluation.evaluate_mjx`` can import it and produce
+    metrics that are directly comparable to training-time returns.
+    """
+    # Positive delta_distance means the brittle star is moving *away* from target.
+    delta_distance = (
+        next_env_state.observations["xy_distance_to_target"]
+        - env_state.observations["xy_distance_to_target"]
+    ).squeeze(-1)
+
+    env_reward = next_env_state.reward
+    clipped_env_reward = jnp.clip(100 * env_reward, -10, 10)
+
+    time_penalty = 0.1
+    distance_penalty = jnp.clip(0.5 * delta_distance, -0.5, 0.5)
+    penalty = time_penalty + distance_penalty
+
+    return jnp.where(next_env_state.terminated, 50.0, clipped_env_reward - penalty)
+
+
+def _step_env_wrapped(
+    episode_stats,
+    env_state,
+    action,
+    reset_rngs,
+    env_step_fn,
+    reset_single_fn,
+    obs_processor,
+):
+    next_env_state_pre_reset = env_step_fn(env_state, action)
+
+    reward = reward_fn(env_state, next_env_state_pre_reset)
+    terminated = next_env_state_pre_reset.terminated
+    truncated = next_env_state_pre_reset.truncated
+    done = terminated | truncated
+
+    new_episode_return = episode_stats.episode_returns + reward
+    new_episode_length = episode_stats.episode_lengths + 1
+
+    episode_stats = episode_stats.replace(
+        episode_returns=new_episode_return * (1 - done),
+        episode_lengths=new_episode_length * (1 - done),
+        returned_episode_returns=jnp.where(
+            done, new_episode_return, episode_stats.returned_episode_returns
+        ),
+        returned_episode_lengths=jnp.where(
+            done, new_episode_length, episode_stats.returned_episode_lengths
+        ),
+    )
+
+    def _maybe_reset(state_i, rng_i, do_reset_i):
+        def _do(_):
+            reset_state = reset_single_fn(rng=rng_i)
+
+            def _cast_leaf(new_leaf, like_leaf):
+                if like_leaf is None or new_leaf is None:
+                    return new_leaf
+
+                # Use jnp.asarray(...) to robustly get dtype for both JAX arrays and Python scalars.
+                like_dtype = jnp.asarray(like_leaf).dtype
+
+                # Avoid unnecessary work when already matching.
+                if hasattr(new_leaf, "dtype") and new_leaf.dtype == like_dtype:
+                    return new_leaf
+
+                return jnp.asarray(new_leaf, dtype=like_dtype)
+
+            # `lax.cond` requires both branches to return identical PyTree types/dtypes.
+            return jax.tree_util.tree_map(_cast_leaf, reset_state, state_i)
+
+        def _dont(_):
+            return state_i
+
+        return jax.lax.cond(do_reset_i, _do, _dont, operand=None)
+
+    # Auto-reset done envs so rollouts continue with fresh episode initial states.
+    next_env_state = jax.vmap(_maybe_reset)(next_env_state_pre_reset, reset_rngs, done)
+
+    return (
+        episode_stats,
+        next_env_state,
+        (obs_processor(next_env_state.observations), reward, done, terminated, truncated),
+    )
+
+
+def apply_shared(net, params, x):
+    # x: (batch, nodes, feat)
+    # If the critic expects a single vector per environment:
+    batch_size = x.shape[0]
+    x_flattened = x.reshape(batch_size, -1)
+    return jax.vmap(lambda xi: net.apply(params, xi))(x_flattened)
+
+
+def _rollout_jit(
+    agent_state,
+    episode_stats,
+    env_state,
+    next_obs,
+    next_done,
+    key,
+    max_steps,
+    step_env_fn,
+    num_envs: int,
+    sensor: nn.Module,
+    feature_extractor: nn.Module,
+    actor: nn.Module,
+    critic: nn.Module,
+    message_passer: Optional[nn.Module],
+    action_low,
+    action_high,
+):
+    terminated_any0 = jnp.zeros((num_envs,), dtype=jnp.bool_)
+    truncated_any0 = jnp.zeros((num_envs,), dtype=jnp.bool_)
+
+    (
+        (
+            agent_state,
+            episode_stats,
+            next_obs,
+            next_done,
+            key,
+            env_state,
+            terminated_any,
+            truncated_any,
+        ),
+        storage,
+    ) = jax.lax.scan(
+        partial(
+            _step_once,
+            sensor=sensor,
+            feature_extractor=feature_extractor,
+            actor=actor,
+            critic=critic,
+            message_passer=message_passer,
+            env_step_fn=step_env_fn,
+            num_envs=num_envs,
+            action_low=action_low,
+            action_high=action_high,
+        ),
+        (
+            agent_state,
+            episode_stats,
+            next_obs,
+            next_done,
+            key,
+            env_state,
+            terminated_any0,
+            truncated_any0,
+        ),
+        (),
+        max_steps,
+    )
+    return (
+        agent_state,
+        episode_stats,
+        next_obs,
+        next_done,
+        storage,
+        key,
+        env_state,
+        terminated_any,
+        truncated_any,
+    )
+
+
+def _compute_gae_once(carry, inp, gamma, gae_lambda):
+    advantages = carry
+    nextdone, nextvalues, curvalues, reward = inp
+    nextnonterminal = 1.0 - nextdone
+    delta = reward + gamma * nextvalues * nextnonterminal - curvalues
+    advantages = delta + gamma * gae_lambda * nextnonterminal * advantages
+    return advantages, advantages
+
+
+def _compute_gae_jit(
+    agent_state,
+    storage,
+    next_obs,
+    next_done,
+    gamma,
+    gae_lambda,
+    num_envs,
+    feature_extractor,
+    critic,
+):
+    next_value = apply_shared(
+        critic,
+        agent_state.params["critic_params"],
+        apply_shared(feature_extractor, agent_state.params["feature_extractor_params"], next_obs),
+    ).squeeze(-1)
+
+    advantages = jnp.zeros((num_envs,))
+    dones = jnp.concatenate([storage.dones, next_done[None, :]], axis=0)
+    values = jnp.concatenate([storage.values, next_value[None, :]], axis=0)
+    _, advantages = jax.lax.scan(
+        partial(_compute_gae_once, gamma=gamma, gae_lambda=gae_lambda),
+        advantages,
+        (dones[1:], values[1:], values[:-1], storage.rewards),
+        reverse=True,
+    )
+    returns = advantages + storage.values
+    advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
+    return storage.replace(advantages=advantages, returns=returns)
+
+
+@dataclass
+class TrainingMeasurements:
+    loss: jnp.ndarray
+    pg_loss: jnp.ndarray
+    v_loss: jnp.ndarray
+    entropy_loss: jnp.ndarray
+    approx_kl: jnp.ndarray
+    avg_episodic_return: float
+    explained_variance: float
+    num_terminated: int
+    num_truncated: int
+    avg_terminated_length: Any
+    avg_truncated_length: Any
+
+
+class PPOTrainer:
+    def __init__(
+        self,
+        cfg: BrittleStarConfig,
+        env: BrittleStarJaxEnvWrapper,
+        run_dir: str,
+        run_name: str,
+    ):
+        self.cfg = cfg
+        self.ppo = cfg.ppo
+        self.experiment = cfg.experiment
+        self.logging_cfg = cfg.logging
+        self.evaluation_cfg = cfg.evaluation
+        self.env = env
+        self.run_dir = run_dir
+        self.run_name = run_name
+        self.logger = get_logger()
+
+        # Derived runtime fields
+        self.batch_size = self.ppo.num_envs * self.ppo.num_steps
+        self.num_iterations = self.ppo.total_timesteps // self.batch_size
+
+        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()
+
+        self.logger.info(f"[INIT]: Used morphology mode {self.morph_mode}")
+        self.adj = build_adjacency(cfg.morphology.segments_per_arm, self.morph_mode)
+
+        (
+            self.sensor,
+            self.message_passer,
+            self.actor,
+            self.feature_extractor,
+            self.critic,
+            self.needed_copies,
+            self.agent_indices,
+        ) = self._init_agent()
+
+        self.sensor.apply = logged_jit(self.sensor.apply)
+        self.feature_extractor.apply = logged_jit(self.feature_extractor.apply)
+        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(),
+            needed_copies=self.needed_copies,
+            num_arms=self.num_arms,
+            morph_mode=self.morph_mode,
+            padding_masks=self.env.padding_masks,
+            segments_per_arm=self.segments_per_arm,
+            agent_indices=self.agent_indices,
+        )
+
+        self.logger.debug(f"needed copies = {self.needed_copies}")
+
+        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)
+        self._action_low = action_low
+        self._action_high = action_high
+
+        self._rollout_jit = logged_jit(
+            partial(
+                _rollout_jit,
+                max_steps=self.ppo.num_steps,
+                step_env_fn=partial(
+                    _step_env_wrapped,
+                    env_step_fn=self.env.step,
+                    reset_single_fn=self.env.raw.reset,
+                    obs_processor=self.obs_processor,
+                ),
+                num_envs=self.ppo.num_envs,
+                sensor=self.sensor,
+                feature_extractor=self.feature_extractor,
+                actor=self.actor,
+                critic=self.critic,
+                message_passer=self.message_passer,
+                action_low=action_low,
+                action_high=action_high,
+            )
+        )
+        self._compute_gae_jit = logged_jit(
+            partial(
+                _compute_gae_jit,
+                num_envs=self.ppo.num_envs,
+                gamma=self.ppo.gamma,
+                gae_lambda=self.ppo.gae_lambda,
+                feature_extractor=self.feature_extractor,
+                critic=self.critic,
+            )
+        )
+
+        def apply_sensor(p, x):
+            return apply_per_node(self.sensor.apply, p, x)
+
+        def apply_actor(p, x):
+            return apply_per_node(self.actor.apply, p, x)
+
+        def apply_critic(p, x):
+            return apply_shared(self.critic, p, x)
+
+        def apply_feature(p, x):
+            return apply_shared(self.feature_extractor, p, x)
+
+        def apply_message_passer(p, x):
+            assert self.message_passer is not None
+            return jax.vmap(lambda x_in: self.message_passer.apply(p, x_in))(x)
+
+        self._ppo = PPO(
+            self.ppo,
+            apply_sensor,
+            apply_actor,
+            apply_critic,
+            apply_feature,
+            apply_message_passer if self.message_passer is not None else None,
+        )
+
+        self.agent_state = self._init_agent_state()
+
+        self.episode_stats = self._init_episode_stats()
+
+        self._init_random()
+        # Lazily-built JIT-compiled MJX eval rollout, created on first evaluation.
+        self._eval_fn = None
+
+    def _init_random(self):
+        self.logger.info(f"[RANDOM]: Setting random seed to {self.experiment.seed}")
+
+        random.seed(self.experiment.seed)
+        np.random.seed(self.experiment.seed)
+
+    def _init_agent(self):
+        self.logger.info("[AGENT]: Initializing agent...")
+        agent_indices = [0, 1, 2, 3, 4]
+        match self.morph_mode:
+            case MorphMode.CENTRALIZED:
+                needed_copies = 1
+            case MorphMode.FULLY_CONNECTED | MorphMode.RING:
+                agent_mask = self.segments_per_arm > 0
+                agent_indices = jnp.where(agent_mask)[0]
+                needed_copies = jnp.where(self.segments_per_arm > 0, 1, 0).sum().item()
+            case MorphMode.SEGMENT:
+                agent_mask = self.segments_per_arm > 0
+                agent_indices = jnp.where(agent_mask)[0]
+                needed_copies = (
+                    self.segments_per_arm.sum() + jnp.where(self.segments_per_arm > 0, 1, 0).sum()
+                ).item()
+
+        # scale actor output with size of model --> more models ==> less actions needed per model
+        actor = Actor(action_dim=self.env.single_action_space.shape[0] // needed_copies)
+        sensor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300])
+        message_passer: Optional[nn.Module] = (
+            MessagePasser(
+                hidden_dim=300,
+                num_propagation_steps=self.cfg.architecture.message_passing_steps or 4,
+                adj_matrix=self.adj,
+            )
+            if self.morph_mode != MorphMode.CENTRALIZED
+            else None
+        )
+
+        feature_extractor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300])
+        critic = OneDenseLayerMLP()
+        return (
+            sensor,
+            message_passer,
+            actor,
+            feature_extractor,
+            critic,
+            needed_copies,
+            agent_indices,
+        )
+
+    def _init_agent_state(self) -> TrainState:
+        self.logger.info("[AGENT STATE]: Initializing agent state...")
+
+        self.key, sensor_key, actor_key, critic_key, feature_extractor_key, message_passer_key = (
+            jax.random.split(self.key, 6)
+        )
+
+        dummy_reset = self.env.reset(seed=0)
+
+        for k, v in dummy_reset.observations.items():
+            self.logger.debug(k, v.shape)
+
+        sample_obs = self.obs_processor(dummy_reset.observations)[0]  # take first env
+
+        self.logger.debug(f"[_init_agent_state] sample_obs: {sample_obs.shape}")
+        self.obs_mean = jnp.zeros((sample_obs.shape[-1],))
+        self.obs_var = jnp.ones((sample_obs.shape[-1],))
+        self.obs_count = 1e-4
+        self.logger.debug(f"[_init_agent_state] obs_mean: {self.obs_mean.shape}")
+        self.logger.debug(f"[_init_agent_state] obs_var: {self.obs_var.shape}")
+
+        self.logger.debug(f"[_init_agent_state]: Needed copies: {self.needed_copies}")
+        sensor_keys = jax.random.split(sensor_key, self.needed_copies)
+        actor_keys = jax.random.split(actor_key, self.needed_copies)
+
+        # (needed_copies, X)
+        sensor_params = jax.vmap(lambda k: self.sensor.init(k, sample_obs))(sensor_keys)
+        self.logger.debug(
+            f"[_init_agent_state] sensor_params: {jax.tree.map(lambda x: x.shape, sensor_params)}"
+        )
+
+        single_sensor_param = jax.tree.map(lambda x: x[0], sensor_params)
+        self.logger.debug(
+            f"[_init_agent_state] single_sensor_param: {
+                jax.tree.map(lambda x: x.shape, single_sensor_param)
+            }"
+        )
+
+        sensor_params_sample = self.sensor.apply(single_sensor_param, sample_obs)
+        self.logger.debug(
+            f"[_init_agent_state] sensor_params_sample shape: {sensor_params_sample.shape}"
+        )
+
+        actor_params = jax.vmap(lambda k: self.actor.init(k, sensor_params_sample))(actor_keys)
+        self.logger.debug(
+            f"[_init_agent_state] actor_params: {jax.tree.map(lambda x: x.shape, actor_params)}"
+        )
+
+        message_passer_params = {}
+        if self.morph_mode != MorphMode.CENTRALIZED:
+            assert self.message_passer is not None, "decentralized modes require a message passer"
+
+            message_passer_params = self.message_passer.init(
+                message_passer_key,
+                self.sensor.apply(single_sensor_param, sample_obs),
+            )
+            self.logger.debug(
+                f"[_init_agent_state] message_passer_params: {
+                    jax.tree.map(lambda x: x.shape, message_passer_params)
+                }"
+            )
+
+        flat_obs = sample_obs.reshape(-1)  # BECAUSE 1 centralized critic
+        self.logger.debug(f"[_init_agent_state] flat_obs: {flat_obs.shape}")
+
+        feature_extractor_params = self.feature_extractor.init(feature_extractor_key, flat_obs)
+        self.logger.debug(
+            f"[_init_agent_state] feature_extractor_params: {
+                jax.tree.map(lambda x: x.shape, feature_extractor_params)
+            }"
+        )
+
+        critic_input = self.feature_extractor.apply(feature_extractor_params, flat_obs)
+        self.logger.debug(f"[_init_agent_state] critic_input: {critic_input.shape}")
+
+        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)}"
+        )
+
+        return TrainState.create(
+            apply_fn=None,
+            params=asdict(
+                AgentParams(
+                    sensor_params,
+                    actor_params,
+                    critic_params,
+                    feature_extractor_params,
+                    message_passer_params,
+                )
+            ),
+            tx=optax.chain(
+                optax.clip_by_global_norm(self.ppo.max_grad_norm),
+                optax.inject_hyperparams(optax.adam)(
+                    learning_rate=partial(
+                        _linear_schedule,
+                        minibatch_count=self.ppo.num_minibatches,
+                        update_epochs=self.ppo.update_epochs,
+                        num_iterations=self.num_iterations,
+                        learning_rate=self.ppo.learning_rate,
+                    )
+                    if self.ppo.anneal_lr
+                    else self.ppo.learning_rate,
+                    eps=1e-5,
+                ),
+            ),
+        )
+
+    def _init_episode_stats(self) -> EpisodeStatistics:
+        self.logger.info("[EPISODE STATS]: Initializing episode stats...")
+
+        return EpisodeStatistics(
+            episode_returns=jnp.zeros(self.ppo.num_envs, dtype=jnp.float32),
+            episode_lengths=jnp.zeros(self.ppo.num_envs, dtype=jnp.int32),
+            returned_episode_returns=jnp.zeros(self.ppo.num_envs, jnp.float32),
+            returned_episode_lengths=jnp.zeros(self.ppo.num_envs, dtype=jnp.int32),
+        )
+
+    def _rollout(self, env_state, next_obs, next_done) -> tuple[Any, ...]:
+        return self._rollout_jit(
+            self.agent_state,
+            self.episode_stats,
+            env_state,
+            next_obs,
+            next_done,
+            self.key,
+        )
+
+    def _compute_gae(self, storage, next_obs, next_done) -> Storage:
+        return self._compute_gae_jit(
+            self.agent_state,
+            storage,
+            next_obs,
+            next_done,
+        )
+
+    def _log(
+        self,
+        global_step,
+        episode_stats,
+        start_time,
+        iteration_time_start,
+        training_measurements,
+        storage,
+    ):
+        data = jax.device_get(
+            {
+                "rewards": storage.rewards,
+                "values": storage.values,
+                "returns": storage.returns,
+                "advantages": storage.advantages,
+            }
+        )
+
+        rollout_metrics = {
+            "rollout/reward_mean": float(np.mean(data["rewards"])),
+            "rollout/return_mean": float(np.mean(data["returns"])),
+            "rollout/value_mean": float(np.mean(data["values"])),
+            "rollout/advantage_mean": float(np.mean(data["advantages"])),
+            "rollout/advantage_std": float(np.std(data["advantages"])),
+            "rollout/value_vs_return_mse": float(np.mean((data["values"] - data["returns"]) ** 2)),
+        }
+
+        metrics = {
+            "charts/episodic_return": training_measurements.avg_episodic_return,
+            "charts/episodic_length": float(
+                np.mean(jax.device_get(episode_stats.returned_episode_lengths))
+            ),
+            "charts/explained_variance": training_measurements.explained_variance,
+            "losses/value_loss": training_measurements.v_loss[-1, -1].item(),
+            "losses/policy_loss": training_measurements.pg_loss[-1, -1].item(),
+            "losses/entropy": training_measurements.entropy_loss[-1, -1].item(),
+            "losses/approx_kl": training_measurements.approx_kl[-1, -1].item(),
+            "charts/learning_rate": self.agent_state.opt_state[1]
+            .hyperparams["learning_rate"]
+            .item(),
+            "charts/SPS": int(global_step / (time.time() - start_time)),
+            "charts/SPS_update": int(
+                self.ppo.num_envs * self.ppo.num_steps / (time.time() - iteration_time_start)
+            ),
+            "termi_trunci/num_terminated": training_measurements.num_terminated,
+            "termi_trunci/num_truncated": training_measurements.num_truncated,
+            "termi_trunci/avg_terminated_ep_length": training_measurements.avg_terminated_length,
+            "termi_trunci/avg_truncated_ep_length": training_measurements.avg_truncated_length,
+            **rollout_metrics,
+        }
+
+        self.logger.log(metrics, step=global_step)
+
+    def _step(self, env_state, next_obs, next_done, iteration: int) -> tuple:
+        if iteration == 1:
+            self.logger.log_non_interactive(f"Starting first rollout (JIT): {time.ctime()}")
+        self.logger.debug(f"[_step] next_obs (in): {next_obs.shape}")
+        (
+            self.agent_state,
+            self.episode_stats,
+            next_obs,
+            next_done,
+            storage,
+            self.key,
+            next_env_state,
+            terminated_any,
+            truncated_any,
+        ) = self._rollout(env_state, next_obs, next_done)
+        self.logger.debug(f"[_step] next_obs (post-rollout): {next_obs.shape}")
+        if iteration == 1:
+            self.logger.log_non_interactive(f"First rollout completed: {time.ctime()}")
+
+        storage = self._compute_gae(storage, next_obs, next_done)
+        self.logger.debug(f"[_step] storage.obs (post-gae): {storage.obs.shape}")
+        if iteration == 1:
+            self.logger.log_non_interactive(f"Starting first PPO update (JIT): {time.ctime()}")
+
+        self.agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, self.key = (
+            self._ppo.update_ppo(self.agent_state, storage, self.key)
+        )
+
+        if iteration == 1:
+            self.logger.log_non_interactive(f"First PPO update completed: {time.ctime()}")
+
+        avg_episodic_return = float(
+            jnp.mean(jax.device_get(self.episode_stats.returned_episode_returns)).item()
+        )
+
+        explained_var = _compute_explained_variance(storage.values, storage.returns)
+
+        terminated = terminated_any
+        truncated = truncated_any
+        episode_lengths = self.episode_stats.returned_episode_lengths
+
+        num_terminated = int(jnp.sum(terminated).item())
+        num_truncated = int(jnp.sum(truncated).item())
+
+        avg_terminated_length = jnp.sum(episode_lengths * terminated) / jnp.maximum(
+            jnp.sum(terminated), 1
+        )
+
+        avg_truncated_length = jnp.sum(episode_lengths * truncated) / jnp.maximum(
+            jnp.sum(truncated), 1
+        )
+
+        return (
+            next_env_state,
+            next_obs,
+            next_done,
+            TrainingMeasurements(
+                loss=loss,
+                pg_loss=pg_loss,
+                v_loss=v_loss,
+                entropy_loss=entropy_loss,
+                approx_kl=approx_kl,
+                avg_episodic_return=avg_episodic_return,
+                explained_variance=explained_var,
+                num_terminated=num_terminated,
+                num_truncated=num_truncated,
+                avg_terminated_length=avg_terminated_length,
+                avg_truncated_length=avg_truncated_length,
+            ),
+            storage,
+        )
+
+    def _close(self):
+        self.env.close()
+
+    def _save_model(self, model_path: str):
+        self.logger.info("[SAVE]: Saving the final model...")
+        self.logger.save_final_model(params=self.agent_state.params, metadata=asdict(self.cfg))
+
+    def _save_checkpoint(self, iteration: int):
+        self.logger.info(f"[SAVE]: Saving checkpoint at iteration {iteration}...")
+        self.logger.save_checkpoint(
+            params=self.agent_state.params, step=iteration, metadata=asdict(self.cfg)
+        )
+
+    def _evaluate_checkpoint(self, iteration: int, *, trained_timesteps: int) -> None:
+        """Evaluate the current checkpoint and persist metrics to CSV.
+
+        Delegates all evaluation logic to `evaluation.evaluate_mjx`.
+        Best-effort: a failure here must never abort training.
+        """
+        if not self.evaluation_cfg.evaluate_checkpoints:
+            return
+
+        max_steps = int(self.evaluation_cfg.eval_max_steps)
+        seed = int(self.evaluation_cfg.eval_seed)
+
+        if max_steps <= 0:
+            self.logger.warning("[EVAL]: eval_max_steps must be > 0; skipping evaluation")
+            return
+
+        if not self.logging_cfg.save_checkpoints or self.logging_cfg.checkpoint_frequency <= 0:
+            self.logger.warning(
+                "[EVAL]: evaluate_checkpoints is enabled but checkpoint saving is disabled; "
+                "skipping evaluation"
+            )
+            return
+
+        try:
+            if self._eval_fn is None:
+                if getattr(self.env, "backend", None) != Backend.MJX:
+                    self.logger.warning(
+                        f"[EVAL]: Training env backend is {self.env.backend}; "
+                        "MJX evaluation may be unavailable/slow."
+                    )
+                self._eval_fn = build_eval_rollout_fn(
+                    env=self.env,
+                    obs_processor=self.obs_processor,
+                    sensor_apply=lambda p, x: apply_per_node(self.sensor.apply, p, x),
+                    actor_apply=lambda p, x: apply_per_node(self.actor.apply, p, x),
+                    message_passer_apply=(
+                        None if self.message_passer is None else self.message_passer.apply
+                    ),
+                    action_low=self._action_low,
+                    action_high=self._action_high,
+                    reward_fn=reward_fn,
+                )
+
+            result = evaluate_checkpoint_mjx(
+                self._eval_fn,
+                self.agent_state.params,
+                seed=seed,
+                max_steps=max_steps,
+            )
+            csv_path = append_checkpoint_eval_row(
+                self.run_dir,
+                iteration=iteration,
+                trained_timesteps=int(trained_timesteps),
+                result=result,
+            )
+            self.logger.sync_file(csv_path)
+        except Exception as e:
+            self.logger.warning(f"[EVAL]: Checkpoint evaluation failed: {e}")
+
+    def train(self):
+        """
+        Train the PPO agent for a specified number of iterations.
+        Closes the environment at the end of training.
+        """
+        self.logger.info(f"running name: {self.run_name}")
+
+        self.logger.info("[TRAIN]: Resetting environment...")
+        self.logger.log_non_interactive(f"Initial reset started: {time.ctime()}")
+
+        env_state = self.env.reset(seed=self.experiment.seed)
+
+        next_obs = self.obs_processor(env_state.observations)
+        self.logger.debug(f"[train] next_obs: {next_obs.shape}")
+
+        next_done = jnp.zeros(self.ppo.num_envs, dtype=jnp.bool_)
+
+        self.logger.log_non_interactive(f"Initial reset completed: {time.ctime()}")
+
+        global_step = 0
+        start_time = time.time()
+
+        iter_bar = self.logger.progress_bar(range(1, self.num_iterations + 1))
+        for iteration in iter_bar:
+            iteration_time_start = time.time()
+
+            env_state, next_obs, next_done, training_measurements, storage = self._step(
+                env_state, next_obs, next_done, iteration=iteration
+            )
+
+            global_step += self.ppo.num_steps * self.ppo.num_envs
+            self._log(
+                global_step,
+                self.episode_stats,
+                start_time,
+                iteration_time_start,
+                training_measurements,
+                storage,
+            )
+
+            sps = int(global_step / (time.time() - start_time))
+            remaining_steps = self.ppo.total_timesteps - global_step
+            eta_seconds = int(remaining_steps / sps) if sps > 0 else 0
+            eta_str = str(datetime.timedelta(seconds=eta_seconds))
+
+            self.logger.log_non_interactive(
+                f"Iteration {iteration}/{self.num_iterations} | "
+                f"Step {global_step}/{self.ppo.total_timesteps} | "
+                f"SPS {sps} | "
+                f"Return {training_measurements.avg_episodic_return:.4f} | "
+                f"ETA {eta_str}"
+            )
+
+            if self.logging_cfg.save_checkpoints and self.logging_cfg.checkpoint_frequency > 0:
+                if iteration % self.logging_cfg.checkpoint_frequency == 0:
+                    self._save_checkpoint(iteration)
+                    self._evaluate_checkpoint(iteration, trained_timesteps=global_step)
+
+            if getattr(self.cfg.experiment, "debug_sanity", False):
+                self.logger.info("\n[SANITY CHECK] Successfully completed 1 epoch")
+                break
+
+        if self.logging_cfg.save_model:
+            model_path = f"{self.run_dir}/{self.experiment.exp_name}.cleanrl_model"
+            self._save_model(model_path=model_path)
+
+        self._close()

src/brittle_star_project/utils/__init__.py

@@ -0,0 +1,3 @@
+from .logged_jit import logged_jit
+
+__all__ = ["logged_jit"]

src/brittle_star_project/utils/logged_jit.py

@@ -0,0 +1,17 @@
+import jax
+from experiment_logger import get_logger
+
+
+def logged_jit(fn, **jit_kwargs):
+    logger = get_logger()
+    name = getattr(fn, "__name__", getattr(fn, "__qualname__", repr(fn)))
+
+    def decorator(func):
+        def traced_func(*args, **kwargs):
+            logger.debug(f"[JIT] Compiling {name}...")
+            return func(*args, **kwargs)
+
+        jitted = jax.jit(traced_func, **jit_kwargs)
+        return jitted
+
+    return decorator(fn)

src/experiment_logger/__init__.py

@@ -0,0 +1,19 @@
+"""Unified logging framework for machine learning experiments.
+
+This package provides a unified interface for logging to multiple backends
+(WandB, disk, stdout) simultaneously, ensuring no data loss.
+"""
+
+from experiment_logger.unified_logger import UnifiedLogger, get_logger, init_logger
+from experiment_logger.simple_logger import SimpleLogger
+from experiment_logger.wandb_utils import finish_wandb, init_wandb
+
+__all__ = [
+    "UnifiedLogger",
+    "SimpleLogger",
+    "get_logger",
+    "init_logger",
+    "init_wandb",
+    "finish_wandb",
+]
+__version__ = "0.1.0"

src/experiment_logger/config_logger.py

@@ -0,0 +1,36 @@
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class LoggingConfig:
+    track: bool = False
+    wandb_project_name: str = "default-project"
+    wandb_entity: Optional[str] = "SEL3-2026-Groep-4"
+    capture_video: bool = False
+
+    # Local Saving
+    save_model: bool = True  # Final model
+    save_checkpoints: bool = True  # Intermediate checkpoints
+    checkpoint_frequency: int = 100
+
+    # Remote Uploading (WandB Artifacts)
+    upload_final_model: bool = False
+    upload_checkpoints: bool = False
+
+    hf_entity: str = ""
+
+    def __post_init__(self):
+        if self.upload_final_model and not (self.track and self.save_model):
+            raise ValueError(
+                "Configuration Error: 'upload_final_model' is True, but it requires "
+                "both 'track' and 'save_model' to also be True."
+            )
+        if self.upload_checkpoints and not (self.track and self.save_checkpoints):
+            raise ValueError(
+                "Configuration Error: 'upload_checkpoints' is True, but it requires "
+                "both 'track' and 'save_checkpoints' to also be True."
+            )
+
+        # NOTE: Checkpoint evaluation settings live under the project's
+        # `evaluation` config group (see brittle_star_project.configs).

src/experiment_logger/simple_logger.py

@@ -0,0 +1,85 @@
+"""Simple terminal logger for running without external backends.
+
+This is used for standalone package usage where WandB or TensorBoard are not desired.
+It preserves the same API as UnifiedLogger but simply prints to stdout.
+"""
+
+import logging
+from typing import Any, Dict, Optional
+
+
+class SimpleLogger:
+    """Simple logger that implements the UnifiedLogger interface via print statements."""
+
+    def __init__(
+        self,
+        run_name: str = "simple_run",
+        full_config: Optional[Dict[str, Any]] = None,
+        logging_cfg: Optional[Any] = None,
+        base_dir: str = "runs",
+        save_code: bool = False,
+        log_level: int = logging.INFO,
+        _set_as_global: bool = False,
+    ):
+        self.is_interactive = True
+        self.run_name = run_name
+        self.full_config = full_config or {}
+        print(f"[INIT] SimpleLogger initialized for run: {run_name}")
+
+    def set_level(self, level: int):
+        pass
+
+    def log_non_interactive(self, msg: str, *args, **kwargs):
+        """In SimpleLogger, we just print everything as we assume interactive use."""
+        self.info(msg, *args, **kwargs)
+
+    def progress_bar(self, iterable=None, *args, **kwargs):
+        """Standard tqdm wrapper that falls back to range if tqdm is missing."""
+        try:
+            import tqdm
+
+            return tqdm.tqdm(iterable, *args, **kwargs)
+        except ImportError:
+            return iterable
+
+    def info(self, msg: str, *args, **kwargs):
+        print(f"[INFO] {msg}")
+
+    def warning(self, msg: str, *args, **kwargs):
+        print(f"[WARNING] {msg}")
+
+    def error(self, msg: str, *args, **kwargs):
+        print(f"[ERROR] {msg}")
+
+    def debug(self, msg: str, *args, **kwargs):
+        print(f"[DEBUG] {msg}")
+
+    def log(self, metrics: Dict[str, Any], step: Optional[int] = None, commit: bool = True):
+        step_str = f"Step {step}" if step is not None else "Log"
+        metric_str = ", ".join(f"{k}: {v}" for k, v in metrics.items())
+        print(f"[{step_str}] {metric_str}")
+
+    def save_checkpoint(
+        self,
+        params: Any,
+        step: int,
+        prefix: str = "checkpoint",
+        metadata: Optional[Dict[str, Any]] = None,
+    ):
+        print(f"[SAVE] Checkpoint '{prefix}' would be saved at step {step} (SimpleLogger: No-Op)")
+
+    def save_final_model(self, params: Any, metadata: Optional[Dict[str, Any]] = None):
+        print("[SAVE] Final model would be saved (SimpleLogger: No-Op)")
+
+    def sync_file(self, path: Any):
+        """No-op for SimpleLogger."""
+        pass
+
+    def finish(self):
+        print(f"[FINISH] SimpleLogger finished for run: {self.run_name}")
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.finish()

src/experiment_logger/unified_logger.py

@@ -0,0 +1,470 @@
+"""Unified logger that writes to multiple backends simultaneously.
+
+This logger ensures all experimental data is preserved by writing to:
+1. Weights & Biases (when available)
+2. Local disk (JSON files, model checkpoints, run.log)
+3. stdout (for real-time monitoring)
+"""
+
+from enum import Enum
+import logging
+import yaml
+import sys
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import flax
+import jax.numpy as jnp
+import numpy as np
+
+from experiment_logger.wandb_utils import finish_wandb, init_wandb
+from experiment_logger.config_logger import LoggingConfig
+
+# Global storage for the active logger and the proxy singleton
+_active_logger: Optional[Any] = None
+_proxy_instance: Optional["LoggerProxy"] = None
+
+
+def _sanitize_for_yaml(obj: Any) -> Any:
+    """Convert non-primitive values into YAML-safe structures.
+
+    In particular, avoids PyYAML serializing Enums as
+    ``!!python/object/apply:...`` which OmegaConf will not load.
+    """
+
+    if isinstance(obj, Enum):
+        return obj.name
+    if isinstance(obj, Path):
+        return str(obj)
+    if isinstance(obj, (np.generic, jnp.ndarray)):
+        try:
+            return obj.item()
+        except Exception:
+            pass
+    if isinstance(obj, np.ndarray):
+        return obj.tolist()
+    if isinstance(obj, dict):
+        return {str(k): _sanitize_for_yaml(v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [_sanitize_for_yaml(v) for v in obj]
+    if isinstance(obj, tuple):
+        return [_sanitize_for_yaml(v) for v in obj]
+    return obj
+
+
+def get_logger() -> "LoggerProxy":
+    """Retrieve the global LoggerProxy.
+
+    This should be used for all logging calls. It returns a proxy that
+    delegates to the active logger (defaulting to a SimpleLogger until
+    init_logger is called).
+    """
+    global _proxy_instance, _active_logger
+    if _proxy_instance is None:
+        if _active_logger is None:
+            # Fallback to SimpleLogger to avoid premature directory creation
+            from experiment_logger.simple_logger import SimpleLogger
+
+            _active_logger = SimpleLogger(run_name="pre_init")
+
+        _proxy_instance = LoggerProxy()
+
+    return _proxy_instance
+
+
+def init_logger(**kwargs) -> "UnifiedLogger":
+    """Initialize the full UnifiedLogger and set it as the active logger.
+
+    This should be called once the configuration is ready. It will create
+    the output directories and set up all logging backends.
+    """
+    global _active_logger
+    logger = UnifiedLogger(**kwargs)
+    _active_logger = logger
+    return logger
+
+
+class LoggerProxy:
+    """Proxy that delegates all method calls to the active logger instance.
+
+    This allows the logger to be swapped out (e.g., from a SimpleLogger to
+    a UnifiedLogger) without any clients needing to update their references.
+    """
+
+    def _get_logger(self) -> Any:
+        global _active_logger
+        if _active_logger is None:
+            # This shouldn't normally happen since get_logger handles it
+            from experiment_logger.simple_logger import SimpleLogger
+
+            _active_logger = SimpleLogger(run_name="pre_init_fallback")
+        return _active_logger
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._get_logger(), name)
+
+    def __enter__(self):
+        return self._get_logger().__enter__()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return self._get_logger().__exit__(exc_type, exc_val, exc_tb)
+
+
+class UnifiedLogger:
+    """Unified logger for scientific experiments with redundant backup."""
+
+    def __init__(
+        self,
+        run_name: str,
+        full_config: Dict[str, Any],
+        logging_cfg: LoggingConfig,
+        base_dir: str = "runs",
+        save_code: bool = True,
+        log_level: int = logging.INFO,
+    ):
+        """Initialize the unified logger.
+
+        Args:
+            run_name: Unique name for this run
+            full_config: Full configuration dictionary with hyperparameters to be saved
+            logging_cfg: Structured logging configuration dataclass
+            base_dir: Base directory for local storage
+            save_code: Whether to save code to WandB
+        """
+        self.run_name = run_name
+        self.full_config = full_config
+        self.use_wandb = logging_cfg.track
+        self.upload_final_model = logging_cfg.upload_final_model
+        self.upload_checkpoints = logging_cfg.upload_checkpoints
+        self.wandb_available = False
+        self.wandb_run = None
+        self.is_interactive = sys.stdout.isatty()
+
+        # Setup local storage
+        self.run_dir = Path(base_dir) / run_name
+        self.run_dir.mkdir(parents=True, exist_ok=True)
+
+        self.checkpoints_dir = self.run_dir / "checkpoints"
+        self.checkpoints_dir.mkdir(exist_ok=True)
+
+        self.metrics_dir = self.run_dir / "metrics"
+        self.metrics_dir.mkdir(exist_ok=True)
+
+        self.config_file = self.run_dir / "config.yaml"
+
+        # Setup standard Python logging mirror
+        self.text_log_file = self.run_dir / "run.log"
+        self._text_logger = logging.getLogger(f"UnifiedLogger_{self.run_name}")
+        self._text_logger.setLevel(log_level)
+        self._text_logger.propagate = False
+
+        # Avoid duplicate handlers if re-instantiated
+        if not self._text_logger.handlers:
+            fh = logging.FileHandler(self.text_log_file)
+            ch = logging.StreamHandler()
+
+            formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
+            fh.setFormatter(formatter)
+            ch.setFormatter(formatter)
+
+            self._text_logger.addHandler(fh)
+            self._text_logger.addHandler(ch)
+
+        # Save config to disk
+        self._save_config()
+
+        # Setup TensorBoard
+        self.writer = None
+        try:
+            from torch.utils.tensorboard import SummaryWriter
+
+            self.writer = SummaryWriter(self.run_dir)
+            self.info("TensorBoard SummaryWriter initialized.")
+        except ImportError:
+            self.warning("tensorboard not installed. Skipping SummaryWriter.")
+
+        # Initialize WandB if requested
+        if self.use_wandb:
+            self._init_wandb(logging_cfg.wandb_project_name, logging_cfg.wandb_entity, save_code)
+
+        # Initialize metrics storage
+        self.metrics_buffer: List[Dict[str, Any]] = []
+        self.step_counter = 0
+
+        self.info(f"Initialized UnifiedLogger for run: {run_name}")
+        self.info(f"Local storage: {self.run_dir.absolute()}")
+        self.info(f"WandB logging: {self.wandb_available}")
+
+    def set_level(self, level: int):
+        """Dynamically update the verbosity of the stdout/text logger."""
+        self._text_logger.setLevel(level)
+
+    def log_non_interactive(self, msg: str, *args, **kwargs):
+        """Log an info message only if running in a non-interactive environment."""
+        if not self.is_interactive:
+            self.info(msg, *args, **kwargs)
+
+    def progress_bar(self, iterable=None, *args, **kwargs):
+        """Wrapper around tqdm that automatically disables in non-interactive environments."""
+        import tqdm
+
+        kwargs.setdefault("disable", not self.is_interactive)
+        return tqdm.tqdm(iterable, *args, **kwargs)
+
+    def info(self, msg: str, *args, **kwargs):
+        """Log an info message to stdout and disk."""
+        self._text_logger.info(msg, *args, **kwargs)
+
+    def warning(self, msg: str, *args, **kwargs):
+        """Log a warning message to stdout and disk."""
+        self._text_logger.warning(msg, *args, **kwargs)
+
+    def error(self, msg: str, *args, **kwargs):
+        """Log an error message to stdout and disk."""
+        self._text_logger.error(msg, *args, **kwargs)
+
+    def debug(self, msg: str, *args, **kwargs):
+        """Log a debug message to stdout and disk."""
+        self._text_logger.debug(msg, *args, **kwargs)
+
+    def _init_wandb(self, project_name: str, entity: Optional[str], save_code: bool):
+        """Initialize Weights & Biases logging."""
+        self.wandb_run = init_wandb(
+            project=project_name,
+            entity=entity,
+            name=self.run_name,
+            config=self.full_config,
+            save_code=save_code,
+            resume="allow",
+        )
+        self.wandb_available = self.wandb_run is not None
+
+    def _save_config(self):
+        """Save configuration to disk."""
+        try:
+            with open(self.config_file, "w") as f:
+                yaml.safe_dump(
+                    _sanitize_for_yaml(self.full_config),
+                    f,
+                    default_flow_style=False,
+                    indent=2,
+                    sort_keys=False,
+                )
+            self.info(f"Config saved to {self.config_file}")
+        except Exception as e:
+            self.error(f"Error saving config: {e}")
+
+    def log(self, metrics: Dict[str, Any], step: Optional[int] = None, commit: bool = True):
+        """Log metrics to all backends.
+
+        Args:
+            metrics: Dictionary of metric name -> value
+            step: Global step counter (auto-incremented if None)
+            commit: Whether to commit to WandB immediately
+        """
+        if step is None:
+            step = self.step_counter
+            self.step_counter += 1
+
+        # Add timestamp
+        metrics_with_metadata = {
+            "step": step,
+            "timestamp": time.time(),
+            **metrics,
+        }
+
+        # Log to stdout
+        self._log_to_stdout(metrics_with_metadata)
+
+        # Log to WandB
+        if self.wandb_run is not None:
+            try:
+                self.wandb_run.log(metrics, step=step, commit=commit)
+            except Exception as e:
+                self.warning(f"WandB logging failed: {e}")
+
+        # Log to TensorBoard
+        if self.writer is not None:
+            for k, v in metrics.items():
+                if isinstance(v, (int, float, np.floating, np.integer)):
+                    self.writer.add_scalar(k, v, step)
+                elif hasattr(v, "item"):
+                    self.writer.add_scalar(k, v.item(), step)
+                elif isinstance(v, (np.ndarray, jnp.ndarray)) and v.size == 1:
+                    self.writer.add_scalar(k, v.item(), step)
+
+        # Buffer for disk storage
+        self.metrics_buffer.append(metrics_with_metadata)
+
+        # Periodically flush to disk
+        if len(self.metrics_buffer) >= 100:
+            self._flush_metrics()
+
+    def _log_to_stdout(self, metrics: Dict[str, Any]):
+        """Log metrics to stdout for real-time monitoring."""
+        step = metrics.get("step", "?")
+        metric_str = ", ".join(
+            f"{k}={v:.6f}" if isinstance(v, (float, np.floating)) else f"{k}={v}"
+            for k, v in metrics.items()
+            if k not in ["step", "timestamp"]
+        )
+        self.info(f"[Step {step}] {metric_str}")
+
+    def _flush_metrics(self):
+        """Flush buffered metrics to disk."""
+        if not self.metrics_buffer:
+            return
+
+        try:
+            metrics_file = self.metrics_dir / "metrics.yaml"
+            with open(metrics_file, "a") as f:
+                for metric in self.metrics_buffer:
+                    # Convert numpy/jax types to native Python types for YAML serialization
+                    serializable_metric = {}
+                    for k, v in metric.items():
+                        if hasattr(v, "item"):  # numpy/jax scalar
+                            serializable_metric[k] = v.item()
+                        elif isinstance(v, (np.ndarray, jnp.ndarray)):
+                            serializable_metric[k] = v.tolist()
+                        else:
+                            serializable_metric[k] = v
+                    f.write("---\n")
+                    yaml.safe_dump(
+                        _sanitize_for_yaml(serializable_metric),
+                        f,
+                        default_flow_style=False,
+                        sort_keys=False,
+                    )
+            self.metrics_buffer.clear()
+        except Exception as e:
+            self.error(f"Error flushing metrics: {e}")
+
+    def save_checkpoint(
+        self,
+        params: Any,
+        step: int,
+        prefix: str = "checkpoint",
+        metadata: Optional[Dict[str, Any]] = None,
+    ):
+        """Save model checkpoint to disk and optionally to WandB."""
+        checkpoint_name = f"{prefix}_step_{step}.flax"
+        checkpoint_path = self.checkpoints_dir / checkpoint_name
+
+        try:
+            # Save to disk using Flax serialization
+            with open(checkpoint_path, "wb") as f:
+                f.write(flax.serialization.to_bytes(params))
+
+            # Save metadata if provided
+            if metadata:
+                metadata_path = self.checkpoints_dir / f"{prefix}_step_{step}_metadata.yaml"
+                with open(metadata_path, "w") as f:
+                    yaml.safe_dump(
+                        _sanitize_for_yaml(metadata),
+                        f,
+                        default_flow_style=False,
+                        indent=2,
+                        sort_keys=False,
+                    )
+
+            self.info(f"Checkpoint saved: {checkpoint_path}")
+
+            # Log to WandB as artifact
+            if self.wandb_run is not None and self.upload_checkpoints:
+                try:
+                    import wandb
+
+                    artifact = wandb.Artifact(
+                        name=f"{self.run_name}_{prefix}",
+                        type="model",
+                        metadata=metadata or {},
+                    )
+                    artifact.add_file(str(checkpoint_path))
+                    if metadata:
+                        artifact.add_file(str(metadata_path))
+                    self.wandb_run.log_artifact(artifact)
+                    self.info("Checkpoint uploaded to WandB")
+                except Exception as e:
+                    self.warning(f"Could not upload checkpoint to WandB: {e}")
+
+        except Exception as e:
+            self.error(f"Error saving checkpoint: {e}")
+
+    def save_final_model(self, params: Any, metadata: Optional[Dict[str, Any]] = None):
+        """Save the final trained model."""
+        final_model_path = self.run_dir / "final_model.flax"
+
+        try:
+            with open(final_model_path, "wb") as f:
+                f.write(flax.serialization.to_bytes(params))
+
+            if metadata:
+                metadata_path = self.run_dir / "final_model_metadata.yaml"
+                with open(metadata_path, "w") as f:
+                    yaml.safe_dump(
+                        _sanitize_for_yaml(metadata),
+                        f,
+                        default_flow_style=False,
+                        indent=2,
+                        sort_keys=False,
+                    )
+
+            self.info(f"Final model saved: {final_model_path}")
+
+            # Log to WandB
+            if self.wandb_run is not None and self.upload_final_model:
+                try:
+                    import wandb
+
+                    artifact = wandb.Artifact(
+                        name=f"{self.run_name}_final_model",
+                        type="model",
+                        metadata=metadata or {},
+                    )
+                    artifact.add_file(str(final_model_path))
+                    if metadata:
+                        artifact.add_file(str(metadata_path))
+                    self.wandb_run.log_artifact(artifact)
+                except Exception as e:
+                    self.warning(f"Could not upload final model to WandB: {e}")
+
+        except Exception as e:
+            self.error(f"Error saving final model: {e}")
+
+    def sync_file(self, path: Path) -> None:
+        """Upload a file to W&B if tracking is enabled.
+
+        Best-effort: logs a warning on failure, never raises.
+        """
+        if self.wandb_run is None:
+            return
+        try:
+            import wandb
+
+            # "Simple sync" behavior: wandb will copy this file into the run.
+            wandb.save(str(path), base_path=str(path.parent))
+        except Exception as e:
+            self.warning(f"Failed to sync file to W&B: {e}")
+
+    def finish(self):
+        """Finalize logging and cleanup."""
+        # Flush remaining metrics
+        self._flush_metrics()
+
+        if self.writer is not None:
+            self.writer.close()
+
+        self.info(f"Run complete. Results saved to: {self.run_dir.absolute()}")
+
+        # Finish WandB run
+        if self.wandb_available:
+            finish_wandb()
+
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit."""
+        self.finish()

src/experiment_logger/wandb_utils.py

@@ -0,0 +1,91 @@
+"""Centralized WandB initialization utilities."""
+
+import logging
+from typing import Any, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+
+def init_wandb(
+    project: str,
+    config: Dict[str, Any],
+    name: Optional[str] = None,
+    entity: Optional[str] = None,
+    sync_tensorboard: bool = False,
+    save_code: bool = True,
+    resume: str = "allow",
+    **kwargs,
+):
+    """Initialize WandB with standardized settings.
+
+    This function provides a centralized way to initialize WandB across different
+    scripts, ensuring consistent configuration and error handling.
+
+    Args:
+        project: WandB project name
+        config: Configuration dictionary to log
+        name: Run name (auto-generated if None)
+        entity: WandB entity (team/user name)
+        sync_tensorboard: Whether to sync tensorboard logs
+        save_code: Whether to save code snapshots
+        resume: Resume strategy ("allow", "must", "never", "auto")
+        **kwargs: Additional arguments to pass to wandb.init()
+
+    Returns:
+        wandb.Run object if successful, None otherwise
+    """
+    try:
+        import wandb
+        import os
+        import sys
+
+        # Robust HPC checking: check for API key
+        has_key = os.environ.get("WANDB_API_KEY") is not None
+        if not has_key:
+            try:
+                # Check if logged in locally via settings/netrc
+                has_key = wandb.setup().settings.api_key is not None
+            except Exception:
+                pass
+
+        is_interactive = sys.stdout.isatty()
+
+        if not has_key and not is_interactive and os.environ.get("WANDB_MODE") != "offline":
+            logger.warning(
+                "WANDB_API_KEY not found and environment is non-interactive. "
+                "Switching to offline mode."
+            )
+            sync_path = f"runs/{name}" if name else "runs"
+            logger.warning(f"WandB is offline. Use 'wandb sync {sync_path}' to upload logs later.")
+            os.environ["WANDB_MODE"] = "offline"
+
+        run = wandb.init(
+            project=project,
+            entity=entity,
+            name=name,
+            config=config,
+            sync_tensorboard=sync_tensorboard,
+            save_code=save_code,
+            resume=resume,
+            **kwargs,
+        )
+        logger.info(f"WandB initialized successfully for project '{project}', run '{run.name}'")
+        return run
+    except ImportError:
+        logger.warning("WandB not installed. Skipping WandB initialization.")
+        return None
+    except Exception as e:
+        logger.error(f"Failed to initialize WandB: {e}")
+        return None
+
+
+def finish_wandb():
+    """Safely finish the current WandB run."""
+    try:
+        import wandb
+
+        if wandb.run is not None:
+            wandb.finish()
+            logger.info("WandB run finished successfully")
+    except Exception as e:
+        logger.warning(f"Error finishing WandB run: {e}")