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