Return indexes of next sample

def sample_batch_transitions(self, batch_size, forward_steps=1):
        """ Return indexes of next sample"""
        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)