code: def testAction(self, batch_size):
      if tf.executing_eagerly():
          self.skipTest(
              "b/123770140: Handling sessions with eager mode is buggy"
          )
  
      single_observation = np.array([1, 2], dtype=np.float32)
      time_steps = ts.restart(single_observation)
      if batch_size is not None:
          time_steps = [time_steps] * batch_size
          time_steps = fast_map_structure(
              lambda *arrays: np.stack(arrays), *time_steps
          )
      policy = py_tf_policy.PyTFPolicy(self._tf_policy)
  
      with self.cached_session():
          policy_state = policy.get_initial_state(batch_size)
          self.evaluate(tf.compat.v1.global_variables_initializer())
          action_steps = policy.action(time_steps, policy_state)
          self.assertEqual(action_steps.action.dtype, np.int32)
          if batch_size is None:
              self.assertEqual(action_steps.action.shape, ())
              self.assertIn(action_steps.action, (0, 1))
              self.assertEqual(action_steps.state, np.zeros([1, 1]))
          else:
              self.assertEqual(action_steps.action.shape, (batch_size,))
              for a in action_steps.action:
                  self.assertIn(a, (0, 1))
              self.assertAllEqual(
                  action_steps.state, np.zeros([batch_size, 1])
              )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  from absl import flags
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tf_agents.networks import network
  from tf_agents.policies import actor_policy
  from tf_agents.policies import ou_noise_policy
  from tf_agents.policies import py_tf_policy
  from tf_agents.policies import q_policy
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.utils import test_utils
label_desc: Use image node_id 6 for calling the PyTFPolicyTest obj's underlying member method code with example usage: obj.testAction(batch_size) without return types