make the game multi sessions

connectfour/app.py

@@ -1,17 +1,94 @@
 import time
 
 import gradio as gr
+import numpy as np
+import onnxruntime as ort
+from pettingzoo.classic import connect_four_v3
 
-from connectfour.connect4 import Connect4
+from connectfour import ERROR_SCREEN
+from models import MODEL_PATH
 
 # poetry export -f requirements.txt --output requirements.txt --without-hashes
 # gradio connectfour/app.py
 
-
+session = ort.InferenceSession(str(MODEL_PATH), None)
 demo = gr.Blocks()
 
+column_count = 7
+game_on_msg = "Game On"
+
+
+def flatten_observation(obs):
+    flatten_action_mask = np.array(obs["action_mask"])
+    flatten_observation = np.reshape(obs["observation"], 2 * 6 * column_count)
+    flatten_obs = np.concatenate([flatten_action_mask, flatten_observation])
+    return flatten_obs[np.newaxis, ...].astype(np.float32)
+
+
+def legal_moves(env, player_id):
+    return np.arange(column_count)[env.observe(player_id)["action_mask"] == 1]
+
+
+def done(env):
+    return np.any(list(env.terminations.values()) + list(env.truncations.values()))
+
+
+def get_state_msg(env, human):
+    if done(env):
+        end_message = "End of the game"
+        if env.rewards[human] > 0:
+            end_message += ": You WIN !!"
+        elif env.rewards[human] < 0:
+            end_message += ": You LOSE !!"
+        return end_message
+
+    return game_on_msg
+
+
+def play(env, human, action=None):
+    try:
+        if human != env.agent_selection:
+            action = session.run(
+                ["output"],
+                {
+                    "obs": flatten_observation(env.observe(env.agent_selection)),
+                    "state_ins": [],
+                },
+            )
+            action = int(np.argmax(action[0]))
+
+        if action not in legal_moves(env, env.agent_selection):
+            action = np.random.choice(legal_moves(env, env.agent_selection))
+
+        env.step(action)
+        return env, get_state_msg(env, human)
+    except Exception as e:
+        return env, f"Restart the Game"
+
+
+def init_env(env, who_plays_first, human):
+    env.reset()
+
+    if who_plays_first != "You":
+        play(env, human)
+
+    return env
+
+
+def error_screen():
+    with open(ERROR_SCREEN, "rb") as f:
+        error_screen = np.load(f)
+    return error_screen
+
+
+def create_env():
+    return init_env(connect_four_v3.env(render_mode="rgb_array"), "You", "player_0")
+
+
 with demo:
-    connect4 = Connect4("You")
+    human = gr.State("player_0")
+    env = gr.State(create_env())
+
     drop_token_btns = []
 
     with gr.Row():
@@ -22,70 +99,83 @@ with demo:
                 label="Who plays first", choices=["You", "Bot"], value="You"
             )
             reinitialize = gr.Button("New Game")
-
-            game_state = gr.Text(value="Game On", interactive=False, label="Status")
+            game_state = gr.Text(value=game_on_msg, interactive=False, label="Status")
 
         with gr.Column(scale=1):
             output = gr.Image(
                 label="Connect Four Grid",
                 type="numpy",
                 show_label=False,
-                value=connect4.env.render(),
+                value=error_screen(),
             )
 
             with gr.Row():
-                for i in range(7):
+                for i in range(column_count):
                     with gr.Column(min_width=20):
                         drop_token_btns.append(gr.Button("X", elem_id=i))
 
-    who_plays_first.change(
-        connect4.init_env, who_plays_first, outputs=[output, game_state]
-    )
-
-    def reinit_game(who_plays_first):
-        output, game_state = connect4.init_env(who_plays_first)
-        return output, game_state, gr.Checkbox.update(interactive=True)
+    def reinit_game(env, who_plays_first, human):
+        env = init_env(env, who_plays_first, human)
+        return [
+            env,
+            env.agent_selection,  # human
+            get_state_msg(env, human),  # state_msg
+            gr.Checkbox.update(interactive=True),  # who_plays_first
+        ]
 
-    reinitialize.click(
-        reinit_game, who_plays_first, outputs=[output, game_state, who_plays_first]
-    )
+    def on_render_change(env):
+        return env.render()
 
     def wait(game_state_value):
-        if game_state_value == "Game On":
+        if game_state_value == game_on_msg:
             time.sleep(0.7)
             return gr.Checkbox.update(interactive=False)
         else:
             return gr.Checkbox.update(interactive=True)
 
-    def bot(game_state_value):
-        if game_state_value == "Game On":
-            rendered_env = connect4.play()
-            return *rendered_env, gr.Checkbox.update(interactive=False) if rendered_env[
-                1
-            ] == "Game On" else gr.Checkbox.update(interactive=True)
+    def bot(env, game_state_value):
+        if game_state_value == game_on_msg:
+            env, state_msg = play(env, human)
+            if state_msg == game_on_msg:
+                return state_msg, gr.Checkbox.update(interactive=False)
+            else:
+                return state_msg, gr.Checkbox.update(interactive=True)
         return (
-            gr.Image.update(),
             game_state_value,
             gr.Checkbox.update(interactive=True),
         )
 
