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","}"]
code: def __init__(
      self,
      ucb_alpha: float = 1.0,
      min_explore_steps: int = 1,
      default_score: float = float("inf"),
      estimate_variance: bool = True,
      min_variance: float = 0.0,
      input_dim: int = 1,  # ununsed argument, keeping it for compatibility with parent class
  ):
      super().__init__()
      self.num_obs_total = 0
      self.obs_count_by_arm_id = {}
      self.sum_reward_by_arm_id = {}
      self.sum_squared_reward_by_arm_id = {}
      self.ucb_alpha = ucb_alpha
      self.min_explore_steps = min_explore_steps
      self.default_score = default_score
      self.estimate_variance = estimate_variance
      self.min_variance = min_variance
  
      # add a dummy parameter so that DDP doesn't compain about lack of parameters with gradient
      self.dummy_param = torch.nn.parameter.Parameter(torch.zeros(1))
  
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 1 to create a new MABBaseModel object from inherited base classes: torch.nn with example: obj = MABBaseModel(ucb_alpha, min_explore_steps, default_score, estimate_variance, min_variance, input_dim)