Update generate_clips.py

generate_clips.py

@@ -1,41 +1,48 @@
+import os
+from PIL import Image
+
 import gym_super_mario_bros
 from gym_super_mario_bros.actions import RIGHT_ONLY
 from nes_py.wrappers import JoypadSpace
-from agent import Agent
 
 from gym import Wrapper
 from gym.wrappers import GrayScaleObservation, ResizeObservation, FrameStack
 
-import os
-from PIL import Image
+from agent import Agent
+
 
-# Modified SkipFrame wrapper to log frames and actions
+# ----- Custom SkipFrame wrapper -----
 class SkipFrame(Wrapper):
-    def __init__(self, env, skip):
+    def __init__(self, env, skip: int):
         super().__init__(env)
         self._skip = skip
-        self.counter = 0
         self.frames_log = []
         self.actions_log = []
-    
+
     def step(self, action):
         total_reward = 0.0
         done = False
+        trunc = False
+        info = {}
+
         for _ in range(self._skip):
             next_state, reward, done, trunc, info = self.env.step(action)
+            # log frames + actions
             self.frames_log.append(next_state.copy())
             self.actions_log.append(action)
             total_reward += reward
-            if done:
+            if done or trunc:
                 break
+
         return next_state, total_reward, done, trunc, info
-    
+
     def reset(self, **kwargs):
         state, info = self.env.reset(**kwargs)
         self.frames_log = [state.copy()]
         self.actions_log = [0]
         return state, info
 
+
 def apply_wrappers(env):
     env = SkipFrame(env, skip=4)
     env = ResizeObservation(env, shape=84)
@@ -44,58 +51,85 @@ def apply_wrappers(env):
     return env
 
 
-ENV_NAME = 'SuperMarioBros-1-1-v0'
-NUM_OF_EPISODES = 1_000
-controllers = [Image.open(f"controllers/{i}.png") for i in range(5)]
-
-env = gym_super_mario_bros.make(ENV_NAME, render_mode='rgb_array', apply_api_compatibility=True)
-env = JoypadSpace(env, RIGHT_ONLY)
-env = apply_wrappers(env)
+# -------- CONFIG --------
+ENV_NAME = "SuperMarioBros-1-1-v0"
+NUM_OF_EPISODES = 5          # just a few episodes to record
+GAMES_ROOT = "games"
+CONTROLLERS_DIR = "controllers"   # folder with 0.png ... 4.png
+SCALING_FACTOR = 10               # upscale frames for nicer video
+# ------------------------
 
-agent = Agent(input_dims=env.observation_space.shape, num_actions=env.action_space.n)
 
-# agent.load_model("models/folder_name/ckpt_name")
+def main():
+    os.makedirs(GAMES_ROOT, exist_ok=True)
 
-for i in range(NUM_OF_EPISODES):
-    done = False
-    state, _ = env.reset()
-    rewards = 0
-    while not done:
-        action = agent.choose_action(state)
-        frame = env.render()
-        new_state, reward, done, truncated, info  = env.step(action)
-        rewards += reward
+    # load controller images (optional, only for separate controller frames)
+    controllers = [Image.open(os.path.join(CONTROLLERS_DIR, f"{i}.png")) for i in range(5)]
 
-        # agent.store_in_memory(state, action, reward, new_state, done)
-        # agent.learn()
+    # ---- create env ----
+    env = gym_super_mario_bros.make(
+        ENV_NAME,
+        render_mode="rgb_array",
+        apply_api_compatibility=True,
+    )
+    env = JoypadSpace(env, RIGHT_ONLY)
+    env = apply_wrappers(env)
 
-        state = new_state
+    # ---- agent ----
+    agent = Agent(input_dims=env.observation_space.shape,
+                  num_actions=env.action_space.n)
 
-        if done:
-            print(f"Episode: {i}, Reward: {rewards}")
-            if info["flag_get"]:
-                os.makedirs(os.path.join("games" f"game_{i}"), exist_ok=True)
-                frame_skip_env = env.env.env.env # Unwrapping the environment to get the SkipFrame wrapper
-                frames_log = frame_skip_env.frames_log
-                actions_log = frame_skip_env.actions_log
-                for j, (frame, action) in enumerate(zip(frames_log, actions_log)):
-                    # upscale frame
-                    scaling_factor = 10
-                    new_dims = (frame.shape[1] * scaling_factor, frame.shape[0] * scaling_factor)
-                    frame = Image.fromarray(frame).resize(new_dims, Image.NEAREST)
+    # TODO: change this path to your trained checkpoint
+    # e.g. "models/best_model.pt" or "models/mario/model_500000.pt"
+    # If you don't have a trained model yet, comment this line.
+    agent.load_model("models/your_model_checkpoint.pt")
 
-                    frame.save(os.path.join("games" f"game_{i}", f"frame_{j}.png"))
-                    controllers[action].save(os.path.join("games" f"game_{i}", f"controller_{j}.png"))
-        
-        # if i % 5000 == 0 and i > 0:
-        #     agent.save_model(os.path.join("models", f"model_{i}_iter.pt"))
+    for i in range(NUM_OF_EPISODES):
+        done = False
+        state, _ = env.reset()
+        rewards = 0
+
+        while not done:
+            action = agent.choose_action(state)
+            frame = env.render()  # rgb_array frame (not used directly, but ok)
+            new_state, reward, done, truncated, info = env.step(action)
+            rewards += reward
+            state = new_state
+
+            if done or truncated:
+                print(f"Episode: {i}, Reward: {rewards}")
+
+                # Only save a game if Mario reached the flag
+                if info.get("flag_get", False):
+                    game_dir = os.path.join(GAMES_ROOT, f"game_{i}")
+                    os.makedirs(game_dir, exist_ok=True)
+
+                    # unwrap env to get SkipFrame wrapper
+                    frame_skip_env = env.env.env.env
+                    frames_log = frame_skip_env.frames_log
+                    actions_log = frame_skip_env.actions_log
+
+                    for j, (frame_array, action_taken) in enumerate(zip(frames_log, actions_log)):
+                        # upscale frame for nicer video
+                        new_dims = (
+                            frame_array.shape[1] * SCALING_FACTOR,
+                            frame_array.shape[0] * SCALING_FACTOR,
+                        )
+                        frame_img = Image.fromarray(frame_array).resize(
+                            new_dims, Image.NEAREST
+                        )
+                        frame_img.save(os.path.join(game_dir, f"frame_{j}.png"))
+
+                        # save controller overlay frames (optional)
+                        controllers[action_taken].save(
+                            os.path.join(game_dir, f"controller_{j}.png")
+                        )
 
-env.close()
+                break
 
-# Bash commands to convert frames to video (Assuming you're in the games folder)
+    env.close()
+    print("Finished generating frame sequences in 'games/'.")
 
-# Video game frames
-# ffmpeg -framerate 60 -i frame_%d.png game.mp4
 
-# Controller frames (For transparent background)
-# ffmpeg -framerate 60 -i controller_%d.png -c:v prores_ks -profile:v 4444 -pix_fmt yuva444p10le -alpha_bits 16 controller.mov
+if __name__ == "__main__":
+    main()