Merge branch 'dev' into simulate-results

src/brittle_star_project/MLPs/mlps.py

@@ -58,6 +58,10 @@ class Storage:
     returns: jnp.array
     rewards: jnp.array
 
+    raw_actions: jnp.ndarray = None  # before clipping
+    means: jnp.ndarray = None  # policy mean
+    stds: jnp.ndarray = 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/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_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,12 @@
+from dataclasses import dataclass
+from typing import Optional
+from brittle_star_project.environment.env_types import Backend
+
+
+@dataclass
+class SimulationSettings:
+    """Settings for the simulation script."""
+
+    model_path: Optional[str] = None
+    model_type: str = "random"
+    backend: Backend = Backend.MJX

src/brittle_star_project/configs/main_config.py

@@ -0,0 +1,28 @@
+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_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
+
+
+@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)
+    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)
+    simulation: SimulationSettings = field(default_factory=SimulationSettings)

src/brittle_star_project/configs/register_configs.py

@@ -0,0 +1,40 @@
+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_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
+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="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="simulation", name="base_simulation", node=SimulationSettings)

src/brittle_star_project/dataclasses/PPOArgs.py

@@ -1,116 +0,0 @@
-from dataclasses import dataclass
-
-import jax
-
-
-@jax.tree_util.register_dataclass
-@dataclass
-class PPOArgs:
-    """
-    source: https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/ppo_atari_envpool_xla_jax_scan.py
-    """
-
-    # path to environment config file, if None, use default config
-    env_config_path: str | None = None
-
-    # path to hyperparameter config file (yaml), if None, use default config
-    hyperparameter_config_path: str | None = None
-
-    # the name of this experiment
-    exp_name: str = "brittle_star_ppo"
-
-    # the directory to save the experiment results
-    run_dir: str | None = None
-
-    # seed of the experiment
-    seed: int = 1
-
-    # if toggled, `torch.backends.cudnn.deterministic=False`
-    torch_deterministic: bool = True
-
-    # if toggled, cuda will be enabled by default
-    cuda: bool = True
-
-    # if toggled, this experiment will be tracked with Weights and Biases
-    track: bool = False
-
-    # the wandb's project name
-    wandb_project_name: str = "PPO-Modularity"
-
-    # the entity (team) of wandb's project
-    wandb_entity: str | None = "SEL3-2026-Groep-4"
-
-    # whether to capture videos of the agent performances (check out `videos` folder)
-    capture_video: bool = False
-
-    # whether to save model into the `runs/{run_name}` folder
-    save_model: bool = True
-
-    # checkpoint frequency (in iterations, 0 = no intermediate checkpoints)
-    checkpoint_frequency: int = 100
-
-    # whether to upload the saved model to huggingface
-    upload_model: bool = False
-
-    # the user or org name of the model repository from the Hugging Face Hub
-    hf_entity: str = ""
-
-    # ==== Algorithm specific dataclasses ====
-
-    # total timesteps of the experiments
-    total_timesteps: int = 10000000
-
-    # the learning rate of the optimizer
-    learning_rate: float = 2.5e-4
-
-    # the number of parallel game environments
-    num_envs: int = 100
-
-    # the number of steps to run in each environment per policy rollout
-    num_steps: int = 128
-
-    # Toggle learning rate annealing for policy and value networks
-    anneal_lr: bool = True
-
-    # the discount factor gamma
-    gamma: float = 0.99
-
-    # the lambda for the general advantage estimation
-    gae_lambda: float = 0.95
-
-    # the number of mini-batches
-    num_minibatches: int = 4
-
-    # the K epochs to update the policy
-    update_epochs: int = 4
-
-    # Toggles advantages normalization
-    norm_adv: bool = True
-
-    # the surrogate clipping coefficient
-    clip_coef: float = 0.1
-
-    # Toggles whether or not to use a clipped loss for the value function, as per the paper.
-    clip_vloss: bool = True
-
-    # coefficient of the entropy
-    ent_coef: float = 0.01
-
-    # coefficient of the value function
-    vf_coef: float = 0.5
-
-    # the maximum norm for the gradient clipping
-    max_grad_norm: float = 0.5
-
-    # the target KL divergence threshold
-    target_kl: float | None = None
-
-    # ==== to be filled in runtime ====
-    # the batch size (computed in runtime)
-    batch_size: int = 0
-
-    # the mini-batch size (computed in runtime)
-    minibatch_size: int = 0
-
-    # the number of iterations (computed in runtime)
-    num_iterations: int = 0

src/brittle_star_project/dataclasses/__init__.py

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

src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py

@@ -1,14 +1,11 @@
 import jax
 import jax.numpy as jnp
 
