Action values for selected actions in the rollout

def model_action_q_dist(self):
        """ Action values for selected actions in the rollout """
        q = self.get('model:q_dist')
        actions = self.get('rollout:actions')
        return q[range(q.size(0)), actions]