scripts/train.py

import argparse
import random
import warnings
import numpy as np
import time
import datetime
import torch

import gym_minigrid.social_ai_envs
import torch_ac
import sys
import json
import utils
from pathlib import Path
from distutils.dir_util import copy_tree
from utils.env import env_args_str_to_dict
from models import *


# Parse arguments

parser = argparse.ArgumentParser()

## General parameters
parser.add_argument("--algo", required=True,
                    help="algorithm to use: ppo (REQUIRED)")
parser.add_argument("--env", required=True,
                    help="name of the environment to train on (REQUIRED)")
parser.add_argument("--model", default=None,
                    help="name of the model (default: {ENV}_{ALGO}_{TIME})")
parser.add_argument("--seed", type=int, default=1,
                    help="random seed (default: 1)")
parser.add_argument("--log-interval", type=int, default=10,
                    help="number of updates between two logs (default: 10)")
parser.add_argument("--save-interval", type=int, default=10,
                    help="number of updates between two saves (default: 10, 0 means no saving)")
parser.add_argument("--procs", type=int, default=16,
                    help="number of processes (default: 16)")
parser.add_argument("--frames", type=int, default=10**7,
                    help="number of frames of training (default: 1e7)")
parser.add_argument("--continue-train", default=None,
                    help="path to the model to finetune", type=str)
parser.add_argument("--finetune-train", default=None,
                    help="path to the model to finetune", type=str)
parser.add_argument("--compact-save", "-cs", action="store_true", default=False,
                    help="Keep only last model save")
parser.add_argument("--lr-schedule-end-frames", type=int, default=0,
                    help="Learning rate will be diminished from --lr to 0 linearly over the period of --lr-schedule-end-frames (default: 0 - no diminsh)")
parser.add_argument("--lr-end", type=float, default=0,
                    help="the final lr that will be reached at 'lr-schedule-end-frames' (default = 0)")

## Periodic test parameters
parser.add_argument("--test-set-name", required=False,
                    help="name of the environment to test on, default use the train env", default="SocialAITestSet")
# parser.add_argument("--test-env", required=False,
#                     help="name of the environment to test on, default use the train env")
# parser.add_argument("--no-test", "-nt", action="store_true", default=False,
#                     help="don't perform periodic testing")
parser.add_argument("--test-seed", type=int, default=0,
                    help="random seed (default: 0)")
parser.add_argument("--test-episodes", type=int, default=50,
                    help="number of episodes to test")
parser.add_argument("--test-interval", type=int, default=-1,
                    help="number of updates between two tests (default: -1, no testing)")
parser.add_argument("--test-env-args", nargs='*', default="like_train_no_acl")

## Parameters for main algorithm
parser.add_argument("--acl", action="store_true", default=False,
                    help="use acl")
parser.add_argument("--acl-type", type=str, default=None,
                    help="acl type")
parser.add_argument("--acl-thresholds", nargs="+", type=float, default=(0.75, 0.75),
                    help="per phase thresholds for expert CL")
parser.add_argument("--acl-minimum-episodes", type=int, default=1000,
                    help="Never go to second phase before this.")
parser.add_argument("--acl-average-interval", type=int, default=500,
                    help="Average the perfromance estimate over this many last episodes")
parser.add_argument("--epochs", type=int, default=4,
                    help="number of epochs for PPO (default: 4)")
parser.add_argument("--exploration-bonus", action="store_true", default=False,
                    help="Use a count based exploration bonus")
parser.add_argument("--exploration-bonus-type", nargs="+", default=["lang"],
                    help="modality on which to use the bonus (lang/grid)")
parser.add_argument("--exploration-bonus-params", nargs="+", type=float, default=(30., 50.),
                    help="parameters for a count based exploration bonus (C, M)")
parser.add_argument("--exploration-bonus-tanh", nargs="+", type=float, default=None,
                    help="tanh expl bonus scale, None means no tanh")
parser.add_argument("--expert-exploration-bonus", action="store_true", default=False,
                    help="Use an expert exploration bonus")
parser.add_argument("--episodic-exploration-bonus", action="store_true", default=False,
                    help="Use the exploration bonus in a episodic setting")
parser.add_argument("--batch-size", type=int, default=256,
                    help="batch size for PPO (default: 256)")
parser.add_argument("--frames-per-proc", type=int, default=None,
                    help="number of frames per process before update (default: 5 for A2C and 128 for PPO)")
parser.add_argument("--discount", type=float, default=0.99,
                    help="discount factor (default: 0.99)")
parser.add_argument("--lr", type=float, default=0.001,
                    help="learning rate (default: 0.001)")
parser.add_argument("--gae-lambda", type=float, default=0.99,
                    help="lambda coefficient in GAE formula (default: 0.99, 1 means no gae)")
parser.add_argument("--entropy-coef", type=float, default=0.01,
                    help="entropy term coefficient (default: 0.01)")
parser.add_argument("--value-loss-coef", type=float, default=0.5,
                    help="value loss term coefficient (default: 0.5)")
