[Minor,Feature] group_optimizers

ghstack-source-id: 81a94ed
Pull Request resolved: #2577

Author: Vincent Moens

docs/source/reference/objectives.rst

@@ -311,11 +311,12 @@ Utils
     :toctree: generated/
     :template: rl_template_noinherit.rst
 
+    HardUpdate
+    SoftUpdate
+    ValueEstimators
+    default_value_kwargs
     distance_loss
+    group_optimizers
     hold_out_net
     hold_out_params
     next_state_value
-    SoftUpdate
-    HardUpdate
-    ValueEstimators
-    default_value_kwargs

torchrl/objectives/__init__.py

@@ -23,6 +23,7 @@
 from .utils import (
     default_value_kwargs,
     distance_loss,
+    group_optimizers,
     HardUpdate,
     hold_out_net,
     hold_out_params,

torchrl/objectives/utils.py

@@ -590,3 +590,21 @@ def _clip_value_loss(
     # Chose the most pessimistic value prediction between clipped and non-clipped
     loss_value = torch.max(loss_value, loss_value_clipped)
     return loss_value, clip_fraction
+
+
+def group_optimizers(*optimizers: torch.optim.Optimizer) -> torch.optim.Optimizer:
+    """Groups multiple optimizers into a single one.
+
+    All optimizers are expected to have the same type.
+    """
+    cls = None
+    params = []
+    for optimizer in optimizers:
+        if optimizer is None:
+            continue
+        if cls is None:
+            cls = type(optimizer)
+        if cls is not type(optimizer):
+            raise ValueError("Cannot group optimizers of different type.")
+        params.extend(optimizer.param_groups)
+    return cls(params)