Move embedder, middleware, and head parameters to framework agnostic modules. (#45)

Part of #28

rl_coach/architectures/tensorflow_components/heads/__init__.py

@@ -0,0 +1,29 @@
+from .categorical_q_head import CategoricalQHead
+from .ddpg_actor_head import DDPGActor
+from .dnd_q_head import DNDQHead
+from .dueling_q_head import DuelingQHead
+from .measurements_prediction_head import MeasurementsPredictionHead
+from .naf_head import NAFHead
+from .policy_head import PolicyHead
+from .ppo_head import PPOHead
+from .ppo_v_head import PPOVHead
+from .q_head import QHead
+from .quantile_regression_q_head import QuantileRegressionQHead
+from .rainbow_q_head import RainbowQHead
+from .v_head import VHead
+
+__all__ = [
+    'CategoricalQHead',
+    'DDPGActor',
+    'DNDQHead',
+    'DuelingQHead',
+    'MeasurementsPredictionHead',
+    'NAFHead',
+    'PolicyHead',
+    'PPOHead',
+    'PPOVHead',
+    'QHead',
+    'QuantileRegressionQHead',
+    'RainbowQHead',
+    'VHead'
+]

rl_coach/architectures/tensorflow_components/heads/categorical_q_head.py

@@ -18,22 +18,12 @@ import tensorflow as tf
 
 from rl_coach.architectures.tensorflow_components.layers import Dense
 
-from rl_coach.architectures.tensorflow_components.heads.head import Head, HeadParameters
+from rl_coach.architectures.tensorflow_components.heads.head import Head
 from rl_coach.base_parameters import AgentParameters
 from rl_coach.core_types import QActionStateValue
 from rl_coach.spaces import SpacesDefinition
 
 
-class CategoricalQHeadParameters(HeadParameters):
-    def __init__(self, activation_function: str ='relu', name: str='categorical_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):
-        super().__init__(parameterized_class=CategoricalQHead, activation_function=activation_function, name=name,
-                         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 CategoricalQHead(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',

rl_coach/architectures/tensorflow_components/heads/cil_head.py

@@ -16,25 +16,14 @@
 
 import tensorflow as tf
 
-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.architectures.tensorflow_components.layers import Dense
+from rl_coach.architectures.tensorflow_components.heads.head import Head
 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',
-                 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, 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',

rl_coach/architectures/tensorflow_components/heads/ddpg_actor_head.py

@@ -17,23 +17,12 @@
 import tensorflow as tf
 
 from rl_coach.architectures.tensorflow_components.layers import batchnorm_activation_dropout, Dense
-from rl_coach.architectures.tensorflow_components.heads.head import Head, HeadParameters
+from rl_coach.architectures.tensorflow_components.heads.head import Head
 from rl_coach.base_parameters import AgentParameters
 from rl_coach.core_types import ActionProbabilities
 from rl_coach.spaces import SpacesDefinition
 
 
-class DDPGActorHeadParameters(HeadParameters):
-    def __init__(self, activation_function: str ='tanh', name: str='policy_head_params', batchnorm: bool=True,
-                 num_output_head_copies: int = 1, rescale_gradient_from_head_by_factor: float = 1.0,
-                 loss_weight: float = 1.0, dense_layer=Dense):
-        super().__init__(parameterized_class=DDPGActor, activation_function=activation_function, name=name,
-                         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)
-        self.batchnorm = batchnorm
-
-
 class DDPGActor(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='tanh',

rl_coach/architectures/tensorflow_components/heads/dnd_q_head.py

@@ -16,23 +16,12 @@
 import tensorflow as tf
 
 from rl_coach.architectures.tensorflow_components.layers import Dense
-from rl_coach.architectures.tensorflow_components.heads.head import HeadParameters
 from rl_coach.architectures.tensorflow_components.heads.q_head import QHead
 from rl_coach.base_parameters import AgentParameters
 from rl_coach.memories.non_episodic import differentiable_neural_dictionary
 from rl_coach.spaces import SpacesDefinition
 
 
-class DNDQHeadParameters(HeadParameters):
-    def __init__(self, activation_function: str ='relu', name: str='dnd_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):
-        super().__init__(parameterized_class=DNDQHead, activation_function=activation_function, name=name,
-                         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 DNDQHead(QHead):
     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',

rl_coach/architectures/tensorflow_components/heads/dueling_q_head.py

@@ -17,21 +17,11 @@
 import tensorflow as tf
 
 from rl_coach.architectures.tensorflow_components.layers import Dense
-from rl_coach.architectures.tensorflow_components.heads.head import HeadParameters
 from rl_coach.architectures.tensorflow_components.heads.q_head import QHead
 from rl_coach.base_parameters import AgentParameters
 from rl_coach.spaces import SpacesDefinition
 
 
-class DuelingQHeadParameters(HeadParameters):
-    def __init__(self, activation_function: str ='relu', name: str='dueling_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):
-        super().__init__(parameterized_class=DuelingQHead, activation_function=activation_function, name=name,
-                         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 DuelingQHead(QHead):
     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',

rl_coach/architectures/tensorflow_components/heads/head.py

@@ -33,19 +33,6 @@ def normalized_columns_initializer(std=1.0):
     return _initializer
 
 
-class HeadParameters(NetworkComponentParameters):
-    def __init__(self, parameterized_class: Type['Head'], activation_function: str = 'relu', name: str= 'head',
-                 num_output_head_copies: int=1, rescale_gradient_from_head_by_factor: float=1.0,
-                 loss_weight: float=1.0, dense_layer=Dense):
-        super().__init__(dense_layer=dense_layer)
-        self.activation_function = activation_function
-        self.name = name
-        self.num_output_head_copies = num_output_head_copies
-        self.rescale_gradient_from_head_by_factor = rescale_gradient_from_head_by_factor
-        self.loss_weight = loss_weight
-        self.parameterized_class_name = parameterized_class.__name__
-
-
 class Head(object):
     """
     A head is the final part of the network. It takes the embedding from the middleware embedder and passes it through
@@ -74,6 +61,8 @@ class Head(object):
         self.return_type = None
         self.activation_function = activation_function
         self.dense_layer = dense_layer
+        if self.dense_layer is None:
+            self.dense_layer = Dense
 
     def __call__(self, input_layer):
         """

rl_coach/architectures/tensorflow_components/heads/measurements_prediction_head.py

@@ -17,23 +17,12 @@
 import tensorflow as tf
 
 from rl_coach.architectures.tensorflow_components.layers import Dense
-
-from rl_coach.architectures.tensorflow_components.heads.head import Head, HeadParameters
+from rl_coach.architectures.tensorflow_components.heads.head import Head
 from rl_coach.base_parameters import AgentParameters
 from rl_coach.core_types import Measurements
 from rl_coach.spaces import SpacesDefinition
 
 
-class MeasurementsPredictionHeadParameters(HeadParameters):
-    def __init__(self, activation_function: str ='relu', name: str='measurements_prediction_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):
-        super().__init__(parameterized_class=MeasurementsPredictionHead, activation_function=activation_function, name=name,
-                         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 MeasurementsPredictionHead(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',

rl_coach/architectures/tensorflow_components/heads/naf_head.py

@@ -17,23 +17,13 @@
 import tensorflow as tf
 
 from rl_coach.architectures.tensorflow_components.layers import Dense
-from rl_coach.architectures.tensorflow_components.heads.head import Head, HeadParameters
+from rl_coach.architectures.tensorflow_components.heads.head import Head
 from rl_coach.base_parameters import AgentParameters
 from rl_coach.core_types import QActionStateValue
 from rl_coach.spaces import BoxActionSpace
 from rl_coach.spaces import SpacesDefinition
 
 
-class NAFHeadParameters(HeadParameters):
-    def __init__(self, activation_function: str ='tanh', name: str='naf_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):
-        super().__init__(parameterized_class=NAFHead, activation_function=activation_function, name=name,
-                         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 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',

rl_coach/architectures/tensorflow_components/heads/policy_head.py

@@ -18,7 +18,7 @@ import numpy as np
 import tensorflow as tf
 
 from rl_coach.architectures.tensorflow_components.layers import Dense
-from rl_coach.architectures.tensorflow_components.heads.head import Head, normalized_columns_initializer, HeadParameters
+from rl_coach.architectures.tensorflow_components.heads.head import Head, normalized_columns_initializer
 from rl_coach.base_parameters import AgentParameters
 from rl_coach.core_types import ActionProbabilities
 from rl_coach.exploration_policies.continuous_entropy import ContinuousEntropyParameters
@@ -27,17 +27,6 @@ from rl_coach.spaces import SpacesDefinition
 from rl_coach.utils import eps, indent_string
 
 
-class PolicyHeadParameters(HeadParameters):
-    def __init__(self, activation_function: str ='tanh', name: str='policy_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):
-        super().__init__(parameterized_class=PolicyHead, activation_function=activation_function, name=name,
-                         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 PolicyHead(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='tanh',

rl_coach/architectures/tensorflow_components/heads/ppo_head.py

@@ -18,7 +18,7 @@ import numpy as np
 import tensorflow as tf
 
 from rl_coach.architectures.tensorflow_components.layers import Dense
-from rl_coach.architectures.tensorflow_components.heads.head import Head, HeadParameters, normalized_columns_initializer
+from rl_coach.architectures.tensorflow_components.heads.head import Head, normalized_columns_initializer
 from rl_coach.base_parameters import AgentParameters
 from rl_coach.core_types import ActionProbabilities
 from rl_coach.spaces import BoxActionSpace, DiscreteActionSpace
@@ -26,16 +26,6 @@ from rl_coach.spaces import SpacesDefinition
 from rl_coach.utils import eps
 
 
-class PPOHeadParameters(HeadParameters):
-    def __init__(self, activation_function: str ='tanh', name: str='ppo_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):
-        super().__init__(parameterized_class=PPOHead, activation_function=activation_function, name=name,
-                         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 PPOHead(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='tanh',

rl_coach/architectures/tensorflow_components/heads/ppo_v_head.py

@@ -17,23 +17,12 @@
 import tensorflow as tf
 
 from rl_coach.architectures.tensorflow_components.layers import Dense
-
-from rl_coach.architectures.tensorflow_components.heads.head import Head, normalized_columns_initializer, HeadParameters
+from rl_coach.architectures.tensorflow_components.heads.head import Head, normalized_columns_initializer
 from rl_coach.base_parameters import AgentParameters
 from rl_coach.core_types import ActionProbabilities
 from rl_coach.spaces import SpacesDefinition
 
 
-class PPOVHeadParameters(HeadParameters):
-    def __init__(self, activation_function: str ='relu', name: str='ppo_v_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):
-        super().__init__(parameterized_class=PPOVHead, activation_function=activation_function, name=name,
-                         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 PPOVHead(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',

rl_coach/architectures/tensorflow_components/heads/q_head.py

@@ -17,23 +17,12 @@
 import tensorflow as tf
 
 from rl_coach.architectures.tensorflow_components.layers import Dense
-
-from rl_coach.architectures.tensorflow_components.heads.head import Head, HeadParameters
+from rl_coach.architectures.tensorflow_components.heads.head import Head
 from rl_coach.base_parameters import AgentParameters
 from rl_coach.core_types import QActionStateValue
 from rl_coach.spaces import SpacesDefinition, BoxActionSpace, DiscreteActionSpace
 
 
-class QHeadParameters(HeadParameters):
-    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):
-        super().__init__(parameterized_class=QHead, activation_function=activation_function, name=name,
-                         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 QHead(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',

rl_coach/architectures/tensorflow_components/heads/quantile_regression_q_head.py

@@ -17,23 +17,12 @@
 import tensorflow as tf
 
 from rl_coach.architectures.tensorflow_components.layers import Dense
-
-from rl_coach.architectures.tensorflow_components.heads.head import Head, HeadParameters
+from rl_coach.architectures.tensorflow_components.heads.head import Head
 from rl_coach.base_parameters import AgentParameters
 from rl_coach.core_types import QActionStateValue
 from rl_coach.spaces import SpacesDefinition
 
 
-class QuantileRegressionQHeadParameters(HeadParameters):
-    def __init__(self, activation_function: str ='relu', name: str='quantile_regression_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):
-        super().__init__(parameterized_class=QuantileRegressionQHead, activation_function=activation_function, name=name,
-                         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 QuantileRegressionQHead(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',

rl_coach/architectures/tensorflow_components/heads/rainbow_q_head.py

@@ -17,22 +17,12 @@
 import tensorflow as tf
 
 from rl_coach.architectures.tensorflow_components.layers import Dense
-from rl_coach.architectures.tensorflow_components.heads.head import HeadParameters, Head
+from rl_coach.architectures.tensorflow_components.heads.head import Head
 from rl_coach.base_parameters import AgentParameters
 from rl_coach.core_types import QActionStateValue
 from rl_coach.spaces import SpacesDefinition
 
 
-class RainbowQHeadParameters(HeadParameters):
-    def __init__(self, activation_function: str ='relu', name: str='rainbow_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):
-        super().__init__(parameterized_class=RainbowQHead, activation_function=activation_function, name=name,
-                         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 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',

rl_coach/architectures/tensorflow_components/heads/v_head.py

@@ -17,23 +17,12 @@
 import tensorflow as tf
 
 from rl_coach.architectures.tensorflow_components.layers import Dense
-
-from rl_coach.architectures.tensorflow_components.heads.head import Head, normalized_columns_initializer, HeadParameters
+from rl_coach.architectures.tensorflow_components.heads.head import Head, normalized_columns_initializer
 from rl_coach.base_parameters import AgentParameters
 from rl_coach.core_types import VStateValue
 from rl_coach.spaces import SpacesDefinition
 
 
-class VHeadParameters(HeadParameters):
-    def __init__(self, activation_function: str ='relu', name: str='v_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):
-        super().__init__(parameterized_class=VHead, activation_function=activation_function, name=name,
-                         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 VHead(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',