parser.add_argument("--max-grad-norm", type=float, default=0.5,
                    help="maximum norm of gradient (default: 0.5)")
parser.add_argument("--optim-eps", type=float, default=1e-8,
                    help="Adam and RMSprop optimizer epsilon (default: 1e-8)")
parser.add_argument("--optim-alpha", type=float, default=0.99,
                    help="RMSprop optimizer alpha (default: 0.99)")
parser.add_argument("--clip-eps", type=float, default=0.2,
                    help="clipping epsilon for PPO (default: 0.2)")
parser.add_argument("--recurrence", type=int, default=1,
                    help="number of time-steps gradient is backpropagated (default: 1). If > 1, a LSTM is added to the model to have memory.")
parser.add_argument("--text", action="store_true", default=False,
                    help="add a GRU to the model to handle text input")
parser.add_argument("--dialogue", action="store_true", default=False,
                    help="add a GRU to the model to handle the history of dialogue input")
parser.add_argument("--current-dialogue-only", action="store_true", default=False,
                    help="add a GRU to the model to handle only the current dialogue input")
parser.add_argument("--multi-headed-agent", action="store_true", default=False,
                    help="add a talking head")
parser.add_argument("--babyai11_agent", action="store_true", default=False,
                    help="use the babyAI 1.1 agent architecture")
parser.add_argument("--multi-headed-babyai11-agent", action="store_true", default=False,
                    help="use the multi headed babyAI 1.1 agent architecture")
parser.add_argument("--custom-ppo", action="store_true", default=False,
                    help="use BabyAI original PPO hyperparameters")
parser.add_argument("--custom-ppo-2", action="store_true", default=False,
                    help="use BabyAI original PPO hyperparameters but with smaller memory")
parser.add_argument("--custom-ppo-3", action="store_true", default=False,
                    help="use BabyAI original PPO hyperparameters but with no memory")
parser.add_argument("--custom-ppo-rnd", action="store_true", default=False,
                    help="rnd reconstruct")
parser.add_argument("--custom-ppo-rnd-reference", action="store_true", default=False,
                    help="rnd reconstruct")
parser.add_argument("--custom-ppo-ride", action="store_true", default=False,
                    help="rnd reconstruct")
parser.add_argument("--custom-ppo-ride-reference", action="store_true", default=False,
                    help="rnd reconstruct")
parser.add_argument("--ppo-hp-tuning", action="store_true", default=False,
                    help="use PPO hyperparameters selected from our HP tuning")
parser.add_argument("--multi-modal-babyai11-agent", action="store_true", default=False,
                    help="use the multi headed babyAI 1.1 agent architecture")

# ride ref
parser.add_argument("--ride-ref-agent", action="store_true", default=False,
                    help="Model from the ride paper")
parser.add_argument("--ride-ref-preprocessor", action="store_true", default=False,
                    help="use ride reference preprocessor (3D images)")

parser.add_argument("--bAI-lang-model", help="lang model type for babyAI models", default="gru")
parser.add_argument("--memory-dim", type=int, help="memory dim (128 is small 2048 is big", default=128)
parser.add_argument("--clipped-rewards", action="store_true", default=False,
                    help="add a talking head")
parser.add_argument("--intrinsic-reward-epochs", type=int, default=0,
                    help="")
parser.add_argument("--balance-moa-training", action="store_true", default=False,
                    help="balance moa training to handle class imbalance.")
parser.add_argument("--moa-memory-dim", type=int, help="memory dim (default=128)", default=128)

# rnd + ride
parser.add_argument("--intrinsic-reward-coef", type=float, default=0.1,
                    help="")
parser.add_argument("--intrinsic-reward-learning-rate", type=float, default=0.0001,
                    help="")
parser.add_argument("--intrinsic-reward-momentum", type=float, default=0,
                    help="")
parser.add_argument("--intrinsic-reward-epsilon", type=float, default=0.01,
                    help="")
parser.add_argument("--intrinsic-reward-alpha", type=float, default=0.99,
                    help="")
parser.add_argument("--intrinsic-reward-max-grad-norm", type=float, default=40,
                    help="")
# rnd + soc_inf
parser.add_argument("--intrinsic-reward-loss-coef", type=float, default=0.1,
                    help="")
# ride
parser.add_argument("--intrinsic-reward-forward-loss-coef", type=float, default=10,
                    help="")
parser.add_argument("--intrinsic-reward-inverse-loss-coef", type=float, default=0.1,
                    help="")

parser.add_argument("--reset-rnd-ride-at-phase", action="store_true", default=False,
                    help="expert knowledge resets rnd ride at acl phase change")

# babyAI1.1 related
parser.add_argument("--arch", default="original_endpool_res",
                  help="image embedding architecture")
parser.add_argument("--num-films", type=int, default=2,
                    help="")

# Put all env related arguments after --env_args, e.g. --env_args nb_foo 1 is_bar True
parser.add_argument("--env-args", nargs='*', default=None)

