DB/run.py

#!/usr/bin/env python
try:
    from OpenGL import GLU
except:
    print("no OpenGL.GLU")
import functools
import os.path as osp
from functools import partial
import os
import gym
import dmc2gym

import utils
import tensorflow as tf
from baselines import logger
from baselines.bench import Monitor
from baselines.common.atari_wrappers import NoopResetEnv, FrameStack
from mpi4py import MPI

from dynamic_bottleneck import DynamicBottleneck    
from cnn_policy import CnnPolicy
from cppo_agent import PpoOptimizer
from utils import random_agent_ob_mean_std
from wrappers import MontezumaInfoWrapper, make_mario_env, make_robo_pong, make_robo_hockey, \
    make_multi_pong, AddRandomStateToInfo, MaxAndSkipEnv, ProcessFrame84, ExtraTimeLimit, StickyActionEnv
import datetime
from wrappers import PixelNoiseWrapper, RandomBoxNoiseWrapper
import json

getsess = tf.get_default_session

def start_experiment(**args):

    make_env = partial(make_env_all_params, add_monitor=True, args=args)

    trainer = Trainer(make_env=make_env,
                      num_timesteps=args['num_timesteps'], hps=args,
                      envs_per_process=args['envs_per_process'])
    log, tf_sess, saver, logger_dir = get_experiment_environment(**args)
    with log, tf_sess:
        logdir = logger.get_dir()
        print("results will be saved to ", logdir)
        trainer.train(saver, logger_dir)

class Trainer(object):
    def __init__(self, make_env, hps, num_timesteps, envs_per_process):
        self.make_env = make_env
        self.hps = hps
        self.envs_per_process = envs_per_process
        self.num_timesteps = num_timesteps
        self._set_env_vars()   

        self.policy = CnnPolicy(scope='pol',
                                ob_space=self.ob_space,
                                ac_space=self.ac_space,
                                hidsize=512,
                                feat_dim=512,
                                ob_mean=self.ob_mean,
                                ob_std=self.ob_std,
                                layernormalize=False,
                                nl=tf.nn.leaky_relu)

        self.dynamic_bottleneck = DynamicBottleneck(
                    policy=self.policy, feat_dim=512, tau=hps['momentum_tau'], loss_kl_weight=hps['loss_kl_weight'],
                    loss_nce_weight=hps['loss_nce_weight'], loss_l2_weight=hps['loss_l2_weight'], aug=hps['aug'])

        self.agent = PpoOptimizer(
            scope='ppo',
            ob_space=self.ob_space,
            ac_space=self.ac_space,
            stochpol=self.policy,
            use_news=hps['use_news'],
            gamma=hps['gamma'],
            lam=hps["lambda"],
            nepochs=hps['nepochs'],
            nminibatches=hps['nminibatches'],
            lr=hps['lr'],
            cliprange=0.1,
            nsteps_per_seg=hps['nsteps_per_seg'],
            nsegs_per_env=hps['nsegs_per_env'],
            ent_coef=hps['ent_coeff'],
            normrew=hps['norm_rew'],
            normadv=hps['norm_adv'],
            ext_coeff=hps['ext_coeff'],
            int_coeff=hps['int_coeff'],
            dynamic_bottleneck=self.dynamic_bottleneck
        )

        self.agent.to_report['db'] = tf.reduce_mean(self.dynamic_bottleneck.loss)
        self.agent.total_loss += self.agent.to_report['db']

        self.agent.db_loss = tf.reduce_mean(self.dynamic_bottleneck.loss)

        self.agent.to_report['feat_var'] = tf.reduce_mean(tf.nn.moments(self.dynamic_bottleneck.features, [0, 1])[1])

    def _set_env_vars(self):
        env = self.make_env(0, add_monitor=False)
        # ob_space.shape=(84, 84, 4)     ac_space.shape=Discrete(4)
        self.ob_space, self.ac_space = env.observation_space, env.action_space
        self.ob_mean, self.ob_std = random_agent_ob_mean_std(env)
        del env
        self.envs = [functools.partial(self.make_env, i) for i in range(self.envs_per_process)]

    def train(self, saver, logger_dir):
        self.agent.start_interaction(self.envs, nlump=self.hps['nlumps'], dynamic_bottleneck=self.dynamic_bottleneck)
        previous_saved_tcount = 0

        # add bai. initialize IB parameters
        print("***Init Momentum Network in Dynamic-Bottleneck.")
        getsess().run(self.dynamic_bottleneck.init_updates)

        while True:
            info = self.agent.step()         # 
            if info['DB_loss_info']:         # add bai. for debug
                logger.logkvs(info['DB_loss_info'])
            if info['update']:
                logger.logkvs(info['update'])
                logger.dumpkvs()
            if self.hps["save_period"] and (int(self.agent.rollout.stats['tcount'] / self.hps["save_freq"]) > previous_saved_tcount):
                previous_saved_tcount += 1
                save_path = saver.save(tf.get_default_session(), os.path.join(logger_dir, "model_"+str(previous_saved_tcount)+".ckpt"))
                print("Periodically model saved in path:", save_path)
            if self.agent.rollout.stats['tcount'] %10000: #self.agent.rollout.stats['tcount'] > self.num_timesteps:
                save_path = saver.save(tf.get_default_session(), os.path.join(logger_dir, "model_last.ckpt"))
                print("Model saved in path:", save_path)
                #break

        self.agent.stop_interaction()


