mirror of
https://github.com/wassname/retro-baselines.git
synced 2026-08-29 11:25:51 +08:00
139 lines
4.5 KiB
Python
139 lines
4.5 KiB
Python
#!/usr/bin/env python
|
|
|
|
"""
|
|
Train an agent on Sonic using PPO2 from OpenAI Baselines.
|
|
|
|
Run with:
|
|
```bash
|
|
python ppo2_agent.py --config configs/sonic.json
|
|
```
|
|
"""
|
|
|
|
import tensorflow as tf
|
|
|
|
import os
|
|
import sys
|
|
import datetime as dt
|
|
import numpy as np
|
|
import logging
|
|
|
|
import baselines.ppo2_rudder.ppo2_rudder as ppo2_rudder
|
|
from baselines import bench, logger
|
|
from sonic_util import make_env as sonic_env
|
|
|
|
from TeLL.config import Config
|
|
from TeLL.utility.plotting import launch_plotting_daemon, save_subplots, save_movie, save_subplots_line_plots
|
|
from TeLL.utility.misc import make_sure_path_exists, Tee
|
|
|
|
# Start subprocess for plotting workers
|
|
# Due to a garbage-collector bug with matplotlib/GPU, launch_plotting_daemon needs so be called before tensorflow
|
|
# import
|
|
launch_plotting_daemon(num_workers=3)
|
|
|
|
|
|
# Log to file and stream
|
|
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
|
|
logger_file = logging.getLogger(__file__)
|
|
|
|
|
|
def train(env_id, num_timesteps, policy, working_dir, config):
|
|
"""Run PPO until the environment throws an exception."""
|
|
# Original modules
|
|
from baselines.common import set_global_seeds
|
|
from baselines.common.vec_env.subproc_vec_env import SubprocVecEnv
|
|
import logging
|
|
import gym
|
|
import os.path as osp
|
|
# Module modified for RUDDER
|
|
from baselines.common.vec_env.vec_frame_stack import VecFrameStackNoZeroPadding
|
|
from baselines.ppo2_rudder.policies import CnnPolicy, LstmPolicy, LstmPolicyDense
|
|
|
|
bl_config = config.bl_config
|
|
|
|
# Set numpy random seed
|
|
rnd_seed = config.get_value('random_seed', 12345)
|
|
rnd_gen = np.random.RandomState(seed=rnd_seed)
|
|
|
|
# Set GPU
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = str(config.get_value("cuda_gpu", "0"))
|
|
|
|
# Tensorflow configuration
|
|
tf_config = tf.ConfigProto(
|
|
allow_soft_placement=True,
|
|
inter_op_parallelism_threads=config.get_value("inter_op_parallelism_threads", 1),
|
|
intra_op_parallelism_threads=config.get_value("intra_op_parallelism_threads", 1),
|
|
log_device_placement=config.get_value("log_device_placement", False)
|
|
)
|
|
tf_config.gpu_options.allow_growth = config.get_value("tf_allow_growth", True)
|
|
|
|
# Start Tensorflow session
|
|
print("Preparing Logger...")
|
|
|
|
gym.logger.setLevel(logging.INFO)
|
|
print("Starting session...")
|
|
tf_session = tf.Session(config=tf_config).__enter__()
|
|
|
|
# Set tensorflow random seed
|
|
tf.set_random_seed(rnd_seed)
|
|
|
|
# Create parallel environments
|
|
print("Preparing Envionments...", end="")
|
|
|
|
def make_env(rank):
|
|
def env_fn():
|
|
np.random.seed(rnd_seed + rank)
|
|
env = sonic_env(scale_rew=True)
|
|
env.unwrapped.rank = rank
|
|
env.seed(rnd_seed + rank)
|
|
env = bench.Monitor(env, logger.get_dir() and osp.join(logger.get_dir(), str(rank)))
|
|
return env
|
|
return env_fn
|
|
|
|
nenvs = bl_config['num_actors']
|
|
print("creating workers...", end="")
|
|
env = SubprocVecEnv([make_env(i) for i in range(nenvs)])
|
|
set_global_seeds(rnd_seed)
|
|
print("stacking frames...", end="")
|
|
env = VecFrameStackNoZeroPadding(env, 4)
|
|
print("Done!")
|
|
|
|
# Enter learning
|
|
policy = {'cnn': CnnPolicy, 'lstmdense': LstmPolicyDense, 'lstm': LstmPolicy}[policy]
|
|
ppo2_rudder.learn(
|
|
policy=policy,
|
|
env=env,
|
|
nsteps=1024,
|
|
nminibatches=4,
|
|
lam=0.95,
|
|
gamma=0.99,
|
|
noptepochs=4,
|
|
log_interval=1,
|
|
ent_coef=bl_config['ent_coef'],
|
|
lr=lambda f: f * 6e-5 * bl_config['lr_coef'],
|
|
cliprange=lambda f: f * 0.2,
|
|
total_timesteps=int(num_timesteps * 1.1), tf_session=tf_session,
|
|
working_dir=working_dir,
|
|
config=config,
|
|
plotting=dict(save_subplots=save_subplots, save_movie=save_movie,
|
|
save_subplots_line_plots=save_subplots_line_plots),
|
|
rnd_gen=rnd_gen
|
|
)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
config = Config()
|
|
working_dir = os.path.join(config.working_dir, config.specs)
|
|
working_dir = os.path.join(working_dir, dt.datetime.now().strftime("%Y-%m-%dT%H-%M-%S"))
|
|
make_sure_path_exists(working_dir)
|
|
|
|
with open(os.path.join(working_dir, 'log.txt'), 'a') as logfile:
|
|
sys.stdout = Tee(sys.stdout, logfile, sys.stdout)
|
|
|
|
bl_config = config.get_value('bl_config')
|
|
|
|
logger.configure(os.path.join(working_dir, 'baselines'), ['tensorboard', 'log', 'stdout'])
|
|
train(env_id=bl_config['env'], num_timesteps=bl_config['num_timesteps'],
|
|
policy=config.get_value('policy'), working_dir=working_dir, config=config)
|
|
|
|
sys.stdout.flush()
|