args = parser.parse_args()

if args.compact_save:
    print("Compact save is deprecated. Don't use it. It doesn't do anything now.")

if args.save_interval != args.log_interval:
    print(f"save_interval ({args.save_interval}) and log_interval ({args.log_interval}) are not the same. This is not ideal for train continuation.")

if args.seed == -1:
    args.seed = np.random.randint(424242)

if args.custom_ppo:
    print("babyAI's ppo config")

    assert not args.custom_ppo_2
    assert not args.custom_ppo_3
    args.frames_per_proc = 40
    args.lr = 1e-4
    args.gae_lambda = 0.99
    args.recurrence = 20
    args.optim_eps = 1e-05
    args.clip_eps = 0.2
    args.batch_size = 1280

elif args.custom_ppo_2:
    print("babyAI's ppo config with smaller memory")

    assert not args.custom_ppo
    assert not args.custom_ppo_3
    args.frames_per_proc = 40
    args.lr = 1e-4
    args.gae_lambda = 0.99
    args.recurrence = 10
    args.optim_eps = 1e-05
    args.clip_eps = 0.2
    args.batch_size = 1280

elif args.custom_ppo_3:
    print("babyAI's ppo config with no memory")

    assert not args.custom_ppo
    assert not args.custom_ppo_2
    args.frames_per_proc = 40
    args.lr = 1e-4
    args.gae_lambda = 0.99
    args.recurrence = 1
    args.optim_eps = 1e-05
    args.clip_eps = 0.2
    args.batch_size = 1280

elif args.custom_ppo_rnd:
    print("RND reconstruct")

    assert not args.custom_ppo
    assert not args.custom_ppo_2
    assert not args.custom_ppo_3
    args.frames_per_proc = 40
    args.lr = 1e-4
    args.recurrence = 1
    # args.recurrence = 5  # use 5 for SocialAI envs
    args.batch_size = 640
    args.epochs = 4

    # args.optim_eps = 1e-05
    # args.entropy_coef = 0.0001
    args.clipped_rewards = True

elif args.custom_ppo_ride:
    print("RIDE reconstruct")

    assert not args.custom_ppo
    assert not args.custom_ppo_2
    assert not args.custom_ppo_3
    assert not args.custom_ppo_rnd

    args.frames_per_proc = 40
    args.lr = 1e-4
    args.recurrence = 1
    # args.recurrence = 5  # use 5 for SocialAI envs
    args.batch_size = 640
    args.epochs = 4

    # args.optim_eps = 1e-05
    # args.entropy_coef = 0.0005
    args.clipped_rewards = True

elif args.custom_ppo_rnd_reference:
    print("RND reconstruct")

    assert not args.custom_ppo
    assert not args.custom_ppo_2
    assert not args.custom_ppo_3

    args.frames_per_proc = 128  # 128 for PPO
    args.lr = 1e-4
    args.recurrence = 64

    args.gae_lambda = 0.99
    args.batch_size = 1280
    args.epochs = 4

    args.optim_eps = 1e-05
    args.clip_eps = 0.2
    args.entropy_coef = 0.0001
    args.clipped_rewards = True


elif args.custom_ppo_ride_reference:
    print("RIDE reference")

    assert not args.custom_ppo
    assert not args.custom_ppo_2
    assert not args.custom_ppo_3
    assert not args.custom_ppo_rnd

    args.frames_per_proc = 128  # 128 for PPO
    args.lr = 1e-4
    args.recurrence = 64

    args.gae_lambda = 0.99
    args.batch_size = 1280
    args.epochs = 4

    args.optim_eps = 1e-05
    args.clip_eps = 0.2
    args.entropy_coef = 0.0005
    args.clipped_rewards = True

elif args.ppo_hp_tuning:

    args.frames_per_proc = 40
    args.lr = 1e-4
    args.recurrence = 5
    args.batch_size = 640
    args.epochs = 4

if args.env not in [
    "MiniGrid-KeyCorridorS3R3-v0",
    "MiniGrid-MultiRoom-N2-S4-v0",
    "MiniGrid-MultiRoom-N4-S5-v0",
    "MiniGrid-MultiRoom-N7-S4-v0",
    "MiniGrid-MultiRoomNoisyTV-N7-S4-v0"
]:
    if args.recurrence <= 1:
        print("You are using recurrence {} with {} env. This is probably unintentional.".format(args.recurrence, args.env))
        # warnings.warn("You are using recurrence {} with {} env. This is probably unintentional.".format(args.recurrence, args.env))


args.mem = args.recurrence > 1

# Set run dir
date = datetime.datetime.now().strftime("%y-%m-%d-%H-%M-%S")
default_model_name = f"{args.env}_{args.algo}_seed{args.seed}_{date}"

model_name = args.model or default_model_name
model_dir = utils.get_model_dir(model_name)

if Path(model_dir).exists() and args.continue_train is None:
    raise ValueError(f"Dir {model_dir} already exists and continue train is None.")

