Multiple improvements and bug fixes (#66)

* Multiple improvements and bug fixes:

    * Using lazy stacking to save on memory when using a replay buffer
    * Remove step counting for evaluation episodes
    * Reset game between heatup and training
    * Major bug fixes in NEC (is reproducing the paper results for pong now)
    * Image input rescaling to 0-1 is now optional
    * Change the terminal title to be the experiment name
    * Observation cropping for atari is now optional
    * Added random number of noop actions for gym to match the dqn paper
    * Fixed a bug where the evaluation episodes won't start with the max possible ale lives
    * Added a script for plotting the results of an experiment over all the atari games

memories/differentiable_neural_dictionary.py

@@ -83,6 +83,11 @@ class AnnoyDictionary(object):
 
     # Returns the stored embeddings and values of the closest embeddings
     def query(self, keys, k):
+        if not self.has_enough_entries(k):
+            # this will only happen when the DND is not yet populated with enough entries, which is only during heatup
+            # these values won't be used and therefore they are meaningless
+            return [0.0], [0.0], [0]
+
         _, indices = self._get_k_nearest_neighbors_indices(keys, k)
 
         embeddings = []
@@ -94,7 +99,7 @@ class AnnoyDictionary(object):
 
         self.current_timestamp += 1
 
-        return embeddings, values
+        return embeddings, values, indices
 
     def has_enough_entries(self, k):
         return self.curr_size > k and (self.built_capacity > k)
@@ -133,9 +138,11 @@ class AnnoyDictionary(object):
 
 
 class QDND:
-    def __init__(self, dict_size, key_width, num_actions, new_value_shift_coefficient=0.1, key_error_threshold=0.01):
+    def __init__(self, dict_size, key_width, num_actions, new_value_shift_coefficient=0.1, key_error_threshold=0.01,
+                 learning_rate=0.01):
         self.num_actions = num_actions
         self.dicts = []
+        self.learning_rate = learning_rate
 
         # create a dict for each action
         for a in range(num_actions):
@@ -155,16 +162,18 @@ class QDND:
             self.dicts[a].add(curr_action_embeddings, curr_action_values)
         return True
 
-    def query(self, embeddings, actions, k):
+    def query(self, embeddings, action, k):
         # query for nearest neighbors to the given embeddings
         dnd_embeddings = []
         dnd_values = []
-        for i, action in enumerate(actions):
-            embedding, value = self.dicts[action].query([embeddings[i]], k)
+        dnd_indices = []
+        for i in range(len(embeddings)):
+            embedding, value, indices = self.dicts[action].query([embeddings[i]], k)
             dnd_embeddings.append(embedding[0])
             dnd_values.append(value[0])
+            dnd_indices.append(indices[0])
 
-        return dnd_embeddings, dnd_values
+        return dnd_embeddings, dnd_values, dnd_indices
 
     def has_enough_entries(self, k):
         # check if each of the action dictionaries has at least k entries
@@ -193,4 +202,5 @@ def load_dnd(model_dir):
                 DND.dicts[a].index.add_item(idx, key)
 
             DND.dicts[a].index.build(50)
+
     return DND

memories/episodic_experience_replay.py

@@ -16,6 +16,7 @@
 
 from memories.memory import *
 import threading
+from typing import Union
 
 
 class EpisodicExperienceReplay(Memory):
@@ -103,7 +104,8 @@ class EpisodicExperienceReplay(Memory):
         if transition.game_over:
             self._num_transitions_in_complete_episodes += last_episode.length()
             self._length += 1
-            self.buffer[-1].update_returns(self.discount, is_bootstrapped=self.return_is_bootstrapped,
+            self.buffer[-1].update_returns(self.discount,
+                                           is_bootstrapped=self.tp.agent.bootstrap_total_return_from_old_policy,
                                            n_step_return=self.tp.agent.n_step)
             self.buffer[-1].update_measurements_targets(self.tp.agent.num_predicted_steps_ahead)
             # self.buffer[-1].update_actions_probabilities() # used for off-policy policy optimization
@@ -146,6 +148,17 @@ class EpisodicExperienceReplay(Memory):
     def get(self, index):
         return self.get_episode(index)
 
+    def get_last_complete_episode(self) -> Union[None, Episode]:
+        """
+        Returns the last complete episode in the memory or None if there are no complete episodes
+        :return: None or the last complete episode
+        """
+        last_complete_episode_index = self.num_complete_episodes()-1
+        if last_complete_episode_index >= 0:
+            return self.get(last_complete_episode_index)
+        else:
+            return None
+
     def update_last_transition_info(self, info):
         episode = self.buffer[-1]
         if episode.length() == 0:

memories/memory.py

@@ -80,9 +80,12 @@ class Episode(object):
             total_return += current_discount * np.pad(rewards[i:], (0, i), 'constant', constant_values=0)
             current_discount *= discount
 
+        # calculate the bootstrapped returns
+        bootstraps = np.array([np.squeeze(t.info['max_action_value']) for t in self.transitions[n_step_return:]])
+        bootstrapped_return = total_return + current_discount * np.pad(bootstraps, (0, n_step_return), 'constant',
+                                                                       constant_values=0)
         if is_bootstrapped:
-            bootstraps = np.array([np.squeeze(t.info['action_value']) for t in self.transitions[n_step_return:]])
-            total_return += current_discount * np.pad(bootstraps, (0, n_step_return), 'constant', constant_values=0)
+            total_return = bootstrapped_return
 
         for transition_idx in range(self.length()):
             self.transitions[transition_idx].total_return = total_return[transition_idx]
@@ -114,7 +117,13 @@ class Episode(object):
         return self.returns_table
 
     def get_returns(self):
-        return [t.total_return for t in self.transitions]
+        return self.get_transitions_attribute('total_return')
+
+    def get_transitions_attribute(self, attribute_name):
+        if hasattr(self.transitions[0], attribute_name):
+            return [t.__dict__[attribute_name] for t in self.transitions]
+        else:
+            raise ValueError("The transitions have no such attribute name")
 
     def to_batch(self):
         batch = []
@@ -141,14 +150,12 @@ class Transition(object):
         :param game_over: A boolean which should be True if the episode terminated after
                           the execution of the action.
         """
-        self.state = copy.deepcopy(state)
-        self.state['observation'] = np.array(self.state['observation'], copy=False)
+        self.state = state
         self.action = action
         self.reward = reward
         self.total_return = None
         if not next_state:
             next_state = state
-        self.next_state = copy.deepcopy(next_state)
-        self.next_state['observation'] = np.array(self.next_state['observation'], copy=False)
+        self.next_state = next_state
         self.game_over = game_over
         self.info = {}