Improved the game state machine