code: def get_deterministic_gaussian_non_stationary_environment(
      observation_shape,
      action_shape,
      batch_size,
      drift_mean=0.0,
      drift_scale=1.0,
  ):
      """Returns a DriftingLinearEnvironment."""
      overall_shape = [batch_size] + observation_shape
      observation_distribution = tfd.Normal(
          loc=tf.zeros(overall_shape), scale=tf.ones(overall_shape)
      )
      observation_to_reward_shape = observation_shape + action_shape
      observation_to_reward_distribution = tfd.Normal(
          loc=tf.zeros(observation_to_reward_shape),
          scale=tf.ones(observation_to_reward_shape),
      )
      drift_distribution = tfd.Normal(loc=drift_mean, scale=drift_scale)
      additive_reward_distribution = tfd.Normal(
          loc=tf.zeros(action_shape), scale=tf.ones(action_shape)
      )
      return dle.DriftingLinearEnvironment(
          observation_distribution,
          observation_to_reward_distribution,
          drift_distribution,
          additive_reward_distribution,
      )
  
file_path: agents/tf_agents/bandits/environments/drifting_linear_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 drifting_linear_environment
label_desc: Use image node_id 2 for calling a global function with example usage: get_deterministic_gaussian_non_stationary_environment(observation_shape, action_shape, batch_size, drift_mean, drift_scale) and returns: dle