code: def test_non_scalar_rewards(self):
      def _context_sampling_fn():
          return np.array([[4, 3], [4, 3], [5, 6]])
  
      # Build a case with 4 arms and 2-dimensional rewards and batch size 3.
      reward_fns = [
          LinearDeterministicMultipleRewards(
              theta
          )  # pylint: disable=g-complex-comprehension
          for theta in [
              np.array([[0, 1], [1, 0]]),
              np.array([[1, 2], [2, 1]]),
              np.array([[2, 3], [3, 2]]),
              np.array([[3, 4], [4, 3]]),
          ]
      ]
      env = sspe.StationaryStochasticPyEnvironment(
          _context_sampling_fn, reward_fns, batch_size=3
      )
      time_step = env.reset()
      self.assertAllEqual(
          time_step.observation, [[4, 3], [4, 3], [5, 6]]
      )
      time_step = env.step([0, 1, 2])
      self.assertAllEqual(
          time_step.reward, [[3.0, 4.0], [10.0, 11.0], [28.0, 27.0]]
      )
      env.reset()
      time_step = env.step([2, 3, 0])
      self.assertAllEqual(
          time_step.reward, [[17.0, 18.0], [24.0, 25.0], [6.0, 5.0]]
      )
      # Check that the reward vectors in the reward spec are 2-dimensional.
      time_step_spec = env.time_step_spec()
      self.assertEqual(time_step_spec.reward.shape[0], 2)
  
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_py_environment
  from tf_agents.policies import random_py_policy
  from tf_agents.specs import array_spec
  from tf_agents.trajectories import time_step
label_desc: Use image node_id 4 for calling the StationaryStochasticBanditPyEnvironmentTest obj's underlying member method code with example usage: obj.test_non_scalar_rewards() and returns: np