code: def testCategoricalActions(self, action_probs):
      action_spec = [
          tensor_spec.BoundedTensorSpec(
              (1,), tf.int32, 0, len(action_probs) - 1
          ),
          tensor_spec.BoundedTensorSpec(
              (), tf.int32, 0, len(action_probs) - 1
          ),
      ]
      wrapped_policy = DistributionPolicy(
          [
              tfp.distributions.Categorical(probs=[action_probs]),
              tfp.distributions.Categorical(probs=action_probs),
          ],
          self._time_step_spec,
          action_spec,
      )
      policy = greedy_policy.GreedyPolicy(wrapped_policy)
  
      self.assertEqual(policy.time_step_spec, self._time_step_spec)
      self.assertEqual(policy.action_spec, action_spec)
  
      observations = tf.constant([[1, 2]], dtype=tf.float32)
      time_step = ts.restart(observations, batch_size=1)
      action_step = policy.action(time_step)
      tf.nest.assert_same_structure(action_spec, action_step.action)
  
      action_ = self.evaluate(action_step.action)
      self.assertEqual(action_[0][0], np.argmax(action_probs))
      self.assertEqual(action_[1], np.argmax(action_probs))
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  import tensorflow
  import tensorflow_probability
  from tf_agents.policies import greedy_policy
  from tf_agents.policies import tf_policy
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.utils import test_utils
label_desc: Use image node_id 2 for calling the GreedyPolicyTest obj's underlying member method code with example usage: obj.testCategoricalActions(action_probs) without return types