train.py

import argparse
import copy
import sys
from datetime import datetime

import numpy as np
import pytz
import ray
from ray import tune
from ray.rllib.agents.registry import get_agent_class
from ray.rllib.models import ModelCatalog
from ray.tune import Experiment
from ray.tune.registry import register_env
from ray.tune.schedulers import PopulationBasedTraining

from algorithms.a3c_baseline import build_a3c_baseline_trainer
from algorithms.a3c_moa import build_a3c_moa_trainer
from algorithms.impala_baseline import build_impala_baseline_trainer
from algorithms.impala_moa import build_impala_moa_trainer
from algorithms.ppo_baseline import build_ppo_baseline_trainer
from algorithms.ppo_moa import build_ppo_moa_trainer
from algorithms.ppo_scm import build_ppo_scm_trainer
from models.baseline_model import BaselineModel
from models.moa_model import MOAModel
from models.scm_model import SocialCuriosityModule
from social_dilemmas.config.default_args import add_default_args
from social_dilemmas.envs.env_creator import get_env_creator
from utility_funcs import update_nested_dict

parser = argparse.ArgumentParser()
add_default_args(parser)


def build_experiment_config_dict(args):
    """
    Create a config dict for a single Experiment object.
    :param args: The parsed arguments.
    :return: An Experiment config dict.
    """
    env_creator = get_env_creator(
        args.env, args.num_agents, args.use_collective_reward, args.num_switches
    )
    env_name = args.env + "_env"
    register_env(env_name, env_creator)

    single_env = env_creator(args.num_agents)
    obs_space = single_env.observation_space
    act_space = single_env.action_space

    model_name = args.model + "_lstm"
    if args.model == "scm":
        ModelCatalog.register_custom_model(model_name, SocialCuriosityModule)
    elif args.model == "moa":
        ModelCatalog.register_custom_model(model_name, MOAModel)
    elif args.model == "baseline":
        ModelCatalog.register_custom_model(model_name, BaselineModel)

    # Each policy can have a different configuration (including custom model)
    def gen_policy():
        return None, obs_space, act_space, {"custom_model": model_name}

    # Create 1 distinct policy per agent
    policy_graphs = {}
    for i in range(args.num_agents):
        policy_graphs["agent-" + str(i)] = gen_policy()

    def policy_mapping_fn(agent_id):
        return agent_id

    agent_cls = get_agent_class(args.algorithm)
    config = copy.deepcopy(agent_cls._default_config)

    config["env"] = env_name
    config["eager"] = args.eager_mode

    # information for replay
    config["env_config"]["func_create"] = env_creator
    config["env_config"]["env_name"] = env_name
    if env_name == "switch_env":
        config["env_config"]["num_switches"] = args.num_switches

    conv_filters = [[6, [3, 3], 1]]
    fcnet_hiddens = [32, 32]
    lstm_cell_size = 128

    train_batch_size = (
        args.train_batch_size
        if args.train_batch_size is not None
        else max(1, args.num_workers) * args.num_envs_per_worker * args.rollout_fragment_length
    )

    lr_schedule = (
        list(zip(args.lr_schedule_steps, args.lr_schedule_weights))
        if args.lr_schedule_steps is not None and args.lr_schedule_weights is not None
        else None
    )

    # hyperparams
    update_nested_dict(
        config,
        {
            "horizon": 1000,
            "gamma": 0.99,
            "lr": args.lr,
            "lr_schedule": lr_schedule,
            "rollout_fragment_length": args.rollout_fragment_length,
            "train_batch_size": train_batch_size,
            "num_workers": args.num_workers,
            "num_envs_per_worker": args.num_envs_per_worker,
            "num_gpus": args.gpus_for_driver,  # The number of GPUs for the driver
            "num_cpus_for_driver": args.cpus_for_driver,
            "num_gpus_per_worker": args.gpus_per_worker,  # Can be a fraction
            "num_cpus_per_worker": args.cpus_per_worker,  # Can be a fraction
            "entropy_coeff": args.entropy_coeff,
            "grad_clip": args.grad_clip,
            "multiagent": {"policies": policy_graphs, "policy_mapping_fn": policy_mapping_fn},
            "callbacks": single_env.get_environment_callbacks(),
            "model": {
                "custom_model": model_name,
                "use_lstm": False,
                "conv_filters": conv_filters,
                "fcnet_hiddens": fcnet_hiddens,
                "custom_options": {
                    "cell_size": lstm_cell_size,
                    "num_other_agents": args.num_agents - 1,
                },
            },
        },
    )

    if args.model != "baseline":
        config["model"]["custom_options"].update(
            {
                "moa_loss_weight": args.moa_loss_weight,
                "influence_reward_clip": 10,
                "influence_reward_weight": args.influence_reward_weight,
                "influence_reward_schedule_steps": args.influence_reward_schedule_steps,
                "influence_reward_schedule_weights": args.influence_reward_schedule_weights,
                "return_agent_actions": True,
                "influence_divergence_measure": "kl",
                "train_moa_only_when_visible": True,
                "influence_only_when_visible": True,
            }
        )

    if args.model == "scm":
        config["model"]["custom_options"].update(
            {
                "scm_loss_weight": args.scm_loss_weight,
                "curiosity_reward_clip": 10,
                "curiosity_reward_weight": args.curiosity_reward_weight,
                "curiosity_reward_schedule_steps": args.curiosity_reward_schedule_steps,
                "curiosity_reward_schedule_weights": args.curiosity_reward_schedule_weights,
                "scm_forward_vs_inverse_loss_weight": args.scm_forward_vs_inverse_loss_weight,
            }
        )

    if args.tune_hparams:
        tune_dict = create_hparam_tune_dict(model=args.model, is_config=True)
        update_nested_dict(config, tune_dict)

    if args.algorithm == "PPO":
        config.update(
            {
                "num_sgd_iter": 10,
                "sgd_minibatch_size": args.ppo_sgd_minibatch_size
                if args.ppo_sgd_minibatch_size is not None
                else train_batch_size / 4,
                "vf_loss_coeff": 1e-4,
                "vf_share_layers": True,
            }
        )
    elif args.algorithm == "A3C" or args.algorithm == "IMPALA":
        config.update({"vf_loss_coeff": 0.1})
    else:
        sys.exit("The only available algorithms are A3C, PPO and IMPALA")

    return config


