diff --git a/overcooked_simulator/gui_2d_vis/overcooked_gui.py b/overcooked_simulator/gui_2d_vis/overcooked_gui.py
index 76f85bb290847ae9f2f62f12f0bb395b13450793..25f9ccde11cf253ba7608d92aba9253998bb3e42 100644
--- a/overcooked_simulator/gui_2d_vis/overcooked_gui.py
+++ b/overcooked_simulator/gui_2d_vis/overcooked_gui.py
@@ -75,7 +75,7 @@ class PyGameGUI:
         port: int,
         manager_ids: list[str],
     ):
-        self.game_screen = None
+        self.game_screen: pygame.Surface = None
         self.FPS = 60
         self.running = True
 
@@ -368,17 +368,19 @@ class PyGameGUI:
         )
 
         border = self.visualization_config["GameWindow"]["black_border_size"]
+        border_rect = pygame.Rect(
+            self.window_width // 2 - (self.game_width // 2) - border,
+            self.window_height // 2 - (self.game_height // 2) - border,
+            self.game_width + 2 * border,
+            self.game_height + 2 * border,
+        )
         pygame.draw.rect(
             self.main_window,
             colors["black"],
-            pygame.Rect(
-                self.screen_margin - border,
-                self.screen_margin - border,
-                self.game_width + 2 * border,
-                self.game_height + 2 * border,
-            ),
+            border_rect,
             width=self.visualization_config["GameWindow"]["black_border_size"],
         )
+
         self.update_score_label(state)
 
     def set_window_size(self):
@@ -672,10 +674,12 @@ class PyGameGUI:
                             self.draw(state)
 
                             game_screen_rect = self.game_screen.get_rect()
+
                             game_screen_rect.center = [
                                 self.window_width // 2,
                                 self.window_height // 2,
                             ]
+
                             self.main_window.blit(self.game_screen, game_screen_rect)
 
                     case MenuStates.End:
@@ -687,6 +691,7 @@ class PyGameGUI:
             except (KeyboardInterrupt, SystemExit):
                 self.running = False
 
+        self.disconnect_websockets()
         pygame.quit()
         sys.exit()