code: def process_experience_for_neural_agents(
      experience: types.NestedTensor,
      accepts_per_arm_features: bool,
      training_data_spec: types.NestedTensorSpec,
  ) -> Tuple[types.NestedTensor, types.Tensor, types.Tensor]:
      """Processes the experience and prepares it for the network of the agent.
  
      First the reward, the action, and the observation are flattened to have only
      one batch dimension. Then, if the experience includes chosen action features
      in the policy info, it gets copied in place of the per-arm observation.
  
      Args:
        experience: The experience coming from the replay buffer.
        accepts_per_arm_features: Whether the agent accepts per-arm features.
        training_data_spec: The data spec describing what the agent expects.
  
      Returns:
        A tuple of (observation, action, reward) tensors to be consumed by the train
          function of the neural agent.
      """
      (
          flattened_experience,
          _,
      ) = nest_utils.flatten_multi_batched_nested_tensors(
          experience, training_data_spec
      )
  
      observation = flattened_experience.observation
      action = flattened_experience.action
      reward = flattened_experience.reward
  
      if not accepts_per_arm_features:
          return observation, action, reward
  
      # The arm observation we train on needs to be copied from the respective
      # policy info field to the per arm observation field. Pretending there was
      # only one action, we fill the action field with zeros.
      chosen_arm_features = (
          flattened_experience.policy_info.chosen_arm_features
      )
      observation[
          bandit_spec_utils.PER_ARM_FEATURE_KEY
      ] = tf.nest.map_structure(
          lambda t: tf.expand_dims(t, axis=1), chosen_arm_features
      )
      action = tf.zeros_like(action)
      if bandit_spec_utils.NUM_ACTIONS_FEATURE_KEY in observation:
          # This change is not crucial but since in training there will be only one
          # action per sample, it's good to follow the convention that the feature
          # value for `num_actions` be less than or equal to the maximum available
          # number of actions.
          observation[
              bandit_spec_utils.NUM_ACTIONS_FEATURE_KEY
          ] = tf.ones_like(
              observation[bandit_spec_utils.NUM_ACTIONS_FEATURE_KEY]
          )
  
      return observation, action, reward
  
file_path: agents/tf_agents/bandits/agents/utils.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Tuple
  import gin
  import numpy
  import tensorflow
  from tf_agents.bandits.specs import utils
  from tf_agents.policies import utils
  from tf_agents.typing import types
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 5 for calling a global function with example usage: process_experience_for_neural_agents(experience, accepts_per_arm_features, training_data_spec) and returns: observation,  action,  reward, observation,  action,  reward