docstr: Test that rewards are being returned correctly for a simple case.
docstr_tok: ["Test","that","rewards","are","being","returned","correctly","for","a","simple","case","."]
code: def testReturnsCorrectRewards(self):
      """Test that rewards are being returned correctly for a simple case."""
      # Reward of 1 is given if action == (context % 3)
      context = tf.reshape(tf.range(128), shape=[128, 1])
      labels = tf.math.mod(context, 3)
      batch_size = 32
      dataset = (
          tf.data.Dataset.from_tensor_slices((context, labels))
          .repeat()
          .shuffle(4 * batch_size)
      )
      reward_distribution = deterministic_reward_distribution(tf.eye(3))
      env = ce.ClassificationBanditEnvironment(
          dataset, reward_distribution, batch_size
      )
      self.evaluate(tf.compat.v1.global_variables_initializer())
      for _ in range(10):
          # Take the 'correct' action
          observation = env.reset().observation
          action = tf.math.mod(observation, 3)
          reward = env.step(action).reward
          np.testing.assert_almost_equal(
              self.evaluate(reward), self.evaluate(tf.ones_like(reward))
          )
  
      for _ in range(10):
          # Take the 'incorrect' action
          observation = env.reset().observation
          action = tf.math.mod(observation + 1, 3)
          reward = env.step(action).reward
          np.testing.assert_almost_equal(
              self.evaluate(reward),
              self.evaluate(tf.zeros_like(reward)),
          )
  
file_path: agents/tf_agents/bandits/environments/classification_environment_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  from absl.testing.absltest import mock
  import numpy
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.environments import classification_environment
label_desc: Use image node_id 3 for calling the ClassificationEnvironmentTest obj's underlying member method code with example usage: obj.testReturnsCorrectRewards() without return types