# Load loggers and Tensorboard writer
txt_logger = utils.get_txt_logger(model_dir)
csv_file, csv_logger = utils.get_csv_logger(model_dir)


# Log command and all script arguments
txt_logger.info("{}\n".format(" ".join(sys.argv)))
txt_logger.info("{}\n".format(args))

# Set seed for all randomness sources
utils.seed(args.seed)

# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
txt_logger.info(f"Device: {device}\n")

# Create env_args dict
env_args = env_args_str_to_dict(args.env_args)

if args.acl:
    # expert_acl = "three_stage_expert"
    expert_acl = args.acl_type
    print(f"Using curriculum: {expert_acl}.")
else:
    expert_acl = None

env_args_no_acl = env_args.copy()
env_args["curriculum"] = expert_acl
env_args["expert_curriculum_thresholds"] = args.acl_thresholds
env_args["expert_curriculum_average_interval"] = args.acl_average_interval
env_args["expert_curriculum_minimum_episodes"] = args.acl_minimum_episodes
env_args["egocentric_observation"] = True

# test env args
if not args.test_env_args:
    test_env_args = {}
elif args.test_env_args == "like_train_no_acl":
    test_env_args = env_args_no_acl
elif args.test_env_args == "like_train":
    test_env_args = env_args
else:
    test_env_args = env_args_str_to_dict(args.test_env_args)


if "SocialAI-" not in args.env:
    env_args = {}
    test_env_args = {}

print("train_env_args:", env_args)
print("test_env_args:", test_env_args)

# Load train environments

envs = []
for i in range(args.procs):
    envs.append(utils.make_env(args.env, args.seed + 10000 * i, env_args=env_args))

txt_logger.info("Environments loaded\n")

if args.continue_train and args.finetune_train:
    raise ValueError(f"Continue path ({args.continue_train}) and finetune path ({args.finetune_train}) can't both be set.")

# Load training status
if args.continue_train:
    if args.continue_train == "auto":
        status_continue_path = Path(model_dir)
        args.continue_train = status_continue_path  # just in case
    else:
        status_continue_path = Path(args.continue_train)

    if status_continue_path.is_dir():
        # if dir, assume experiment dir so append the seed
        # status_continue_path = Path(status_continue_path) / str(args.seed)
        status_continue_path = utils.get_status_path(status_continue_path)

    else:
        if not status_continue_path.is_file():
            raise ValueError(f"{status_continue_path} is not a file")

        if "status" not in status_continue_path.name:
            raise UserWarning(f"{status_continue_path} is does not contain status, is this the correct file? ")

    status = utils.load_status(status_continue_path)

    txt_logger.info("Training status loaded\n")
    txt_logger.info(f"{model_name} continued from {status_continue_path}")

    # copy everything from model_dir to backup_dir
    assert Path(status_continue_path).is_file()

elif args.finetune_train:

    status_finetune_path = Path(args.finetune_train)

    if status_finetune_path.is_dir():
        # if dir, assume experiment dir so append the seed
        status_finetune_seed_path = Path(status_finetune_path) / str(args.seed)
        if status_finetune_seed_path.exists():
            # if a seed folder exists assume that you use that one
            status_finetune_path = utils.get_status_path(status_finetune_seed_path)

        else:
            # if not assume that no seed folder exists
            status_finetune_path = utils.get_status_path(status_finetune_path)

    else:
        if not status_finetune_path.is_file():
            raise ValueError(f"{status_finetune_path} is not dir or a file")

        if "status" not in status_finetune_path.name:
            raise UserWarning(f"{status_finetune_path} is does not contain status, is this the correct file? ")

    status = utils.load_status(status_finetune_path)

    txt_logger.info("Training status loaded\n")
    txt_logger.info(f"{model_name} finetuning from {status_finetune_path}")

    # copy everything from model_dir to backup_dir
    assert Path(status_finetune_path).is_file()

    # reset parameters for finetuning
    status["num_frames"] = 0
    status["update"] = 0
    del status["optimizer_state"]
    del status["lr_scheduler_state"]
    del status["env_args"]

else:
    status = {"num_frames": 0, "update": 0}

# Parameter sanity checks
if args.dialogue and args.current_dialogue_only:
        raise ValueError("Either use dialogue or current-dialogue-only")

if not args.dialogue and not args.current_dialogue_only:
    warnings.warn("Not using dialogue")

if args.text:
    raise ValueError("Text should not be used. Use dialogue instead.")


# Load observations preprocessor
obs_space, preprocess_obss = utils.get_obss_preprocessor(
    obs_space=envs[0].observation_space,
    text=args.text,
    dialogue_current=args.current_dialogue_only,
    dialogue_history=args.dialogue,
    custom_image_preprocessor=utils.ride_ref_image_preprocessor if args.ride_ref_preprocessor else None,
    custom_image_space_preprocessor=utils.ride_ref_image_space_preprocessor if args.ride_ref_preprocessor else None,
)

