adding dueling support for rainbow dqn (now only missing n-step)

rl_coach/architectures/tensorflow_components/heads/categorical_q_head.py

@@ -42,9 +42,6 @@ class CategoricalQHead(Head):
         self.return_type = QActionStateValue
 
     def _build_module(self, input_layer):
-        self.actions = tf.placeholder(tf.int32, [None], name="actions")
-        self.input = [self.actions]
-
         values_distribution = self.dense_layer(self.num_actions * self.num_atoms)(input_layer, name='output')
         values_distribution = tf.reshape(values_distribution, (tf.shape(values_distribution)[0], self.num_actions,
                                                                self.num_atoms))

rl_coach/architectures/tensorflow_components/heads/rainbow_q_head.py

@@ -17,11 +17,10 @@
 import tensorflow as tf
 
 from rl_coach.architectures.tensorflow_components.architecture import Dense
+from rl_coach.architectures.tensorflow_components.heads.head import HeadParameters, Head
 from rl_coach.base_parameters import AgentParameters
-from rl_coach.spaces import SpacesDefinition
-
-from rl_coach.architectures.tensorflow_components.heads.head import Head, HeadParameters
 from rl_coach.core_types import QActionStateValue
+from rl_coach.spaces import SpacesDefinition
 
 
 class RainbowQHeadParameters(HeadParameters):
@@ -30,15 +29,41 @@ class RainbowQHeadParameters(HeadParameters):
                          dense_layer=dense_layer)
 
 
-class RainbowQHead():
+class RainbowQHead(Head):
     def __init__(self, agent_parameters: AgentParameters, spaces: SpacesDefinition, network_name: str,
-                 head_idx: int = 0, loss_weight: float = 1., is_local: bool = True, activation_function: str ='relu',
+                 head_idx: int = 0, loss_weight: float = 1., is_local: bool = True, activation_function: str='relu',
                  dense_layer=Dense):
         super().__init__(agent_parameters, spaces, network_name, head_idx, loss_weight, is_local, activation_function,
                          dense_layer=dense_layer)
-        self.name = 'rainbow_dqn_head'
         self.num_actions = len(self.spaces.action.actions)
+        self.num_atoms = agent_parameters.algorithm.atoms
         self.return_type = QActionStateValue
+        self.name = 'rainbow_q_values_head'
 
     def _build_module(self, input_layer):
-        pass
+        # state value tower - V
+        with tf.variable_scope("state_value"):
+            state_value = self.dense_layer(self.num_atoms)(input_layer, name='fc1')
+            state_value = tf.expand_dims(state_value, axis=1)
+
+        # action advantage tower - A
+        with tf.variable_scope("action_advantage"):
+            action_advantage = self.dense_layer(self.num_actions * self.num_atoms)(input_layer, name='fc1')
+            action_advantage = tf.reshape(action_advantage, (tf.shape(input_layer)[0], self.num_actions,
+                                                             self.num_atoms))
+            action_mean = tf.reduce_mean(action_advantage, axis=1, keepdims=True)
+            action_advantage = action_advantage - action_mean
+
+        # merge to state-action value function Q
+        values_distribution = tf.add(state_value, action_advantage, name='output')
+
+        # softmax on atoms dimension
+        self.output = tf.nn.softmax(values_distribution)
+
+        # calculate cross entropy loss
+        self.distributions = tf.placeholder(tf.float32, shape=(None, self.num_actions, self.num_atoms),
+                                            name="distributions")
+        self.target = self.distributions
+        self.loss = tf.nn.softmax_cross_entropy_with_logits(labels=self.target, logits=values_distribution)
+        tf.losses.add_loss(self.loss)
+