code: def testSampleTransitionBatchExtra(self):
      replay_capacity = 10
      memory = circular_replay_buffer.ReplayBuffer(
          stack_size=1, replay_capacity=replay_capacity, batch_size=2
      )
      num_adds = 50  # The number of transitions to add to the memory.
      for i in range(num_adds):
          memory.add(
              observation=np.full(
                  OBSERVATION_SHAPE, i, dtype=OBS_DTYPE
              ),
              action=0,
              reward=0,
              terminal=i % 4,
              extra1=i % 2,
              extra2=[i % 2, 0],
          )  # Every 4 transitions is terminal.
      # Test sampling with default batch size.
      for _i in range(1000):
          batch = memory.sample_transition_batch()
          self.assertEqual(batch[0].shape[0], 2)
      # Test changing batch sizes.
      for _i in range(1000):
          batch = memory.sample_transition_batch(BATCH_SIZE)
          self.assertEqual(batch[0].shape[0], BATCH_SIZE)
      # Verify we revert to default batch size.
      for _i in range(1000):
          batch = memory.sample_transition_batch()
          self.assertEqual(batch[0].shape[0], 2)
  
      # Verify we can specify what indices to sample.
      indices = [1, 2, 3, 5, 8]
      expected_states = np.array(
          [
              np.full(OBSERVATION_SHAPE, i, dtype=OBS_DTYPE)
              for i in indices
          ]
      )
      expected_next_states = (expected_states + 1) % replay_capacity
      # Because the replay buffer is circular, we can exactly compute what the
      # states will be at the specified indices by doing a little mod math:
      expected_states += num_adds - replay_capacity
      expected_next_states += num_adds - replay_capacity
      # This is replicating the formula that was used above to determine what
      # transitions are terminal when adding observation (i % 4).
      expected_terminal = np.expand_dims(
          np.array(
              [
                  min((x + num_adds - replay_capacity) % 4, 1)
                  for x in indices
              ]
          ),
          1,
      ).astype(bool)
      expected_extra1 = np.expand_dims(
          np.array(
              [(x + num_adds - replay_capacity) % 2 for x in indices]
          ),
          1,
      )
      expected_next_extra1 = np.expand_dims(
          np.array(
              [
                  (x + 1 + num_adds - replay_capacity) % 2
                  for x in indices
              ]
          ),
          1,
      )
      expected_extra2 = np.stack(
          [
              [(x + num_adds - replay_capacity) % 2 for x in indices],
              np.zeros((len(indices),)),
          ],
          axis=1,
      )
      expected_next_extra2 = np.stack(
          [
              [
                  (x + 1 + num_adds - replay_capacity) % 2
                  for x in indices
              ],
              np.zeros((len(indices),)),
          ],
          axis=1,
      )
      batch = memory.sample_transition_batch(
          batch_size=len(indices), indices=torch.tensor(indices)
      )
      npt.assert_array_equal(batch.state, expected_states)
      npt.assert_array_equal(batch.action, np.zeros((len(indices), 1)))
      npt.assert_array_equal(batch.reward, np.zeros((len(indices), 1)))
      npt.assert_array_equal(
          batch.next_action, np.zeros((len(indices), 1))
      )
      npt.assert_array_equal(
          batch.next_reward, np.zeros((len(indices), 1))
      )
      npt.assert_array_equal(batch.next_state, expected_next_states)
      npt.assert_array_equal(batch.terminal, expected_terminal)
      npt.assert_array_equal(
          batch.indices, np.expand_dims(np.array(indices), 1)
      )
      npt.assert_array_equal(batch.extra1, expected_extra1)
      npt.assert_array_equal(batch.next_extra1, expected_next_extra1)
      npt.assert_array_equal(batch.extra2, expected_extra2)
      npt.assert_array_equal(batch.next_extra2, expected_next_extra2)
  
imports: import tempfile
  import unittest
  import numpy
  import numpy.testing
  import torch
  from reagent.replay_memory import circular_replay_buffer
label_desc: Use image node_id 9 for calling the ReplayBufferTest obj's underlying member method code with example usage: obj.testSampleTransitionBatchExtra() without return types