parameter noise exploration - using Noisy Nets

rl_coach/architectures/tensorflow_components/heads/naf_head.py

@@ -16,6 +16,7 @@
 
 import tensorflow as tf
 
+from rl_coach.architectures.tensorflow_components.architecture import Dense
 from rl_coach.architectures.tensorflow_components.heads.head import Head, HeadParameters
 from rl_coach.base_parameters import AgentParameters
 from rl_coach.core_types import QActionStateValue
@@ -24,14 +25,17 @@ from rl_coach.spaces import SpacesDefinition
 
 
 class NAFHeadParameters(HeadParameters):
-    def __init__(self, activation_function: str ='tanh', name: str='naf_head_params'):
-        super().__init__(parameterized_class=NAFHead, activation_function=activation_function, name=name)
+    def __init__(self, activation_function: str ='tanh', name: str='naf_head_params', dense_layer=Dense):
+        super().__init__(parameterized_class=NAFHead, activation_function=activation_function, name=name,
+                         dense_layer=dense_layer)
 
 
 class NAFHead(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'):
-        super().__init__(agent_parameters, spaces, network_name, head_idx, loss_weight, is_local, activation_function)
+                 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)
         if not isinstance(self.spaces.action, BoxActionSpace):
             raise ValueError("NAF works only for continuous action spaces (BoxActionSpace)")
 
@@ -50,15 +54,15 @@ class NAFHead(Head):
         self.input = self.action
 
         # V Head
-        self.V = tf.layers.dense(input_layer, 1, name='V')
+        self.V = self.dense_layer(1)(input_layer, name='V')
 
         # mu Head
-        mu_unscaled = tf.layers.dense(input_layer, self.num_actions, activation=self.activation_function, name='mu_unscaled')
+        mu_unscaled = self.dense_layer(self.num_actions)(input_layer, activation=self.activation_function, name='mu_unscaled')
         self.mu = tf.multiply(mu_unscaled, self.output_scale, name='mu')
 
         # A Head
         # l_vector is a vector that includes a lower-triangular matrix values
-        self.l_vector = tf.layers.dense(input_layer, (self.num_actions * (self.num_actions + 1)) / 2, name='l_vector')
+        self.l_vector = self.dense_layer((self.num_actions * (self.num_actions + 1)) / 2)(input_layer, name='l_vector')
 
         # Convert l to a lower triangular matrix and exponentiate its diagonal