-from brittle_star_project import (
-    EnvConfig,
-    BrittleStarEnvFactory,
-    MorphologyConfig,
-    ArenaConfig,
-    Backend,
-)
-from brittle_star_project.environment import from_file
+from experiment_logger import get_logger
+from .env_config import EnvConfig, MorphologyConfig, ArenaConfig
+from .env_types import Backend
+from .factory import BrittleStarEnvFactory
+from .padded_obs_wrapper import compute_padding_masks, pad_observations_batched
 
 
 class BrittleStarJaxEnvWrapper:
@@ -29,14 +26,15 @@ class BrittleStarJaxEnvWrapper:
             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
 
-        from experiment_logger import get_logger
-
         self.logger = get_logger()
         self.logger.info(
             f"Initialized BrittleStarJaxEnvWrapper with {num_envs} envs on {backend.value}"
@@ -62,7 +60,12 @@ class BrittleStarJaxEnvWrapper:
         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))
-        return self._vectorized_reset(rng=env_rngs)
+        state = self._vectorized_reset(rng=env_rngs)
+
+        state = state.replace(
+            observations=pad_observations_batched(state.observations, self._padding_masks)
+        )
+        return state
 
     def sample_actions(self):
         assert self._action_rng is not None, "Call reset() before sample_actions()"
@@ -72,7 +75,12 @@ class BrittleStarJaxEnvWrapper:
         return self._vectorized_action_sample(rng=jnp.array(sub_rngs))
 
     def step(self, state, action):
-        return self._vectorized_step(state=state, action=action)
+        next_state = self._vectorized_step(state=state, action=action)
+
+        next_state = next_state.replace(
+            observations=pad_observations_batched(next_state.observations, self._padding_masks)
+        )
+        return next_state
 
     def close(self):
         self._env.close()
@@ -86,15 +94,6 @@ class BrittleStarJaxEnvWrapper:
             morphology, arena, env_config, num_envs=num_envs, backend=backend
         )
 
-    @staticmethod
-    def from_config(
-        config_path: str, num_envs: int, backend: Backend = Backend.MJX
-    ) -> "BrittleStarJaxEnvWrapper":
-        morphology_cfg, arena_cfg, env_cfg = from_file(config_path)
-        return BrittleStarJaxEnvWrapper(
-            morphology_cfg, arena_cfg, env_cfg, num_envs=num_envs, backend=backend
-        )
-
     def __str__(self):
         morphology_str = str(self._morphology)
         arena_str = str(self._arena)

src/brittle_star_project/environment/__init__.py

@@ -1,4 +1,4 @@
-from .env_config import ArenaConfig, EnvConfig, MorphologyConfig, from_file
+from .env_config import ArenaConfig, EnvConfig, MorphologyConfig
 from .env_types import Backend, Task
 from .env_wrapper import BrittleStarEnv, StepResult
 from .factory import BrittleStarEnvFactory
@@ -12,5 +12,4 @@ __all__ = [
     "BrittleStarEnv",
     "StepResult",
     "BrittleStarEnvFactory",
-    "from_file",
 ]

src/brittle_star_project/environment/env_config.py

@@ -5,24 +5,37 @@ from dataclasses import dataclass, field
 from .env_types import Task
 
 
-@dataclass(frozen=True, slots=True)
+@dataclass
 class MorphologyConfig:
-    num_arms: int = 5
-    num_segments_per_arm: int = 4
+    """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
 
+    @property
+    def num_arms(self) -> int:
+        return len(self.segments_per_arm)
 
-@dataclass(frozen=True, slots=True)
+
+@dataclass
 class ArenaConfig:
-    size: tuple[float, float] = (10.0, 5.0)
+    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(frozen=True, slots=True)
+@dataclass
 class EnvConfig:
     """Shared environment settings.
 
@@ -31,13 +44,13 @@ class EnvConfig:
 
     task: Task = Task.DIRECTED_LOCOMOTION
 
-    simulation_time: float = 5.0
+    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: tuple[int, int] = (480, 640)
+    render_size: list[int] = field(default_factory=lambda: [480, 640])
 
     joint_randomization_noise_scale: float = 0.0
 
@@ -47,16 +60,3 @@ class EnvConfig:
     # Light escape
     # Per docs in upstream env config: integer factors of 200.
     light_perlin_noise_scale: int = 0
-
-
-def from_file(path: str) -> tuple[MorphologyConfig, ArenaConfig, EnvConfig]:
-    """Load configurations from a YAML file."""
-    import yaml
-
-    with open(path, "r") as f:
-        config_dict = yaml.safe_load(f)
-
-        morphology = MorphologyConfig(**config_dict.get("morphology", {}))
-        arena = ArenaConfig(**config_dict.get("arena", {}))
-        env = EnvConfig(**config_dict.get("env", {}))
-        return morphology, arena, env

