Allow for publishing of reward network in discrete CRR

Summary: Allow for publishing of reward network in discrete_crr.py

Differential Revision: D32711991

fbshipit-source-id: 7959d630ddce1e3c57688d9b449ffbeef7b9d8be

Author: David Vengerov

reagent/model_managers/discrete/discrete_crr.py

@@ -15,6 +15,12 @@
     param_hash,
 )
 from reagent.evaluation.evaluator import get_metrics_to_score
+from reagent.fb.prediction.cfeval.bandit_reward_network_predictor import (
+    FbBanditRewardNetPredictorWrapper,
+)
+from reagent.fb.prediction.cfeval.bandit_reward_network_predictor import (
+    FbBanditRewardNetPredictorWrapper,
+)
 from reagent.gym.policies.policy import Policy
 from reagent.gym.policies.predictor_policies import create_predictor_policy_from_model
 from reagent.model_managers.discrete_dqn_base import DiscreteDQNBase
@@ -196,7 +202,7 @@ def get_reporter(self):
     # in utils.py
 
     def serving_module_names(self):
-        module_names = ["default_model", "dqn", "actor_dqn"]
+        module_names = ["default_model", "dqn", "actor_dqn", "reward"]
         if len(self.action_names) == 2:
             module_names.append("binary_difference_scorer")
         return module_names
@@ -219,6 +225,7 @@ def build_serving_modules(
             "dqn": self._build_dqn_module(
                 trainer_module.q1_network, normalization_data_map
             ),
+            "reward": self.build_reward_module(trainer_module, normalization_data_map),
             "actor_dqn": self._build_dqn_module(
                 ActorDQN(trainer_module.actor_network), normalization_data_map
             ),
@@ -286,6 +293,23 @@ def build_actor_module(
             action_feature_ids=list(range(len(self.action_names))),
         )
 
+    def build_reward_module(
+        self,
+        trainer_module: DiscreteCRRTrainer,
+        normalization_data_map: Dict[str, NormalizationData],
+    ) -> torch.nn.Module:
+        """
+        Returns a TorchScript predictor module
+        """
+        net_builder = self.cpe_net_builder.value
+        return net_builder.build_serving_module(
+            trainer_module.reward_network,
+            normalization_data_map[NormalizationKey.STATE],
+            action_names=self.action_names,
+            state_feature_config=self.state_feature_config,
+            predictor_wrapper_type=FbBanditRewardNetPredictorWrapper,
+        )
+
 
 class ActorDQN(ModelBase):
     def __init__(self, actor):

reagent/training/discrete_crr_trainer.py

@@ -135,6 +135,8 @@ def __init__(
                 # pyre-fixme[16]: Optional type has no attribute `__getitem__`.
                 self.reward_boosts[0, i] = rl.reward_boost[k]
 
+        # The function below adds reward_network as a member object to DQNTrainerBaseLightning,
+        # from which DiscreteCRRTrainer is derived.
         self._initialize_cpe(
             reward_network,
             q_network_cpe,