exploitability.py

from distutils.util import strtobool
import argparse, os, yaml
import ray

from rl_games.torch_runner import Runner

os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-t", "--train", required=False, help="train network", action="store_true"
    )
    parser.add_argument(
        "-p", "--play", required=False, help="play(test) network", action="store_true"
    )
    parser.add_argument("-c", "--checkpoint", required=False, help="path to checkpoint")
    parser.add_argument("-f", "--file", required=True, help="path to config")
    os.makedirs("runs", exist_ok=True)

    args = vars(parser.parse_args())
    config_name = args["file"]

    print("Loading config: ", config_name)
    with open(config_name, "r") as stream:
        config = yaml.safe_load(stream)

    ray.init(object_store_memory=1024 * 1024 * 1000)

    runner = Runner()
    try:
        runner.load(config)
    except yaml.YAMLError as exc:
        print(exc)

    runner.run(args)

    ray.shutdown()