docstr: Test the multiple-rewards case. Ensure correct shapes are returned.
docstr_tok: ["Test","the","multiple-rewards","case",".","Ensure","correct","shapes","are","returned","."]
code: def testMultipleRewardsEnvironment(
      self, batch_size, observation_shape, num_rewards
  ):
      """Test the multiple-rewards case. Ensure correct shapes are returned."""
  
      env = MultipleRewardsEnvironment(
          observation_shape=observation_shape,
          batch_size=batch_size,
          num_rewards=num_rewards,
      )
  
      observation = env.reset().observation
      reward = env.step(tf.zeros(batch_size)).reward
  
      expected_observation = np.zeros([batch_size] + observation_shape)
      expected_reward = np.zeros([batch_size, num_rewards])
  
      np.testing.assert_array_almost_equal(
          expected_observation, self.evaluate(observation)
      )
      np.testing.assert_array_almost_equal(
          expected_reward, self.evaluate(reward)
      )
  
      self.assertEqual(env.reward_spec().shape, num_rewards)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tf_agents.bandits.environments import bandit_tf_environment
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.utils import common
label_desc: Use image node_id 3 for calling the BanditTFEnvironmentTest obj's underlying member method code with example usage: obj.testMultipleRewardsEnvironment(batch_size, observation_shape, num_rewards) without return types