class_docstr: Base class to build policies based on reward predictions.
class_docstr_tok: ["Base","class","to","build","policies","based","on","reward","predictions","."]
docstr: Extracts and saves the chosen arm features in the policy info.
  
  If the policy accepts arm features, this method extracts the arm features
  from `time_step.observation` corresponding to the input `action` tensor,
  saves it in the policy info of the input `step` and returns the modified
  step. Otherwise, the method returns the input `step`.
  
  Args:
    time_step: A `TimeStep` tuple corresponding to `time_step_spec()`.
    action: A `types.Tensor` of chosen actions.
    step: A `PolicyStep`.
  
  Returns:
    A `PolicyStep` resulting from saving the chosen arm features in the policy
    info of the input `step` if the policy accepts arm features, or the input
    `step` otherwise.
docstr_tok: ["Extracts","and","saves","the","chosen","arm","features","in","the","policy","info",".","If","the","policy","accepts","arm","features",",","this","method","extracts","the","arm","features","from","`","time_step.observation","`","corresponding","to","the","input","`","action","`","tensor",",","saves","it","in","the","policy","info","of","the","input","`","step","`","and","returns","the","modified","step",".","Otherwise",",","the","method","returns","the","input","`","step","`",".","Args",":","time_step",":","A","`","TimeStep","`","tuple","corresponding","to","`","time_step_spec","(",")","`",".","action",":","A","`","types.Tensor","`","of","chosen","actions",".","step",":","A","`","PolicyStep","`",".","Returns",":","A","`","PolicyStep","`","resulting","from","saving","the","chosen","arm","features","in","the","policy","info","of","the","input","`","step","`","if","the","policy","accepts","arm","features",",","or","the","input","`","step","`","otherwise","."]
code: def _maybe_save_chosen_arm_features(
      self,
      time_step: ts.TimeStep,
      action: types.Tensor,
      step: policy_step.PolicyStep,
  ) -> policy_step.PolicyStep:
      """Extracts and saves the chosen arm features in the policy info.
  
      If the policy accepts arm features, this method extracts the arm features
      from `time_step.observation` corresponding to the input `action` tensor,
      saves it in the policy info of the input `step` and returns the modified
      step. Otherwise, the method returns the input `step`.
  
      Args:
        time_step: A `TimeStep` tuple corresponding to `time_step_spec()`.
        action: A `types.Tensor` of chosen actions.
        step: A `PolicyStep`.
  
      Returns:
        A `PolicyStep` resulting from saving the chosen arm features in the policy
        info of the input `step` if the policy accepts arm features, or the input
        `step` otherwise.
      """
      if self.accepts_per_arm_features:
          # Saving the features for the chosen action to the policy_info.
          chosen_arm_features = tf.nest.map_structure(
              lambda obs: tf.gather(
                  params=obs, indices=action, batch_dims=1
              ),
              time_step.observation[
                  bandit_spec_utils.PER_ARM_FEATURE_KEY
              ],
          )
          step = step._replace(
              info=step.info._replace(
                  chosen_arm_features=chosen_arm_features
              )
          )
      return step
  
imports: import abc
  from typing import Iterable, Optional, Text, Tuple
  import gin
  import six
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.networks import heteroscedastic_q_network
  from tf_agents.bandits.policies import constraints
  from tf_agents.bandits.specs import utils
  from tf_agents.policies import tf_policy
  from tf_agents.policies import utils
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
label_desc: Use image node_id 5 for calling the RewardPredictionBasePolicy obj's underlying member method code with example usage: obj._maybe_save_chosen_arm_features(time_step, action, step) and returns: step