docstr: Exercise `reset` and `step`. Ensure correct shapes are returned.
docstr_tok: ["Exercise","`","reset","`","and","`","step","`",".","Ensure","correct","shapes","are","returned","."]
code: def testObservationAndRewardShapes(
      self, observation_shape, action_shape, batch_size
  ):
      """Exercise `reset` and `step`. Ensure correct shapes are returned."""
      env = get_gaussian_random_environment(
          observation_shape, action_shape, batch_size
      )
      observation = env.reset().observation
      reward = env.step(tf.zeros(batch_size)).reward
  
      expected_observation_shape = np.array(
          [batch_size] + observation_shape
      )
      expected_reward_shape = np.array([batch_size])
  
      self.assertAllEqual(
          expected_observation_shape,
          self.evaluate(tf.shape(observation)),
      )
      self.assertAllEqual(
          expected_reward_shape, self.evaluate(tf.shape(reward))
      )
  
file_path: agents/tf_agents/bandits/environments/random_bandit_environment_test.py
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.bandits.environments import random_bandit_environment
  from tf_agents.specs import tensor_spec
label_desc: Use image node_id 5 for calling the RandomBanditEnvironmentTest obj's underlying member method code with example usage: obj.testObservationAndRewardShapes(observation_shape, action_shape, batch_size) without return types