diff --git a/main.py b/main.py index 646c27d..d1b5b38 100644 --- a/main.py +++ b/main.py @@ -18,40 +18,73 @@ plot_evaluation, plot_train) +# ============================================================================== +# ARGUMENT PARSING +# ============================================================================== def parse_args(): + """ + Parses command-line arguments for training or evaluation. + Returns the parsed arguments object. + """ default_base_dir = '/Users/tchu/Documents/rl_test/deeprl_dist/ia2c_grid_0.9' default_config_dir = './config/config_ia2c_grid.ini' + parser = argparse.ArgumentParser() parser.add_argument('--base-dir', type=str, required=False, default=default_base_dir, help="experiment base dir") + + # Subparsers for different execution modes (train vs evaluate) subparsers = parser.add_subparsers(dest='option', help="train or evaluate") + + # Training arguments sp = subparsers.add_parser('train', help='train a single agent under base dir') sp.add_argument('--config-dir', type=str, required=False, default=default_config_dir, help="experiment config path") + + # Evaluation arguments sp = subparsers.add_parser('evaluate', help="evaluate and compare agents under base dir") sp.add_argument('--evaluation-seeds', type=str, required=False, default=','.join([str(i) for i in range(2000, 2500, 10)]), help="random seeds for evaluation, split by ,") sp.add_argument('--demo', action='store_true', help="shows SUMO gui") + args = parser.parse_args() + + # Ensure an option is provided, else show help if not args.option: parser.print_help() exit(1) return args +# ============================================================================== +# ENVIRONMENT SETUP +# ============================================================================== def init_env(config, port=0): + """ + Initializes the simulation environment based on the configuration scenario. + """ scenario = config.get('scenario') + + # Check if the scenario is an Adaptive Traffic Signal Control (ATSC) environment if scenario.startswith('atsc'): if scenario.endswith('large_grid'): return LargeGridEnv(config, port=port) else: return RealNetEnv(config, port=port) + # Otherwise, default to Cooperative Adaptive Cruise Control (CACC) environment else: return CACCEnv(config) +# ============================================================================== +# AGENT INITIALIZATION +# ============================================================================== def init_agent(env, config, total_step, seed): + """ + Instantiates the multi-agent reinforcement learning (MARL) model + based on the specified environment configuration. + """ if env.agent == 'ia2c': return IA2C(env.n_s_ls, env.n_a_ls, env.neighbor_mask, env.distance_mask, env.coop_gamma, total_step, config, seed=seed) @@ -75,88 +108,121 @@ def init_agent(env, config, total_step, seed): return None +# ============================================================================== +# TRAINING PROCEDURE +# ============================================================================== def train(args): + """ + Main loop for training the RL agents. + Sets up directories, configures logging, initializes the environment and agents, + runs the trainer, and finally saves the trained model. + """ + # 1. Setup directories and logging base_dir = args.base_dir dirs = init_dir(base_dir) init_log(dirs['log']) + + # 2. Load and copy configuration files config_dir = args.config_dir copy_file(config_dir, dirs['data']) config = configparser.ConfigParser() config.read(config_dir) - # init env + # 3. Initialize environment env = init_env(config['ENV_CONFIG']) logging.info('Training: a dim %r, agent dim: %d' % (env.n_a_ls, env.n_agent)) - # init step counter + # 4. Initialize step counters total_step = int(config.getfloat('TRAIN_CONFIG', 'total_step')) test_step = int(config.getfloat('TRAIN_CONFIG', 'test_interval')) log_step = int(config.getfloat('TRAIN_CONFIG', 'log_interval')) global_counter = Counter(total_step, test_step, log_step) - # init centralized or multi agent + # 5. Initialize agents (centralized or multi-agent) seed = config.getint('ENV_CONFIG', 'seed') model = init_agent(env, config['MODEL_CONFIG'], total_step, seed) + # 6. Execute training # disable multi-threading for safe SUMO implementation summary_writer = tf.summary.FileWriter(dirs['log']) trainer = Trainer(env, model, global_counter, summary_writer, output_path=dirs['data']) trainer.run() - # save model + # 7. Save model final_step = global_counter.cur_step logging.info('Training: save final model at step %d ...' % final_step) model.save(dirs['model'], final_step) +# ============================================================================== +# EVALUATION PROCEDURES +# ============================================================================== def evaluate_fn(agent_dir, output_dir, seeds, port, demo): + """ + Helper function to evaluate a single trained agent model. + """ agent = agent_dir.split('/')[-1] if not check_dir(agent_dir): logging.error('Evaluation: %s does not exist!' % agent) return - # load config file + + # Load configuration file config_dir = find_file(agent_dir + '/data/') if not config_dir: return config = configparser.ConfigParser() config.read(config_dir) - # init env + # Initialize environment with testing seeds env = init_env(config['ENV_CONFIG'], port=port) env.init_test_seeds(seeds) - # load model for agent + # Load previously trained model for agent model = init_agent(env, config['MODEL_CONFIG'], 0, 0) if model is None: return model_dir = agent_dir + '/model/' if not model.load(model_dir): return - # collect evaluation data + + # Execute evaluation and collect data evaluator = Evaluator(env, model, output_dir, gui=demo) evaluator.run() def evaluate(args): + """ + Main evaluation entry point. Sets up log directories and determines seeds, + then evaluates the agent's performance. + """ base_dir = args.base_dir + + # Setup directories for evaluation logs if not args.demo: dirs = init_dir(base_dir, pathes=['eva_data', 'eva_log']) init_log(dirs['eva_log']) output_dir = dirs['eva_data'] else: output_dir = None - # enforce the same evaluation seeds across agents + + # Enforce the same evaluation seeds across different agents for fairness seeds = args.evaluation_seeds logging.info('Evaluation: random seeds: %s' % seeds) if not seeds: seeds = [] else: seeds = [int(s) for s in seeds.split(',')] + evaluate_fn(base_dir, output_dir, seeds, 1, args.demo) +# ============================================================================== +# MAIN EXECUTION +# ============================================================================== if __name__ == '__main__': + # Parse arguments and decide whether to train or evaluate args = parse_args() + if args.option == 'train': train(args) else: