comments added for beter understanding

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: