ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more

tf_agents.specs.bandit_spec_utils.create_per_arm_observation_spec

Creates an observation spec with per-arm features and possibly action mask.

global_dim (int) The global feature dimension.
per_arm_dim (int) The per-arm feature dimension.
max_num_actions If specified (int), this is the maximum number of actions in any sample, and the num_actions dimension of the per-arm features will be set to this number. The actual number of actions for a given sample can be lower than this parameter: it can be specified via the NUM_ACTIONS_FEATURE_KEY, or an action mask.
add_num_actions_feature (bool) whether to use the num_actions feature key to encode the number of actions per sample.

A nested structure of observation spec.