src/brittle_star_project/environment/factory.py

@@ -22,7 +22,7 @@ class BrittleStarEnvFactory:
 
         spec = default_brittle_star_morphology_specification(
             num_arms=config.num_arms,
-            num_segments_per_arm=config.num_segments_per_arm,
+            num_segments_per_arm=list(config.segments_per_arm),
             use_p_control=config.use_p_control,
             use_torque_control=config.use_torque_control,
         )

src/brittle_star_project/environment/padded_obs_wrapper.py

@@ -0,0 +1,108 @@
+"""Observation padding wrapper for amputated brittle star morphologies.
+
+When using a centralized controller, the global observation vector must remain
+a constant size regardless of how many segments are amputated. This wrapper pads
+the observation dictionary values with zeros using spatial insertion so that the
+flattened observation maintains the correct physical mapping to the neural network.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+import jax.numpy as jnp
+
+# Observation keys whose size scales with the number of joints (2 per segment).
+_JOINT_SCALED_KEYS = frozenset(
+    {
+        "joint_position",
+        "joint_velocity",
+        "joint_actuator_force",
+        "actuator_force",
+    }
+)
+
+# Observation keys whose size scales with the number of segments (1 per segment).
+_SEGMENT_SCALED_KEYS = frozenset(
+    {
+        "segment_contact",
+    }
+)
+
+
+def compute_padding_masks(
+    segments_per_arm: tuple[int, ...],
+    reference_segments_per_arm: tuple[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 (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)
+        # 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,
+    }
+
+
+def pad_observation(
+    obs: dict[str, Any],
+    masks: dict[str, Any],
+) -> dict[str, Any]:
+    """Pad an observation dict using spatial insertion."""
+    padded = {}
+    for key, value in obs.items():
+        if key in _JOINT_SCALED_KEYS:
+            out = jnp.zeros(masks["target_size_2x"], dtype=value.dtype)
+            padded[key] = out.at[masks["mask_2x"]].set(value)
+        elif key in _SEGMENT_SCALED_KEYS:
+            out = jnp.zeros(masks["target_size_1x"], dtype=value.dtype)
+            padded[key] = out.at[masks["mask_1x"]].set(value)
+        else:
+            padded[key] = value
+    return padded
+
+
+def pad_observations_batched(
+    obs: dict[str, Any],
+    masks: dict[str, Any],
+) -> dict[str, Any]:
+    """Pad a batched observation dict (leading batch dimension) using spatial insertion."""
+    padded = {}
+    for key, value in obs.items():
+        batch_size = value.shape[0]
+        if key in _JOINT_SCALED_KEYS:
+            out = jnp.zeros((batch_size, masks["target_size_2x"]), dtype=value.dtype)
+            padded[key] = out.at[:, masks["mask_2x"]].set(value)
+        elif key in _SEGMENT_SCALED_KEYS:
+            out = jnp.zeros((batch_size, masks["target_size_1x"]), dtype=value.dtype)
+            padded[key] = out.at[:, masks["mask_1x"]].set(value)
+        else:
+            padded[key] = value
+    return padded

src/brittle_star_project/ppo.py

@@ -99,6 +99,7 @@ def get_action_and_value(
     hidden_critic = feature_extractor_apply(params["feature_extractor_params"], x)
     hidden_sensor = message_passer(hidden_sensor)
     mean, log_std = actor_apply(params["actor_params"], hidden_sensor)
+    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)).sum(-1)
@@ -142,6 +143,7 @@ def ppo_loss(
     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

src/brittle_star_project/trainers/PPOTrainer.py

@@ -13,7 +13,8 @@ from flax.training.train_state import TrainState
 
 from experiment_logger import get_logger
 
-from brittle_star_project.dataclasses import EpisodeStatistics, PPOArgs
+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.MLPs.mlps import (
     Actor,
@@ -24,6 +25,33 @@ from brittle_star_project.MLPs.mlps import (
 )
 from brittle_star_project.ppo import PPO
 
+# TODO: move to config
+_ALLOWED_OBS_KEYS = {
+    "joint_position",
+    "joint_velocity",
+    "joint_actuator_force",
+    "actuator_force",
+    "disk_position",
+    "disk_rotation",
+    "disk_linear_velocity",
+    "disk_angular_velocity",
+    "unit_xy_direction_to_target",
+    "xy_distance_to_target",
+}
+# TODO: clip scaled reward?
+
+
+@jax.jit
+def _get_xy_distance_to_target(obs_dict: dict) -> jnp.ndarray:
+    """Extract xy_distance_to_target for all environments."""
+    # obs_dict is a dict of arrays with leading batch dimension (num_envs, ...)
+    return obs_dict["xy_distance_to_target"].squeeze(-1)  # shape: (num_envs,)
+
+
+@jax.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)
@@ -37,11 +65,25 @@ def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, lear
     return learning_rate * frac
 
 
+@jax.jit
+def _normalize_obs(obs, mean, var, eps=1e-8):
+    return jnp.clip((obs - mean) / jnp.sqrt(var + eps), -10.0, 10.0)
+
+
 @jax.jit
 def _convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray:
-    return jax.vmap(lambda o: jnp.concatenate([v.flatten() for v in o.values() if v.size > 0]))(
-        obs_dict
-    )
+    """Convert the raw observation dict → flat array, filtering unwanted keys."""
+
+    def _filter_and_flatten(o: dict) -> jnp.ndarray:
+        values = []
+        for key in sorted(o.keys()):
+            if key in _ALLOWED_OBS_KEYS:  # TODO: NORMALIZATION or .. of observations??
+                v = o[key]
+                if v.size > 0:
+                    values.append(jnp.asarray(v).flatten())
+        return jnp.concatenate(values)
+
+    return jax.vmap(_filter_and_flatten)(obs_dict)
 
 
 def _get_action_and_value_noise(
@@ -52,6 +94,8 @@ def _get_action_and_value_noise(
     agent_state: TrainState,
     next_obs: jnp.ndarray,
     key: jax.random.PRNGKey,
+    action_low,
+    action_high,
 ):
     hidden = sensor.apply(agent_state.params["sensor_params"], next_obs)
     hidden_critic = feature_extractor.apply(
@@ -59,13 +103,16 @@ def _get_action_and_value_noise(
     )
 
     mean, log_std = 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)
-    action = mean + noise * std
-    logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
+    raw_action = mean + noise * std
+    clipped_action = _clip_action(raw_action, action_low, action_high)
+    logprob = -0.5 * (((raw_action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
     value = critic.apply(agent_state.params["critic_params"], hidden_critic)
-    return action, logprob, value.squeeze(-1), key
+
+    return clipped_action, raw_action, logprob, value.squeeze(-1), mean, std, key
 
 
 def _step_once(
@@ -76,23 +123,28 @@ def _step_once(
     feature_extractor: GenericDenseLayersWithActivation,
     actor: Actor,
     critic: OneDenseLayerMLP,
+    action_low,
+    action_high,
 ):
     agent_state, episode_stats, obs, done, key, env_state = carry
-    action, logprob, value, key = _get_action_and_value_noise(
-        sensor, feature_extractor, actor, critic, agent_state, obs, key
+    clipped_action, raw_action, logprob, value, mean, std, key = _get_action_and_value_noise(
+        sensor, feature_extractor, actor, critic, agent_state, obs, key, action_low, action_high
     )
 
     episode_stats, env_state, (next_obs, reward, next_done) = env_step_fn(
-        episode_stats, env_state, action
+        episode_stats, env_state, clipped_action
     )
 
     storage = Storage(
         obs=obs,
-        actions=action,
+        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),
     )
@@ -103,6 +155,8 @@ def _step_env_wrapped(episode_stats, env_state, action, env_step_fn):
     next_env_state = env_step_fn(env_state, action)
 
     reward = next_env_state.reward
+    reward *= 20000
+    reward = jnp.clip(reward, -10, 10)
     terminated = next_env_state.terminated
     truncated = next_env_state.truncated
     done = terminated | truncated
@@ -140,6 +194,8 @@ def _rollout_jit(
     feature_extractor: GenericDenseLayersWithActivation,
     actor: Actor,
     critic: OneDenseLayerMLP,
+    action_low,
+    action_high,
 ):
     (agent_state, episode_stats, next_obs, next_done, key, env_state), storage = jax.lax.scan(
         partial(
@@ -149,6 +205,8 @@ def _rollout_jit(
             actor=actor,
             critic=critic,
             env_step_fn=step_env_fn,
+            action_low=action_low,
+            action_high=action_high,
         ),
         (agent_state, episode_stats, next_obs, next_done, key, env_state),
         (),
@@ -191,7 +249,9 @@ def _compute_gae_jit(
         (dones[1:], values[1:], values[:-1], storage.rewards),
         reverse=True,
     )
-    return storage.replace(advantages=advantages, returns=advantages + storage.values)
+    returns = advantages + storage.values
+    advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
+    return storage.replace(advantages=advantages, returns=returns)
 
 
 @dataclass
@@ -210,14 +270,23 @@ class TrainingMeasurements:
 
 
 class PPOTrainer:
-    def __init__(self, args: PPOArgs, env: BrittleStarJaxEnvWrapper, run_dir: str, run_name: str):
-        self.args = args
+    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.env = env
         self.run_dir = run_dir
         self.run_name = run_name
         self.logger = get_logger()
 
-        self.key = jax.random.PRNGKey(args.seed)
+        # 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.sensor, self.feature_extractor, self.actor, self.critic = self._init_agent()
         self.sensor.apply = jax.jit(self.sensor.apply)
@@ -225,29 +294,34 @@ class PPOTrainer:
         self.actor.apply = jax.jit(self.actor.apply)
         self.critic.apply = jax.jit(self.critic.apply)
 
+        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._rollout_jit = jax.jit(
             partial(
                 _rollout_jit,
-                max_steps=self.args.num_steps,
+                max_steps=self.ppo.num_steps,
                 step_env_fn=partial(_step_env_wrapped, env_step_fn=self.env.step),
                 sensor=self.sensor,
                 feature_extractor=self.feature_extractor,
                 actor=self.actor,
                 critic=self.critic,
+                action_low=action_low,
+                action_high=action_high,
             )
         )
         self._compute_gae_jit = jax.jit(
             partial(
                 _compute_gae_jit,
-                num_envs=self.args.num_envs,
-                gamma=self.args.gamma,
-                gae_lambda=self.args.gae_lambda,
+                num_envs=self.ppo.num_envs,
+                gamma=self.ppo.gamma,
+                gae_lambda=self.ppo.gae_lambda,
                 feature_extractor=self.feature_extractor,
                 critic=self.critic,
             )
         )
 
-        self._ppo = PPO(self.args, self.sensor, self.actor, self.critic, self.feature_extractor)
+        self._ppo = PPO(self.ppo, self.sensor, self.actor, self.critic, self.feature_extractor)
 
         self.agent_state = self._init_agent_state()
 
@@ -256,16 +330,16 @@ class PPOTrainer:
         self._init_random()
 
     def _init_random(self):
-        self.logger.info(f"[RANDOM]: Setting random seed to {self.args.seed}")
+        self.logger.info(f"[RANDOM]: Setting random seed to {self.experiment.seed}")
 
-        random.seed(self.args.seed)
-        np.random.seed(self.args.seed)
+        random.seed(self.experiment.seed)
+        np.random.seed(self.experiment.seed)
 
     def _init_agent(self):
         self.logger.info("[AGENT]: Initializing agent...")
 
-        sensor = GenericDenseLayersWithActivation()
-        feature_extractor = GenericDenseLayersWithActivation()
+        sensor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300])
+        feature_extractor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300])
         actor = Actor(action_dim=self.env.single_action_space.shape[0])
         critic = OneDenseLayerMLP()
         return sensor, feature_extractor, actor, critic
@@ -277,15 +351,11 @@ class PPOTrainer:
             self.key, 5
         )
 
-        sample_obs = jnp.concatenate(
-            [
-                v.flatten()
-                for v in self.env.single_observation_space.sample(
-                    rng=jax.random.PRNGKey(0)
-                ).values()
-                if v.size > 0
-            ]
-        )
+        dummy_reset = self.env.reset(seed=0)
+        sample_obs = _convert_obs_dict_to_array(dummy_reset.observations)[0]  # take first env
+        self.obs_mean = jnp.zeros((len(sample_obs),))
+        self.obs_var = jnp.ones((len(sample_obs),))
+        self.obs_count = 1e-4
         sensor_params = self.sensor.init(sensor_key, sample_obs)
         feature_extractor_params = self.feature_extractor.init(feature_extractor_key, sample_obs)
         actor_params = self.actor.init(actor_key, self.sensor.apply(sensor_params, sample_obs))
@@ -299,17 +369,17 @@ class PPOTrainer:
                 AgentParams(sensor_params, actor_params, critic_params, feature_extractor_params)
             ),
             tx=optax.chain(
-                optax.clip_by_global_norm(self.args.max_grad_norm),
+                optax.clip_by_global_norm(self.ppo.max_grad_norm),
                 optax.inject_hyperparams(optax.adam)(
                     learning_rate=partial(
                         _linear_schedule,
-                        minibatch_count=self.args.num_minibatches,
-                        update_epochs=self.args.update_epochs,
-                        num_iterations=self.args.num_iterations,
-                        learning_rate=self.args.learning_rate,
+                        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.args.anneal_lr
-                    else self.args.learning_rate,
+                    if self.ppo.anneal_lr
+                    else self.ppo.learning_rate,
                     eps=1e-5,
                 ),
             ),
@@ -319,12 +389,31 @@ class PPOTrainer:
         self.logger.info("[EPISODE STATS]: Initializing episode stats...")
 
         return EpisodeStatistics(
-            episode_returns=jnp.zeros(self.args.num_envs, dtype=jnp.float32),
-            episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32),
-            returned_episode_returns=jnp.zeros(self.args.num_envs, jnp.float32),
-            returned_episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32),
+            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 _update_obs_stats(self, obs: jnp.ndarray):
+        batch_mean = jnp.mean(obs, axis=0)
+        batch_var = jnp.var(obs, axis=0)
+        batch_count = obs.shape[0]
+
+        delta = batch_mean - self.obs_mean
+        total_count = self.obs_count + batch_count
+
+        new_mean = self.obs_mean + delta * batch_count / total_count
+
+        m_a = self.obs_var * self.obs_count
+        m_b = batch_var * batch_count
+        M2 = m_a + m_b + delta**2 * self.obs_count * batch_count / total_count
+        new_var = M2 / total_count
+
+        self.obs_mean = new_mean
+        self.obs_var = new_var
+        self.obs_count = total_count
+
     def _rollout(self, env_state, next_obs, next_done) -> tuple[Any, ...]:
         return self._rollout_jit(
             self.agent_state,
@@ -350,7 +439,39 @@ class PPOTrainer:
         start_time,
         iteration_time_start,
         training_measurements,
+        storage,
+        next_obs,
+        xy_distance,
     ):
+        data = jax.device_get(
+            {
+                "rewards": storage.rewards[0],
+                "values": storage.values[0],
+                "returns": storage.returns[0],
+                "advantages": storage.advantages[0],
+                "actions": storage.actions[0],
+                "raw_actions": storage.raw_actions[0],
+                "means": storage.means[0],
+                "stds": storage.stds[0],
+                "logprobs": storage.logprobs[0],
+            }
+        )
+
+        storage_metrics = {
+            "rollout/env0/return_mean": float(np.mean(data["returns"])),
+            "rollout/env0/advantage_mean": float(np.mean(data["advantages"])),
+            "rollout/env0/value_mean": float(np.mean(data["values"])),
+            "rollout/env0/value_vs_return_diff": float(np.mean(data["values"] - data["returns"])),
+            "rollout/env0/reward_mean": float(np.mean(data["rewards"])),
+            "rollout/env0/mean_mean": float(np.mean(data["means"])),
+            "rollout/env0/logprob_mean": float(np.mean(data["logprobs"])),
+            "rollout/env0/action_mean": float(np.mean(data["actions"])),
+            "rollout/env0/raw_action_mean": float(np.mean(data["raw_actions"])),
+        }
+
+        for i in range(len(xy_distance)):
+            storage_metrics[f"env_data/env{i}_xy_dist_target"] = float(xy_distance[i])
+
         metrics = {
             "charts/avg_episodic_return": training_measurements.avg_episodic_return,
             "charts/avg_episodic_length": np.mean(
@@ -371,8 +492,9 @@ class PPOTrainer:
             "losses/loss": training_measurements.loss[-1, -1].item(),
             "charts/SPS": int(global_step / (time.time() - start_time)),
             "charts/SPS_update": int(
-                self.args.num_envs * self.args.num_steps / (time.time() - iteration_time_start)
+                self.ppo.num_envs * self.ppo.num_steps / (time.time() - iteration_time_start)
             ),
+            **storage_metrics,
         }
         self.logger.log(metrics, step=global_step)
 
@@ -443,6 +565,7 @@ class PPOTrainer:
                 avg_terminated_length=avg_terminated_length,
                 avg_truncated_length=avg_truncated_length,
             ),
+            storage,
         )
 
     def _close(self):
@@ -451,8 +574,14 @@ class PPOTrainer:
     def _save_model(self, model_path: str):
         self.logger.info("[SAVE]: Saving the final model...")
 
+        from dataclasses import asdict as _asdict
+
+        config_dict = {
+            "experiment": _asdict(self.experiment),
+            "ppo": _asdict(self.ppo),
+        }
         params = [
-            vars(self.args),
+            config_dict,
             [
                 self.agent_state.params["sensor_params"],
                 self.agent_state.params["actor_params"],
@@ -464,8 +593,7 @@ class PPOTrainer:
 
     def train(self):
         """
