code: def sample_batch_transitions(self, batch_size, forward_steps=1):
          
          results = []
  
          for i in range(self.num_envs):
              results.append(self.sample_frame_single_env(batch_size, forward_steps=forward_steps))
  
          return np.stack(results, axis=-1)
docs: Return indexes of next sample