class_docstr: The trainer for a Contextual Bandit model, where deep represent layer serves as feature processor,
  and then processed features are fed to LinUCB layer to produce UCB score.
  This is extension of LinUCBTrainer. More details refer to docstring of LinUCBTrainer.
  
  Reference:
  - LinUCB : https://arxiv.org/pdf/1003.0146.pdf
  - DeepRepresentLinUCB : https://arxiv.org/pdf/2012.01780.pdf
class_docstr_tok: ["The","trainer","for","a","Contextual","Bandit","model",",","where","deep","represent","layer","serves","as","feature","processor",",","and","then","processed","features","are","fed","to","LinUCB","layer","to","produce","UCB","score",".","This","is","extension","of","LinUCBTrainer",".","More","details","refer","to","docstring","of","LinUCBTrainer",".","Reference",":","-","LinUCB",":","https",":","\/\/arxiv.org\/pdf\/1003.0146.pdf","-","DeepRepresentLinUCB",":","https",":","\/\/arxiv.org\/pdf\/2012.01780.pdf"]
docstr: Here scorer (DeepRepresentLinearRegressionUCB) outputs Dict
      {
          "pred_label": pred_label,
          "pred_sigma": pred_sigma,
          "ucb": ucb,
          "mlp_out_with_ones": mlp_out_with_ones,
      }
  The pred_label is useful for calculating MSE loss.
  The mlp_out_with_ones is useful for updating LinUCB parameters (A, b, etc)
docstr_tok: ["Here","scorer","(","DeepRepresentLinearRegressionUCB",")","outputs","Dict","{","``","pred_label","''",":","pred_label",",","``","pred_sigma","''",":","pred_sigma",",","``","ucb","''",":","ucb",",","``","mlp_out_with_ones","''",":","mlp_out_with_ones",",","}","The","pred_label","is","useful","for","calculating","MSE","loss",".","The","mlp_out_with_ones","is","useful","for","updating","LinUCB","parameters","(","A",",","b",",","etc",")"]
code: def cb_training_step(
      self, batch: CBInput, batch_idx: int, optimizer_idx: int = 0
  ) -> torch.Tensor:
      """
      Here scorer (DeepRepresentLinearRegressionUCB) outputs Dict
          {
              "pred_label": pred_label,
              "pred_sigma": pred_sigma,
              "ucb": ucb,
              "mlp_out_with_ones": mlp_out_with_ones,
          }
      The pred_label is useful for calculating MSE loss.
      The mlp_out_with_ones is useful for updating LinUCB parameters (A, b, etc)
      """
      model_output = self.scorer(
          inp=batch.features_of_chosen_arm
      )  # noqa
      # this calls scorer.forward() so as to update pred_u, and to grad descent on deep_represent module
  
      pred_label, mlp_out_with_ones = (
          model_output["pred_label"],
          model_output["mlp_out_with_ones"],
      )
  
      assert batch.label is not None  # to satisfy Pyre
      label = batch.label.squeeze(-1)
      assert (
          pred_label.shape == label.shape
      ), f"Shapes of model prediction {pred_label.shape} and label {label.shape} have to match"
      # compute the NN loss
      # weighted average loss
      losses = self.loss_fn(pred_label, label, reduction="none")
      weight = batch.effective_weight
      loss = (losses * weight.squeeze(-1)).sum() / losses.shape[0]
  
      # update LinUCB parameters
      self.update_params(
          mlp_out_with_ones.detach(), batch.label, weight
      )
  
      return loss
  
file_path: ReAgent/reagent/training/cb/deep_represent_linucb_trainer.py
imports: import logging
  import torch
  from reagent.core.configuration import resolve_defaults
  from reagent.core.types import CBInput
  from reagent.gym.policies.policy import Policy
  from reagent.models.deep_represent_linucb import DeepRepresentLinearRegressionUCB
  from reagent.training.cb.linucb_trainer import LinUCBTrainer
  from reagent.training.cb.supervised_trainer import LOSS_TYPES
label_desc: Use image node_id 3 for calling the DeepRepresentLinUCBTrainer obj's underlying member method code with example usage: obj.cb_training_step(batch, batch_idx, optimizer_idx) and returns: loss