if args.continue_train is not None or args.finetune_train is not None:
    assert "vocab" in status
    preprocess_obss.vocab.load_vocab(status["vocab"])
    txt_logger.info("Observations preprocessor loaded")

if args.exploration_bonus:
    if args.expert_exploration_bonus:
        warnings.warn("You are using expert exploration bonus.")

# Load model
assert sum(map(int, [
    args.multi_modal_babyai11_agent,
    args.multi_headed_babyai11_agent,
    args.babyai11_agent,
    args.multi_headed_agent,
])) <= 1

if args.multi_modal_babyai11_agent:
    acmodel = MultiModalBaby11ACModel(
        obs_space=obs_space,
        action_space=envs[0].action_space,
        arch=args.arch,
        use_text=args.text,
        use_dialogue=args.dialogue,
        use_current_dialogue_only=args.current_dialogue_only,
        use_memory=args.mem,
        lang_model=args.bAI_lang_model,
        memory_dim=args.memory_dim,
        num_films=args.num_films
    )
elif args.ride_ref_agent:
    assert args.mem
    assert not args.text
    assert not args.dialogue

    acmodel = RefACModel(
        obs_space=obs_space,
        action_space=envs[0].action_space,
        use_memory=args.mem,
        use_text=args.text,
        use_dialogue=args.dialogue,
        input_size=obs_space['image'][-1],
    )
    if args.current_dialogue_only: raise NotImplementedError("current dialogue only")

else:
    acmodel = ACModel(
        obs_space=obs_space,
        action_space=envs[0].action_space,
        use_memory=args.mem,
        use_text=args.text,
        use_dialogue=args.dialogue,
        input_size=obs_space['image'][-1],
    )
    if args.current_dialogue_only: raise NotImplementedError("current dialogue only")

# if args.continue_train is not None:
#     assert "model_state" in status
#     acmodel.load_state_dict(status["model_state"])

acmodel.to(device)
txt_logger.info("Model loaded\n")
txt_logger.info("{}\n".format(acmodel))

# Load algo
assert args.algo == "ppo"
algo = torch_ac.PPOAlgo(
    envs=envs,
    acmodel=acmodel,
    device=device,
    num_frames_per_proc=args.frames_per_proc,
    discount=args.discount,
    lr=args.lr,
    gae_lambda=args.gae_lambda,
    entropy_coef=args.entropy_coef,
    value_loss_coef=args.value_loss_coef,
    max_grad_norm=args.max_grad_norm,
    recurrence=args.recurrence,
    adam_eps=args.optim_eps,
    clip_eps=args.clip_eps,
    epochs=args.epochs,
    batch_size=args.batch_size,
    preprocess_obss=preprocess_obss,
    exploration_bonus=args.exploration_bonus,
    exploration_bonus_tanh=args.exploration_bonus_tanh,
    exploration_bonus_type=args.exploration_bonus_type,
    exploration_bonus_params=args.exploration_bonus_params,
    expert_exploration_bonus=args.expert_exploration_bonus,
    episodic_exploration_bonus=args.episodic_exploration_bonus,
    clipped_rewards=args.clipped_rewards,
    # for rnd, ride, and social influence
    intrinsic_reward_coef=args.intrinsic_reward_coef,
    # for rnd and ride
    intrinsic_reward_epochs=args.intrinsic_reward_epochs,
    intrinsic_reward_learning_rate=args.intrinsic_reward_learning_rate,
    intrinsic_reward_momentum=args.intrinsic_reward_momentum,
    intrinsic_reward_epsilon=args.intrinsic_reward_epsilon,
    intrinsic_reward_alpha=args.intrinsic_reward_alpha,
    intrinsic_reward_max_grad_norm=args.intrinsic_reward_max_grad_norm,
    # for rnd and social influence
    intrinsic_reward_loss_coef=args.intrinsic_reward_loss_coef,
    # for ride
    intrinsic_reward_forward_loss_coef=args.intrinsic_reward_forward_loss_coef,
    intrinsic_reward_inverse_loss_coef=args.intrinsic_reward_inverse_loss_coef,
    # for social influence
    balance_moa_training=args.balance_moa_training,
    moa_memory_dim=args.moa_memory_dim,
    lr_schedule_end_frames=args.lr_schedule_end_frames,
    end_lr=args.lr_end,
    reset_rnd_ride_at_phase=args.reset_rnd_ride_at_phase,
)

if args.continue_train or args.finetune_train:
    algo.load_status_dict(status)
    # txt_logger.info(f"Model + Algo loaded from {args.continue_train or args.finetune_train}\n")
    if args.continue_train:
        txt_logger.info(f"Model + Algo loaded from {status_continue_path} \n")
    elif args.finetune_train:
        txt_logger.info(f"Model + Algo loaded from {status_finetune_path} \n")


