def compute_expectation(self, beta): return F.sum(F.softmax(beta * self.q_values) * self.q_values, axis=1)