From ce02009d6494d97965339f8bbdb949f3a7178844 Mon Sep 17 00:00:00 2001
From: Mark <mark@betalupi.com>
Date: Thu, 16 Feb 2023 12:11:04 -0800
Subject: [PATCH] Fixed frame timings

---
 celeste/celeste.py | 14 ++++++++------
 celeste/main.py    | 29 ++++++++++++++++++++++-------
 2 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/celeste/celeste.py b/celeste/celeste.py
index 28b530a..1168ab1 100755
--- a/celeste/celeste.py
+++ b/celeste/celeste.py
@@ -20,10 +20,7 @@ class Celeste:
 		"dash-lu"	# dash left-up
 	]
 
-	def __init__(self, on_get_state):
-
-		self.on_get_state = on_get_state
-
+	def __init__(self):
 		# 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)
\ No newline at end of file
+			if before_out is not None:
+				after(self, before_out)
+			before_out = before(self)
\ No newline at end of file
diff --git a/celeste/main.py b/celeste/main.py
index 6379b8d..aa972b0 100644
--- a/celeste/main.py
+++ b/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(
-	on_state
-)
+c = Celeste()
 
-c.update_loop()
+c.update_loop(
+	on_state_before,
+	on_state_after
+)