def get_trainer(args, config):
    """
    Creates a trainer depending on what args are specified.
    :param args: The parsed arguments.
    :param config: The config dict that is provided to the trainer.
    :return: A new trainer.
    """
    if args.model == "baseline":
        if args.algorithm == "A3C":
            trainer = build_a3c_baseline_trainer(config)
        if args.algorithm == "PPO":
            trainer = build_ppo_baseline_trainer(config)
        if args.algorithm == "IMPALA":
            trainer = build_impala_baseline_trainer(config)
    elif args.model == "moa":
        if args.algorithm == "A3C":
            trainer = build_a3c_moa_trainer(config)
        if args.algorithm == "PPO":
            trainer = build_ppo_moa_trainer(config)
        if args.algorithm == "IMPALA":
            trainer = build_impala_moa_trainer(config)
    elif args.model == "scm":
        if args.algorithm == "A3C":
            # trainer = build_a3c_scm_trainer(config)
            raise NotImplementedError
        if args.algorithm == "PPO":
            trainer = build_ppo_scm_trainer(config)
        if args.algorithm == "IMPALA":
            # trainer = build_impala_scm_trainer(config)
            raise NotImplementedError
    if trainer is None:
        raise NotImplementedError("The provided combination of model and algorithm was not found.")
    return trainer


def initialize_ray(args):
    """
    Initialize ray and automatically turn on local mode when debugging.
    :param args: The parsed arguments.
    """
    if sys.gettrace() is not None:
        print(
            "Debug mode detected through sys.gettrace(), turning on ray local mode. Saving"
            " experiment under ray_results/debug_experiment"
        )
        args.local_mode = True
    if args.multi_node and args.local_mode:
        sys.exit("You cannot have both local mode and multi node on at the same time")
    ray.init(
        address=args.address,
        local_mode=args.local_mode,
        memory=args.memory,
        object_store_memory=args.object_store_memory,
        redis_max_memory=args.redis_max_memory,
        include_webui=False,
    )


def get_experiment_name(args):
    """
    Build an experiment name based on environment, model and algorithm.
    :param args: The parsed arguments.
    :return: The experiment name.
    """
    if sys.gettrace() is not None:
        exp_name = "debug_experiment"
    elif args.exp_name is None:
        exp_name = args.env + "_" + args.model + "_" + args.algorithm
    else:
        exp_name = args.exp_name
    return exp_name


