Fixed frame timings

celeste/celeste.py

@@ -20,10 +20,7 @@ class Celeste:
 		"dash-lu"	# dash left-up
 	]
 
-	def __init__(self, on_get_state):
+	def __init__(self):
-
-		self.on_get_state = on_get_state
-
 		# Start pico-8
 		self.process = subprocess.Popen(
 			"bin/pico-8/linux/pico8",
@@ -182,11 +179,14 @@ class Celeste:
 			if k == "!RESTART":
 				break
 
-	def update_loop(self):
+
+	def update_loop(self, before, after):
 
 		# Get state, call callback, wait for state
 		# One line => one frame.
 
+		before_out = None
+
 		it = iter(self.process.stdout.readline, "")
 
 
@@ -232,4 +232,6 @@ class Celeste:
 			self.dist = dist
 			
 			# Call step callback
-			self.on_get_state(self)
+			if before_out is not None:
+				after(self, before_out)
+			before_out = before(self)

celeste/main.py

@@ -303,7 +303,7 @@ def optimize_model():
 	optimizer.step()
 
 
-def on_state(celeste):
+def on_state_before(celeste):
 	global steps_done
 
 	# Conversion to pytorch
@@ -324,14 +324,28 @@ def on_state(celeste):
 
 	# Turn number into action string
 	str_action = Celeste.action_space[action]
+
+
+	celeste.act(str_action)
+
+	return state, action
+
+def on_state_after(celeste, before_out):
+
+	state, action = before_out
+
+	pt_state = torch.tensor(
+		[state[x] for x in state_number_map],
+		dtype = torch.float32,
+		device = compute_device
+	).unsqueeze(0)
+
 	pt_action = torch.tensor(
 		[[ action ]],
 		device = compute_device,
 		dtype = torch.long
 	)
 
-	celeste.act(str_action)
-
 	next_state = celeste.status
 
 	if next_state["deaths"] != 0:
@@ -390,8 +404,9 @@ def on_state(celeste):
 
 
 
-c = Celeste(
+c = Celeste()
-	on_state
-)
 
-c.update_loop()
+c.update_loop(
+	on_state_before,
+	on_state_after
+)