"""
Train the environment using the selected algorithm.
The original code is from https://github.com/PKU-MARL/HARL
Several modifications are made to adapt to the SustainDC environment.
"""
import os
import sys
import warnings
import argparse
import json
warnings.filterwarnings('ignore')
# sys.path.insert(0, os.getcwd())
from harl.utils.configs_tools import get_defaults_yaml_args, update_args
[docs]
def main():
"""Main function to train the environment using the selected algorithm."""
# Create an argument parser
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
# Add command-line arguments
parser.add_argument(
"--algo",
type=str,
default="happo",
choices=[
"happo",
"hatrpo",
"haa2c",
"haddpg",
"hatd3",
"hasac",
"had3qn",
"maddpg",
"matd3",
"mappo",
],
help="Algorithm name. Choose from: happo, hatrpo, haa2c, haddpg, hatd3, hasac, had3qn, maddpg, matd3, mappo."
)
parser.add_argument(
"--env",
type=str,
default="sustaindc",
choices=["sustaindc"],
help="Environment name. Choose from: sustaindc."
)
parser.add_argument(
"--exp_name",
type=str,
default="installtest",
help="Experiment name."
)
parser.add_argument(
"--load_config",
type=str,
default="",
help="If set, load existing experiment config file instead of reading from yaml config file."
)
# Parse known arguments and process unknown arguments
args, unparsed_args = parser.parse_known_args()
def process(arg):
"""Evaluate the argument if possible, otherwise return the argument as is."""
try:
return eval(arg)
except:
return arg
# Process unparsed arguments to a dictionary
keys = [k[2:] for k in unparsed_args[0::2]] # remove -- from argument
values = [process(v) for v in unparsed_args[1::2]]
unparsed_dict = {k: v for k, v in zip(keys, values)}
# Convert args to dictionary
args = vars(args)
# Load configuration
if args["load_config"] != "":
# Load config from existing config file
with open(args["load_config"], encoding="utf-8") as file:
all_config = json.load(file)
args["algo"] = all_config["main_args"]["algo"]
args["env"] = all_config["main_args"]["env"]
algo_args = all_config["algo_args"]
env_args = all_config["env_args"]
else:
# Load config from corresponding yaml file
algo_args, env_args = get_defaults_yaml_args(args["algo"], args["env"])
# Update args from command line
update_args(unparsed_dict, algo_args, env_args)
# Start training
from harl.runners import RUNNER_REGISTRY
# Initialize and run the selected algorithm
runner = RUNNER_REGISTRY[args["algo"]](args, algo_args, env_args)
runner.run()
runner.close()
if __name__ == "__main__":
main()