feat: simulation config

configs/main_config.yaml

@@ -11,6 +11,7 @@ defaults:
   - morphology: 5_arms_full
   - arena: default
   - environment: directed_locomotion
+  - simulation: default
   - _self_
 
 hydra:

configs/simulation/default.yaml

@@ -0,0 +1,11 @@
+# Default Simulation Settings
+# These values are used by scripts/simulate.py
+
+# Path to the trained model (optional)
+model_path: null
+
+# Type of model to use if no path is provided (e.g., random)
+model_type: "random"
+
+# Execution backend (MJX or BRAX)
+backend: "MJX"

scripts/simulate.py

@@ -29,19 +29,23 @@ MODEL_OPTIONS = sorted(MODEL_BY_NAME)
 
 @hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
 def main(dict_cfg: DictConfig) -> None:
-    cfg: BrittleStarConfig = OmegaConf.to_object(dict_cfg)
+    # 1. Convert DictConfig to structured dataclass, ensuring the root schema is applied correctly.
+    config: BrittleStarConfig = OmegaConf.to_object(
+        OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg)
+    )
 
-    # Allow overriding these via Hydra CLI or a dedicated simulate config group
+    # Use the configurable settings from the simulation group
-    # For now, defaults matching the old argparse behavior
+    backend = config.simulation.backend
-    backend = Backend.MJX
+    model_type = config.simulation.model_type
-    model_type = "random"
+    model_path = config.simulation.model_path
-    model_path = None
+    seed = config.experiment.seed
-    seed = cfg.experiment.seed
 
     # ======= ENVIRONMENT SETUP =======
     factory = BrittleStarEnvFactory()
-    raw_env = factory.create_environment(backend, cfg.morphology, cfg.arena, cfg.environment)
+    raw_env = factory.create_environment(
-    env = BrittleStarEnv(raw_env, backend=backend, config=cfg.environment)
+        backend, config.morphology, config.arena, config.environment
+    )
+    env = BrittleStarEnv(raw_env, backend=backend, config=config.environment)
 
     state = env.reset(seed=seed)
 

scripts/train.py

@@ -22,8 +22,10 @@ def make_env(cfg: BrittleStarConfig) -> BrittleStarJaxEnvWrapper:
 
 @hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
 def main(dict_cfg: DictConfig):
-    # 1. Convert DictConfig to structured dataclass
+    # 1. Convert DictConfig to structured dataclass, ensuring the root schema is applied correctly.
-    config: BrittleStarConfig = OmegaConf.to_object(dict_cfg)
+    config: BrittleStarConfig = OmegaConf.to_object(
+        OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg)
+    )
 
     # 2. Setup run metadata
     # Hydra changes CWD to the output directory by default.

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

@@ -4,6 +4,7 @@ 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
 
 
@@ -24,3 +25,4 @@ class BrittleStarConfig:
     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

@@ -7,6 +7,7 @@ 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
 
@@ -36,3 +37,4 @@ def register_configs() -> None:
     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)

tests/test_configs.py

@@ -13,9 +13,10 @@ def test_config_composition_centralized():
         # We compose the config; it follows main_config.yaml
         cfg = compose(config_name="main_config", overrides=["architecture=centralized"])
 
-        # Merge with the dataclass class to get a structured DictConfig,
+        # Merge with the structured schema and convert to a real dataclass instance to verify validation.
-        # then convert to a real dataclass instance to verify validation.
+        structured_cfg = OmegaConf.to_object(
-        structured_cfg = OmegaConf.to_object(OmegaConf.merge(BrittleStarConfig, cfg))
+            OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), cfg)
+        )
 
         # Basic assertions
         assert structured_cfg.architecture.name == "centralized"
@@ -30,7 +31,9 @@ def test_config_composition_decentralized():
         cfg = compose(config_name="main_config", overrides=["architecture=decentralized"])
 
         # Merge and convert to dataclass instance
-        structured_cfg = OmegaConf.to_object(OmegaConf.merge(BrittleStarConfig, cfg))
+        structured_cfg = OmegaConf.to_object(
+            OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), cfg)
+        )
 
         # Basic assertions
         assert structured_cfg.architecture.name == "decentralized"