chore: import upstream snapshot with attribution

rllib/core/learner/utils.py

@@ -0,0 +1,58 @@
+import copy
+
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.typing import NetworkType
+from ray.util import PublicAPI
+
+torch, _ = try_import_torch()
+
+
+def make_target_network(main_net: NetworkType) -> NetworkType:
+    """Creates a (deep) copy of `main_net` (including synched weights) and returns it.
+
+    Args:
+        main_net: The main network to return a target network for
+
+    Returns:
+        The copy of `main_net` that can be used as a target net. Note that the weights
+        of the returned net are already synched (identical) with `main_net`.
+    """
+    # Deepcopy the main net (this should already take care of synching all weights).
+    target_net = copy.deepcopy(main_net)
+    # Make the target net not trainable.
+    if isinstance(main_net, torch.nn.Module):
+        target_net.requires_grad_(False)
+    else:
+        raise ValueError(f"Unsupported framework for given `main_net` {main_net}!")
+
+    return target_net
+
+
+@PublicAPI(stability="beta")
+def update_target_network(
+    *,
+    main_net: NetworkType,
+    target_net: NetworkType,
+    tau: float,
+) -> None:
+    """Updates a target network (from a "main" network) using Polyak averaging.
+
+    Thereby:
+    new_target_net_weight = (
+        tau * main_net_weight + (1.0 - tau) * current_target_net_weight
+    )
+
+    Args:
+        main_net: The nn.Module to update from.
+        target_net: The target network to update.
+        tau: The tau value to use in the Polyak averaging formula. Use 1.0 for a
+            complete sync of the weights (target and main net will be the exact same
+            after updating).
+    """
+    if isinstance(main_net, torch.nn.Module):
+        from ray.rllib.utils.torch_utils import update_target_network as _update_target
+
+    else:
+        raise ValueError(f"Unsupported framework for given `main_net` {main_net}!")
+
+    _update_target(main_net=main_net, target_net=target_net, tau=tau)