def _setup_pgops_mock(sequence_length=3, batch_size=2, num_policies=3):
  """Setup ops using mock distribution for numerical tests."""
  t, b = sequence_length, batch_size
  policies = [MockDistribution((t, b), i + 1)  for i in xrange(num_policies)]
  actions = [tf.constant(np.arange(t * b).reshape((t, b)))
             for i in xrange(num_policies)]
  if num_policies == 1:
    policies, actions = policies[0], actions[0]
  entropy_scale_op = lambda policies: len(nest.flatten(policies))
  return policies, actions, entropy_scale_op
Setup ops using mock distribution for numerical tests.
trfl/policy_gradient_ops_test.py