[RLlib] Issue 9071 A3C w/ RNN not working due to VF assuming no RNN. (#13238)

rllib/utils/tf_ops.py

@@ -92,7 +92,7 @@ def minimize_and_clip(optimizer, objective, var_list, clip_val=10.0):
     variable is clipped to `clip_val`
     """
     # Accidentally passing values < 0.0 will break all gradients.
-    assert clip_val > 0.0, clip_val
+    assert clip_val is None or clip_val > 0.0, clip_val
 
     if tf.executing_eagerly():
         tape = optimizer.tape
@@ -102,10 +102,8 @@ def minimize_and_clip(optimizer, objective, var_list, clip_val=10.0):
         grads_and_vars = optimizer.compute_gradients(
             objective, var_list=var_list)
 
-    for i, (grad, var) in enumerate(grads_and_vars):
-        if grad is not None:
-            grads_and_vars[i] = (tf.clip_by_norm(grad, clip_val), var)
-    return grads_and_vars
+    return [(tf.clip_by_norm(g, clip_val) if clip_val is not None else g, v)
+            for (g, v) in grads_and_vars if g is not None]
 
 
 def make_tf_callable(session_or_none, dynamic_shape=False):

rllib/utils/torch_ops.py

@@ -14,6 +14,30 @@ FLOAT_MIN = -3.4e38
 FLOAT_MAX = 3.4e38
 
 
+def apply_grad_clipping(policy, optimizer, loss):
+    """Applies gradient clipping to already computed grads inside `optimizer`.
+
+    Args:
+        policy (TorchPolicy): The TorchPolicy, which calculated `loss`.
+        optimizer (torch.optim.Optimizer): A local torch optimizer object.
+        loss (torch.Tensor): The torch loss tensor.
+    """
+    info = {}
+    if policy.config["grad_clip"]:
+        for param_group in optimizer.param_groups:
+            # Make sure we only pass params with grad != None into torch
+            # clip_grad_norm_. Would fail otherwise.
+            params = list(
+                filter(lambda p: p.grad is not None, param_group["params"]))
+            if params:
+                grad_gnorm = nn.utils.clip_grad_norm_(
+                    params, policy.config["grad_clip"])
+                if isinstance(grad_gnorm, torch.Tensor):
+                    grad_gnorm = grad_gnorm.cpu().numpy()
+                info["grad_gnorm"] = grad_gnorm
+    return info
+
+
 def atanh(x):
     return 0.5 * torch.log((1 + x) / (1 - x))