Add support for specifying loss type in ReAgent Neural LinUCB

Summary: Add support for specifying which loss function to use in Neural LinUCB. MSE, MAE and binary cross-entropy (BCE) are supported. The default is MSE

Reviewed By: BerenLuthien

Differential Revision: D50992907

fbshipit-source-id: 997c08a39c487ac68c62dea561e4b0d43d0e821f

Author: Alex Nikulkov

reagent/test/training/cb/test_deep_represent_linucb.py

@@ -23,7 +23,9 @@ class TestDeepRepresentLinUCB(unittest.TestCase):
 
     def setUp(self):
 
-        self.params = DeepRepresentLinUCBTrainerParameters(lr=1e-1)
+        self.params = DeepRepresentLinUCBTrainerParameters(
+            lr=1e-1, loss_type="cross_entropy"
+        )
 
         input_dim = 100
         sizes = [20]
@@ -43,14 +45,15 @@ def setUp(self):
             sizes=sizes + [linucb_inp_dim],
             activations=activations,
             mlp_layers=customized_layers,
+            output_activation="sigmoid",
         )
 
         self.policy = Policy(scorer=policy_network, sampler=GreedyActionSampler())
         self.trainer = DeepRepresentLinUCBTrainer(self.policy, **self.params.asdict())
         self.batch = CBInput(
             context_arm_features=torch.rand(2, 2, input_dim),
             action=torch.tensor([[0], [1]], dtype=torch.long),
-            reward=torch.tensor([[1.5], [-2.3]]),
+            reward=torch.tensor([[0.3], [0.1]]),
         )  # random Gaussian features
 
     def test_linucb_training_step(self):

reagent/training/cb/deep_represent_linucb_trainer.py

@@ -8,6 +8,7 @@
 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
 
 logger = logging.getLogger(__name__)
 
@@ -29,6 +30,7 @@ def __init__(
         policy: Policy,
         lr: float = 1e-3,
         weight_decay: float = 0.0,
+        loss_type: str = "mse",  # one of the LOSS_TYPES names
         **kwargs,
     ):
         super().__init__(
@@ -40,7 +42,7 @@ def __init__(
             policy.scorer, DeepRepresentLinearRegressionUCB
         ), "Trainer requires the policy scorer to be DeepRepresentLinearRegressionUCB"
         self.scorer = policy.scorer
-        self.loss_fn = torch.nn.functional.mse_loss
+        self.loss_fn = LOSS_TYPES[loss_type]
         self.lr = lr
         self.weight_decay = weight_decay
 

reagent/training/cb/supervised_trainer.py

@@ -28,7 +28,7 @@ class SupervisedTrainer(BaseCBTrainerWithEval):
     def __init__(
         self,
         policy: Policy,
-        loss_type: str = "mse",  # one of the LossTypes names
+        loss_type: str = "mse",  # one of the LOSS_TYPES names
         lr: float = 1e-3,
         weight_decay: float = 0.0,
         *args,