class_docstr: Tests for policy_gradient op.
class_docstr_tok: ["Tests","for","policy_gradient","op","."]
code: def _setUp_loss(self, sequence_length, batch_size, action_dim):
      # Use default single MVN action setup
      policies, self._actions, self._policy_vars, _ = _setup_pgops(
          sequence_length=sequence_length,
          batch_size=batch_size,
          num_mvn_actions=action_dim,
      )
      self._action_values = tf.placeholder(
          tf.float32, shape=(sequence_length, batch_size)
      )
      self._loss = pg_ops.policy_gradient(
          policies, self._actions, self._action_values
      )
  
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 1 for calling the PolicyGradientTest obj's underlying member method code with example usage: obj._setUp_loss(sequence_length, batch_size, action_dim) without return types