From e160a55d95c6ec829fc6acf4c9428b417ecd0477 Mon Sep 17 00:00:00 2001
From: Tibo De Peuter <tibo.depeuter@ugent.be>
Date: Tue, 31 Mar 2026 23:05:43 +0200
Subject: [PATCH] fix(train): integrate YAML config

---
 pyproject.toml                                  | 1 +
 src/brittle_star_project/dataclasses/PPOArgs.py | 2 +-
 src/train.py                                    | 9 ++++++++-
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 43c977b..9ba5a9a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -17,6 +17,7 @@ dependencies = [
     "optax>=0.2.6",
     "pyopengl>=3.1.10",
     "pyopengl-accelerate>=3.1.10",
+    "pyyaml>=6.0",
     "tyro>=1.0.10",
     "wandb==0.24.2",
 ]
diff --git a/src/brittle_star_project/dataclasses/PPOArgs.py b/src/brittle_star_project/dataclasses/PPOArgs.py
index b27265f..3a6836b 100644
--- a/src/brittle_star_project/dataclasses/PPOArgs.py
+++ b/src/brittle_star_project/dataclasses/PPOArgs.py
@@ -26,7 +26,7 @@ class PPOArgs:
     wandb_project_name: str = "PPO-Modularity"
 
     # the entity (team) of wandb's project
-    wandb_entity: str | None = "tdpeuter-ghent-university"
+    wandb_entity: str | None = None
 
     # whether to capture videos of the agent performances (check out `videos` folder)
     capture_video: bool = False
diff --git a/src/train.py b/src/train.py
index 00e8c60..ad31b56 100644
--- a/src/train.py
+++ b/src/train.py
@@ -21,6 +21,7 @@ from brittle_star_project.dataclasses.EpisodeStatistics import EpisodeStatistics
 from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
 from brittle_star_project.rl import Actor, AgentParams, Critic, Network, Storage
 from experiment_logger import UnifiedLogger
+from experiment_logger.config_utils import merge_config_with_cli, print_config
 
 log = logging.getLogger(__name__)
 
@@ -418,7 +419,13 @@ def main() -> None:
         level=logging.INFO,
         format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
     )
-    args = tyro.cli(PPOArgs)
+    
+    # Enhanced argument parsing with YAML config support
+    args = merge_config_with_cli(PPOArgs)
+    
+    # Print final configuration
+    print_config(args, "Final Training Configuration")
+    
     train(args)