fix(train): integrate YAML config
This commit is contained in:
parent
2001a92e75
commit
e160a55d95
3 changed files with 10 additions and 2 deletions
|
|
@ -17,6 +17,7 @@ dependencies = [
|
||||||
"optax>=0.2.6",
|
"optax>=0.2.6",
|
||||||
"pyopengl>=3.1.10",
|
"pyopengl>=3.1.10",
|
||||||
"pyopengl-accelerate>=3.1.10",
|
"pyopengl-accelerate>=3.1.10",
|
||||||
|
"pyyaml>=6.0",
|
||||||
"tyro>=1.0.10",
|
"tyro>=1.0.10",
|
||||||
"wandb==0.24.2",
|
"wandb==0.24.2",
|
||||||
]
|
]
|
||||||
|
|
|
||||||
|
|
@ -26,7 +26,7 @@ class PPOArgs:
|
||||||
wandb_project_name: str = "PPO-Modularity"
|
wandb_project_name: str = "PPO-Modularity"
|
||||||
|
|
||||||
# the entity (team) of wandb's project
|
# 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)
|
# whether to capture videos of the agent performances (check out `videos` folder)
|
||||||
capture_video: bool = False
|
capture_video: bool = False
|
||||||
|
|
|
||||||
|
|
@ -21,6 +21,7 @@ from brittle_star_project.dataclasses.EpisodeStatistics import EpisodeStatistics
|
||||||
from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
|
from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
|
||||||
from brittle_star_project.rl import Actor, AgentParams, Critic, Network, Storage
|
from brittle_star_project.rl import Actor, AgentParams, Critic, Network, Storage
|
||||||
from experiment_logger import UnifiedLogger
|
from experiment_logger import UnifiedLogger
|
||||||
|
from experiment_logger.config_utils import merge_config_with_cli, print_config
|
||||||
|
|
||||||
log = logging.getLogger(__name__)
|
log = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
@ -418,7 +419,13 @@ def main() -> None:
|
||||||
level=logging.INFO,
|
level=logging.INFO,
|
||||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
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)
|
train(args)
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
Reference in a new issue