[BugFix, Feature] Vmap randomness in losses

torchrl/objectives/cql.py

@@ -335,8 +335,12 @@ def __init__(
                 torch.nn.Parameter(torch.tensor(math.log(1.0), device=device)),
             )
 
-        self._vmap_qvalue_networkN0 = _vmap_func(self.qvalue_network, (None, 0))
-        self._vmap_qvalue_network00 = _vmap_func(self.qvalue_network)
+        self._vmap_qvalue_networkN0 = _vmap_func(
+            self.qvalue_network, (None, 0), randomness="different"
+        )
+        self._vmap_qvalue_network00 = _vmap_func(
+            self.qvalue_network, randomness="different"
+        )
 
     @property
     def target_entropy(self):

torchrl/objectives/iql.py

@@ -286,7 +286,9 @@ def __init__(
         if gamma is not None:
             warnings.warn(_GAMMA_LMBDA_DEPREC_WARNING, category=DeprecationWarning)
             self.gamma = gamma
-        self._vmap_qvalue_networkN0 = _vmap_func(self.qvalue_network, (None, 0))
+        self._vmap_qvalue_networkN0 = _vmap_func(
+            self.qvalue_network, (None, 0), randomness="different"
+        )
 
     @property
     def device(self) -> torch.device:

torchrl/objectives/redq.py

@@ -318,7 +318,9 @@ def __init__(
             warnings.warn(_GAMMA_LMBDA_DEPREC_WARNING, category=DeprecationWarning)
             self.gamma = gamma
 
-        self._vmap_qvalue_network00 = _vmap_func(self.qvalue_network)
+        self._vmap_qvalue_network00 = _vmap_func(
+            self.qvalue_network, randomness="different"
+        )
         self._vmap_getdist = _vmap_func(self.actor_network, func="get_dist_params")
 
     @property

torchrl/objectives/sac.py

@@ -376,9 +376,11 @@ def __init__(
         if gamma is not None:
             warnings.warn(_GAMMA_LMBDA_DEPREC_WARNING, category=DeprecationWarning)
             self.gamma = gamma
-        self._vmap_qnetworkN0 = _vmap_func(self.qvalue_network, (None, 0))
+        self._vmap_qnetworkN0 = _vmap_func(
+            self.qvalue_network, (None, 0), randomness="different"
+        )
         if self._version == 1:
-            self._vmap_qnetwork00 = _vmap_func(qvalue_network)
+            self._vmap_qnetwork00 = _vmap_func(qvalue_network, randomness="different")
 
     @property
     def target_entropy_buffer(self):

torchrl/objectives/td3.py

@@ -26,6 +26,25 @@
 )
 from torchrl.objectives.value import TD0Estimator, TD1Estimator, TDLambdaEstimator
 
+RANDOM_MODULE_LIST = (
+    torch.nn.Dropout,
+    torch.nn.Dropout2d,
+    torch.nn.Dropout3d,
+    torch.nn.AlphaDropout,
+    torch.nn.FeatureAlphaDropout,
+    torch.nn.GaussianDropout,
+    torch.nn.GaussianNoise,
+    torch.nn.SyncBatchNorm,
+    torch.nn.GroupNorm,
+    torch.nn.LayerNorm,
+    torch.nn.LocalResponseNorm,
+    torch.nn.InstanceNorm1d,
+    torch.nn.InstanceNorm2d,
+    torch.nn.InstanceNorm3d,
+    torch.nn.SpatialDropout,
+    torch.nn.SpatialCrossMapLRN,
+)
+
 
 class TD3Loss(LossModule):
     """TD3 Loss module.
@@ -201,6 +220,7 @@ class _AcceptedKeys:
         "next_state_value",
         "target_value",
     ]
+    _vmap_randomness = None
 
     def __init__(
         self,
@@ -219,7 +239,6 @@ def __init__(
         priority_key: str = None,
         separate_losses: bool = False,
     ) -> None:
-
         super().__init__()
         self._in_keys = None
         self._set_deprecated_ctor_keys(priority=priority_key)
@@ -296,8 +315,12 @@ def __init__(
         if gamma is not None:
             warnings.warn(_GAMMA_LMBDA_DEPREC_WARNING, category=DeprecationWarning)
             self.gamma = gamma
-        self._vmap_qvalue_network00 = _vmap_func(self.qvalue_network)
-        self._vmap_actor_network00 = _vmap_func(self.actor_network)
+        self._vmap_qvalue_network00 = _vmap_func(
+            self.qvalue_network, randomness=self.vmap_randomness
+        )
+        self._vmap_actor_network00 = _vmap_func(
+            self.actor_network, randomness=self.vmap_randomness
+        )
 
     def _forward_value_estimator_keys(self, **kwargs) -> None:
         if self._value_estimator is not None:
@@ -343,6 +366,28 @@ def _cached_stack_actor_params(self):
             [self.actor_network_params, self.target_actor_network_params], 0
         )
 
+    @property
+    def vmap_randomness(self):
+        if self._vmap_randomness is None:
+            do_break = False
+            for val in self.__dict__.values():
+                if isinstance(val, torch.nn.Module):
+                    for module in val.modules():
+                        if isinstance(module, RANDOM_MODULE_LIST):
+                            self._vmap_randomness = "different"
+                            do_break = True
+                            break
+                if do_break:
+                    # double break
+                    break
+            else:
+                self._vmap_randomness = "error"
+
+        return self._vmap_randomness
+
+    def set_vmap_randomness(self, value):
+        self._vmap_randomness = value
+
     def actor_loss(self, tensordict):
         tensordict_actor_grad = tensordict.select(*self.actor_network.in_keys)
         with self.actor_network_params.to_module(self.actor_network):

torchrl/objectives/utils.py

@@ -478,13 +478,23 @@ def new_fun(self, netname=None):
 
 
 def _vmap_func(module, *args, func=None, **kwargs):
-    def decorated_module(*module_args_params):
-        params = module_args_params[-1]
-        module_args = module_args_params[:-1]
-        with params.to_module(module):
-            if func is None:
-                return module(*module_args)
-            else:
-                return getattr(module, func)(*module_args)
+    try:
 
-    return vmap(decorated_module, *args, **kwargs)  # noqa: TOR101
+        def decorated_module(*module_args_params):
+            params = module_args_params[-1]
+            module_args = module_args_params[:-1]
+            with params.to_module(module):
+                if func is None:
+                    return module(*module_args)
+                else:
+                    return getattr(module, func)(*module_args)
+
+        return vmap(decorated_module, *args, **kwargs)  # noqa: TOR101
+
+    except RuntimeError as err:
+        if "vmap: called random operation while in randomness error mode" in str(
+            err
+        ):  # better to use re.match here but anyway
+            raise RuntimeError(
+                "Please use loss_module.set_vmap_randomness to handle random operations during vmap."
+            ) from err