Combine figures 3 and 4

.gitignore

@@ -0,0 +1,4 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class

RED/configs/example/Figure_3_RT3D_chemostat.yaml

@@ -4,6 +4,8 @@ defaults:
   - _self_
 
 policy_delay: 2
+wandb_team: rl-oed
+wandb_project_name: figure3-example
 initial_explore_rate: 1
 explore_rate_mul: 1
 test_episode: False

RED/configs/example/Figure_4_RT3D_chemostat.yaml

@@ -4,6 +4,9 @@ defaults:
   - _self_
 
 policy_delay: 2
+number_of_trials: 10
+wandb_team: rl-oed
+wandb_project_name: figure4-example
 initial_explore_rate: 1
 explore_rate_mul: 1
 test_episode: False

examples/Figure_4_RT3D_chemostat/train_RT3D.py → examples/Figure_3_and_4_RT3D_chemostat/train_RT3D.py

@@ -1,14 +1,15 @@
-
 import json
 import math
 import os
 import sys
+from datetime import datetime
 
+# TODO: Can we just run using python -m?
 IMPORT_PATH = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 sys.path.append(IMPORT_PATH)
 
 import multiprocessing
-
+import wandb
 import hydra
 import numpy as np
 from casadi import *
@@ -35,6 +36,30 @@ def train_RT3D(cfg : DictConfig):
         sep="\n\n"
     )
 
+    group_name = None
+    number_of_trials = 1
+    if hasattr(cfg, "number_of_trials"):
+        group_name = datetime.now().strftime("%d/%m/%Y_%H:%M")
+        number_of_trials = cfg.number_of_trials
+
+    original_save_path = cfg.save_path     
+    for _ in range(number_of_trials):
+        # Append datetime to differentiate trials
+        cfg.save_path = os.path.join(original_save_path, group_name)
+        run_single_experiment(cfg, group_name=group_name)
+
+
+
+def run_single_experiment(cfg : DictConfig, group_name : str):
+
+    # start a new wandb run to track this script
+    if group_name:
+        wandb.init(reinit=True, project=cfg.wandb_project_name,
+                entity=cfg.wandb_team, group=group_name, config=dict(cfg))
+    else:
+        wandb.init(reinit=True, project=cfg.wandb_project_name,
+                entity=cfg.wandb_team, config=dict(cfg))
+
     ### prepare save path
     os.makedirs(cfg.save_path, exist_ok=True)
     print("Results will be saved in: ", cfg.save_path)
@@ -81,7 +106,7 @@ def train_RT3D(cfg : DictConfig):
         actual_params = np.random.uniform(
             low=cfg.environment.lb,
             high=cfg.environment.ub,
-            size=(cfg.environment.n_parallel_experiments, 3)
+            size=(cfg.environment.n_parallel_experiments, n_params)
         )
         env.param_guesses = DM(actual_params)
 
@@ -173,6 +198,11 @@ def train_RT3D(cfg : DictConfig):
         history["us"].extend(e_us)
         history["explore_rate"].append(explore_rate)
 
+        ### log results to w and b
+        for i in range(len(e_returns)):
+            wandb.log({"returns": e_returns[i], "actions": np.array(e_actions).transpose(1, 0, 2)[i],
+                       "us": e_us[i], "explore_rate": explore_rate})
+
         print(
             f"\nEPISODE: [{episode}/{total_episodes}] ({episode * cfg.environment.n_parallel_experiments} experiments)",
             f"explore rate:\t{explore_rate:.2f}",
@@ -211,6 +241,8 @@ def train_RT3D(cfg : DictConfig):
         conv_window=25,
     )
 
+    wandb.finish()
+
 
 def setup_env(cfg):
     n_cores = multiprocessing.cpu_count()