Create get_data_module() on OSS WorldModelBase AND on FB FbWorldModel

Summary: As titled. See T83887308 & T83886520 for more details.

Reviewed By: kaiwenw

Differential Revision: D26498062

fbshipit-source-id: 2202b8a3dcde22bc97d03231228d518e947ca7db

Author: Ban Kawas

reagent/model_managers/world_model_base.py

@@ -5,11 +5,15 @@
 from reagent.core.dataclasses import dataclass
 from reagent.core.parameters import NormalizationData, NormalizationKey
 from reagent.data.data_fetcher import DataFetcher
+from reagent.data.manual_data_module import ManualDataModule
 from reagent.data.reagent_data_module import ReAgentDataModule
 from reagent.gym.policies.policy import Policy
 from reagent.preprocessing.batch_preprocessor import BatchPreprocessor
+from reagent.preprocessing.types import InputColumn
+from reagent.workflow.identify_types_flow import identify_normalization_parameters
 from reagent.workflow.types import (
     Dataset,
+    PreprocessingOptions,
     ReaderOptions,
     ResourceOptions,
     RewardOptions,
@@ -40,7 +44,7 @@ def create_policy(self) -> Policy:
 
     @property
     def should_generate_eval_dataset(self) -> bool:
-        return False
+        raise RuntimeError
 
     @property
     def required_normalization_keys(self) -> List[str]:
@@ -49,7 +53,7 @@ def required_normalization_keys(self) -> List[str]:
     def run_feature_identification(
         self, input_table_spec: TableSpec
     ) -> Dict[str, NormalizationData]:
-        raise NotImplementedError()
+        raise RuntimeError
 
     def query_data(
         self,
@@ -58,10 +62,30 @@ def query_data(
         reward_options: RewardOptions,
         data_fetcher: DataFetcher,
     ) -> Dataset:
-        raise NotImplementedError()
+        raise RuntimeError
 
-    def build_batch_preprocessor(self, use_gpu: bool) -> BatchPreprocessor:
-        raise NotImplementedError()
+    def build_batch_preprocessor(self) -> BatchPreprocessor:
+        raise RuntimeError
+
+    def get_data_module(
+        self,
+        *,
+        input_table_spec: Optional[TableSpec] = None,
+        reward_options: Optional[RewardOptions] = None,
+        reader_options: Optional[ReaderOptions] = None,
+        setup_data: Optional[Dict[str, bytes]] = None,
+        saved_setup_data: Optional[Dict[str, bytes]] = None,
+        resource_options: Optional[ResourceOptions] = None,
+    ) -> Optional[ReAgentDataModule]:
+        return WorldModelDataModule(
+            input_table_spec=input_table_spec,
+            reward_options=reward_options,
+            setup_data=setup_data,
+            saved_setup_data=saved_setup_data,
+            reader_options=reader_options,
+            resource_options=resource_options,
+            model_manager=self,
+        )
 
     def train(
         self,
@@ -84,3 +108,49 @@ def train(
         - validation_output
         """
         raise NotImplementedError()
+
+
+class WorldModelDataModule(ManualDataModule):
+    @property
+    def should_generate_eval_dataset(self) -> bool:
+        return True
+
+    @property
+    def required_normalization_keys(self) -> List[str]:
+        return [NormalizationKey.STATE]
+
+    def run_feature_identification(
+        self, input_table_spec: TableSpec
+    ) -> Dict[str, NormalizationData]:
+        # Run state feature identification
+        state_preprocessing_options = PreprocessingOptions()
+        state_features = [
+            ffi.feature_id
+            for ffi in self.model_manager.state_feature_config.float_feature_infos
+        ]
+        logger.info(f"state allowedlist_features: {state_features}")
+        state_preprocessing_options = state_preprocessing_options._replace(
+            allowedlist_features=state_features
+        )
+
+        state_normalization_parameters = identify_normalization_parameters(
+            input_table_spec, InputColumn.STATE_FEATURES, state_preprocessing_options
+        )
+
+        return {
+            NormalizationKey.STATE: NormalizationData(
+                dense_normalization_parameters=state_normalization_parameters
+            )
+        }
+
+    def query_data(
+        self,
+        input_table_spec: TableSpec,
+        sample_range: Optional[Tuple[float, float]],
+        reward_options: RewardOptions,
+        data_fetcher: DataFetcher,
+    ) -> Dataset:
+        raise NotImplementedError()
+
+    def build_batch_preprocessor(self) -> BatchPreprocessor:
+        raise NotImplementedError()