Removed tensorflow specific code in presets (#59)

* Add generic layer specification for using in presets

* Modify presets to use the generic scheme

rl_coach/architectures/tensorflow_components/layers.py

@@ -1,9 +1,11 @@
 import math
-from typing import List, Union
+from types import FunctionType
+from typing import Any
 
 import tensorflow as tf
 
-from rl_coach.utils import force_list
+from rl_coach.architectures import layers
+from rl_coach.architectures.tensorflow_components import utils
 
 
 def batchnorm_activation_dropout(input_layer, batchnorm, activation_function, dropout, dropout_rate, is_training, name):
@@ -17,6 +19,8 @@ def batchnorm_activation_dropout(input_layer, batchnorm, activation_function, dr
 
     # activation
     if activation_function:
+        if isinstance(activation_function, str):
+            activation_function = utils.get_activation_function(activation_function)
         layers.append(
             activation_function(layers[-1], name="{}_activation".format(name))
         )
@@ -33,11 +37,35 @@ def batchnorm_activation_dropout(input_layer, batchnorm, activation_function, dr
     return layers
 
 
-class Conv2d(object):
+# define global dictionary for storing layer type to layer implementation mapping
+tf_layer_dict = dict()
+
+
+def reg_to_tf(layer_type) -> FunctionType:
+    """ function decorator that registers layer implementation
+    :return: decorated function
+    """
+    def reg_impl_decorator(func):
+        assert layer_type not in tf_layer_dict
+        tf_layer_dict[layer_type] = func
+        return func
+    return reg_impl_decorator
+
+
+def convert_layer(layer):
+    """
+    If layer is callable, return layer, otherwise convert to TF type
+    :param layer: layer to be converted
+    :return: converted layer if not callable, otherwise layer itself
+    """
+    if callable(layer):
+        return layer
+    return tf_layer_dict[type(layer)](layer)
+
+
+class Conv2d(layers.Conv2d):
     def __init__(self, num_filters: int, kernel_size: int, strides: int):
-        self.num_filters = num_filters
-        self.kernel_size = kernel_size
-        self.strides = strides
+        super(Conv2d, self).__init__(num_filters=num_filters, kernel_size=kernel_size, strides=strides)
 
     def __call__(self, input_layer, name: str=None, is_training=None):
         """
@@ -49,16 +77,19 @@ class Conv2d(object):
         return tf.layers.conv2d(input_layer, filters=self.num_filters, kernel_size=self.kernel_size,
                                 strides=self.strides, data_format='channels_last', name=name)
 
-    def __str__(self):
-        return "Convolution (num filters = {}, kernel size = {}, stride = {})"\
-            .format(self.num_filters, self.kernel_size, self.strides)
+    @staticmethod
+    @reg_to_tf(layers.Conv2d)
+    def to_tf(base: layers.Conv2d):
+        return Conv2d(
+            num_filters=base.num_filters,
+            kernel_size=base.kernel_size,
+            strides=base.strides)
 
 
-class BatchnormActivationDropout(object):
+class BatchnormActivationDropout(layers.BatchnormActivationDropout):
     def __init__(self, batchnorm: bool=False, activation_function=None, dropout_rate: float=0):
-        self.batchnorm = batchnorm
-        self.activation_function = activation_function
-        self.dropout_rate = dropout_rate
+        super(BatchnormActivationDropout, self).__init__(
+            batchnorm=batchnorm, activation_function=activation_function, dropout_rate=dropout_rate)
 
     def __call__(self, input_layer, name: str=None, is_training=None):
         """
@@ -72,20 +103,18 @@ class BatchnormActivationDropout(object):
                                             dropout=self.dropout_rate > 0, dropout_rate=self.dropout_rate,
                                             is_training=is_training, name=name)
 
-    def __str__(self):
-        result = []
-        if self.batchnorm:
-            result += ["Batch Normalization"]
-        if self.activation_function:
-            result += ["Activation (type = {})".format(self.activation_function.__name__)]
-        if self.dropout_rate > 0:
-            result += ["Dropout (rate = {})".format(self.dropout_rate)]
-        return "\n".join(result)
+    @staticmethod
+    @reg_to_tf(layers.BatchnormActivationDropout)
+    def to_tf(base: layers.BatchnormActivationDropout):
+        return BatchnormActivationDropout(
+            batchnorm=base.batchnorm,
+            activation_function=base.activation_function,
+            dropout_rate=base.dropout_rate)
 
 
-class Dense(object):
+class Dense(layers.Dense):
     def __init__(self, units: int):
-        self.units = units
+        super(Dense, self).__init__(units=units)
 
     def __call__(self, input_layer, name: str=None, kernel_initializer=None, activation=None, is_training=None):
         """
@@ -97,11 +126,13 @@ class Dense(object):
         return tf.layers.dense(input_layer, self.units, name=name, kernel_initializer=kernel_initializer,
                                activation=activation)
 
-    def __str__(self):
-        return "Dense (num outputs = {})".format(self.units)
+    @staticmethod
+    @reg_to_tf(layers.Dense)
+    def to_tf(base: layers.Dense):
+        return Dense(units=base.units)
 
 
-class NoisyNetDense(object):
+class NoisyNetDense(layers.NoisyNetDense):
     """
     A factorized Noisy Net layer
 
@@ -109,8 +140,7 @@ class NoisyNetDense(object):
     """
 
     def __init__(self, units: int):
-        self.units = units
-        self.sigma0 = 0.5
+        super(NoisyNetDense, self).__init__(units=units)
 
     def __call__(self, input_layer, name: str, kernel_initializer=None, activation=None, is_training=None):
         """
@@ -125,6 +155,16 @@ class NoisyNetDense(object):
         #      forward (either act or train, both for online and target networks).
         #      A3C, on the other hand, should sample noise only when policy changes (i.e. after every t_max steps)
 
+        def _f(values):
+            return tf.sqrt(tf.abs(values)) * tf.sign(values)
+
+        def _factorized_noise(inputs, outputs):
+            # TODO: use factorized noise only for compute intensive algos (e.g. DQN).
+            #      lighter algos (e.g. DQN) should not use it
+            noise1 = _f(tf.random_normal((inputs, 1)))
+            noise2 = _f(tf.random_normal((1, outputs)))
+            return tf.matmul(noise1, noise2)
+
         num_inputs = input_layer.get_shape()[-1].value
         num_outputs = self.units
 
@@ -145,23 +185,14 @@ class NoisyNetDense(object):
                                             initializer=kernel_stddev_initializer)
             bias_stddev = tf.get_variable('bias_stddev', shape=(num_outputs,),
                                           initializer=kernel_stddev_initializer)
-            bias_noise = self.f(tf.random_normal((num_outputs,)))
-            weight_noise = self.factorized_noise(num_inputs, num_outputs)
+            bias_noise = _f(tf.random_normal((num_outputs,)))
+            weight_noise = _factorized_noise(num_inputs, num_outputs)
 
         bias = bias_mean + bias_stddev * bias_noise
         weight = weight_mean + weight_stddev * weight_noise
         return activation(tf.matmul(input_layer, weight) + bias)
 
-    def factorized_noise(self, inputs, outputs):
-        # TODO: use factorized noise only for compute intensive algos (e.g. DQN).
-        #      lighter algos (e.g. DQN) should not use it
-        noise1 = self.f(tf.random_normal((inputs, 1)))
-        noise2 = self.f(tf.random_normal((1, outputs)))
-        return tf.matmul(noise1, noise2)
-
     @staticmethod
-    def f(values):
-        return tf.sqrt(tf.abs(values)) * tf.sign(values)
-
-    def __str__(self):
-        return "Noisy Dense (num outputs = {})".format(self.units)
+    @reg_to_tf(layers.NoisyNetDense)
+    def to_tf(base: layers.NoisyNetDense):
+        return NoisyNetDense(units=base.units)