-        Train the PPO agent for a specified number of iterations
-        (passed through PPOArgs in constructor).
+        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}")
@@ -473,47 +601,58 @@ class PPOTrainer:
         self.logger.info("[TRAIN]: Resetting environment...")
         self.logger.log_non_interactive(f"Initial reset started: {time.ctime()}")
 
-        env_state = self.env.reset(seed=self.args.seed)
+        env_state = self.env.reset(seed=self.experiment.seed)
         next_obs = _convert_obs_dict_to_array(env_state.observations)
-        next_done = jnp.zeros(self.args.num_envs, dtype=jnp.bool_)
+        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.args.num_iterations + 1))
+        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 = self._step(
+            env_state, next_obs, next_done, training_measurements, storage = self._step(
                 env_state, next_obs, next_done, iteration=iteration
             )
+            self._update_obs_stats(next_obs)
+            next_obs = _normalize_obs(next_obs, self.obs_mean, self.obs_var)
 
-            global_step += self.args.num_steps * self.args.num_envs
+            xy_distance = _get_xy_distance_to_target(env_state.observations)
+
+            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,
+                next_obs,
+                xy_distance,
             )
 
             sps = int(global_step / (time.time() - start_time))
-            remaining_steps = self.args.total_timesteps - global_step
+            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.args.num_iterations} | "
-                f"Step {global_step}/{self.args.total_timesteps} | "
+                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.args.save_model:
-            model_path = f"{self.run_dir}/{self.args.exp_name}.cleanrl_model"
+            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/experiment_logger/__init__.py

