code: def populate_policy_info(
      arm_observations: types.Tensor,
      chosen_actions: types.Tensor,
      rewards_for_argmax: types.Tensor,
      est_rewards: types.Tensor,
      emit_policy_info: Sequence[Text],
      accepts_per_arm_features: bool,
  ) -> PolicyInfo:
      """Populates policy info given all needed input.
  
      Args:
        arm_observations: In case the policy accepts per-arm feautures, this is a
          Tensor with the per-arm features. Otherwise its value is unused.
        chosen_actions: A Tensor with the indices of the chosen actions.
        rewards_for_argmax: The sampled or optimistically boosted reward estimates
          based on which the policy chooses the action greedily.
        est_rewards: A Tensor with the rewards estimated by the model.
        emit_policy_info: A set of policy info keys, specifying wich info fields to
          populate
        accepts_per_arm_features: (bool) Whether the policy accepts per-arm
          features.
  
      Returns:
        A policy info.
      """
      if accepts_per_arm_features:
          # Saving the features for the chosen action to the policy_info.
          chosen_arm_features = tf.nest.map_structure(
              lambda t: tf.gather(
                  params=t, indices=chosen_actions, batch_dims=1
              ),
              arm_observations,
          )
          policy_info = PerArmPolicyInfo(
              predicted_rewards_optimistic=(
                  rewards_for_argmax
                  if InfoFields.PREDICTED_REWARDS_OPTIMISTIC
                  in emit_policy_info
                  else ()
              ),
              predicted_rewards_sampled=(
                  rewards_for_argmax
                  if InfoFields.PREDICTED_REWARDS_SAMPLED
                  in emit_policy_info
                  else ()
              ),
              predicted_rewards_mean=(
                  est_rewards
                  if InfoFields.PREDICTED_REWARDS_MEAN
                  in emit_policy_info
                  else ()
              ),
              chosen_arm_features=chosen_arm_features,
          )
      else:
          policy_info = PolicyInfo(
              predicted_rewards_optimistic=(
                  rewards_for_argmax
                  if InfoFields.PREDICTED_REWARDS_OPTIMISTIC
                  in emit_policy_info
                  else ()
              ),
              predicted_rewards_sampled=(
                  rewards_for_argmax
                  if InfoFields.PREDICTED_REWARDS_SAMPLED
                  in emit_policy_info
                  else ()
              ),
              predicted_rewards_mean=(
                  est_rewards
                  if InfoFields.PREDICTED_REWARDS_MEAN
                  in emit_policy_info
                  else ()
              ),
          )
      return policy_info
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  from typing import Optional, Sequence, Text, Union
  import tensorflow
  from tf_agents.specs import bandit_spec_utils
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.typing import types
  from tf_agents.utils import common
label_desc: Use image node_id 2 for calling a global function with example usage: populate_policy_info(arm_observations, chosen_actions, rewards_for_argmax, est_rewards, emit_policy_info, accepts_per_arm_features) and returns: policy_info