# todo: make nicer
# Set and load test environment
if args.test_set_name:
    if args.test_set_name == "SocialAITestSet":
        # "SocialAI-AskEyeContactLanguageBoxesInformationSeekingParamEnv-v1",
        # "SocialAI-NoIntroPointingBoxesInformationSeekingParamEnv-v1"
        test_env_names = [
            "SocialAI-TestLanguageColorBoxesInformationSeekingEnv-v1",
            "SocialAI-TestLanguageFeedbackBoxesInformationSeekingEnv-v1",
            "SocialAI-TestPointingBoxesInformationSeekingEnv-v1",
            "SocialAI-TestEmulationBoxesInformationSeekingEnv-v1",
            "SocialAI-TestLanguageColorSwitchesInformationSeekingEnv-v1",
            "SocialAI-TestLanguageFeedbackSwitchesInformationSeekingEnv-v1",
            "SocialAI-TestPointingSwitchesInformationSeekingEnv-v1",
            "SocialAI-TestEmulationSwitchesInformationSeekingEnv-v1",
            "SocialAI-TestLanguageColorMarbleInformationSeekingEnv-v1",
            "SocialAI-TestLanguageFeedbackMarbleInformationSeekingEnv-v1",
            "SocialAI-TestPointingMarbleInformationSeekingEnv-v1",
            "SocialAI-TestEmulationMarbleInformationSeekingEnv-v1",
            "SocialAI-TestLanguageColorGeneratorsInformationSeekingEnv-v1",
            "SocialAI-TestLanguageFeedbackGeneratorsInformationSeekingEnv-v1",
            "SocialAI-TestPointingGeneratorsInformationSeekingEnv-v1",
            "SocialAI-TestEmulationGeneratorsInformationSeekingEnv-v1",
            "SocialAI-TestLanguageColorLeversInformationSeekingEnv-v1",
            "SocialAI-TestLanguageFeedbackLeversInformationSeekingEnv-v1",
            "SocialAI-TestPointingLeversInformationSeekingEnv-v1",
            "SocialAI-TestEmulationLeversInformationSeekingEnv-v1",
            "SocialAI-TestLanguageColorDoorsInformationSeekingEnv-v1",
            "SocialAI-TestLanguageFeedbackDoorsInformationSeekingEnv-v1",
            "SocialAI-TestPointingDoorsInformationSeekingEnv-v1",
            "SocialAI-TestEmulationDoorsInformationSeekingEnv-v1",

            "SocialAI-TestLeverDoorCollaborationEnv-v1",
            "SocialAI-TestMarblePushCollaborationEnv-v1",
            "SocialAI-TestMarblePassCollaborationEnv-v1",

            "SocialAI-TestAppleStealingPerspectiveTakingEnv-v1"
        ]
    elif args.test_set_name == "SocialAIGSTestSet":
        test_env_names = [
            "SocialAI-GridSearchParamEnv-v1",
            "SocialAI-GridSearchPointingParamEnv-v1",
            "SocialAI-GridSearchLangColorParamEnv-v1",
            "SocialAI-GridSearchLangFeedbackParamEnv-v1",
        ]
    elif args.test_set_name == "SocialAICuesGSTestSet":
        test_env_names = [
            "SocialAI-CuesGridSearchParamEnv-v1",
            "SocialAI-CuesGridSearchPointingParamEnv-v1",
            "SocialAI-CuesGridSearchLangColorParamEnv-v1",
            "SocialAI-CuesGridSearchLangFeedbackParamEnv-v1",
        ]
    elif args.test_set_name == "BoxesPointingTestSet":
        test_env_names = [
            "SocialAI-TestPointingBoxesInformationSeekingParamEnv-v1",
        ]
    elif args.test_set_name == "PointingTestSet":
        test_env_names = gym_minigrid.social_ai_envs.pointing_test_set
    elif args.test_set_name == "LangColorTestSet":
        test_env_names = gym_minigrid.social_ai_envs.language_color_test_set
    elif args.test_set_name == "LangFeedbackTestSet":
        test_env_names = gym_minigrid.social_ai_envs.language_feedback_test_set
    # joint attention
    elif args.test_set_name == "JAPointingTestSet":
        test_env_names = gym_minigrid.social_ai_envs.ja_pointing_test_set
    elif args.test_set_name == "JALangColorTestSet":
        test_env_names = gym_minigrid.social_ai_envs.ja_language_color_test_set
    elif args.test_set_name == "JALangFeedbackTestSet":
        test_env_names = gym_minigrid.social_ai_envs.ja_language_feedback_test_set
    # emulation
    elif args.test_set_name == "DistrEmulationTestSet":
        test_env_names = gym_minigrid.social_ai_envs.distr_emulation_test_set
    elif args.test_set_name == "NoDistrEmulationTestSet":
        test_env_names = gym_minigrid.social_ai_envs.no_distr_emulation_test_set
    # formats
    elif args.test_set_name == "NFormatsTestSet":
        test_env_names = gym_minigrid.social_ai_envs.N_formats_test_set
    elif args.test_set_name == "EFormatsTestSet":
        test_env_names = gym_minigrid.social_ai_envs.E_formats_test_set
    elif args.test_set_name == "AFormatsTestSet":
        test_env_names = gym_minigrid.social_ai_envs.A_formats_test_set
    elif args.test_set_name == "AEFormatsTestSet":
        test_env_names = gym_minigrid.social_ai_envs.AE_formats_test_set

    elif args.test_set_name == "RoleReversalTestSet":
        test_env_names = gym_minigrid.social_ai_envs.role_reversal_test_set

    else:
        raise ValueError("Undefined test set name.")