def make_env_all_params(rank, add_monitor, args):
    env = dmc2gym.make(
        domain_name='cartpole',
        task_name='swingup',
        seed=args["seed"],
        visualize_reward=False,
        from_pixels='pixel',
        height=84,
        width=84,
        frame_skip=4,
        img_source=args["img_source"],
        resource_files=args["resource_files"],
        total_frames=args["total_frames"]
    )
    env.seed(args["seed"])

    env = utils.FrameStack(env, k=4)

    """
    if args["env_kind"] == 'atari':
        env = gym.make(args['env'])
        assert 'NoFrameskip' in env.spec.id
        if args["stickyAtari"]:               # 
            env._max_episode_steps = args['max_episode_steps'] * 4
            env = StickyActionEnv(env)
        else:
            env = NoopResetEnv(env, noop_max=args['noop_max'])
        env = MaxAndSkipEnv(env, skip=4)            # 
        if args['pixelNoise']:                      # add pixel noise
            env = PixelNoiseWrapper(env)
        if args['randomBoxNoise']:
            env = RandomBoxNoiseWrapper(env)
        env = ProcessFrame84(env, crop=False)       #
        env = FrameStack(env, 4)                    #
        # env = ExtraTimeLimit(env, args['max_episode_steps'])
        if not args["stickyAtari"]:
            env = ExtraTimeLimit(env, args['max_episode_steps'])  # 
        if 'Montezuma' in args['env']:              # 
            env = MontezumaInfoWrapper(env)
        env = AddRandomStateToInfo(env)
    elif args["env_kind"] == 'mario':               # 
        env = make_mario_env()
    elif args["env_kind"] == "retro_multi":         # 
        env = make_multi_pong()
    elif args["env_kind"] == 'robopong':
        if args["env"] == "pong":
            env = make_robo_pong()
        elif args["env"] == "hockey":
            env = make_robo_hockey()
    """
    if add_monitor:
        env = Monitor(env, osp.join(logger.get_dir(), '%.2i' % rank))
    return env


def get_experiment_environment(**args):
    from utils import setup_mpi_gpus, setup_tensorflow_session
    from baselines.common import set_global_seeds
    from gym.utils.seeding import hash_seed
    process_seed = args["seed"] + 1000 * MPI.COMM_WORLD.Get_rank()
    process_seed = hash_seed(process_seed, max_bytes=4)
    set_global_seeds(process_seed)
    setup_mpi_gpus()

    # log dir name
    logger_dir = './logs/' + args["env"].replace("NoFrameskip-v4", "")
    # logger_dir += "-KLloss-"+str(args["loss_kl_weight"])     
    # logger_dir += "-NCEloss-" + str(args["loss_nce_weight"]) 
    # logger_dir += "-L2loss-" + str(args["loss_l2_weight"])
    if args['pixelNoise'] is True:
        logger_dir += "-pixelNoise"
    if args['randomBoxNoise'] is True:
        logger_dir += "-randomBoxNoise"
    if args['stickyAtari'] is True:
        logger_dir += "-stickyAtari"
    if args["comments"] != "":
        logger_dir += '-' + args["comments"]
    logger_dir += datetime.datetime.now().strftime("-%m-%d-%H-%M-%S")

    # write config
    logger.configure(dir=logger_dir)
    with open(os.path.join(logger_dir, 'parameters.txt'), 'w') as f:
        f.write("\n".join([str(x[0]) + ": " + str(x[1]) for x in args.items()]))

    logger_context = logger.scoped_configure(
        dir=logger_dir,
        format_strs=['stdout', 'log', 'csv'] if MPI.COMM_WORLD.Get_rank() == 0 else ['log'])
    tf_context = setup_tensorflow_session()

    # saver 
    saver = tf.train.Saver()
    return logger_context, tf_context, saver, logger_dir


