ctm-dqn/main.py

45 lines
987 B
Python

"""
Main entry point for CTM-DQN speed limit control system.
"""
import argparse
from train import train
from test import test
def main():
parser = argparse.ArgumentParser(
description="DQN-based Dynamic Speed Limit Control with CTM"
)
parser.add_argument(
"--mode",
type=str,
choices=["train", "test"],
default="train",
help="Mode: train or test",
)
parser.add_argument(
"--config",
type=str,
default="config.yaml",
help="Path to configuration file",
)
parser.add_argument(
"--model",
type=str,
default=None,
help="Path to model checkpoint (for testing)",
)
args = parser.parse_args()
if args.mode == "train":
print("Starting training mode...")
train(args.config)
elif args.mode == "test":
print("Starting testing mode...")
test(args.config, args.model)
if __name__ == "__main__":
main()