UCI unitest fix

Summary: Fix of the UCI unitest Reviewed By: alexnikulkov Differential Revision: D53361355 fbshipit-source-id: 2ecf0ff8ae7c42a71a9510a05fb634bbba561c05
facebookresearch · Feb 5, 2024 · 8a41fd5 · 8a41fd5
1 parent 27ebc21
commit 8a41fd5
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 9 deletions.
diff --git a/pearl/utils/scripts/cb_benchmark/cb_benchmark_config.py b/pearl/utils/scripts/cb_benchmark/cb_benchmark_config.py
@@ -35,7 +35,7 @@
 """
 Experiment config
 """
-run_config: Dict[str, Any] = {
+run_config_def: Dict[str, Any] = {
     "T": 10000,
     "training_rounds": 10,
     "num_of_experiments": 5,
@@ -87,7 +87,7 @@
 
 def return_neural_squarecb_config(
     env: SLCBEnvironment,
-    run_config: Dict[str, Any] = run_config,
+    run_config: Dict[str, Any] = run_config_def,
 ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
 
     feature_dim: int = env.observation_dim
@@ -123,7 +123,7 @@ def return_neural_squarecb_config(
 
 def return_neural_lin_ucb_config(
     env: SLCBEnvironment,
-    run_config: Dict[str, Any] = run_config,
+    run_config: Dict[str, Any] = run_config_def,
     alpha: float = 1.0,
 ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
 
@@ -156,7 +156,7 @@ def return_neural_lin_ucb_config(
 
 def return_neural_lin_ts_config(
     env: SLCBEnvironment,
-    run_config: Dict[str, Any] = run_config,
+    run_config: Dict[str, Any] = run_config_def,
     alpha: float = 0.25,
 ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
 
@@ -208,7 +208,7 @@ def return_offline_eval_config(
 
 def return_neural_fastcb_config(
     env: SLCBEnvironment,
-    run_config: Dict[str, Any] = run_config,
+    run_config: Dict[str, Any] = run_config_def,
 ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
 
     feature_dim: int = env.observation_dim

diff --git a/pearl/utils/scripts/cb_benchmark/run_cb_benchmarks.py b/pearl/utils/scripts/cb_benchmark/run_cb_benchmarks.py
@@ -41,7 +41,7 @@
     return_neural_lin_ucb_config,
     return_neural_squarecb_config,
     return_offline_eval_config,
-    run_config,
+    run_config_def,
     satimage_uci_dict,
     yeast_uci_dict,
 )
@@ -181,6 +181,7 @@ def run_experiments(
     num_of_experiments: int,
     policy_learner_dict: Dict[str, Any],
     exploration_module_dict: Dict[str, Any],
+    run_config: Dict[str, Any],
     save_results_path: str,
     dataset_name: str,
     run_offline: bool = False,
@@ -296,6 +297,7 @@ def run_cb_benchmarks(
                 num_of_experiments=run_config["num_of_experiments"],
                 policy_learner_dict=policy_learner_dict,
                 exploration_module_dict=exploration_module_dict,
+                run_config=run_config,
                 save_results_path=save_results_path,
                 dataset_name=dataset_name,
                 run_offline=algorithm == "OfflineEval",
@@ -323,5 +325,5 @@ def run_cb_benchmarks(
     run_cb_benchmarks(
         cb_algorithms_config=cb_algorithms_config,
         test_environments_config=test_environments_config,
-        run_config=run_config,
+        run_config=run_config_def,
     )
diff --git a/test/unit/with_pytorch/test_agent.py b/test/unit/with_pytorch/test_agent.py
@@ -260,7 +260,7 @@ def test_contextual_bandit_on_uci_datasets(self) -> None:
         # CB Algorithms are the neural versions of LinUCB, LinTS, and SquareCB with shared models.
 
         # set number of time steps to be small, just for unit testing purposes
-        run_config: Dict[str, Any] = {
+        run_config_test: Dict[str, Any] = {
             "T": 300,
             "training_rounds": 1,
             "num_of_experiments": 1,
@@ -281,5 +281,5 @@ def test_contextual_bandit_on_uci_datasets(self) -> None:
         run_cb_benchmarks(
             cb_algorithms_config=cb_algorithms_config,
             test_environments_config=test_environments_config,
-            run_config=run_config,
+            run_config=run_config_test,
         )