diff --git a/agents/rainbow.docker b/agents/rainbow.docker new file mode 100644 index 0000000..3ac27b6 --- /dev/null +++ b/agents/rainbow.docker @@ -0,0 +1,21 @@ +FROM agent:tensorflow + +# Needed for OpenCV. +RUN apt-get update && \ + apt-get install -y libgtk2.0-dev && \ + rm -rf /var/lib/apt/lists/* + +# Baselines has some unneeded and cumbersome dependencies, +# so we manually fetch the deps we need. +RUN . ~/venv/bin/activate && \ + pip install scipy tqdm joblib zmq dill progressbar2 cloudpickle opencv-python && \ + pip install --no-deps git+https://github.com/openai/baselines.git + +# Use the anyrl open source RL framework. +RUN . ~/venv/bin/activate && \ + pip install anyrl==0.11.17 + +ADD rainbow_agent.py ./agent.py +ADD sonic_util.py . + +CMD ["python", "-u", "/root/compo/agent.py"] diff --git a/agents/rainbow_agent.py b/agents/rainbow_agent.py new file mode 100644 index 0000000..c18a3db --- /dev/null +++ b/agents/rainbow_agent.py @@ -0,0 +1,48 @@ +#!/usr/bin/env python + +""" +Train an agent on Sonic using an open source Rainbow DQN +implementation. +""" + +import tensorflow as tf + +from anyrl.algos import DQN +from anyrl.envs import BatchedGymEnv +from anyrl.envs.wrappers import BatchedFrameStack +from anyrl.models import rainbow_models +from anyrl.rollouts import BatchedPlayer, PrioritizedReplayBuffer, NStepPlayer +from anyrl.spaces import gym_space_vectorizer +import gym_remote.exceptions as gre + +from sonic_util import AllowBacktracking, make_env + +def main(): + """Run PPO until the environment throws an exception.""" + env = AllowBacktracking(make_env(stack=False, scale_rew=False)) + env = BatchedFrameStack(BatchedGymEnv([[env]]), num_images=4, concat=False) + config = tf.ConfigProto() + config.gpu_options.allow_growth = True # pylint: disable=E1101 + with tf.Session(config=config) as sess: + dqn = DQN(*rainbow_models(sess, + env.action_space.n, + gym_space_vectorizer(env.observation_space), + min_val=-200, + max_val=200)) + player = NStepPlayer(BatchedPlayer(env, dqn.online_net), 3) + optimize = dqn.optimize(learning_rate=1e-4) + sess.run(tf.global_variables_initializer()) + dqn.train(num_steps=2000000, # Make sure an exception arrives before we stop. + player=player, + replay_buffer=PrioritizedReplayBuffer(500000, 0.5, 0.4, epsilon=0.1), + optimize_op=optimize, + train_interval=1, + target_interval=8192, + batch_size=32, + min_buffer_size=20000) + +if __name__ == '__main__': + try: + main() + except gre.GymRemoteError as exc: + print('exception', exc) diff --git a/agents/sonic_util.py b/agents/sonic_util.py index 699ac47..d2a245b 100644 --- a/agents/sonic_util.py +++ b/agents/sonic_util.py @@ -8,15 +8,17 @@ import numpy as np from baselines.common.atari_wrappers import WarpFrame, FrameStack import gym_remote.client as grc -def make_env(): +def make_env(stack=True, scale_rew=True): """ Create an environment with some standard wrappers. """ env = grc.RemoteEnv('tmp/sock') env = SonicDiscretizer(env) - env = RewardScaler(env) + if scale_rew: + env = RewardScaler(env) env = WarpFrame(env) - env = FrameStack(env, 4) + if stack: + env = FrameStack(env, 4) return env class SonicDiscretizer(gym.ActionWrapper): @@ -49,3 +51,27 @@ class RewardScaler(gym.RewardWrapper): """ def reward(self, reward): return reward * 0.01 + +class AllowBacktracking(gym.Wrapper): + """ + Use deltas in max(X) as the reward, rather than deltas + in X. This way, agents are not discouraged too heavily + from exploring backwards if there is no way to advance + head-on in the level. + """ + def __init__(self, env): + super(AllowBacktracking, self).__init__(env) + self._cur_x = 0 + self._max_x = 0 + + def reset(self, **kwargs): # pylint: disable=E0202 + self._cur_x = 0 + self._max_x = 0 + return self.env.reset(**kwargs) + + def step(self, action): # pylint: disable=E0202 + obs, rew, done, info = self.env.step(action) + self._cur_x += rew + rew = max(0, self._cur_x - self._max_x) + self._max_x = max(self._max_x, self._cur_x) + return obs, rew, done, info