code: def testRewardDistribution(self):
      reward_distr = dataset_utilities.mushroom_reward_distribution(
          r_noeat=0.0,
          r_eat_safe=5.0,
          r_eat_poison_bad=-35.0,
          r_eat_poison_good=5.0,
          prob_poison_bad=0.5,
      )
      self.assertAllEqual(reward_distr.mean(), [[0, -15.0], [0, 5.0]])
  
file_path: agents/tf_agents/bandits/environments/dataset_utilities_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  import tensorflow
  from tf_agents.bandits.environments import dataset_utilities
label_desc: Use image node_id 2 for calling the DatasetUtilitiesTest obj's underlying member method code with example usage: obj.testRewardDistribution() without return types