network_imporvements branch merge

rl_coach/architectures/tensorflow_components/heads/cil_head.py

@@ -16,27 +16,34 @@
 
 import tensorflow as tf
 
-from rl_coach.architectures.tensorflow_components.architecture import Dense
+from rl_coach.architectures.tensorflow_components.layers import Dense, batchnorm_activation_dropout
 
 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
 from rl_coach.spaces import SpacesDefinition, BoxActionSpace, DiscreteActionSpace
+from rl_coach.utils import force_list
 
 
 class RegressionHeadParameters(HeadParameters):
-    def __init__(self, activation_function: str ='relu', name: str='q_head_params', dense_layer=Dense):
+    def __init__(self, activation_function: str ='relu', name: str='q_head_params',
+                 num_output_head_copies: int = 1, rescale_gradient_from_head_by_factor: float = 1.0,
+                 loss_weight: float = 1.0, dense_layer=Dense, scheme=[Dense(256), Dense(256)]):
         super().__init__(parameterized_class=RegressionHead, activation_function=activation_function, name=name,
-                         dense_layer=dense_layer)
+                         dense_layer=dense_layer, num_output_head_copies=num_output_head_copies,
+                         rescale_gradient_from_head_by_factor=rescale_gradient_from_head_by_factor,
+                         loss_weight=loss_weight)
 
 
 class RegressionHead(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',
-                 dense_layer=Dense):
+                 dense_layer=Dense, scheme=[Dense(256), Dense(256)]):
         super().__init__(agent_parameters, spaces, network_name, head_idx, loss_weight, is_local, activation_function,
                          dense_layer=dense_layer)
         self.name = 'regression_head'
+        self.scheme = scheme
+        self.layers = []
         if isinstance(self.spaces.action, BoxActionSpace):
             self.num_actions = self.spaces.action.shape[0]
         elif isinstance(self.spaces.action, DiscreteActionSpace):
@@ -48,9 +55,18 @@ class RegressionHead(Head):
             self.loss_type = tf.losses.mean_squared_error
 
     def _build_module(self, input_layer):
-        self.fc1 = self.dense_layer(256)(input_layer)
-        self.fc2 = self.dense_layer(256)(self.fc1)
-        self.output = self.dense_layer(self.num_actions)(self.fc2, name='output')
+        self.layers.append(input_layer)
+        for idx, layer_params in enumerate(self.scheme):
+            self.layers.extend(force_list(
+                layer_params(input_layer=self.layers[-1], name='{}_{}'.format(layer_params.__class__.__name__, idx))
+            ))
 
+        self.layers.append(self.dense_layer(self.num_actions)(self.layers[-1], name='output'))
+        self.output = self.layers[-1]
 
+    def __str__(self):
+        result = []
+        for layer in self.layers:
+            result.append(str(layer))
+        return '\n'.join(result)