diff --git a/Assets/Scripts/Control&Input/Peer2PeerTransport.cs b/Assets/Scripts/Control&Input/Peer2PeerTransport.cs
index d73f69c5e..ff1a233f4 100644
--- a/Assets/Scripts/Control&Input/Peer2PeerTransport.cs
+++ b/Assets/Scripts/Control&Input/Peer2PeerTransport.cs
@@ -502,6 +502,15 @@ private void AddAiPlayers()
         }
     }
 
+    private void RemoveAiPlayers()
+    {
+        foreach (var ai in players.Values.Where(p => p.type is PlayerType.AI))
+        {
+            NetworkServer.SendToAll(new PlayerLeftMessage(ai.id));
+        }
+    }
+
+
     // Called after shooting rounds (TODO just use the same matchcontroller stuff???)
     private static void UpdatePlayerDetailsAfterShootingRound()
     {
@@ -543,7 +552,8 @@ private static void UpdatePlayerInventoryForIdentity(PlayerIdentity identity)
 
     public override void OnServerChangeScene(string newSceneName)
     {
-        var needsExtraAiPlayers = PlayerInputManagerController.Singleton.MatchHasAI && !MatchController.Singleton;
+        var isAiEnabled = PlayerInputManagerController.Singleton.MatchHasAI;
+        var needsExtraAiPlayers = isAiEnabled && !MatchController.Singleton;
         if (needsExtraAiPlayers)
         {
             AddAiPlayers();
@@ -560,6 +570,8 @@ public override void OnServerChangeScene(string newSceneName)
                 break;
             case Scenes.Menu:
                 isInMatch = false;
+                if (isAiEnabled)
+                    RemoveAiPlayers();
                 NetworkServer.RegisterHandler<PlayerConnectedMessage>(OnSpawnPlayerInput);
                 break;
             default:
diff --git a/Assets/Scripts/UI/MainMenu/MainMenuController.cs b/Assets/Scripts/UI/MainMenu/MainMenuController.cs
index 8b44b9deb..c08af0ccd 100644
--- a/Assets/Scripts/UI/MainMenu/MainMenuController.cs
+++ b/Assets/Scripts/UI/MainMenu/MainMenuController.cs
@@ -132,6 +132,8 @@ private void GoBackToLobby()
             SelectControl(playerSelectMenuDefaultButton);
             SetStartButtonState();
             playerSelectManager.UpdateLobby();
+            if (playerInputManagerController.MatchHasAI)
+                aIButton.Toggle();
         }
         else
         {