test: additional testing and visual confirmations

src/brittle_star_project/configs/config_experiment.py

@@ -7,3 +7,4 @@ class ExperimentConfig:
     seed: int = 1
     torch_deterministic: bool = True
     cuda: bool = True
+    debug_sanity: bool = False

src/brittle_star_project/configs/main_config.py

@@ -18,8 +18,9 @@ class BrittleStarConfig:
     experiment: ExperimentConfig = field(default_factory=ExperimentConfig)
     logging: LoggingConfig = field(default_factory=LoggingConfig)
     ppo: PPOConfig = field(default_factory=PPOConfig)
-    # Default to centralized; swap with architecture=decentralized on the CLI.
-    architecture: ArchitectureConfig = field(default_factory=CentralizedConfig)
+    # 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)

src/brittle_star_project/environment/env_config.py

@@ -5,7 +5,7 @@ from dataclasses import dataclass, field
 from .env_types import Task
 
 
-@dataclass(frozen=True, slots=True)
+@dataclass
 class MorphologyConfig:
     """Brittle star morphology configuration.
 
@@ -17,7 +17,7 @@ class MorphologyConfig:
     The upstream biorobot library natively supports per-arm segment counts.
     """
 
-    segments_per_arm: tuple[int, ...] = (4, 4, 4, 4, 4)
+    segments_per_arm: list[int] = field(default_factory=lambda: [4, 4, 4, 4, 4])
     use_p_control: bool = True
     use_torque_control: bool = False
 
@@ -26,16 +26,16 @@ class MorphologyConfig:
         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.
 
@@ -50,7 +50,7 @@ class EnvConfig:
 
     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
 

src/brittle_star_project/trainers/PPOTrainer.py

@@ -527,6 +527,13 @@ class PPOTrainer:
                 f"ETA {eta_str}"
             )
 
+            if getattr(self.config.experiment, "debug_sanity", False):
+                self.logger.log(
+                    "\n[SANITY CHECK] Successfully completed 1 epoch of data collection and gradient updates."
+                )
+                self.logger.log("[SANITY CHECK] Gradients flowed without NaN. Exiting gracefully.")
+                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)