@@ -4,8 +4,8 @@ This package provides a unified interface for logging to multiple backends
 (WandB, disk, stdout) simultaneously, ensuring no data loss.
 """
 
-from experiment_logger.config_utils import load_yaml_config, merge_config_with_cli
-from experiment_logger.unified_logger import UnifiedLogger, get_logger
+from experiment_logger.config_utils import load_yaml_config
+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
 
@@ -13,9 +13,9 @@ __all__ = [
     "UnifiedLogger",
     "SimpleLogger",
     "get_logger",
+    "init_logger",
     "init_wandb",
     "finish_wandb",
     "load_yaml_config",
-    "merge_config_with_cli",
 ]
 __version__ = "0.1.0"

src/experiment_logger/config_logger.py

@@ -0,0 +1,14 @@
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class LoggingConfig:
+    track: bool = False
+    wandb_project_name: str = "PPO-Modularity"
+    wandb_entity: Optional[str] = "SEL3-2026-Groep-4"
+    capture_video: bool = False
+    save_model: bool = True
+    checkpoint_frequency: int = 100
+    upload_model: bool = False
+    hf_entity: str = ""

src/experiment_logger/config_utils.py

@@ -1,15 +1,12 @@
 """Configuration utilities for loading YAML configs and merging with CLI args."""
 
 import os
-import sys
 from typing import Dict, Any, Type, TypeVar
 import yaml
 from dataclasses import fields, is_dataclass
 
 from experiment_logger.unified_logger import get_logger
 
-log = get_logger()
-
 T = TypeVar("T")
 
 
@@ -24,7 +21,7 @@ def load_yaml_config(config_path: str) -> Dict[str, Any]:
     if config is None:
         return {}
 
-    log.info(f"Loaded configuration from: {config_path}")
+    get_logger().info(f"Loaded configuration from: {config_path}")
     return config
 
 
@@ -35,7 +32,7 @@ def save_yaml_config(config: Dict[str, Any], config_path: str):
     with open(config_path, "w") as f:
         yaml.dump(config, f, default_flow_style=False, indent=2, sort_keys=False)
 
-    log.info(f"Saved configuration to: {config_path}")
+    get_logger().info(f"Saved configuration to: {config_path}")
 
 
 def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T:
@@ -66,82 +63,21 @@ def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T:
                     else:
                         filtered_config[key] = field.type(value) if value is not None else None  # type: ignore
             except (ValueError, TypeError) as e:
-                log.warning(f"Could not convert {key}={value} to {field.type}: {e}")
+                get_logger().warning(f"Could not convert {key}={value} to {field.type}: {e}")
                 filtered_config[key] = value
         else:
-            log.warning(f"Unknown configuration parameter: {key}")
+            get_logger().warning(f"Unknown configuration parameter: {key}")
 
     return cls(**filtered_config)
 
 
-def merge_config_with_cli(config_class: Type[T], config_file: str | None = None) -> T:
-    """Merge YAML config with CLI arguments, with CLI taking precedence.
-
-    Args:
-        config_class: Dataclass type to create
-        config_file: Path to YAML config file (optional)
-
-    Returns:
-        Instance of config_class with merged configuration
-    """
-    # Parse CLI args first to get the default/CLI values
-    import tyro
-
-    # Check if --config is in sys.argv and extract it
-    extracted_config_file = config_file
-    if "--config" in sys.argv:
-        config_idx = sys.argv.index("--config")
-        if config_idx + 1 < len(sys.argv):
-            extracted_config_file = sys.argv[config_idx + 1]
-            # Remove from sys.argv so tyro doesn't see it
-            sys.argv.pop(config_idx)  # Remove --config
-            sys.argv.pop(config_idx)  # Remove config file path
-
-    # Load YAML config if available
-    yaml_config = {}
-    if extracted_config_file and os.path.exists(extracted_config_file):
-        yaml_config = load_yaml_config(extracted_config_file)
-        log.info(f"Merging YAML config from {extracted_config_file} with CLI args")
-    elif extracted_config_file:
-        log.warning(f"Config file not found: {extracted_config_file}, using CLI args only")
-
-    # Create default instance to know what the defaults are
-    default_instance = config_class()
-    default_dict = {f.name: getattr(default_instance, f.name) for f in fields(config_class)}  # type: ignore
-
-    # Parse CLI args
-    cli_instance = tyro.cli(config_class)
-    cli_dict = {f.name: getattr(cli_instance, f.name) for f in fields(config_class)}  # type: ignore
-
-    # Merge configs: YAML as base, CLI overrides non-default values
-    final_config = {}
-
-    for field in fields(config_class):  # type: ignore
-        field_name = field.name
-        default_value = default_dict[field_name]
-        yaml_value = yaml_config.get(field_name, default_value)
-        cli_value = cli_dict[field_name]
-
-        # Use CLI value if it's different from default, otherwise use YAML value
-        if cli_value != default_value:
-            final_config[field_name] = cli_value
-            if yaml_value != default_value and yaml_value != cli_value:
-                log.info(f"CLI override: {field_name}={cli_value} (YAML had {yaml_value})")
-        else:
-            final_config[field_name] = yaml_value
-            if yaml_value != default_value:
-                log.info(f"YAML config: {field_name}={yaml_value}")
-
-    return config_class(**final_config)
-
-
 def print_config(config: Any, title: str = "Configuration"):
     """Pretty print configuration."""
-    log.info(f"{title}:")
+    get_logger().info(f"{title}:")
     if is_dataclass(config):
         for field in fields(config):
             value = getattr(config, field.name)
-            log.info(f"  {field.name}: {value}")
+            get_logger().info(f"  {field.name}: {value}")
     else:
         for key, value in vars(config).items():
-            log.info(f"  {key}: {value}")
+            get_logger().info(f"  {key}: {value}")

src/experiment_logger/unified_logger.py

@@ -6,9 +6,7 @@ This logger ensures all experimental data is preserved by writing to:
 3. stdout (for real-time monitoring)
 """
 