def add_environments_params(parser):
    parser.add_argument('--env', help='environment ID', default='BreakoutNoFrameskip-v4', type=str)
    parser.add_argument('--max-episode-steps', help='maximum number of timesteps for episode', default=4500, type=int)
    parser.add_argument('--env_kind', type=str, default="atari")
    parser.add_argument('--noop_max', type=int, default=30)
    parser.add_argument('--stickyAtari', action='store_true', default=False)
    parser.add_argument('--pixelNoise', action='store_true', default=False)
    parser.add_argument('--randomBoxNoise', action='store_true', default=False)
    parser.add_argument('--img_source', default=None, type=str, choices=['color', 'noise', 'images', 'video', 'none'])
    parser.add_argument('--resource_files', type=str)
    parser.add_argument('--total_frames', default=100, type=int)


def add_optimization_params(parser):
    parser.add_argument('--lambda', type=float, default=0.95)
    parser.add_argument('--gamma', type=float, default=0.99)                  # lambda, gamma 用于计算 GAE advantage
    parser.add_argument('--nminibatches', type=int, default=8)
    parser.add_argument('--norm_adv', type=int, default=1)                    # 
    parser.add_argument('--norm_rew', type=int, default=1)                    # 
    parser.add_argument('--lr', type=float, default=1e-4)                     # 
    parser.add_argument('--ent_coeff', type=float, default=0.001)             # 
    parser.add_argument('--nepochs', type=int, default=3)                     # 
    parser.add_argument('--num_timesteps', type=int, default=int(1e8))
    parser.add_argument('--save_period', action='store_true', default=False)  # 1e7
    parser.add_argument('--save_freq', type=int, default=int(1e7))            # 1e7
    # Parameters of Dynamic-Bottleneck
    parser.add_argument('--loss_kl_weight', type=float, default=0.1)          # KL loss weight
    parser.add_argument('--loss_l2_weight', type=float, default=0.1)        # l2 loss weight
    parser.add_argument('--loss_nce_weight', type=float, default=0.01)         # nce loss weight
    parser.add_argument('--momentum_tau', type=float, default=0.001)          # momentum tau
    parser.add_argument('--aug', action='store_true', default=False)          # data augmentation (bottleneck)
    parser.add_argument('--comments', type=str, default="")


def add_rollout_params(parser):
    parser.add_argument('--nsteps_per_seg', type=int, default=128)
    parser.add_argument('--nsegs_per_env', type=int, default=1)
    parser.add_argument('--envs_per_process', type=int, default=128)
    parser.add_argument('--nlumps', type=int, default=1)


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    add_environments_params(parser)
    add_optimization_params(parser)
    add_rollout_params(parser)

    parser.add_argument('--exp_name', type=str, default='')
    parser.add_argument('--seed', help='RNG seed', type=int, default=0)
    parser.add_argument('--dyn_from_pixels', type=int, default=0)     
    parser.add_argument('--use_news', type=int, default=0)
    parser.add_argument('--ext_coeff', type=float, default=0.)
    parser.add_argument('--int_coeff', type=float, default=1.)
    parser.add_argument('--layernorm', type=int, default=0)

    args = parser.parse_args()

    # load paramets
    with open("para.json") as f:
        d = json.load(f) 
    env_name_para = args.env.replace("NoFrameskip-v4", "")
    if env_name_para not in list(d["standard"].keys()):
        env_name_para = "other"
    
    if args.pixelNoise is True:
        print("pixel noise")
        args.loss_kl_weight = d["pixelNoise"][env_name_para]["kl"]
        args.loss_nce_weight = d["pixelNoise"][env_name_para]["nce"]
    elif args.randomBoxNoise is True:
        print("random box noise")
        args.loss_kl_weight = d["randomBox"][env_name_para]["kl"]
        args.loss_nce_weight = d["randomBox"][env_name_para]["nce"]
    elif args.stickyAtari is True:
        print("sticky noise")
        args.loss_kl_weight = d["stickyAtari"][env_name_para]["kl"]
        args.loss_nce_weight = d["stickyAtari"][env_name_para]["nce"]
    else:
        print("standard atari")
        args.loss_kl_weight = d["standard"][env_name_para]["kl"]
        args.loss_nce_weight = d["standard"][env_name_para]["nce"]

    print("env_name:", env_name_para, "kl:", args.loss_kl_weight, ", nce:", args.loss_nce_weight)
    start_experiment(**args.__dict__)