def METHOD_NAME(
    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