-import datetime
 import logging
-import subprocess
 import yaml
 import sys
 import time
@@ -21,38 +19,67 @@ import numpy as np
 
 from experiment_logger.wandb_utils import finish_wandb, init_wandb
 
-# Global singleton storage
-_global_logger = None
+# Global storage for the active logger and the proxy singleton
+_active_logger: Optional[Any] = None
+_proxy_instance: Optional["LoggerProxy"] = None
 
 
-def get_logger() -> "UnifiedLogger":
-    """Retrieve the global UnifiedLogger. If not initialized, fallback to auto-initialization."""
-    global _global_logger
-    if _global_logger is None:
-        try:
-            commit_hash = (
-                subprocess.check_output(
-                    ["git", "rev-parse", "--short", "HEAD"], stderr=subprocess.STDOUT
-                )
-                .decode("utf-8")
-                .strip()
-            )
-        except Exception:
-            commit_hash = "unknown"
+def get_logger() -> "LoggerProxy":
+    """Retrieve the global LoggerProxy.
 
-        timestamp_str = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-        generic_name = f"{timestamp_str}_{commit_hash}_brittle_star"
+    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
 
-        # Initialize generic fallback logger without WandB
-        _global_logger = UnifiedLogger(
-            run_name=generic_name,
-            config={"auto_initialized": True},
-            use_wandb=False,
-            _set_as_global=False,  # Prevent recursive call inside __init__
-        )
-        _global_logger.warning(f"UnifiedLogger auto-initialized with name: {generic_name}")
+            _active_logger = SimpleLogger(run_name="pre_init")
 
-    return _global_logger
+        _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:
@@ -68,7 +95,6 @@ class UnifiedLogger:
         use_wandb: bool = True,
         save_code: bool = True,
         log_level: int = logging.INFO,
-        _set_as_global: bool = True,
     ):
         """Initialize the unified logger.
 
@@ -80,7 +106,6 @@ class UnifiedLogger:
             base_dir: Base directory for local storage
             use_wandb: Whether to use WandB logging
             save_code: Whether to save code to WandB
-            _set_as_global: Internal flag to override the global singleton
         """
         self.run_name = run_name
         self.config = config
@@ -119,11 +144,6 @@ class UnifiedLogger:
             self._text_logger.addHandler(fh)
             self._text_logger.addHandler(ch)
 
-        # Set as global singleton
-        global _global_logger
-        if _set_as_global:
-            _global_logger = self
-
         # Save config to disk
         self._save_config()