Add gradient norm tracking

taylorhansen · Oct 4, 2023 · aa9b815 · aa9b815
1 parent 78c25cb
commit aa9b815
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/src/py/agents/dqn_agent.py b/src/py/agents/dqn_agent.py
@@ -427,7 +427,11 @@ def _learn_step_impl(
                 **{f"{n}/activation": a for n, a in activations.items()},
                 **{
                     f"{w.name}/grads": g
-                    for w, g in zip(gradients, self.model.trainable_weights)
+                    for g, w in zip(gradients, self.model.trainable_weights)
+                },
+                **{
+                    f"{w.name}/grad_norm": tf.sqrt(tf.reduce_sum(tf.square(g)))
+                    for g, w in zip(gradients, self.model.trainable_weights)
                 },
                 **{
                     f"{w.name}/weights": w for w in self.model.trainable_weights

diff --git a/src/py/agents/drqn_agent.py b/src/py/agents/drqn_agent.py
@@ -471,7 +471,11 @@ def _learn_step_impl(
                 **{f"{n}/activation": a for n, a in activations.items()},
                 **{
                     f"{w.name}/grads": g
-                    for w, g in zip(gradients, self.model.trainable_weights)
+                    for g, w in zip(gradients, self.model.trainable_weights)
+                },
+                **{
+                    f"{w.name}/grad_norm": tf.sqrt(tf.reduce_sum(tf.square(g)))
+                    for g, w in zip(gradients, self.model.trainable_weights)
                 },
                 **{
                     f"{w.name}/weights": w for w in self.model.trainable_weights