else:
    test_env_names = [args.env]

# test_envs = []
testers = []
if args.test_interval > 0:
    for test_env_name in test_env_names:
        make_env_args = {
            "env_key": test_env_name,
            "seed": args.test_seed,
            "env_args": test_env_args,
        }
        testers.append(utils.Tester(
            make_env_args, args.test_seed, args.test_episodes, model_dir, acmodel, preprocess_obss, device)
        )

        # test_env = utils.make_env(test_env_name, args.test_seed, env_args=test_env_args)
        # test_envs.append(test_env)

        # init tester
        # testers.append(utils.Tester(test_env, args.test_seed, args.test_episodes, model_dir, acmodel, preprocess_obss, device))

if args.continue_train:
    for tester in testers:
        tester.load()


# Save config
env_args_ = {k: v.__repr__() if k == "curriculum" else v for k, v in env_args.items()}
test_env_args_ = {k: v.__repr__() if k == "curriculum" else v for k, v in test_env_args.items()}
config_dict = {
    "seed": args.seed,
    "env": args.env,
    "env_args": env_args_,
    "test_seed": args.test_seed,
    "test_env": args.test_set_name,
    "test_env_args": test_env_args_
}
config_dict.update(algo.get_config_dict())
config_dict.update(acmodel.get_config_dict())
with open(model_dir+'/config.json', 'w') as fp:
    json.dump(config_dict, fp)


# Train model

num_frames = status["num_frames"]
update = status["update"]
start_time = time.time()

log_add_headers = num_frames == 0 or not args.continue_train

long_term_save_interval = 5000000

