diff --git a/configs/main_config.yaml b/configs/main_config.yaml index a4efbb7..745755a 100644 --- a/configs/main_config.yaml +++ b/configs/main_config.yaml @@ -11,6 +11,7 @@ defaults: - morphology: 5_arms_full - arena: default - environment: directed_locomotion + - simulation: default - _self_ hydra: diff --git a/configs/simulation/default.yaml b/configs/simulation/default.yaml new file mode 100644 index 0000000..00c4e44 --- /dev/null +++ b/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" diff --git a/scripts/simulate.py b/scripts/simulate.py index 63b1cfe..8fb597d 100644 --- a/scripts/simulate.py +++ b/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 - # For now, defaults matching the old argparse behavior - backend = Backend.MJX - model_type = "random" - model_path = None - seed = cfg.experiment.seed + # Use the configurable settings from the simulation group + backend = config.simulation.backend + model_type = config.simulation.model_type + model_path = config.simulation.model_path + seed = config.experiment.seed # ======= ENVIRONMENT SETUP ======= factory = BrittleStarEnvFactory() - raw_env = factory.create_environment(backend, cfg.morphology, cfg.arena, cfg.environment) - env = BrittleStarEnv(raw_env, backend=backend, config=cfg.environment) + raw_env = factory.create_environment( + backend, config.morphology, config.arena, config.environment + ) + env = BrittleStarEnv(raw_env, backend=backend, config=config.environment) state = env.reset(seed=seed) diff --git a/scripts/train.py b/scripts/train.py index 409a196..409f050 100644 --- a/scripts/train.py +++ b/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 - config: 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) + ) # 2. Setup run metadata # Hydra changes CWD to the output directory by default. diff --git a/src/brittle_star_project/configs/config_simulation.py b/src/brittle_star_project/configs/config_simulation.py new file mode 100644 index 0000000..56747ae --- /dev/null +++ b/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 diff --git a/src/brittle_star_project/configs/main_config.py b/src/brittle_star_project/configs/main_config.py index 2164a30..5a937a7 100644 --- a/src/brittle_star_project/configs/main_config.py +++ b/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) diff --git a/src/brittle_star_project/configs/register_configs.py b/src/brittle_star_project/configs/register_configs.py index 74360a3..b677b31 100644 --- a/src/brittle_star_project/configs/register_configs.py +++ b/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) diff --git a/tests/test_configs.py b/tests/test_configs.py index c8cad63..43ef95e 100644 --- a/tests/test_configs.py +++ b/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, - # then convert to a real dataclass instance to verify validation. - structured_cfg = OmegaConf.to_object(OmegaConf.merge(BrittleStarConfig, cfg)) + # Merge with the structured schema and convert to a real dataclass instance to verify validation. + structured_cfg = OmegaConf.to_object( + 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"