From 76ba835ee64b32e44f239fe15eef161af3931181 Mon Sep 17 00:00:00 2001
From: Jona Reynaert <jona.reynaert@ugent.be>
Date: Wed, 15 Apr 2026 14:37:08 +0200
Subject: [PATCH] feat: added cli config path to simulate script

---
 scripts/simulate.py                           | 30 +++++++++++++++++--
 .../environment/env_wrapper.py                | 16 ++++++++--
 2 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/scripts/simulate.py b/scripts/simulate.py
index 90dc65f..b38d9a4 100644
--- a/scripts/simulate.py
+++ b/scripts/simulate.py
@@ -12,7 +12,7 @@ import numpy as np
 from brittle_star_project import (
     Backend,
 )
-from brittle_star_project.environment import from_file
+from brittle_star_project.environment import ArenaConfig, EnvConfig, MorphologyConfig, from_file
 
 def _flatten_obs_dict(obs_dict: dict[str, Any]) -> jnp.ndarray:
     """Flatten the env's observation dict into a 1D vector.
@@ -279,6 +279,16 @@ def parse_args() -> argparse.Namespace:
     p = argparse.ArgumentParser(
         description="Run a trained policy for exactly one episode (viewer or headless)."
     )
+    p.add_argument(
+        "--config-path",
+        type=str,
+        default=None,
+        help=(
+            "Path to an environment JSON config (morphology/arena/env). "
+            "If omitted, uses the environment defaults. "
+            "Relative paths are resolved from the repository root."
+        ),
+    )
     p.add_argument(
         "--model",
         type=str,
@@ -317,7 +327,16 @@ def main() -> None:
 
     args = parse_args()
 
-    morphology_cfg, arena_cfg, env_cfg = from_file("../configs/test.yaml")
+    if args.config_path is None:
+        morphology_cfg = MorphologyConfig()
+        arena_cfg = ArenaConfig()
+        env_cfg = EnvConfig()
+    else:
+        repo_root = Path(__file__).resolve().parents[1]
+        config_path = Path(args.config_path)
+        if not config_path.is_absolute():
+            config_path = repo_root / config_path
+        morphology_cfg, arena_cfg, env_cfg = from_file(str(config_path))
 
     # ======= ENVIRONMENT SETUP =======
 
@@ -325,7 +344,12 @@ def main() -> None:
 
     factory = BrittleStarEnvFactory()
     raw_env = factory.create_environment(backend, morphology_cfg, arena_cfg, env_cfg)
-    env = BrittleStarEnv(raw_env, backend=backend, config=env_cfg)
+    env = BrittleStarEnv(
+        raw_env,
+        backend=backend,
+        config=env_cfg,
+        morphology_config=morphology_cfg,
+    )
 
     seed_for_env = int(args.seed) if args.seed is not None else 0
     state = env.reset(seed=seed_for_env)
diff --git a/src/brittle_star_project/environment/env_wrapper.py b/src/brittle_star_project/environment/env_wrapper.py
index f74a1ca..8b3e51f 100644
--- a/src/brittle_star_project/environment/env_wrapper.py
+++ b/src/brittle_star_project/environment/env_wrapper.py
@@ -6,7 +6,7 @@ from typing import Any
 
 import numpy as np
 
-from .env_config import EnvConfig
+from .env_config import EnvConfig, MorphologyConfig
 from .env_types import Backend
 
 
@@ -25,10 +25,18 @@ class BrittleStarEnv:
     Goal: hide backend-specific RNG setup and provide a stable place to plug in RL.
     """
 
-    def __init__(self, env: Any, *, backend: Backend, config: EnvConfig) -> None:
+    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:
@@ -42,6 +50,10 @@ class BrittleStarEnv:
     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)