Spaces:
Runtime error
Runtime error
| import argparse | |
| import random | |
| import ray | |
| import ray.rllib.algorithms.ppo as ppo | |
| from pettingzoo.classic import connect_four_v3 | |
| from ray import air, tune | |
| from ray.rllib.policy.policy import PolicySpec | |
| from ray.rllib.utils.framework import try_import_torch | |
| from ray.tune import CLIReporter, register_env | |
| from connectfour.training.callbacks import create_self_play_callback | |
| from connectfour.training.dummy_policies import ( | |
| AlwaysSameHeuristic, | |
| BeatLastHeuristic, | |
| LinearHeuristic, | |
| RandomHeuristic, | |
| ) | |
| from connectfour.training.models import Connect4MaskModel | |
| from connectfour.training.wrappers import Connect4Env | |
| torch, nn = try_import_torch() | |
| def get_cli_args(): | |
| """ | |
| Create CLI parser and return parsed arguments | |
| python connectfour/training/train.py --num-cpus 4 --num-gpus 1 --stop-iters 10 --win-rate-threshold 0.50 | |
| python connectfour/training/train.py --num-gpus 1 --stop-iters 10 --win-rate-threshold 0.50 | |
| python connectfour/training/train.py --num-cpus 5 --num-gpus 1 --stop-iters 200 | |
| """ | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--num-cpus", type=int, default=0) | |
| parser.add_argument("--num-gpus", type=int, default=0) | |
| parser.add_argument("--num-workers", type=int, default=2) | |
| parser.add_argument( | |
| "--stop-iters", type=int, default=200, help="Number of iterations to train." | |
| ) | |
| parser.add_argument( | |
| "--stop-timesteps", | |
| type=int, | |
| default=10000000, | |
| help="Number of timesteps to train.", | |
| ) | |
| parser.add_argument( | |
| "--win-rate-threshold", | |
| type=float, | |
| default=0.95, | |
| help="Win-rate at which we setup another opponent by freezing the " | |
| "current main policy and playing against a uniform distribution " | |
| "of previously frozen 'main's from here on.", | |
| ) | |
| args = parser.parse_args() | |
| print(f"Running with following CLI args: {args}") | |
| return args | |
| def select_policy(agent_id, episode, **kwargs): | |
| if episode.episode_id % 2 == int(agent_id[-1:]): | |
| return "learned" | |
| else: | |
| return random.choice(["always_same", "beat_last", "random", "linear"]) | |
| if __name__ == "__main__": | |
| args = get_cli_args() | |
| ray.init( | |
| num_cpus=args.num_cpus or None, num_gpus=args.num_gpus, include_dashboard=False | |
| ) | |
| # define how to make the environment | |
| env_creator = lambda config: connect_four_v3.env(render_mode="rgb_array") | |
| # register that way to make the environment under an rllib name | |
| register_env("connect4", lambda config: Connect4Env(env_creator(config))) | |
| config = ( | |
| ppo.PPOConfig() | |
| .environment("connect4") | |
| .framework("torch") | |
| .training(model={"custom_model": Connect4MaskModel}) | |
| .callbacks( | |
| create_self_play_callback( | |
| win_rate_thr=args.win_rate_threshold, | |
| opponent_policies=["always_same", "beat_last", "random", "linear"], | |
| ) | |
| ) | |
| .rollouts( | |
| num_rollout_workers=args.num_workers, | |
| num_envs_per_worker=5, | |
| ) | |
| .multi_agent( | |
| policies={ | |
| "learned": PolicySpec(), | |
| "always_same": PolicySpec(policy_class=AlwaysSameHeuristic), | |
| "linear": PolicySpec(policy_class=LinearHeuristic), | |
| "beat_last": PolicySpec(policy_class=BeatLastHeuristic), | |
| "random": PolicySpec(policy_class=RandomHeuristic), | |
| }, | |
| policy_mapping_fn=select_policy, | |
| policies_to_train=["learned"], | |
| ) | |
| ) | |
| stop = { | |
| "timesteps_total": args.stop_timesteps, | |
| "training_iteration": args.stop_iters, | |
| } | |
| results = tune.Tuner( | |
| "PPO", | |
| param_space=config.to_dict(), | |
| run_config=air.RunConfig( | |
| stop=stop, | |
| verbose=2, | |
| progress_reporter=CLIReporter( | |
| metric_columns={ | |
| "training_iteration": "iter", | |
| "time_total_s": "time_total_s", | |
| "timesteps_total": "ts", | |
| "episodes_this_iter": "train_episodes", | |
| "policy_reward_mean/learned": "reward", | |
| "win_rate": "win_rate", | |
| "league_size": "league_size", | |
| }, | |
| sort_by_metric=True, | |
| ), | |
| checkpoint_config=air.CheckpointConfig( | |
| checkpoint_at_end=True, | |
| checkpoint_frequency=10, | |
| ), | |
| ), | |
| ).fit() | |
| print("Best checkpoint", results.get_best_result().checkpoint) | |
| ray.shutdown() | |