Removed "can_dash" input value

celeste/celeste_ai/plotting/plot_best_action.py

@@ -1,6 +1,7 @@
 import torch
 import numpy as np
 from pathlib import Path
+import matplotlib as mpl
 import matplotlib.pyplot as plt
 
 # All of the following are required to load
@@ -34,7 +35,7 @@ def best_action(
 
 
 	# Compute preditions
-	p = np.zeros((128, 128, 2), dtype=np.float32)
+	p = np.zeros((128, 128), dtype=np.float32)
 	with torch.no_grad():
 		for r in range(len(p)):
 			for c in range(len(p[r])):
@@ -43,26 +44,31 @@ def best_action(
 
 				k = np.asarray(policy_net(
 					torch.tensor(
-						[x, y, 0],
+						[x, y],
 						dtype = torch.float32,
 						device = device
 					).unsqueeze(0)
 				)[0])
-				p[r][c][0] = np.argmax(k)
+				p[r][c] = np.argmax(k)
 
-				k = np.asarray(policy_net(
-					torch.tensor(
-						[x, y, 1],
-						dtype = torch.float32,
-						device = device
-					).unsqueeze(0)
-				)[0])
-				p[r][c][1] = np.argmax(k)
 
+	cmap = mpl.colors.ListedColormap(
+		[
+			"forestgreen",
+			"firebrick",
+			"lightgreen",
+			"salmon",
+			"darkturquoise",
+			"sandybrown",
+			"olive",
+			"darkorchid",
+			"mediumvioletred"
+		]
+	)
 
 	# Plot predictions
-	fig, axs = plt.subplots(1, 2, figsize = (10, 10))
-	ax = axs[0]
+	fig, axs = plt.subplots(1, 1, figsize = (20, 20))
+	ax = axs
 	ax.set(
 		adjustable = "box",
 		aspect = "equal",
@@ -70,30 +76,16 @@ def best_action(
 	)
 
 	plot = ax.pcolor(
-		p[:,:,0],
-		cmap = "Set1",
+		p,
+		cmap = cmap,
 		vmin = 0,
 		vmax = 8
 	)
 	ax.invert_yaxis()
-	fig.colorbar(plot)
+	cbar = fig.colorbar(plot, ticks = list(range(0, 9)))
+	cbar.ax.set_yticklabels(Celeste.action_space)
 
-	ax = axs[1]
-	ax.set(
-		adjustable = "box",
-		aspect = "equal",
-		title = "Best Action"
-	)
-
-	plot = ax.pcolor(
-		p[:,:,0],
-		cmap = "Set1",
-		vmin = 0,
-		vmax = 8
-	)
-
-	ax.invert_yaxis()
-	fig.colorbar(plot)
+	
 
 	fig.savefig(out_filename)
 	plt.close()

celeste/celeste_ai/plotting/plot_predicted_reward.py

@@ -43,7 +43,7 @@ def predicted_reward(
 
 				k = np.asarray(policy_net(
 					torch.tensor(
-						[x, y, 0],
+						[x, y],
 						dtype = torch.float32,
 						device = device
 					).unsqueeze(0)