refactor: refinements in utils and associated functions

pearl/neural_networks/common/epistemic_neural_networks.py

@@ -17,7 +17,7 @@
 
 import torch
 import torch.nn as nn
-from pearl.neural_networks.common.utils import init_weights, mlp_block
+from pearl.neural_networks.common.utils import xavier_init_weights, mlp_block
 from torch import Tensor
 
 
@@ -155,12 +155,12 @@ def __init__(
         for _ in range(self.index_dim):
             model = mlp_block(self.input_dim, self.hidden_dims, self.output_dim)
             # Xavier uniform initalization
-            model.apply(init_weights)
+            model.apply(xavier_init_weights)
             models.append(model)
         self.base_model: nn.Module = mlp_block(
             self.input_dim, self.hidden_dims, self.output_dim
         )
-        self.base_model.apply(init_weights)
+        self.base_model.apply(xavier_init_weights)
         self.base_model = self.base_model.to("meta")
         self.models: nn.ModuleList = nn.ModuleList(models)
 
@@ -228,7 +228,7 @@ def __init__(
         self.epinet: nn.Module = mlp_block(
             epinet_input_dim, self.epi_hiddens, self.index_dim * self.output_dim
         )
-        self.epinet.apply(init_weights)
+        self.epinet.apply(xavier_init_weights)
         # Priornet
         self.priornet = Priornet(
             self.input_dim, self.prior_hiddens, self.output_dim, self.index_dim

pearl/neural_networks/common/utils.py

@@ -9,6 +9,7 @@
 
 import logging
 from typing import Any
+from enum import Enum
 
 import torch
 import torch.nn as nn
@@ -24,29 +25,54 @@ def __init__(self, dim: int = -1) -> None:
         self.dim: int = dim
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return nn.Softplus()(x) / torch.sum(
-            nn.Softplus()(x), dim=self.dim, keepdim=True
-        )
+        values = nn.functional.softplus(x)
+        return values / torch.sum(values, dim=self.dim, keepdim=True)
 
 
 # Activations and loss functions
-# TODO: Make these into Enums
-ACTIVATION_MAP = {
-    "tanh": nn.Tanh,
-    "relu": nn.ReLU,
-    "leaky_relu": nn.LeakyReLU,
-    "linear": nn.Identity,
-    "sigmoid": nn.Sigmoid,
-    "softplus": nn.Softplus,
-    "softmax": nn.Softmax,
-    "normalized_softplus": NormalizedSoftplus,
-}
-
-LOSS_TYPES = {
-    "mse": nn.functional.mse_loss,
-    "mae": nn.functional.l1_loss,
-    "cross_entropy": nn.functional.binary_cross_entropy,
-}
+class ActivationType(Enum):
+    TANH = "tanh"
+    RELU = "relu"
+    LEAKY_RELU = "leaky_relu"
+    LINEAR = "linear"
+    SIGMOID = "sigmoid"
+    SOFTPLUS = "softplus"
+    SOFTMAX = "softmax"
+    NORMALIZED_SOFTPLUS = "normalized_softplus"
+
+    def module(self) -> nn.Module:
+        if self is ActivationType.TANH:
+            return nn.Tanh()
+        if self is ActivationType.RELU:
+            return nn.ReLU()
+        if self is ActivationType.LEAKY_RELU:
+            return nn.LeakyReLU()
+        if self is ActivationType.LINEAR:
+            return nn.Identity()
+        if self is ActivationType.SIGMOID:
+            return nn.Sigmoid()
+        if self is ActivationType.SOFTPLUS:
+            return nn.Softplus()
+        if self is ActivationType.SOFTMAX:
+            return nn.Softmax(dim=-1)
+        if self is ActivationType.NORMALIZED_SOFTPLUS:
+            return NormalizedSoftplus()
+        raise ValueError(f"Unhandled activation type {self}")
+
+
+class LossType(Enum):
+    MSE = "mse"
+    MAE = "mae"
+    CROSS_ENTROPY = "cross_entropy"
+
+    def function(self) -> Any:
+        if self is LossType.MSE:
+            return nn.functional.mse_loss
+        if self is LossType.MAE:
+            return nn.functional.l1_loss
+        if self is LossType.CROSS_ENTROPY:
+            return nn.functional.binary_cross_entropy
+        raise ValueError(f"Unhandled loss type {self}")
 
 
 def mlp_block(
@@ -90,7 +116,7 @@ def mlp_block(
             single_layers.append(nn.LayerNorm(output_dim_current_layer))
         if dropout_ratio > 0:
             single_layers.append(nn.Dropout(p=dropout_ratio))
-        single_layers.append(ACTIVATION_MAP[hidden_activation]())
+        single_layers.append(ActivationType(hidden_activation).module())
         if use_batch_norm:
             single_layers.append(nn.BatchNorm1d(output_dim_current_layer))
         single_layer_model = nn.Sequential(*single_layers)
@@ -109,7 +135,7 @@ def mlp_block(
     last_layer = []
     last_layer.append(nn.Linear(dims[-2], dims[-1]))
     if last_activation is not None:
-        last_layer.append(ACTIVATION_MAP[last_activation]())
+        last_layer.append(ActivationType(last_activation).module())
     last_layer_model = nn.Sequential(*last_layer)
     if use_skip_connections:
         if dims[-2] == dims[-1]:
@@ -158,19 +184,18 @@ def conv_block(
             padding=padding,
         )
         layers.append(conv_layer)
-        if use_batch_norm and input_channels_count > 1:
-            layers.append(
-                nn.BatchNorm2d(input_channels_count)
-            )  # input to Batchnorm 2d is the number of input channels
+        if use_batch_norm:
+            # batch norm should normalize the output of the convolutional layer
+            layers.append(nn.BatchNorm2d(out_channels))
         layers.append(nn.ReLU())
         # number of input channels to next layer is number of output channels of previous layer:
         input_channels_count = out_channels
 
     return nn.Sequential(*layers)
 
 
-# TODO: the name of this function needs to be revised to xavier_init_weights
-def init_weights(m: nn.Module) -> None:
+def xavier_init_weights(m: nn.Module) -> None:
+    """Initialize Linear layer weights with Xavier uniform initializer."""
     if isinstance(m, nn.Linear):
         nn.init.xavier_uniform_(m.weight)
         m.bias.data.fill_(0.01)
@@ -267,10 +292,19 @@ def update_target_networks(
 
 
 def compute_output_dim_model_cnn(
-    input_channels: int, input_width: int, input_height: int, model_cnn: nn.Module
+    input_channels: int,
+    input_width: int,
+    input_height: int,
+    model_cnn: nn.Module,
 ) -> int:
-    dummy_input = torch.zeros(1, input_channels, input_width, input_height)
-    dummy_output_flattened = torch.flatten(
-        model_cnn(dummy_input), start_dim=1, end_dim=-1
-    )
+    """Return flattened output dimension of a CNN block.
+
+    Parameters correspond to the input tensor layout
+    ``(batch_size, input_channels, input_height, input_width)``.
+    """
+    dummy_input = torch.zeros(1, input_channels, input_height, input_width)
+    with torch.no_grad():
+        dummy_output_flattened = torch.flatten(
+            model_cnn(dummy_input), start_dim=1, end_dim=-1
+        )
     return dummy_output_flattened.shape[1]

pearl/neural_networks/contextual_bandit/neural_linear_regression.py

@@ -12,7 +12,7 @@
 
 import torch
 import torch.nn as nn
-from pearl.neural_networks.common.utils import ACTIVATION_MAP
+from pearl.neural_networks.common.utils import ActivationType
 from pearl.neural_networks.common.value_networks import VanillaValueNetwork
 from pearl.neural_networks.contextual_bandit.base_cb_model import MuSigmaCBModel
 from pearl.neural_networks.contextual_bandit.linear_regression import LinearRegression
@@ -51,7 +51,7 @@ def __init__(
             force_pinv: If True, we will always use pseudo inverse to invert the `A` matrix. If
                 False, we will first try to use regular matrix inversion. If it fails, we will
                 fallback to pseudo inverse.
-            output_activation_name: output activation function name (see ACTIVATION_MAP)
+            output_activation_name: output activation function name (see ActivationType)
             use_batch_norm: whether to use batch normalization
             use_layer_norm: whether to use layer normalization
             hidden_activation: activation function for hidden layers
@@ -80,7 +80,7 @@ def __init__(
             gamma=gamma,
             force_pinv=force_pinv,
         )
-        self.output_activation: nn.Module = ACTIVATION_MAP[output_activation_name]()
+        self.output_activation: nn.Module = ActivationType(output_activation_name).module()
         self.linear_layer_e2e = nn.Linear(
             in_features=hidden_dims[-1], out_features=1, bias=False
         )  # used only if nn_e2e is True

pearl/neural_networks/sequential_decision_making/twin_critic.py

@@ -13,7 +13,7 @@
 
 import torch
 import torch.nn as nn
-from pearl.neural_networks.common.utils import init_weights
+from pearl.neural_networks.common.utils import xavier_init_weights
 from pearl.neural_networks.sequential_decision_making.q_value_networks import (
     QValueNetwork,
     VanillaQValueNetwork,
@@ -34,7 +34,7 @@ def __init__(
         state_dim: Optional[int] = None,
         action_dim: Optional[int] = None,
         hidden_dims: Optional[Iterable[int]] = None,
-        init_fn: Callable[[nn.Module], None] = init_weights,
+        init_fn: Callable[[nn.Module], None] = xavier_init_weights,
         network_type: type[QValueNetwork] = VanillaQValueNetwork,
         output_dim: int = 1,
         network_instance_1: Optional[QValueNetwork] = None,

pearl/policy_learners/contextual_bandits/neural_bandit.py

@@ -21,7 +21,7 @@
     HistorySummarizationModule,
     SubjectiveState,
 )
-from pearl.neural_networks.common.utils import LOSS_TYPES
+from pearl.neural_networks.common.utils import LossType
 from pearl.neural_networks.common.value_networks import VanillaValueNetwork
 from pearl.policy_learners.contextual_bandits.contextual_bandit_base import (
     ContextualBanditBase,
@@ -58,7 +58,7 @@ def __init__(
         batch_size: int = 128,
         learning_rate: float = 0.001,
         state_features_only: bool = False,
-        loss_type: str = "mse",  # one of the LOSS_TYPES names, e.g., mse, mae, xentropy
+        loss_type: LossType | str = LossType.MSE,  # one of the LossType names, e.g., MSE, MAE, CROSS_ENTROPY
         action_representation_module: ActionRepresentationModule | None = None,
     ) -> None:
         super().__init__(
@@ -77,7 +77,7 @@ def __init__(
             self.model.parameters(), lr=learning_rate, amsgrad=True
         )
         self._state_features_only = state_features_only
-        self.loss_type = loss_type
+        self.loss_type = LossType(loss_type) if isinstance(loss_type, str) else loss_type
 
     def learn_batch(self, batch: TransitionBatch) -> dict[str, Any]:
         expected_values = batch.reward
@@ -94,7 +94,7 @@ def learn_batch(self, batch: TransitionBatch) -> dict[str, Any]:
         # forward pass
         predicted_values = self.model(input_features)
 
-        criterion = LOSS_TYPES[self.loss_type]
+        criterion = self.loss_type.function()
 
         # don't reduce the loss, so that we can calculate weighted loss
         loss = criterion(

pearl/policy_learners/contextual_bandits/neural_linear_bandit.py

@@ -21,7 +21,7 @@
     HistorySummarizationModule,
     SubjectiveState,
 )
-from pearl.neural_networks.common.utils import LOSS_TYPES
+from pearl.neural_networks.common.utils import LossType
 from pearl.neural_networks.contextual_bandit.neural_linear_regression import (
     NeuralLinearRegression,
 )
@@ -80,7 +80,7 @@ def __init__(
         apply_discounting_interval: float = 0.0,  # set to 0 to disable
         force_pinv: bool = False,  # If True, use pseudo inverse instead of regular inverse for `A`
         state_features_only: bool = True,
-        loss_type: str = "mse",  # one of the LOSS_TYPES names: [mse, mae, cross_entropy]
+        loss_type: LossType | str = LossType.MSE,  # one of the LossType names: [MSE, MAE, CROSS_ENTROPY]
         output_activation_name: str = "linear",
         use_batch_norm: bool = False,
         use_layer_norm: bool = False,
@@ -120,7 +120,7 @@ def __init__(
             self.model.parameters(), lr=learning_rate, amsgrad=True
         )
         self._state_features_only = state_features_only
-        self.loss_type = loss_type
+        self.loss_type = LossType(loss_type) if isinstance(loss_type, str) else loss_type
         self.apply_discounting_interval = apply_discounting_interval
         self.last_sum_weight_when_discounted = 0.0
         self.separate_uncertainty = separate_uncertainty
@@ -177,8 +177,8 @@ def learn_batch(self, batch: TransitionBatch) -> dict[str, Any]:
         else:
             # criterion = mae, mse, Xentropy
             # Xentropy loss apply Sigmoid, MSE or MAE apply Identiy
-            criterion = LOSS_TYPES[self.loss_type]
-            if self.loss_type == "cross_entropy":
+            criterion = self.loss_type.function()
+            if self.loss_type is LossType.CROSS_ENTROPY:
                 assert torch.all(expected_values >= 0) and torch.all(
                     expected_values <= 1
                 )

pearl/policy_learners/sequential_decision_making/actor_critic_base.py

@@ -24,7 +24,7 @@
 from pearl.history_summarization_modules.history_summarization_module import (
     HistorySummarizationModule,
 )
-from pearl.neural_networks.common.utils import init_weights, update_target_network
+from pearl.neural_networks.common.utils import xavier_init_weights, update_target_network
 from pearl.neural_networks.common.value_networks import ValueNetwork
 from pearl.neural_networks.sequential_decision_making.actor_networks import (
     ActorNetwork,
@@ -154,7 +154,7 @@ def __init__(
                 ),
                 action_space=action_space,
             )
-        self._actor.apply(init_weights)
+        self._actor.apply(xavier_init_weights)
         if actor_optimizer is not None:
             self._actor_optimizer: optim.Optimizer = actor_optimizer
         else:

pearl/utils/functional_utils/learning/critic_utils.py

@@ -13,7 +13,7 @@
 import torch.nn as nn
 
 from pearl.neural_networks.common.utils import (
-    init_weights,
+    xavier_init_weights,
     update_target_network,
     update_target_networks,
 )
@@ -81,7 +81,7 @@ class `TwinCritic` will be instantiated with the specified `network_type`. Note
             action_dim=action_dim,
             hidden_dims=hidden_dims,
             network_type=network_type,
-            init_fn=init_weights,
+            init_fn=xavier_init_weights,
         )
     else:
         if network_type == VanillaQValueNetwork:

test/integration/test_integration.py

@@ -7,7 +7,7 @@
 
 # pyre-strict
 
-from pearl.neural_networks.common.utils import init_weights
+from pearl.neural_networks.common.utils import xavier_init_weights
 from pearl.neural_networks.common.value_networks import VanillaValueNetwork
 from pearl.neural_networks.sequential_decision_making.actor_networks import (
     GaussianActorNetwork,
@@ -602,7 +602,7 @@ def test_sac_network_instance(self) -> None:
             action_dim=action_representation_module.max_number_actions,
             hidden_dims=[64, 64, 64],
             network_type=VanillaQValueNetwork,
-            init_fn=init_weights,
+            init_fn=xavier_init_weights,
         )
 
         agent = PearlAgent(
@@ -692,7 +692,7 @@ def test_continuous_sac_network_instance(self) -> None:
             action_dim=env.action_space.action_dim,
             hidden_dims=[64, 64],
             network_type=VanillaQValueNetwork,
-            init_fn=init_weights,
+            init_fn=xavier_init_weights,
         )
 
         agent = PearlAgent(
@@ -861,7 +861,7 @@ def test_td3_network_instance(self) -> None:
             action_dim=env.action_space.action_dim,
             hidden_dims=[400, 300],
             network_type=VanillaQValueNetwork,
-            init_fn=init_weights,
+            init_fn=xavier_init_weights,
         )
 
         agent = PearlAgent(