code: def test_reward_info(self):
      env = deepmind_lab.Lab(
          "tests/entity_info_test",
          ["DEBUG.PICKUPS"],
          config={"fps": "60", "width": "32", "height": "32"},
      )
  
      action_spec = env.action_spec()
      action_index = {
          action["name"]: i for i, action in enumerate(action_spec)
      }
  
      env.reset()
      obs = env.observations()
      pickups_state = np.array(
          [
              [150.0, 50.0, 16.125, 1.0, 1.0],
              [250.0, 50.0, 16.125, 1.0, 1.0],
              [350.0, 50.0, 16.125, 1.0, 1.0],
              [450.0, 50.0, 16.125, 1.0, -1.0],
              [550.0, 50.0, 16.125, 1.0, -1.0],
              [650.0, 50.0, 16.125, 1.0, -1.0],
          ]
      )
      self.assertTrue(np.allclose(obs["DEBUG.PICKUPS"], pickups_state))
  
      action = np.zeros([len(action_spec)], dtype=np.intc)
      action[action_index["MOVE_BACK_FORWARD"]] = 1
      reward = 0
      for _ in six.moves.range(0, 60):
          reward += env.step(action, 1)
          if reward == 3:
              break
      else:
          self.fail("Failed to all positive rewards.")
  
      pickups_state = np.array(
          [
              [150.0, 50.0, 16.125, 0.0, 1.0],
              [250.0, 50.0, 16.125, 0.0, 1.0],
              [350.0, 50.0, 16.125, 0.0, 1.0],
              [450.0, 50.0, 16.125, 1.0, -1.0],
              [550.0, 50.0, 16.125, 1.0, -1.0],
              [650.0, 50.0, 16.125, 1.0, -1.0],
          ]
      )
      obs = env.observations()
      self.assertTrue(np.allclose(obs["DEBUG.PICKUPS"], pickups_state))
      for _ in six.moves.range(0, 600):
          reward += env.step(action, 1)
          if reward == 0:
              break
      else:
          self.fail("Failed to pickup last negative reward")
  
      pickups_state = np.array(
          [
              [150.0, 50.0, 16.125, 0.0, 1.0],
              [250.0, 50.0, 16.125, 0.0, 1.0],
              [350.0, 50.0, 16.125, 0.0, 1.0],
              [450.0, 50.0, 16.125, 0.0, -1.0],
              [550.0, 50.0, 16.125, 0.0, -1.0],
              [650.0, 50.0, 16.125, 0.0, -1.0],
          ]
      )
  
      obs = env.observations()
      self.assertTrue(np.allclose(obs["DEBUG.PICKUPS"], pickups_state))
  
file_path: lab/python/tests/entity_info_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  from absl.testing import absltest
  import numpy
  import six
  import deepmind_lab
label_desc: Use image node_id 1 for calling the EntityInfo obj's underlying member method code with example usage: obj.test_reward_info() without return types