[Refactor] refactor noisy linear (#3082)

Author: vmoens

sota-implementations/redq/utils.py

@@ -5,8 +5,8 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Sequence
-
 from copy import copy
+from functools import partial
 
 import torch
 from omegaconf import OmegaConf
@@ -417,7 +417,11 @@ def make_redq_model(
     if qvalue_net_kwargs is None:
         qvalue_net_kwargs = {}
 
-    linear_layer_class = torch.nn.Linear if not cfg.exploration.noisy else NoisyLinear
+    linear_layer_class = (
+        torch.nn.Linear
+        if not cfg.exploration.noisy
+        else partial(NoisyLinear, use_exploration_type=True)
+    )
 
     out_features_actor = (2 - gSDE) * action_spec.shape[-1]
     if cfg.env.from_pixels: