Fixed attention mechanism and added further progress

Author: hedrox

.gitignore

@@ -1 +1,2 @@
 .idea
+__pycache__/

README.md

@@ -18,3 +18,13 @@ python run_atari.py --gamma_ext 0.999
 To use more than one gpu/machine, use MPI (e.g. `mpiexec -n 8 python run_atari.py --num_env 128 --gamma_ext 0.999` should use 1024 parallel environments to collect experience on an 8 gpu machine). 
 
 ### [Blog post and videos](https://blog.openai.com/reinforcement-learning-with-prediction-based-rewards/)
+
+### Installation Guide
+First install the conda environment
+```bash
+conda create --name <env_name> --file conda_requirements.txt
+```
+Then install dependencies that cannot be installed with conda
+```bash
+pip install -r pip_requirements.txt
+```

atari_wrappers.py

@@ -67,6 +67,61 @@ def observation(self, frame):
         frame = cv2.resize(frame, (self.width, self.height), interpolation=cv2.INTER_AREA)
         return frame[:, :, None]
 
+class WarpEgo(gym.ObservationWrapper):
+    def __init__(self, env):
+        """Warp frames to 84x84 as done in the Nature paper and later work."""
+        # check that env is montezuma not something else
+        gym.ObservationWrapper.__init__(self, env)
+        # self.width = 84
+        # self.height = 84
+        self.width = 51
+        self.height = 30
+
+        self.observation_space = spaces.Box(low=0, high=255,
+            shape=(self.height, self.width, 1), dtype=np.uint8)
+
+        self.lower_color = np.array([199, 71, 71], dtype="uint8")
+        self.upper_color = np.array([201, 73, 73], dtype="uint8")
+
+    def find_character_in_frame(self, frame):
+        mask = cv2.inRange(frame, self.lower_color, self.upper_color)
+        output = cv2.bitwise_and(frame, frame, mask=mask)
+
+        pix_x, pix_y, _ = np.where(output > 0)
+        if pix_x.size != 0:
+            pix_x = pix_x[np.where(pix_x > 19)]
+            pix_y = pix_y[-pix_x.size:]
+
+            # If array is even then median doesn't exist in the array, because it's the average
+            # between the middle twos
+            median_x = int(np.median(pix_x))
+            while median_x not in pix_x:
+                median_x += 1
+
+            median_y = int(pix_y[np.where(pix_x == median_x)[0][0]])
+        else:
+            median_x = output.shape[0] // 2
+            median_y = output.shape[1] // 2
+
+        low_x = median_x-self.height
+        high_x = median_x+self.height
+        low_y = median_y-self.width
+        high_y = median_y+self.width
+
+        low_x = low_x if low_x > 0 else 0
+        high_x = high_x if high_x < frame.shape[0] else frame.shape[0]
+        low_y = low_y if low_y > 0 else 0
+        high_y = high_y if high_y < frame.shape[1] else frame.shape[1]
+
+        roi = frame[low_x:high_x, low_y:high_y]
+        return roi
+
+    def observation(self, frame):
+        frame = self.find_character_in_frame(frame)
+        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
+        frame = cv2.resize(frame, (self.width, self.height), interpolation=cv2.INTER_AREA)
+        return frame[:, :, None]
+
 class FrameStack(gym.Wrapper):
     def __init__(self, env, k):
         """Stack k last frames.
@@ -212,10 +267,15 @@ def make_atari(env_id, max_episode_steps=4500):
     env = AddRandomStateToInfo(env)
     return env
 
-def wrap_deepmind(env, clip_rewards=True, frame_stack=False, scale=False):
+
+def wrap_deepmind(env, clip_rewards=True, frame_stack=False, scale=False, ego=False):
     """Configure environment for DeepMind-style Atari.
     """
-    env = WarpFrame(env)
+    if ego:
+        env = WarpEgo(env)
+    else:
+        env = WarpFrame(env)
+
     if scale:
         env = ScaledFloatFrame(env)
     if clip_rewards:

conda_requirements.txt

@@ -0,0 +1,119 @@
+# This file may be used to create an environment using:
+# $ conda create --name <env> --file <this file>
+# platform: linux-64
+_tflow_select=2.1.0=gpu
+absl-py=0.7.1=py36_0
+astor=0.7.1=py36_0
+backcall=0.1.0=py36_0
+blas=1.0=mkl
+bzip2=1.0.6=h14c3975_5
+c-ares=1.15.0=h7b6447c_1
+ca-certificates=2019.1.23=0
+cairo=1.14.12=h8948797_3
+certifi=2019.3.9=py36_0
+cudatoolkit=10.0.130=0
+cudnn=7.3.1=cuda10.0_0
+cupti=10.0.130=0
+cycler=0.10.0=py36_0
+dbus=1.13.6=h746ee38_0
+decorator=4.4.0=py36_1
+expat=2.2.6=he6710b0_0
+ffmpeg=4.0=hcdf2ecd_0
+fontconfig=2.13.0=h9420a91_0
+freeglut=3.0.0=hf484d3e_5
+freetype=2.9.1=h8a8886c_1
+gast=0.2.2=py36_0
+glib=2.56.2=hd408876_0
+graphite2=1.3.13=h23475e2_0
+grpcio=1.16.1=py36hf8bcb03_1
+gst-plugins-base=1.14.0=hbbd80ab_1
+gstreamer=1.14.0=hb453b48_1
+h5py=2.8.0=py36h989c5e5_3
+harfbuzz=1.8.8=hffaf4a1_0
+hdf5=1.10.2=hba1933b_1
+icu=58.2=h9c2bf20_1
+imageio=2.5.0=py36_0
+intel-openmp=2019.3=199
+ipython=7.5.0=py36h39e3cac_0
+ipython_genutils=0.2.0=py36_0
+jasper=2.0.14=h07fcdf6_1
+jedi=0.13.3=py36_0
+jpeg=9b=h024ee3a_2
+keras=2.2.4=0
+keras-applications=1.0.7=py_0
+keras-base=2.2.4=py36_0
+keras-preprocessing=1.0.9=py_0
+kiwisolver=1.1.0=py36he6710b0_0
+libedit=3.1.20181209=hc058e9b_0
+libffi=3.2.1=hd88cf55_4
+libgcc-ng=8.2.0=hdf63c60_1
+libgfortran-ng=7.3.0=hdf63c60_0
+libglu=9.0.0=hf484d3e_1
+libopencv=3.4.2=hb342d67_1
+libopus=1.3=h7b6447c_0
+libpng=1.6.37=hbc83047_0
+libprotobuf=3.7.1=hd408876_0
+libstdcxx-ng=8.2.0=hdf63c60_1
+libtiff=4.0.10=h2733197_2
+libuuid=1.0.3=h1bed415_2
+libvpx=1.7.0=h439df22_0
+libxcb=1.13=h1bed415_1
+libxml2=2.9.9=he19cac6_0
+markdown=3.1=py36_0
+matplotlib=3.0.3=py36h5429711_0
+mkl=2019.3=199
+mkl_fft=1.0.12=py36ha843d7b_0
+mkl_random=1.0.2=py36hd81dba3_0
+mock=2.0.0=py36_0
+mpi4py=2.0.0=py36_2
+mpich2=1.4.1p1=0
+ncurses=6.1=he6710b0_1
+numpy=1.16.3=py36h7e9f1db_0
+numpy-base=1.16.3=py36hde5b4d6_0
+olefile=0.46=py36_0
+opencv=3.4.2=py36h6fd60c2_1
+openssl=1.1.1b=h7b6447c_1
+pandas=0.24.2=py36he6710b0_0
+parso=0.4.0=py_0
+pbr=5.1.3=py_0
+pcre=8.43=he6710b0_0
+pexpect=4.7.0=py36_0
+pickleshare=0.7.5=py36_0
+pillow=6.0.0=py36h34e0f95_0
+pip=19.0.3=py36_0
+pixman=0.38.0=h7b6447c_0
+prompt_toolkit=2.0.9=py36_0
+protobuf=3.7.1=py36he6710b0_0
+psutil=5.6.2=py36h7b6447c_0
+ptyprocess=0.6.0=py36_0
+py-opencv=3.4.2=py36hb342d67_1
+pygments=2.4.0=py_0
+pyparsing=2.4.0=py_0
+pyqt=5.9.2=py36h05f1152_2
+python=3.6.8=h0371630_0
+python-dateutil=2.8.0=py36_0
+pytz=2019.1=py_0
+pyyaml=5.1=py36h7b6447c_0
+qt=5.9.7=h5867ecd_1
+readline=7.0=h7b6447c_5
+scipy=1.2.1=py36h7c811a0_0
+setuptools=40.8.0=py36_0
+sip=4.19.8=py36hf484d3e_0
+six=1.12.0=py36_0
+sqlite=3.27.2=h7b6447c_0
+tensorboard=1.13.1=py36hf484d3e_0
+tensorflow=1.13.1=gpu_py36h3991807_0
+tensorflow-base=1.13.1=gpu_py36h8d69cac_0
+tensorflow-estimator=1.13.0=py_0
+tensorflow-gpu=1.13.1=h0d30ee6_0
+termcolor=1.1.0=py36_1
+tk=8.6.8=hbc83047_0
+tornado=6.0.2=py36h7b6447c_0
+traitlets=4.3.2=py36_0
+wcwidth=0.1.7=py36_0
+werkzeug=0.15.2=py_0
+wheel=0.33.1=py36_0
+xz=5.2.4=h14c3975_4
+yaml=0.1.7=had09818_2
+zlib=1.2.11=h7b6447c_3
+zstd=1.3.7=h0b5b093_0

pip_requirements.txt

@@ -0,0 +1,3 @@
+gym==0.12.1
+atari-py==0.1.7
+git+https://github.com/openai/baselines.git@0182fe1877e95b2ef0a82747c20bed1523fb5a3f

policies/cnn_policy_param_matched.py

@@ -103,7 +103,7 @@ def apply_policy(self, ph_ob, reuse, scope, hidsize, memsize, extrahid, sy_nenvs
         yes_gpu = any(get_available_gpus())
         with tf.variable_scope(scope, reuse=reuse), tf.device('/gpu:0' if yes_gpu else '/cpu:0'):
             X = activ(conv(X, 'c1', nf=32, rf=8, stride=4, init_scale=np.sqrt(2), data_format=data_format))
-            #X = activ(conv(X, 'c2', nf=64, rf=4, stride=2, init_scale=np.sqrt(2), data_format=data_format))
+            X = activ(conv(X, 'c2', nf=64, rf=4, stride=2, init_scale=np.sqrt(2), data_format=data_format))
             #X = activ(conv(X, 'c3', nf=64, rf=4, stride=1, init_scale=np.sqrt(2), data_format=data_format))
 
             with tf.variable_scope("augmented1"):
@@ -150,7 +150,8 @@ def define_self_prediction_rew(self, convfeat, rep_size, enlargement):
 
                 xr = tf.nn.leaky_relu(conv(xr, 'c1r', nf=convfeat * 1, rf=8, stride=4, init_scale=np.sqrt(2)))
                 xr = tf.nn.leaky_relu(conv(xr, 'c2r', nf=convfeat * 2 * 1, rf=4, stride=2, init_scale=np.sqrt(2)))
-                xr = tf.nn.leaky_relu(conv(xr, 'c3r', nf=convfeat * 2 * 1, rf=3, stride=1, init_scale=np.sqrt(2)))
+                # rf=3 in case of 84x84 image else might need to be changed to 2 or 1
+                xr = tf.nn.leaky_relu(conv(xr, 'c3r', nf=convfeat * 2 * 1, rf=2, stride=1, init_scale=np.sqrt(2)))
                 rgbr = [to2d(xr)]
                 X_r = fc(rgbr[0], 'fc1r', nh=rep_size, init_scale=np.sqrt(2))
 
@@ -165,7 +166,8 @@ def define_self_prediction_rew(self, convfeat, rep_size, enlargement):
 
                 xrp = tf.nn.leaky_relu(conv(xrp, 'c1rp_pred', nf=convfeat, rf=8, stride=4, init_scale=np.sqrt(2)))
                 xrp = tf.nn.leaky_relu(conv(xrp, 'c2rp_pred', nf=convfeat * 2, rf=4, stride=2, init_scale=np.sqrt(2)))
-                xrp = tf.nn.leaky_relu(conv(xrp, 'c3rp_pred', nf=convfeat * 2, rf=3, stride=1, init_scale=np.sqrt(2)))
+                # rf=3 in case of 84x84 image else might need to be changed to 2 or 1
+                xrp = tf.nn.leaky_relu(conv(xrp, 'c3rp_pred', nf=convfeat * 2, rf=2, stride=1, init_scale=np.sqrt(2)))
                 rgbrp = to2d(xrp)
                 # X_r_hat = tf.nn.relu(fc(rgb[0], 'fc1r_hat1', nh=256 * enlargement, init_scale=np.sqrt(2)))
                 X_r_hat = tf.nn.relu(fc(rgbrp, 'fc1r_hat1_pred', nh=256 * enlargement, init_scale=np.sqrt(2)))

ppo_agent.py

@@ -207,6 +207,7 @@ def __init__(self, *, scope,
         if restore_model:
             tf_util.load_state(model_path)
         else:
+            #self.activate_graph_debugging()
             tf.get_default_session().run(tf.variables_initializer(allvars))
 
         #Syncs initialization across mpi workers.
@@ -354,12 +355,12 @@ def update(self):
                      'ret_int': rets_int,
                      'ret_ext': rets_ext,
                      }
+
         if self.I.venvs[0].record_obs:
             to_record['obs'] = self.I.buf_obs[None]
         self.recorder.record(bufs=to_record,
                              infos=self.I.buf_epinfos)
 
-
         #Create feeddict for optimization.
         envsperbatch = self.I.nenvs // self.nminibatches
         ph_buf = [
@@ -386,6 +387,16 @@ def update(self):
             logger.info(" "*6 + fmt_row(13, self.loss_names))
 
 
+        to_record_attention = None
+        attention_output = None
+        try:
+            #attention_output = tf.get_default_graph().get_tensor_by_name("ppo/pol/augmented2/attention_output_combined:0")
+            #attention_output = tf.get_default_graph().get_tensor_by_name("ppo/pol/augmented2/attention_output_combined/kernel:0")
+            attention_output = tf.get_default_graph().get_tensor_by_name("ppo/pol/augmented2/attention_output_combined/Conv2D:0")
+        except Exception as e:
+            logger.error("Exception in attention_output: {}".format(e))
+            attention_output = None
+
         epoch = 0
         start = 0
         #Optimizes on current data for several epochs.
@@ -402,7 +413,45 @@ def update(self):
 
             fd.update({self.stochpol.ph_mean:self.stochpol.ob_rms.mean, self.stochpol.ph_std:self.stochpol.ob_rms.var**0.5})
 
-            ret = tf.get_default_session().run(self._losses+[self._train], feed_dict=fd)[:-1]
+            if attention_output is not None:
+                _train_losses = [attention_output, self._train]
+            else:
+                _train_losses = [self._train]
+
+            ret = tf.get_default_session().run(self._losses + _train_losses, feed_dict=fd)[:-1]
+
+            if attention_output is not None:
+                attn_output = ret[-1]
+                ret = ret[:-1]
+                outshape = list(self.I.buf_obs[None][mbenvinds].shape[:2])+list(attn_output.shape[1:])
+                attn_output = np.reshape(attn_output, outshape)
+                attn_output = attn_output[:,:,:,:,:64]
+
+                # attn_output = attn_output[:,:,:,:,:1]
+                # for x in range(attn_output.shape[0]):
+                #     for y in range(attn_output.shape[1]):
+                #         attn_min = np.stack([attn_output[x,y,...,0].min()])
+                #         attn_max = np.stack([attn_output[x,y,...,0].max()])
+                #         attn_output[x,y,...] = (1 * ((attn_output[x,y,...] - attn_min)/(attn_max-attn_min)))
+
+                #         #attn_output[x,y,...] = (255 * ((attn_output[x,y,...] - attn_min)/(attn_max-attn_min)))
+
+                # #((oldval - Min) * (255/(Max-Min)))
+                # for x in range(attn_output.shape[0]):
+                #     for y in range(attn_output.shape[1]):
+                #         attn_min = np.stack([attn_output[x,y,...,0].min(),
+                #                              attn_output[x,y,...,1].min(),
+                #                              attn_output[x,y,...,2].min(),
+                #                              attn_output[x,y,...,3].min()])
+
+                #         attn_max = np.stack([attn_output[x,y,...,0].max(),
+                #                              attn_output[x,y,...,1].max(),
+                #                              attn_output[x,y,...,2].max(),
+                #                              attn_output[x,y,...,3].max()])
+
+                #         #attn_output[x,y,...] = (attn_output[x,y,...] - attn_min) * (255/(attn_max-attn_min))
+                #         attn_output[x,y,...] = (attn_output[x,y,...] - attn_min) / (attn_max-attn_min)
+
             if not self.testing:
                 lossdict = dict(zip([n for n in self.loss_names], ret), axis=0)
             else:
@@ -419,6 +468,20 @@ def update(self):
                 epoch += 1
                 start = 0
 
+                if attention_output is not None:
+                    if to_record_attention is None:
+                        to_record_attention = attn_output
+                    else:
+                        to_record_attention = np.concatenate([to_record_attention,
+                                                              attn_output])
+
+        if to_record_attention is not None:
+            to_record['obs'] = self.I.buf_obs[None]
+            to_record['attention'] = to_record_attention
+        self.recorder.record(bufs=to_record,
+                             infos=self.I.buf_epinfos)
+        to_record_attention = None
+
         if self.is_train_leader:
             self.I.stats["n_updates"] += 1
             info.update([('opt_'+n, lossdict[n]) for n in self.loss_names])
@@ -567,6 +630,18 @@ def step(self):
 
         return {'update' : update_info}
 
+    def activate_graph_debugging(self):
+        """
+        Necessary in order to debug tensorflow using the CLI Tensorflow 1.0 Debugger
+        (before the 2.0 eager execution)
+        """
+        sess = tf.get_default_session()
+        from tensorflow.python import debug as tf_debug
+
+        sess_debug = tf_debug.LocalCLIDebugWrapperSession(sess)
+        sess._default_session = sess_debug.as_default()
+        sess._default_session.__enter__()
+
 
 class RewardForwardFilter(object):
     def __init__(self, gamma):