def choose_best_action(observation, planner_index):
      """Choose the best action, update best Monte Carlo values."""
      best_mc_values = self._best_mc_values[planner_index]
      action_probs = self._rollout_agent.action_distribution(
          np.array([observation] * self._rollout_agent.batch_size)
      )[0, :]
      sums_and_counts = run_batches_from(observation, planner_index)

      def monte_carlo_value(action):
        (value_sum, count) = sums_and_counts[action]
        if count > 0:
          mean_value = value_sum / count
        else:
          mean_value = -np.inf
        return mean_value

      mc_values = np.array(
          [monte_carlo_value(action) for action in range(self.action_space.n)]
      )
      best_mc_values.append(mc_values.max())

      normalizer = max(
          np.std(best_mc_values[-self._normalizer_window_size:]),
          self._normalizer_epsilon
      )
      normalized_mc_values = mc_values / normalizer

      uct_bonuses = np.array(
          [self._uct_bonus(sums_and_counts[action][1], action_probs[action])
           for action in range(self.action_space.n)]
      )
      values = normalized_mc_values + uct_bonuses
      return np.argmax(values)
Choose the best action, update best Monte Carlo values.
tensor2tensor/rl/rl_utils.py