code: def get_gaussian_random_environment(
      observation_shape, action_shape, batch_size
  ):
      """Returns a RandomBanditEnvironment with Gaussian observation and reward."""
      overall_shape = [batch_size] + observation_shape
      observation_distribution = tfd.Independent(
          tfd.Normal(
              loc=tf.zeros(overall_shape), scale=tf.ones(overall_shape)
          )
      )
      reward_distribution = tfd.Normal(
          loc=tf.zeros(batch_size), scale=tf.ones(batch_size)
      )
      action_spec = tensor_spec.TensorSpec(
          shape=action_shape, dtype=tf.float32
      )
      return random_bandit_environment.RandomBanditEnvironment(
          observation_distribution, reward_distribution, action_spec
      )
  
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 1 for calling a global function with example usage: get_gaussian_random_environment(observation_shape, action_shape, batch_size) and returns: random_bandit_environment