def get_model_index(arm_index: int, accepts_per_arm_features: bool) -> int:
  """Returns the model index for a specific arm.

  The number of models depends on the observation format: If the policy accepts
  per-arm features, there is only one single model used for every arm. Otherwise
  there is a model for every arm.

  Args:
    arm_index: The index of the arm for which the model index is needed.
    accepts_per_arm_features: (bool) Whether the policy works with per-arm
      features.

  Returns:
    The index of the model for the arm requested.
  """
  return 0 if accepts_per_arm_features else arm_index