c5d9c4a1b2
* wrap retro envs correctly for other (non-deepq) algorithms * flake and csh comments * flake and csh comments
224 lines
6.1 KiB
Python
224 lines
6.1 KiB
Python
import sys
|
|
import multiprocessing
|
|
import os.path as osp
|
|
import gym
|
|
from collections import defaultdict
|
|
import tensorflow as tf
|
|
import numpy as np
|
|
|
|
from baselines.common.vec_env.vec_frame_stack import VecFrameStack
|
|
from baselines.common.cmd_util import common_arg_parser, parse_unknown_args, make_vec_env, make_env
|
|
from baselines.common.tf_util import get_session
|
|
from baselines import logger
|
|
from importlib import import_module
|
|
|
|
from baselines.common.vec_env.vec_normalize import VecNormalize
|
|
|
|
try:
|
|
from mpi4py import MPI
|
|
except ImportError:
|
|
MPI = None
|
|
|
|
try:
|
|
import pybullet_envs
|
|
except ImportError:
|
|
pybullet_envs = None
|
|
|
|
try:
|
|
import roboschool
|
|
except ImportError:
|
|
roboschool = None
|
|
|
|
_game_envs = defaultdict(set)
|
|
for env in gym.envs.registry.all():
|
|
# TODO: solve this with regexes
|
|
env_type = env._entry_point.split(':')[0].split('.')[-1]
|
|
_game_envs[env_type].add(env.id)
|
|
|
|
# reading benchmark names directly from retro requires
|
|
# importing retro here, and for some reason that crashes tensorflow
|
|
# in ubuntu
|
|
_game_envs['retro'] = {
|
|
'BubbleBobble-Nes',
|
|
'SuperMarioBros-Nes',
|
|
'TwinBee3PokoPokoDaimaou-Nes',
|
|
'SpaceHarrier-Nes',
|
|
'SonicTheHedgehog-Genesis',
|
|
'Vectorman-Genesis',
|
|
'FinalFight-Snes',
|
|
'SpaceInvaders-Snes',
|
|
}
|
|
|
|
|
|
def train(args, extra_args):
|
|
env_type, env_id = get_env_type(args.env)
|
|
print('env_type: {}'.format(env_type))
|
|
|
|
total_timesteps = int(args.num_timesteps)
|
|
seed = args.seed
|
|
|
|
learn = get_learn_function(args.alg)
|
|
alg_kwargs = get_learn_function_defaults(args.alg, env_type)
|
|
alg_kwargs.update(extra_args)
|
|
|
|
env = build_env(args)
|
|
|
|
if args.network:
|
|
alg_kwargs['network'] = args.network
|
|
else:
|
|
if alg_kwargs.get('network') is None:
|
|
alg_kwargs['network'] = get_default_network(env_type)
|
|
|
|
print('Training {} on {}:{} with arguments \n{}'.format(args.alg, env_type, env_id, alg_kwargs))
|
|
|
|
model = learn(
|
|
env=env,
|
|
seed=seed,
|
|
total_timesteps=total_timesteps,
|
|
**alg_kwargs
|
|
)
|
|
|
|
return model, env
|
|
|
|
|
|
def build_env(args):
|
|
ncpu = multiprocessing.cpu_count()
|
|
if sys.platform == 'darwin': ncpu //= 2
|
|
nenv = args.num_env or ncpu
|
|
alg = args.alg
|
|
seed = args.seed
|
|
|
|
env_type, env_id = get_env_type(args.env)
|
|
|
|
if env_type in {'atari', 'retro'}:
|
|
if alg == 'acer':
|
|
env = make_vec_env(env_id, env_type, nenv, seed)
|
|
elif alg == 'deepq':
|
|
env = make_env(env_id, env_type, seed=seed, wrapper_kwargs={'frame_stack': True})
|
|
elif alg == 'trpo_mpi':
|
|
env = make_env(env_id, env_type, seed=seed)
|
|
else:
|
|
frame_stack_size = 4
|
|
env = make_vec_env(env_id, env_type, nenv, seed, gamestate=args.gamestate, reward_scale=args.reward_scale)
|
|
env = VecFrameStack(env, frame_stack_size)
|
|
|
|
else:
|
|
config = tf.ConfigProto(allow_soft_placement=True,
|
|
intra_op_parallelism_threads=1,
|
|
inter_op_parallelism_threads=1)
|
|
config.gpu_options.allow_growth = True
|
|
get_session(config=config)
|
|
|
|
env = make_vec_env(env_id, env_type, args.num_env or 1, seed, reward_scale=args.reward_scale)
|
|
|
|
if env_type == 'mujoco':
|
|
env = VecNormalize(env)
|
|
|
|
return env
|
|
|
|
|
|
def get_env_type(env_id):
|
|
if env_id in _game_envs.keys():
|
|
env_type = env_id
|
|
env_id = [g for g in _game_envs[env_type]][0]
|
|
else:
|
|
env_type = None
|
|
for g, e in _game_envs.items():
|
|
if env_id in e:
|
|
env_type = g
|
|
break
|
|
assert env_type is not None, 'env_id {} is not recognized in env types'.format(env_id, _game_envs.keys())
|
|
|
|
return env_type, env_id
|
|
|
|
|
|
def get_default_network(env_type):
|
|
if env_type == 'atari':
|
|
return 'cnn'
|
|
else:
|
|
return 'mlp'
|
|
|
|
def get_alg_module(alg, submodule=None):
|
|
submodule = submodule or alg
|
|
try:
|
|
# first try to import the alg module from baselines
|
|
alg_module = import_module('.'.join(['baselines', alg, submodule]))
|
|
except ImportError:
|
|
# then from rl_algs
|
|
alg_module = import_module('.'.join(['rl_' + 'algs', alg, submodule]))
|
|
|
|
return alg_module
|
|
|
|
|
|
def get_learn_function(alg):
|
|
return get_alg_module(alg).learn
|
|
|
|
|
|
def get_learn_function_defaults(alg, env_type):
|
|
try:
|
|
alg_defaults = get_alg_module(alg, 'defaults')
|
|
kwargs = getattr(alg_defaults, env_type)()
|
|
except (ImportError, AttributeError):
|
|
kwargs = {}
|
|
return kwargs
|
|
|
|
|
|
|
|
def parse_cmdline_kwargs(args):
|
|
'''
|
|
convert a list of '='-spaced command-line arguments to a dictionary, evaluating python objects when possible
|
|
'''
|
|
def parse(v):
|
|
|
|
assert isinstance(v, str)
|
|
try:
|
|
return eval(v)
|
|
except (NameError, SyntaxError):
|
|
return v
|
|
|
|
return {k: parse(v) for k,v in parse_unknown_args(args).items()}
|
|
|
|
|
|
|
|
def main():
|
|
# configure logger, disable logging in child MPI processes (with rank > 0)
|
|
|
|
arg_parser = common_arg_parser()
|
|
args, unknown_args = arg_parser.parse_known_args()
|
|
extra_args = parse_cmdline_kwargs(unknown_args)
|
|
|
|
if MPI is None or MPI.COMM_WORLD.Get_rank() == 0:
|
|
rank = 0
|
|
logger.configure()
|
|
else:
|
|
logger.configure(format_strs=[])
|
|
rank = MPI.COMM_WORLD.Get_rank()
|
|
|
|
model, env = train(args, extra_args)
|
|
env.close()
|
|
|
|
if args.save_path is not None and rank == 0:
|
|
save_path = osp.expanduser(args.save_path)
|
|
model.save(save_path)
|
|
|
|
if args.play:
|
|
logger.log("Running trained model")
|
|
env = build_env(args)
|
|
obs = env.reset()
|
|
def initialize_placeholders(nlstm=128,**kwargs):
|
|
return np.zeros((args.num_env or 1, 2*nlstm)), np.zeros((1))
|
|
state, dones = initialize_placeholders(**extra_args)
|
|
while True:
|
|
actions, _, state, _ = model.step(obs,S=state, M=dones)
|
|
obs, _, done, _ = env.step(actions)
|
|
env.render()
|
|
done = done.any() if isinstance(done, np.ndarray) else done
|
|
|
|
if done:
|
|
obs = env.reset()
|
|
|
|
env.close()
|
|
|
|
if __name__ == '__main__':
|
|
main()
|