ctm-dqn/training/train_madqn.py

15 lines
411 B
Python

"""MA-DQN training wrapper."""
from agents.madqn_agent import MADQNAgent
from training.train_value_based import train_sumo_value_based
def train_sumo_madqn(log_dir=None, checkpoint_dir=None, run_timestamp=None):
return train_sumo_value_based(
"madqn",
"MA-DQN",
MADQNAgent,
log_dir=log_dir,
checkpoint_dir=checkpoint_dir,
run_timestamp=run_timestamp,
)