tf_agents.bandits.agents.utils.process_experience_for_neural_agents

Processes the experience and prepares it for the network of the agent.

tf_agents.bandits.agents.utils.process_experience_for_neural_agents(
    experience: tf_agents.typing.types.NestedTensor,
    accepts_per_arm_features: bool,
    training_data_spec: tf_agents.typing.types.NestedTensorSpec
) -> Tuple[tf_agents.typing.types.NestedTensor, tf_agents.typing.types.Tensor, tf_agents.typing.types.Tensor]

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.

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.

A tuple of (observation, action, reward) tensors to be consumed by the train function of the neural agent.