code: def mushroom_reward_distribution(
      r_noeat,
      r_eat_safe,
      r_eat_poison_bad,
      r_eat_poison_good,
      prob_poison_bad,
  ):
      """Creates a distribution for rewards for the mushroom environment.
  
      Args:
        r_noeat: (float) Reward value for not eating the mushroom.
        r_eat_safe: (float) Reward value for eating an edible mushroom.
        r_eat_poison_bad: (float) Reward value for eating and getting poisoned from
          a poisonous mushroom.
        r_eat_poison_good: (float) Reward value for surviving after eating a
          poisonous mushroom.
        prob_poison_bad: Probability of getting poisoned by a poisonous mushroom.
  
      Returns:
        A reward distribution table, instance of `tfd.Distribution`.
      """
  
      # The function works by first creating a 2x2 Bernoulli with all but one having
      # parameter 0. The fourth one, that corresponds to eating a poisonous mushroom
      # has parameter `prob_poison_bad`. Then, the whole table is shifted and scaled
      # to the desired values.
  
      distr = tfd.Bernoulli(
          probs=[[0, prob_poison_bad], [0, 0]], dtype=tf.float32
      )
      reward_distr = tfp.bijectors.Shift(
          [[r_noeat, r_eat_poison_bad], [r_noeat, r_eat_safe]]
      )(
          tfp.bijectors.Scale(
              [[1, r_eat_poison_good - r_eat_poison_bad], [1, 1]]
          )(distr)
      )
      return tfd.Independent(reward_distr, reinterpreted_batch_ndims=2)
  
file_path: agents/tf_agents/bandits/environments/dataset_utilities.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import csv
  import gin
  import numpy
  from six.moves import range
  from six.moves import zip
  import tensorflow
  import tensorflow_probability
label_desc: Use image node_id 4 for calling a global function with example usage: mushroom_reward_distribution(r_noeat, r_eat_safe, r_eat_poison_bad, r_eat_poison_good, prob_poison_bad) and returns: tfd