Allow for publishing of reward network in discrete CRR (#588)

Summary:
Pull Request resolved: #588

Allow for publishing of reward network in discrete_crr.py

Differential Revision: D32711991

fbshipit-source-id: d13fcf724cd5de0c04609378a86b779c07db9efb

Author: David Vengerov

reagent/model_managers/discrete/discrete_crr.py

@@ -196,7 +196,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 +219,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 +287,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
+        assert trainer_module.reward_network is not None
+        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,
+        )
+
 
 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,