docstr: Given avaialble actions(bandit arms) and features of these actions, the agent chooses the best action.
  Args:
      features_all_arms: 3D tensor of shape (batch_size, num_arms_per_episode, feature_dim)
  What this method does:
      1) calculate the UCB score for each arm
      2) choose the best arm for each arm by finding max UCB
docstr_tok: ["Given","avaialble","actions","(","bandit","arms",")","and","features","of","these","actions",",","the","agent","chooses","the","best","action",".","Args",":","features_all_arms",":","3D","tensor","of","shape","(","batch_size",",","num_arms_per_episode",",","feature_dim",")","What","this","method","does",":","1",")","calculate","the","UCB","score","for","each","arm","2",")","choose","the","best","arm","for","each","arm","by","finding","max","UCB"]
code: def choose_action(
      self, features_all_arms
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """
      Given avaialble actions(bandit arms) and features of these actions, the agent chooses the best action.
      Args:
          features_all_arms: 3D tensor of shape (batch_size, num_arms_per_episode, feature_dim)
      What this method does:
          1) calculate the UCB score for each arm
          2) choose the best arm for each arm by finding max UCB
      """
      actor_output = self.policy.act(obs=features_all_arms)
      chosen_action = torch.argmax(actor_output.action, dim=1)
      log_prob = actor_output.log_prob
      return torch.unsqueeze(chosen_action, 1), log_prob
  
imports: import random
  from dataclasses import replace
  from typing import Any, Optional, Tuple
  import reagent.core.types
  import torch
  from reagent.core.types import CBInput
  from reagent.gym.policies.policy import Policy
  from reagent.gym.policies.samplers.discrete_sampler import GreedyActionSampler
  from reagent.gym.types import Sampler, Scorer
  from reagent.models.linear_regression import LinearRegressionUCB
  from reagent.training.cb.base_trainer import BaseCBTrainerWithEval
  from reagent.training.cb.linucb_trainer import LinUCBTrainer
  from reagent.training.parameters import LinUCBTrainerParameters
label_desc: Use image node_id 4 for calling the DynamicBanditAgent obj's underlying member method code with example usage: obj.choose_action(features_all_arms) and returns: torch,  log_prob