refactor: migrate configs from JSON to YAML

configs/example.json

@@ -1,8 +0,0 @@
-{
-  "morphology": {
-    "num_arms": 2,
-    "num_segments_per_arm": 4,
-    "use_p_control": true,
-    "use_torque_control": false
-  }
-}

configs/example.yaml

@@ -0,0 +1,5 @@
+morphology:
+  num_arms: 2
+  num_segments_per_arm: 4
+  use_p_control: true
+  use_torque_control: false

scripts/simulate.py

@@ -10,7 +10,7 @@ from brittle_star_project import (
     SimulationConfig,
     simulate_policy,
 )
-from brittle_star_project.environment import from_json
+from brittle_star_project.environment import from_file
 from brittle_star_project.rl import RLModel  # imports concrete models via rl.__init__
 from brittle_star_project.rl.base import get_rl_model_registry
 
@@ -44,7 +44,7 @@ def parse_args() -> argparse.Namespace:
 def main() -> None:
     args = parse_args()
 
-    morphology_cfg, arena_cfg, env_cfg = from_json("../configs/test.json")
+    morphology_cfg, arena_cfg, env_cfg = from_file("../configs/test.yaml")
 
     # ======= ENVIRONMENT SETUP =======
 

src/brittle_star_project/environment/env_config.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-import json
 
 from .env_types import Task
 
@@ -51,14 +50,11 @@ class EnvConfig:
 
 
 def from_file(path: str) -> tuple[MorphologyConfig, ArenaConfig, EnvConfig]:
-    """Load configurations from a JSON or YAML file."""
-    with open(path, "r") as f:
-        if path.endswith(".yaml") or path.endswith(".yml"):
-            import yaml
+    """Load configurations from a YAML file."""
+    import yaml
 
-            config_dict = yaml.safe_load(f)
-        else:
-            config_dict = json.load(f)
+    with open(path, "r") as f:
+        config_dict = yaml.safe_load(f)
 
         morphology = MorphologyConfig(**config_dict.get("morphology", {}))
         arena = ArenaConfig(**config_dict.get("arena", {}))

src/experiment_logger/unified_logger.py

@@ -7,9 +7,9 @@ This logger ensures all experimental data is preserved by writing to:
 """
 
 import datetime
-import json
 import logging
 import subprocess
+import yaml
 import sys
 import time
 from pathlib import Path
@@ -99,7 +99,7 @@ class UnifiedLogger:
         self.metrics_dir = self.run_dir / "metrics"
         self.metrics_dir.mkdir(exist_ok=True)
 
-        self.config_file = self.run_dir / "config.json"
+        self.config_file = self.run_dir / "config.yaml"
 
         # Setup standard Python logging mirror
         self.text_log_file = self.run_dir / "run.log"
@@ -196,7 +196,7 @@ class UnifiedLogger:
         """Save configuration to disk."""
         try:
             with open(self.config_file, "w") as f:
-                json.dump(self.config, f, indent=2)
+                yaml.dump(self.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}")
@@ -263,10 +263,10 @@ class UnifiedLogger:
             return
 
         try:
-            metrics_file = self.metrics_dir / "metrics.jsonl"
+            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 JSON serialization
+                    # 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
@@ -275,7 +275,8 @@ class UnifiedLogger:
                             serializable_metric[k] = v.tolist()
                         else:
                             serializable_metric[k] = v
-                    f.write(json.dumps(serializable_metric) + "\n")
+                    f.write("---\n")
+                    yaml.dump(serializable_metric, f, default_flow_style=False)
             self.metrics_buffer.clear()
         except Exception as e:
             self.error(f"Error flushing metrics: {e}")
@@ -298,9 +299,9 @@ class UnifiedLogger:
 
             # Save metadata if provided
             if metadata:
-                metadata_path = self.checkpoints_dir / f"{prefix}_step_{step}_metadata.json"
+                metadata_path = self.checkpoints_dir / f"{prefix}_step_{step}_metadata.yaml"
                 with open(metadata_path, "w") as f:
-                    json.dump(metadata, f, indent=2)
+                    yaml.dump(metadata, f, default_flow_style=False)
 
             self.info(f"Checkpoint saved: {checkpoint_path}")
 
@@ -334,9 +335,9 @@ class UnifiedLogger:
                 f.write(flax.serialization.to_bytes(params))
 
             if metadata:
-                metadata_path = self.run_dir / "final_model_metadata.json"
+                metadata_path = self.run_dir / "final_model_metadata.yaml"
                 with open(metadata_path, "w") as f:
-                    json.dump(metadata, f, indent=2)
+                    yaml.dump(metadata, f, default_flow_style=False)
 
             self.info(f"Final model saved: {final_model_path}")