N-step returns for rainbow (#67)

* n_step returns for rainbow
* Rename CartPole_PPO -> CartPole_ClippedPPO

rl_coach/agents/categorical_dqn_agent.py

@@ -17,14 +17,11 @@
 from typing import Union
 
 import numpy as np
-
 from rl_coach.agents.dqn_agent import DQNNetworkParameters, DQNAlgorithmParameters, DQNAgentParameters
 from rl_coach.agents.value_optimization_agent import ValueOptimizationAgent
 from rl_coach.architectures.head_parameters import CategoricalQHeadParameters
-from rl_coach.base_parameters import AgentParameters
 from rl_coach.core_types import StateType
 from rl_coach.exploration_policies.e_greedy import EGreedyParameters
-from rl_coach.memories.non_episodic.experience_replay import ExperienceReplayParameters
 from rl_coach.memories.non_episodic.prioritized_experience_replay import PrioritizedExperienceReplay
 from rl_coach.schedules import LinearSchedule
 
@@ -85,28 +82,47 @@ class CategoricalDQNAgent(ValueOptimizationAgent):
 
         # for the action we actually took, the error is calculated by the atoms distribution
         # for all other actions, the error is 0
-        distributed_q_st_plus_1, TD_targets = self.networks['main'].parallel_prediction([
+        distributional_q_st_plus_1, TD_targets = self.networks['main'].parallel_prediction([
             (self.networks['main'].target_network, batch.next_states(network_keys)),
             (self.networks['main'].online_network, batch.states(network_keys))
         ])
 
-        # only update the action that we have actually done in this transition
-        target_actions = np.argmax(self.distribution_prediction_to_q_values(distributed_q_st_plus_1), axis=1)
+        # select the optimal actions for the next state
+        target_actions = np.argmax(self.distribution_prediction_to_q_values(distributional_q_st_plus_1), axis=1)
         m = np.zeros((self.ap.network_wrappers['main'].batch_size, self.z_values.size))
 
         batches = np.arange(self.ap.network_wrappers['main'].batch_size)
+
+        # an alternative to the for loop. 3.7x perf improvement vs. the same code done with for looping.
+        # only 10% speedup overall - leaving commented out as the code is not as clear.
+
+        # tzj_ = np.fmax(np.fmin(batch.rewards() + (1.0 - batch.game_overs()) * self.ap.algorithm.discount *
+        #                        np.transpose(np.repeat(self.z_values[np.newaxis, :], batch.size, axis=0), (1, 0)),
+        #                     self.z_values[-1]),
+        #             self.z_values[0])
+        #
+        # bj_ = (tzj_ - self.z_values[0]) / (self.z_values[1] - self.z_values[0])
+        # u_ = (np.ceil(bj_)).astype(int)
+        # l_ = (np.floor(bj_)).astype(int)
+        # m_ = np.zeros((self.ap.network_wrappers['main'].batch_size, self.z_values.size))
+        # np.add.at(m_, [batches, l_],
+        #           np.transpose(distributional_q_st_plus_1[batches, target_actions], (1, 0)) * (u_ - bj_))
+        # np.add.at(m_, [batches, u_],
+        #           np.transpose(distributional_q_st_plus_1[batches, target_actions], (1, 0)) * (bj_ - l_))
+
         for j in range(self.z_values.size):
             tzj = np.fmax(np.fmin(batch.rewards() +
                                   (1.0 - batch.game_overs()) * self.ap.algorithm.discount * self.z_values[j],
-                                  self.z_values[self.z_values.size - 1]),
+                                  self.z_values[-1]),
                           self.z_values[0])
             bj = (tzj - self.z_values[0])/(self.z_values[1] - self.z_values[0])
             u = (np.ceil(bj)).astype(int)
             l = (np.floor(bj)).astype(int)
-            m[batches, l] = m[batches, l] + (distributed_q_st_plus_1[batches, target_actions, j] * (u - bj))
-            m[batches, u] = m[batches, u] + (distributed_q_st_plus_1[batches, target_actions, j] * (bj - l))
+            m[batches, l] += (distributional_q_st_plus_1[batches, target_actions, j] * (u - bj))
+            m[batches, u] += (distributional_q_st_plus_1[batches, target_actions, j] * (bj - l))
 
         # total_loss = cross entropy between actual result above and predicted result for the given action
+        # only update the action that we have actually done in this transition
         TD_targets[batches, batch.actions()] = m
 
         # update errors in prioritized replay buffer
@@ -120,7 +136,7 @@ class CategoricalDQNAgent(ValueOptimizationAgent):
         # TODO: fix this spaghetti code
         if isinstance(self.memory, PrioritizedExperienceReplay):
             errors = losses[0][np.arange(batch.size), batch.actions()]
-            self.memory.update_priorities(batch.info('idx'), errors)
+            self.call_memory('update_priorities', (batch.info('idx'), errors))
 
         return total_loss, losses, unclipped_grads