Sample experience from replay buffer and return a batch

def sample(self, batch_info: BatchInfo, model: RlModel, number_of_steps: int) -> Rollout:
        """ Sample experience from replay buffer and return a batch """
        # Sample trajectories
        rollout = self.replay_buffer.sample_trajectories(rollout_length=number_of_steps, batch_info=batch_info)

        last_observations = rollout.transition_tensors['observations_next'][-1].to(self.device)
        final_values = model.value(last_observations).cpu()

        # Add 'final_values' to the rollout
        rollout.rollout_tensors['final_values'] = final_values

        return rollout