main.py

from __future__ import annotations

import argparse
import sys


AGENT_MODULES = {
    "ppo": "baselines.agents.ppo_rnn",
    "ppo_rnn": "baselines.agents.ppo_rnn",
    "rnd": "baselines.agents.ppo_rnd",
    "ppo_rnd": "baselines.agents.ppo_rnd",
    "icm": "baselines.agents.ppo_icm",
    "ppo_icm": "baselines.agents.ppo_icm",
}


def main(argv: list[str] | None = None) -> None:
    argv = list(sys.argv[1:] if argv is None else argv)
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument("--algorithm", choices=sorted(AGENT_MODULES), default="ppo")
    args, remaining = parser.parse_known_args(argv)

    module_name = AGENT_MODULES[args.algorithm]
    module = __import__(module_name, fromlist=["train", "make_parser"])
    result = module.train(module.make_parser().parse_args(remaining))
    print(f"compile_time_s: {result.compile_time_s:.6f}")
    print(f"training_time_s: {result.training_time_s:.6f}")
    print(f"steps_per_second_excluding_compile: {result.steps_per_second:.2f}")
    print(f"aggregate_completed_timesteps: {result.aggregate_completed_timesteps}")
    print(f"Saved results to {result.results_path}")


if __name__ == "__main__":
    main()