-    def click_column(evt: gr.EventData):
-        output, game_state = connect4.play(int(evt.target.elem_id))
-        return output, game_state
+    def click_column(env, human, evt: gr.EventData):
+        env, state_msg = play(env, human, int(evt.target.elem_id))
+        return env, state_msg
 
     def game_state_change(value):
-        if value == "Game On":
-            return [gr.Button.update(interactive=True)] * 7
-        else:
-            return [gr.Button.update(interactive=False)] * 7
+        return [gr.Button.update(interactive=value == game_on_msg)] * column_count
+
+    who_plays_first.change(
+        reinit_game,
+        [env, who_plays_first, human],
+        outputs=[env, human, game_state, who_plays_first],
+    ).then(on_render_change, inputs=[env], outputs=[output])
 
-    for i in range(7):
+    reinitialize.click(
+        reinit_game,
+        [env, who_plays_first, human],
+        outputs=[env, human, game_state, who_plays_first],
+    ).then(on_render_change, inputs=[env], outputs=[output])
+
+    for i in range(column_count):
         drop_token_btns[i].click(
             click_column,
-            outputs=[output, game_state],
+            inputs=[env, human],
+            outputs=[env, game_state],
+        ).then(on_render_change, inputs=[env], outputs=[output]).then(
+            wait, inputs=[game_state], outputs=[who_plays_first]
+        ).then(
+            bot, inputs=[env, game_state], outputs=[game_state, who_plays_first]
         ).then(
-            wait, inputs=[game_state], outputs=who_plays_first
-        ).then(bot, inputs=[game_state], outputs=[output, game_state, who_plays_first])
+            on_render_change, inputs=[env], outputs=[output]
+        )
 
     game_state.change(
         game_state_change,
@@ -93,5 +183,4 @@ with demo:
         outputs=drop_token_btns,
     )
 
-
 demo.launch()

connectfour/connect4.py

@@ -1,92 +0,0 @@
-import numpy as np
-import onnxruntime as ort
-from pettingzoo.classic import connect_four_v3
-from connectfour import ERROR_SCREEN
-
-from models import MODEL_PATH
-
-
-class Connect4:
-    def __init__(self, who_plays_first) -> None:
-        self.init_env(who_plays_first)
-        self.session = ort.InferenceSession(str(MODEL_PATH), None)
-
-    def init_env(self, who_plays_first):
-        self.env = connect_four_v3.env(render_mode="rgb_array")
-        self.env.reset()
-
-        if who_plays_first == "You":
-            self.human = self.current_player_id
-        else:
-            self.play()
-            self.human = self.current_player_id
-
-        return self.render_and_state
-
-    def flatten_observation(self, obs):
-        flatten_action_mask = np.array(obs["action_mask"])
-        flatten_observation = np.reshape(obs["observation"], 2 * 6 * 7)
-        flatten_obs = np.concatenate([flatten_action_mask, flatten_observation])
-        return flatten_obs[np.newaxis, ...].astype(np.float32)
-
-    def play(self, action=None):
-        try:
-            if self.human != self.current_player_id:
-                action = self.session.run(
-                    ["output"],
-                    {
-                        "obs": self.flatten_observation(
-                            self.env.observe(self.current_player_id)
-                        ),
-                        "state_ins": [],
-                    },
-                )
-                action = int(np.argmax(action[0]))
-
-            if action not in self.legal_moves:
-                action = np.random.choice(self.legal_moves)
-
-            self.env.step(action)
-
-            return self.render_and_state
-        except:
-            return self.blue_screen()
-
-    @property
-    def current_player_id(self):
-        return self.env.agent_selection
-
-    @property
-    def current_observation(self):
-        return self.env.observe(self.current_player_id)
-
-    @property
-    def legal_moves(self):
-        return np.arange(7)[self.current_observation["action_mask"] == 1]
-
-    @property
-    def done(self):
-        return np.any(
-            list(self.env.terminations.values()) + list(self.env.truncations.values())
-        )
-
-    @property
-    def render_and_state(self):
-        if self.done:
-            if self.human not in self.env.rewards:
-                return self.blue_screen()
-
-            end_message = "End of the game"
-            if self.env.rewards[self.human] > 0:
-                end_message += ": You WIN !!"
-            elif self.env.rewards[self.human] < 0:
-                end_message += ": You LOSE !!"
-            return self.env.render(), end_message
-
-        return self.env.render(), "Game On"
-
-    def blue_screen(self):
-        with open(ERROR_SCREEN, "rb") as f:
-            error_screen = np.load(f)
-
-        return (error_screen, "Restart the Game")