Add with_bias parameter to create_mlp

docs/misc/changelog.rst

@@ -17,6 +17,7 @@ Breaking Changes:
 New Features:
 ^^^^^^^^^^^^^
 - Introduced mypy type checking
+- Added ``with_bias`` argument to ``create_mlp``
 
 SB3-Contrib
 ^^^^^^^^^^^

stable_baselines3/common/torch_layers.py

@@ -98,6 +98,7 @@ def create_mlp(
     output_dim: int,
     net_arch: List[int],
     activation_fn: Type[nn.Module] = nn.ReLU,
+    with_bias: bool = True,
     squash_output: bool = False,
 ) -> List[nn.Module]:
     """
@@ -111,23 +112,24 @@ def create_mlp(
         The length of this list is the number of layers.
     :param activation_fn: The activation function
         to use after each layer.
+    :param with_bias: If set to False, the layers will not learn an additive bias
     :param squash_output: Whether to squash the output using a Tanh
         activation function
     :return:
     """
 
     if len(net_arch) > 0:
-        modules = [nn.Linear(input_dim, net_arch[0]), activation_fn()]
+        modules = [nn.Linear(input_dim, net_arch[0], bias=with_bias), activation_fn()]
     else:
         modules = []
 
     for idx in range(len(net_arch) - 1):
-        modules.append(nn.Linear(net_arch[idx], net_arch[idx + 1]))
+        modules.append(nn.Linear(net_arch[idx], net_arch[idx + 1], bias=with_bias))
         modules.append(activation_fn())
 
     if output_dim > 0:
         last_layer_dim = net_arch[-1] if len(net_arch) > 0 else input_dim
-        modules.append(nn.Linear(last_layer_dim, output_dim))
+        modules.append(nn.Linear(last_layer_dim, output_dim, bias=with_bias))
     if squash_output:
         modules.append(nn.Tanh())
     return modules