class_docstr: Tests for policy_gradient op.
class_docstr_tok: ["Tests","for","policy_gradient","op","."]
code: def testGradients(self, sequence_length, batch_size, action_dim):
      self._setUp_loss(sequence_length, batch_size, action_dim)
      total_loss = tf.reduce_sum(self._loss)
      for policy_var in self._policy_vars:
          gradients = tf.gradients(total_loss, policy_var)
          self.assertEqual(
              gradients[0].get_shape().as_list(),
              policy_var.get_shape().as_list(),
          )
      gradients = tf.gradients(
          [total_loss], [self._actions, self._action_values]
      )
      self.assertEqual(gradients, [None, None])
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  from six.moves import xrange
  import tensorflow.compat.v1
  import tensorflow_probability
  import tree
  from trfl import policy_gradient_ops
label_desc: Use image node_id 3 for calling the PolicyGradientTest obj's underlying member method code with example usage: obj.testGradients(sequence_length, batch_size, action_dim) without return types