class_docstr: A base class for Multi-Armed Bandits (MAB). These are non-contextual bandit models.
  IMPORTANT USAGE GUIDILINES:
  1. Use manual optimization as there is no loss function to optimize.
  2. Distributed training is not supported yet, so only run this with a single process.
      Distributed training could be added if needed using `torch.distributed.all_gather_object()`
  3. A minimal implementation of a specific MAB algorithm needs to implement only `get_score_single_arm_id()` method
  4. The base class keeps track of: (1) number of observations per arm (2) sum of rewards per arm (3) sum of squared rewards per arm.
      If a specific MAB algorithm wants to keep track of other metrics, it needs to override the `learn()` method.
  5. This model is intended to be used only for Offline Evaluation (e.g. hyperparameter search, algorithm comparison).
      DO NOT TRY TO SERVE THIS MODEL IN PRODUCTION.
  
  Args:
      ucb_alpha: The coefficient on the standard deviation in UCB formula.
          Set it to 0 to predict the expected value instead of UCB.
      min_explore_steps: Minimum number of observations (per arm) necessary before we
          start returning non-default UCB scores.
      default_score: The default score returned instead of UCB score during
          first `min_explore_steps` for each arm. The detault is `inf`, which should prioritize
          this arm ahead of any arms that already have more than `min_explore_steps` observations.
      estimate_variance: Whether to estimate reward variance using sample variance or not. If not,
          default value of 1.0 is used
      min_variance: If estimating variance, this value is applied to the estimate as the minimum.
  
  Outputs:
  Dict {"pred_label": pred_label, "pred_sigma": pred_sigma, "ucb": ucb}
class_docstr_tok: ["A","base","class","for","Multi-Armed","Bandits","(","MAB",")",".","These","are","non-contextual","bandit","models",".","IMPORTANT","USAGE","GUIDILINES",":","1",".","Use","manual","optimization","as","there","is","no","loss","function","to","optimize",".","2",".","Distributed","training","is","not","supported","yet",",","so","only","run","this","with","a","single","process",".","Distributed","training","could","be","added","if","needed","using","`","torch.distributed.all_gather_object","(",")","`","3",".","A","minimal","implementation","of","a","specific","MAB","algorithm","needs","to","implement","only","`","get_score_single_arm_id","(",")","`","method","4",".","The","base","class","keeps","track","of",":","(","1",")","number","of","observations","per","arm","(","2",")","sum","of","rewards","per","arm","(","3",")","sum","of","squared","rewards","per","arm",".","If","a","specific","MAB","algorithm","wants","to","keep","track","of","other","metrics",",","it","needs","to","override","the","`","learn","(",")","`","method",".","5",".","This","model","is","intended","to","be","used","only","for","Offline","Evaluation","(","e.g",".","hyperparameter","search",",","algorithm","comparison",")",".","DO","NOT","TRY","TO","SERVE","THIS","MODEL","IN","PRODUCTION",".","Args",":","ucb_alpha",":","The","coefficient","on","the","standard","deviation","in","UCB","formula",".","Set","it","to","0","to","predict","the","expected","value","instead","of","UCB",".","min_explore_steps",":","Minimum","number","of","observations","(","per","arm",")","necessary","before","we","start","returning","non-default","UCB","scores",".","default_score",":","The","default","score","returned","instead","of","UCB","score","during","first","`","min_explore_steps","`","for","each","arm",".","The","detault","is","`","inf","`",",","which","should","prioritize","this","arm","ahead","of","any","arms","that","already","have","more","than","`","min_explore_steps","`","observations",".","estimate_variance",":","Whether","to","estimate","reward","variance","using","sample","variance","or","not",".","If","not",",","default","value","of","1.0","is","used","min_variance",":","If","estimating","variance",",","this","value","is","applied","to","the","estimate","as","the","minimum",".","Outputs",":","Dict","{","``","pred_label","''",":","pred_label",",","``","pred_sigma","''",":","pred_sigma",",","``","ucb","''",":","ucb","}"]
docstr: Get the scores (predicted mu, sigma and UCB) for a batch of arm_ids.
docstr_tok: ["Get","the","scores","(","predicted","mu",",","sigma","and","UCB",")","for","a","batch","of","arm_ids","."]
code: def forward(
      self, arm_ids: torch.Tensor, ucb_alpha: Optional[float] = None
  ) -> Dict[str, torch.Tensor]:
      """
      Get the scores (predicted mu, sigma and UCB) for a batch of arm_ids.
      """
      start = time.time()
      if ucb_alpha is None:
          ucb_alpha = self.ucb_alpha
      all_scores = [
          [
              self.get_score_single_arm_id_if_enough_data(
                  arm_id, ucb_alpha
              )
              for arm_id in single_obs_arm_ids.tolist()
          ]
          for single_obs_arm_ids in arm_ids
      ]
      check1 = time.time()
      logger.info(f"Check1 forward time: {check1 - start}")
      ret = {
          k: torch.tensor(
              [
                  [
                      single_arm_scores[k]
                      for single_arm_scores in slate_score
                  ]
                  for slate_score in all_scores
              ]
          )
          for k in ["pred_label", "pred_sigma", "ucb"]
      }
      end = time.time()
      logger.info(f"Batch forward time: {end - start}")
      # from reagent.core.utils import ForkedPdb
      # ForkedPdb().set_trace()
      return ret
  
file_path: ReAgent/reagent/models/mab.py
imports: import logging
  import math
  import time
  from abc import abstractmethod
  from typing import Dict, Optional
  import torch
  from reagent.core.types import CBInput
label_desc: Use image node_id 6 for calling the MABBaseModel obj's underlying member method code with example usage: obj.forward(arm_ids, ucb_alpha) and returns: ret