43 lines
1.4 KiB
Python
43 lines
1.4 KiB
Python
"""Central registry for training entry functions."""
|
|
from typing import Callable, Dict, List
|
|
|
|
from training.train_appo import train_sumo_appo
|
|
from training.train_dcmappo import train_sumo_dcmappo
|
|
from training.train_ddpg import train_sumo_ddpg
|
|
from training.train_dqn import train_sumo_dqn
|
|
from training.train_mappo import train_sumo_mappo
|
|
from training.train_ppo import train_sumo_ppo
|
|
from training.train_sctd3 import train_sumo_sctd3
|
|
from training.train_tcamappo import train_sumo_tcamappo
|
|
from training.train_td3 import train_sumo_td3
|
|
|
|
|
|
DEFAULT_MODELS: List[str] = ["ppo", "mappo", "tcamappo", "dcmappo", "dqn", "ddpg", "td3", "sctd3"]
|
|
ALL_MODELS: List[str] = ["ppo", "appo", "mappo", "tcamappo", "dcmappo", "dqn", "ddpg", "td3", "sctd3"]
|
|
|
|
|
|
TRAINERS: Dict[str, Callable] = {
|
|
"ppo": train_sumo_ppo,
|
|
"appo": train_sumo_appo,
|
|
"mappo": train_sumo_mappo,
|
|
"tcamappo": train_sumo_tcamappo,
|
|
"dcmappo": train_sumo_dcmappo,
|
|
"dqn": train_sumo_dqn,
|
|
"ddpg": train_sumo_ddpg,
|
|
"td3": train_sumo_td3,
|
|
"sctd3": train_sumo_sctd3,
|
|
}
|
|
|
|
|
|
def normalize_model_name(name: str) -> str:
|
|
model = name.strip().lower()
|
|
if model not in TRAINERS:
|
|
raise ValueError(f"Unsupported model name: {name}")
|
|
return model
|
|
|
|
|
|
def normalize_model_list(models: List[str] | None) -> List[str]:
|
|
if not models:
|
|
return list(DEFAULT_MODELS)
|
|
return [normalize_model_name(model) for model in models]
|