Update

[ghstack-poisoned]

torchrl/objectives/common.py

@@ -592,6 +592,7 @@ def _make_meta_params(param):
         return pd
 
     def _make_vmap(self):
+        """Caches the the vmap callers to reduce the overhead at runtime."""
         raise NotImplementedError(
             f"_make_vmap has been called but is not implemented for loss of type {type(self).__name__}."
         )

torchrl/objectives/cql.py

@@ -373,7 +373,7 @@ def __init__(
                 "log_alpha_prime",
                 torch.nn.Parameter(torch.tensor(math.log(1.0), device=device)),
             )
-
+        self._make_vmap()
         self.reduction = reduction
 
     def _make_vmap(self):

torchrl/objectives/deprecated.py

@@ -221,6 +221,7 @@ def __init__(
         self._action_spec = action_spec
         self.target_entropy_buffer = None
         self.gSDE = gSDE
+        self._make_vmap()
         self.reduction = reduction
 
         if gamma is not None: