Unittest of ModelManager (BehavioralCloning)

Summary: - Model Manager for BehaviorCloning - UnitTest of the ModelManager - DataModule for UnitTest Reviewed By: czxttkl Differential Revision: D33829752 fbshipit-source-id: 9d1d6af293f652e095b914608108fc0d215ff257
alexnikulkov · Mar 3, 2022 · 79a1891 · 79a1891
1 parent 3eaf1cf
commit 79a1891
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/reagent/training/__init__.py b/reagent/training/__init__.py
@@ -35,6 +35,7 @@
     CRRTrainerParameters,
     ReinforceTrainerParameters,
     PPOTrainerParameters,
+    BehavioralCloningTrainerParameters,
 )
 
 
@@ -69,6 +70,7 @@
     "ReinforceTrainerParameters",
     "PPOTrainer",
     "PPOTrainerParameters",
+    "BehavioralCloningTrainerParameters",
 ]
 
 if IS_FB_ENVIRONMENT:

diff --git a/reagent/training/parameters.py b/reagent/training/parameters.py
@@ -4,6 +4,7 @@
 from reagent.core.configuration import make_config_class
 from reagent.core.types import BaseDataClass
 
+from .behavioral_cloning_trainer import BehavioralCloningTrainer
 from .c51_trainer import C51Trainer
 from .cb.linucb_trainer import LinUCBTrainer
 from .discrete_crr_trainer import DiscreteCRRTrainer
@@ -128,6 +129,11 @@ class RewardNetworkTrainerParameters:
     pass
 
 
+@make_config_class(BehavioralCloningTrainer.__init__, blocklist=["bc_net"])
+class BehavioralCloningTrainerParameters:
+    pass
+
+
 @make_config_class(
     Seq2SlateTrainer.__init__,
     blocklist=[