def build_experiment_dict(args, experiment_name, trainer, config):
    """
    Creates all parameters needed to create an Experiment object and puts them into a dict.
    :param args: The parsed arguments .
    :param experiment_name: The experiment name.
    :param trainer: The trainer used for the experiment.
    :param config: The config dict with experiment parameters.
    :return: A dict that can be unpacked to create an Experiment object.
    """
    experiment_dict = {
        "name": experiment_name,
        "run": trainer,
        "stop": {},
        "checkpoint_freq": args.checkpoint_frequency,
        "config": config,
        "num_samples": args.num_samples,
        "max_failures": -1,
    }
    if args.stop_at_episode_reward_min is not None:
        experiment_dict["stop"]["episode_reward_min"] = args.stop_at_episode_reward_min
    if args.stop_at_timesteps_total is not None:
        experiment_dict["stop"]["timesteps_total"] = args.stop_at_timesteps_total

    if args.use_s3:
        date = datetime.now(tz=pytz.utc)
        date = date.astimezone(pytz.timezone("US/Pacific")).strftime("%m-%d-%Y")
        s3_string = "s3://ssd-reproduce/" + date + "/" + experiment_name
        experiment_dict["upload_dir"] = s3_string

    return experiment_dict


def create_experiment(args):
    """
    Create a single experiment from arguments.
    :param args: The parsed arguments.
    :return: A new experiment with its own trainer.
    """
    experiment_name = get_experiment_name(args)
    config = build_experiment_config_dict(args)
    trainer = get_trainer(args=args, config=config)
    experiment_dict = build_experiment_dict(args, experiment_name, trainer, config)
    return Experiment(**experiment_dict)


def create_hparam_tune_dict(model, is_config=False):
    """
    Create a hyperparameter tuning dict for population-based training.
    :param is_config: Whether these hyperparameters are being used in the config dict or not.
    When used for the config dict, all hyperparameter-generating functions need to be wrapped with
    tune.sample_from, so we do this automatically here.
    When it is not used for the config dict, it is for PBT initialization, where a lambda is needed
    as a function wrapper.
    :return: The hyperparameter tune dict.
    """

    def wrapper(fn):
        if is_config:
            return tune.sample_from(lambda spec: fn)
        else:
            return lambda: fn

    baseline_options = {}
    model_options = {}
    if model == "baseline":
        baseline_options = {
            "entropy_coeff": wrapper(np.random.exponential(1 / 1000)),
            "lr": wrapper(np.random.uniform(0.00001, 0.01)),
        }
    if model == "moa":
        model_options = {
            "moa_loss_weight": wrapper(np.random.exponential(1 / 15)),
            "influence_reward_weight": wrapper(np.random.exponential(1)),
        }
    elif model == "scm":
        model_options = {
            "scm_loss_weight": wrapper(np.random.exponential(1 / 2)),
            "curiosity_reward_weight": wrapper(np.random.exponential(1)),
            "scm_forward_vs_inverse_loss_weight": wrapper(np.random.uniform(0, 1)),
        }

    hparam_dict = {
        **baseline_options,
        "model": {"custom_options": model_options},
    }
    return hparam_dict


def create_pbt_scheduler(model):
    """
    Create a population-based training (PBT) scheduler.
    :return: A new PBT scheduler.
    """
    hyperparam_mutations = create_hparam_tune_dict(model=model, is_config=False)

    pbt = PopulationBasedTraining(
        time_attr="training_iteration",
        perturbation_interval=10,
        metric="episode_reward_mean",
        mode="max",
        hyperparam_mutations=hyperparam_mutations,
    )
    return pbt


def run(args, experiments):
    """
    Run one or more experiments, with ray settings contained in args.
    :param args: The args to initialize ray with
    :param experiments: A list of experiments to run
    """
    initialize_ray(args)
    scheduler = create_pbt_scheduler(args.model) if args.tune_hparams else None
    tune.run_experiments(
        experiments,
        queue_trials=args.use_s3,
        resume=args.resume,
        scheduler=scheduler,
        reuse_actors=args.tune_hparams,
    )


if __name__ == "__main__":
    parsed_args = parser.parse_args()
    experiment = create_experiment(parsed_args)
    run(parsed_args, experiment)