code: def test_with_uniform_context_and_normal_mu_reward(self):
      def _global_context_sampling_fn():
          return np.random.randint(-10, 10, [4])
  
      def _arm_context_sampling_fn():
          return np.random.randint(-2, 3, [5])
  
      reward_fn = LinearNormalReward([0, 1, 2, 3, 4, 5, 6, 7, 8])
  
      env = sspe.StationaryStochasticPerArmPyEnvironment(
          _global_context_sampling_fn,
          _arm_context_sampling_fn,
          6,
          reward_fn,
          batch_size=2,
      )
      time_step_spec = env.time_step_spec()
      action_spec = array_spec.BoundedArraySpec(
          shape=(), minimum=0, maximum=5, dtype=np.int32
      )
  
      random_policy = random_py_policy.RandomPyPolicy(
          time_step_spec=time_step_spec, action_spec=action_spec
      )
  
      for _ in range(5):
          time_step = env.reset()
          print("time_step", time_step)
          print("time_step_spec", time_step_spec)
          self.assertTrue(
              check_unbatched_time_step_spec(
                  time_step=time_step,
                  time_step_spec=time_step_spec,
                  batch_size=env.batch_size,
              )
          )
  
          action = random_policy.action(time_step).action
          self.assertAllEqual(action.shape, [2])
          self.assertAllGreaterEqual(action, 0)
          self.assertAllLess(action, 6)
          time_step = env.step(action)
  
file_path: agents/tf_agents/bandits/environments/stationary_stochastic_per_arm_py_environment_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  import tensorflow
  from tf_agents.bandits.environments import stationary_stochastic_per_arm_py_environment
  from tf_agents.policies import random_py_policy
  from tf_agents.specs import array_spec
label_desc: Use image node_id 1 for calling the StationaryStochasticPerArmBanditPyEnvironmentTest obj's underlying member method code with example usage: obj.test_with_uniform_context_and_normal_mu_reward() and returns: np, np