Updated main plot script

celeste/plot.py

@@ -8,6 +8,7 @@ from multiprocessing import Pool
 
 m = Path("model_data/current")
 
+scaled = True
 
 # Make "predicted reward" plots
 def plot_pred(src_model):
@@ -15,9 +16,19 @@ def plot_pred(src_model):
 		src_model,
 		m / f"plots/predicted/{src_model.stem}.png",
 
-		device = torch.device("cpu")
+		device = torch.device("cpu"),
+		scaled = scaled
 	)
 
+# Make "best action" plots
+def plot_best(src_model):
+	plotting.best_action(
+		src_model,
+		m / f"plots/best_action/{src_model.stem}.png",
+
+		device = torch.device("cpu"),
+		scaled = scaled
+	)
 
 # Make "actual reward" plots
 def plot_act(src_model):
@@ -30,8 +41,17 @@ def plot_act(src_model):
 	)
 
 
+# Which plots should we make?
+plots = {
+	"prediction": True,
+	"actual": False,
+	"best": True
+}
+
 
 if __name__ == "__main__":
+
+	if plots["prediction"]:
 		print("Making prediction plots...")
 		with Pool(5) as p:
 			p.map(
@@ -39,6 +59,15 @@ if __name__ == "__main__":
 				list((m / "model_archive").iterdir())
 			)
 
+	if plots["best"]:
+		print("Making best-action plots...")
+		with Pool(5) as p:
+			p.map(
+				plot_best,
+				list((m / "model_archive").iterdir())
+			)
+
+	if plots["actual"]:
 		print("Making actual plots...")
 		with Pool(5) as p:
 			p.map(