drafted rainbow baseline

This is untested at present.

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"]

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)

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