if args.continue_train:
    # set next long term save interval
    next_long_term_save = (1 + num_frames // long_term_save_interval) * long_term_save_interval

else:
    next_long_term_save = 0  # for long term logging


while num_frames < args.frames:
    # Update model parameters

    update_start_time = time.time()
    # print("current_seed_pre_train:", np.random.get_state()[1][0])
    exps, logs1 = algo.collect_experiences()
    logs2 = algo.update_parameters(exps)
    logs = {**logs1, **logs2}
    update_end_time = time.time()

    num_frames += logs["num_frames"]
    update += 1

    NPC_intro = np.mean(logs["NPC_introduced_to"])

    # Print logs

    if update % args.log_interval == 0:
        fps = logs["num_frames"]/(update_end_time - update_start_time)
        duration = int(time.time() - start_time)
        return_per_episode = utils.synthesize(logs["return_per_episode"])
        extrinsic_return_per_episode = utils.synthesize(logs["extrinsic_return_per_episode"])
        exploration_bonus_per_episode = utils.synthesize(logs["exploration_bonus_per_episode"])
        success_rate = utils.synthesize(logs["success_rate_per_episode"])
        curriculum_max_success_rate = utils.synthesize(logs["curriculum_max_mean_perf_per_episode"])
        curriculum_param = utils.synthesize(logs["curriculum_param_per_episode"])
        rreturn_per_episode = utils.synthesize(logs["reshaped_return_per_episode"])
        num_frames_per_episode = utils.synthesize(logs["num_frames_per_episode"])

        # intrinsic_reward_perf = utils.synthesize(logs["intr_reward_perf"])
        # intrinsic_reward_perf_ = utils.synthesize(logs["intr_reward_perf_"])

        intrinsic_reward_perf = logs["intr_reward_perf"]
        intrinsic_reward_perf_ = logs["intr_reward_perf_"]

        lr_ = logs["lr"]

        time_now = int(datetime.datetime.now().strftime("%d%m%Y%H%M%S"))

        header = ["update", "frames", "FPS", "duration", "time"]
        data = [update, num_frames, fps, duration, time_now]
        data_to_print = [update, num_frames, fps, duration, time_now]

        header += ["success_rate_" + key for key in success_rate.keys()]
        data += success_rate.values()
        data_to_print += success_rate.values()

        header += ["curriculum_max_success_rate_" + key for key in curriculum_max_success_rate.keys()]
        data += curriculum_max_success_rate.values()
        if args.acl:
            data_to_print += curriculum_max_success_rate.values()

        header += ["curriculum_param_" + key for key in curriculum_param.keys()]
        data += curriculum_param.values()
        if args.acl:
            data_to_print += curriculum_param.values()

        header += ["extrinsic_return_" + key for key in extrinsic_return_per_episode.keys()]
        data += extrinsic_return_per_episode.values()
        data_to_print += extrinsic_return_per_episode.values()

        # turn on
        header += ["exploration_bonus_" + key for key in exploration_bonus_per_episode.keys()]
        data += exploration_bonus_per_episode.values()
        data_to_print += exploration_bonus_per_episode.values()

        header += ["rreturn_" + key for key in rreturn_per_episode.keys()]
        data += rreturn_per_episode.values()
        data_to_print += rreturn_per_episode.values()


        header += ["intrinsic_reward_perf_"]
        data += [intrinsic_reward_perf]
        # data_to_print += [intrinsic_reward_perf]

        header += ["intrinsic_reward_perf2_"]
        data += [intrinsic_reward_perf_]
        # data_to_print += [intrinsic_reward_perf_]

        # header += ["num_frames_" + key for key in num_frames_per_episode.keys()]
        # data += num_frames_per_episode.values()

        header += ["NPC_intro"]
        data += [NPC_intro]
        data_to_print += [NPC_intro]

        header += ["lr"]
        data += [lr_]
        data_to_print += [lr_]

        # header += ["entropy", "value", "policy_loss", "value_loss", "grad_norm"]
        # data += [logs["entropy"], logs["value"], logs["policy_loss"], logs["value_loss"], logs["grad_norm"]]

        # curr_history_len = len(algo.env.envs[0].curriculum.performance_history)
        # header += ["curr_history_len"]
        # data += [curr_history_len]

        txt_logger.info("".join([
            "U {} | F {:06} | FPS {:04.0f} | D {} | T {} ",
            "| SR:μσmM {:.2f} {:.1f} {:.1f} {:.1f} ",
            "| CurMaxSR:μσmM {:.2f} {:.1f} {:.1f} {:.1f} " if args.acl else "",
            "| CurPhase:μσmM {:.2f} {:.1f} {:.1f} {:.1f} " if args.acl else "",
            "| ExR:μσmM {:.2f} {:.1f} {:.1f} {:.1f} ",
            "| InR:μσmM {:.2f} {:.1f} {:.1f} {:.1f} ",
            "| rR:μσmM {:.6f} {:.1f} {:.1f} {:.1f} ",
            # "| irp:μσmM {:.6f} {:.2f} {:.2f} {:.2f} ",
            # "| irp_:μσmM {:.6f} {:.2f} {:.2f} {:.2f} ",
            # "| F:μσmM {:.1f} {:.1f} {} {} ",
            "| NPC_intro: {:.3f}",
            "| lr: {:.5f}",
            # "| cur_his_len: {:.5f}" if args.acl else "",
            # "| H {:.3f} | V {:.3f} | pL {:.3f} | vL {:.3f} | ∇ {:.3f}"
        ]).format(*data_to_print))

        header += ["return_" + key for key in return_per_episode.keys()]
        data += return_per_episode.values()

        if log_add_headers:
            csv_logger.writerow(header)
            log_add_headers = False
        csv_logger.writerow(data)
        csv_file.flush()

    # Save status
    long_term_save = False
    if num_frames >= next_long_term_save:
        next_long_term_save += long_term_save_interval
        long_term_save = True

    if (args.save_interval > 0 and update % args.save_interval == 0) or long_term_save:
        # continuing train works best when save_interval == log_interval, the csv is cleaner wo redundancies
        status = {"num_frames": num_frames, "update": update}

        algo_status = algo.get_status_dict()
        status = {**status, **algo_status}

        if hasattr(preprocess_obss, "vocab"):
            status["vocab"] = preprocess_obss.vocab.vocab
            status["env_args"] = env_args

        if long_term_save:
            utils.save_status(status, model_dir, num_frames=num_frames)
            utils.save_model(acmodel, model_dir, num_frames=num_frames)
            txt_logger.info("Status and Model saved for {} frames".format(num_frames))

        else:
            utils.save_status(status, model_dir)
            utils.save_model(acmodel, model_dir)
            txt_logger.info("Status and Model saved")

    if args.test_interval > 0 and (update % args.test_interval == 0 or update == 1):
        txt_logger.info(f"Testing at update {update}.")
        test_success_rates = []
        for tester in testers:
            mean_success_rate, mean_rewards = tester.test_agent(num_frames)
            test_success_rates.append(mean_success_rate)
            txt_logger.info(f"\t{tester.envs[0].spec.id} -> {mean_success_rate} (SR)")
            tester.dump()

        if len(testers):
            txt_logger.info(f"Test set SR: {np.array(test_success_rates).mean()}")


# save at the end
status = {"num_frames": num_frames, "update": update}
algo_status = algo.get_status_dict()
status = {**status, **algo_status}

if hasattr(preprocess_obss, "vocab"):
    status["vocab"] = preprocess_obss.vocab.vocab
    status["env_args"] = env_args

utils.save_status(status, model_dir)
utils.save_model(acmodel, model_dir)
txt_logger.info("Status and Model saved at the end")