Merge pull request Ziktofel#335 from Salzkorn/sc2-next

Mission Race Balancing option

worlds/sc2/__init__.py

@@ -55,7 +55,7 @@ class Starcraft2WebWorld(WebWorld):
 
     custom_mission_orders_en = Tutorial(
         "Custom Mission Order Usage Guide",
-        "Documentation for the `custom_mission_order` YAML option",
+        "Documentation for the custom_mission_order YAML option",
         "English",
         "en_Custom Mission Orders.md",
         "custom_mission_orders/en",

worlds/sc2/mission_order/mission_pools.py

@@ -1,7 +1,7 @@
 from enum import IntEnum
 from typing import TYPE_CHECKING, Dict, Set, List
 
-from ..mission_tables import SC2Mission, lookup_id_to_mission, MissionPools, MissionFlag, SC2Campaign
+from ..mission_tables import SC2Mission, lookup_id_to_mission, MissionFlag, SC2Campaign
 from worlds.AutoWorld import World
 
 if TYPE_CHECKING:
@@ -48,6 +48,8 @@ class SC2MOGenMissionPools:
     _used_flags: Dict[MissionFlag, int]
     _used_missions: List[SC2Mission]
     _updated_difficulties: Dict[int, Difficulty]
+    _flag_ratios: Dict[MissionFlag, float]
+    _flag_weights: Dict[MissionFlag, int]
 
     def __init__(self) -> None:
         self.master_list = {mission.id for mission in SC2Mission}
@@ -58,6 +60,8 @@ def __init__(self) -> None:
         self._used_flags = {}
         self._used_missions = []
         self._updated_difficulties = {}
+        self._flag_ratios = {}
+        self._flag_weights = {}
 
     def set_exclusions(self, excluded: List[SC2Mission], unexcluded: List[SC2Mission]) -> None:
         """Prevents all the missions that appear in the `excluded` list, but not in the `unexcluded` list,
@@ -102,6 +106,62 @@ def get_used_missions(self) -> List[SC2Mission]:
         """Returns a set of all missions used in the mission order."""
         return self._used_missions
 
+    def set_flag_balances(self, flag_ratios: Dict[MissionFlag, int], flag_weights: Dict[MissionFlag, int]):
+        # Ensure the ratios are percentages
+        ratio_sum = sum(ratio for ratio in flag_ratios.values())
+        self._flag_ratios = {flag: ratio / ratio_sum for flag, ratio in flag_ratios.items()}
+        self._flag_weights = flag_weights
+
+    def pick_balanced_mission(self, world: World, pool: List[int]) -> int:
+        """Applies ratio-based and weight-based balancing to pick a preferred mission from a given mission pool."""
+        # Currently only used for race balancing
+        # Untested for flags that may overlap or not be present at all, but should at least generate
+        balanced_pool = pool
+        if len(self._flag_ratios) > 0:
+            relevant_used_flag_count = max(sum(self._used_flags.get(flag, 0) for flag in self._flag_ratios), 1)
+            current_ratios = {
+                flag: self._used_flags.get(flag, 0) / relevant_used_flag_count
+                for flag in self._flag_ratios
+            }
+            # Desirability of missions is the difference between target and current ratios for relevant flags
+            flag_scores = {
+                flag: self._flag_ratios[flag] - current_ratios[flag]
+                for flag in self._flag_ratios
+            }
+            mission_scores = [
+                sum(
+                    flag_scores[flag] for flag in self._flag_ratios
+                    if flag in lookup_id_to_mission[mission].flags
+                )
+                for mission in balanced_pool
+            ]
+            # Only keep the missions that create the best balance
+            best_score = max(mission_scores)
+            balanced_pool = [mission for idx, mission in enumerate(balanced_pool) if mission_scores[idx] == best_score]
+
+        balanced_weights = [1 for _ in balanced_pool]
+        if len(self._flag_weights) > 0:
+            relevant_used_flag_count = max(sum(self._used_flags.get(flag, 0) for flag in self._flag_weights), 1)
+            # Higher usage rate of relevant flags means lower desirability
+            flag_scores = {
+                flag: (relevant_used_flag_count - self._used_flags.get(flag, 0)) * self._flag_weights[flag]
+                for flag in self._flag_weights
+            }
+            # Mission scores are averaged across the mission's flags,
+            # else flags that aren't always present will inflate weights
+            mission_scores = [
+                sum(
+                    flag_scores[flag] for flag in self._flag_weights
+                    if flag in lookup_id_to_mission[mission].flags
+                ) / sum(flag in lookup_id_to_mission[mission].flags for flag in self._flag_weights)
+                for mission in balanced_pool
+            ]
+            balanced_weights = mission_scores
+
+        if sum(balanced_weights) == 0.0:
+            balanced_weights = [1.0 for _ in balanced_weights]
+        return world.random.choices(balanced_pool, balanced_weights, k=1)[0]
+
     def pull_specific_mission(self, mission: SC2Mission) -> None:
         """Marks the given mission as present in the mission order."""
         # Remove the mission from the master list and whichever difficulty pool it is in
@@ -169,7 +229,7 @@ def pull_random_mission(self, world: World, slot: 'SC2MOGenMission', *, prefer_c
                 difficulty_offset += 1
 
         # Remove the mission from the master list
-        mission = lookup_id_to_mission[world.random.choice(final_pool)]
+        mission = lookup_id_to_mission[self.pick_balanced_mission(world, final_pool)]
         self.master_list.remove(mission.id)
         self.difficulty_pools[self.get_modified_mission_difficulty(mission)].remove(mission.id)
         self._add_mission_stats(mission)

worlds/sc2/mission_order/options.py

@@ -83,6 +83,7 @@
 class CustomMissionOrder(OptionDict):
     """
     Used to generate a custom mission order. Please see documentation to understand usage.
+    Will do nothing unless `mission_order` is set to `custom`.
     """
     display_name = "Custom Mission Order"
     visibility = Visibility.template

worlds/sc2/options.py

@@ -303,6 +303,22 @@ class EnableRaceSwapVariants(Choice):
     default = option_disabled
 
 
+class EnableMissionRaceBalancing(Choice):
+    """
+    If enabled, picks missions in such a way that the appearance rate of races is roughly equal.
+    The final rates may deviate if there are not enough missions enabled to accomodate each race.
+
+    Disabled: Pick missions at random.
+    Semi Balanced: Use a weighting system to pick missions in a random, but roughly equal ratio.
+    Fully Balanced: Pick missions to preserve equal race counts whenever possible.
+    """
+    display_name = "Enable Mission Race Balancing"
+    option_disabled = 0
+    option_semi_balanced = 1
+    option_fully_balanced = 2
+    default = option_semi_balanced
+
+
 class ShuffleCampaigns(DefaultOnToggle):
     """
     Shuffles the missions between campaigns if enabled.
@@ -1002,6 +1018,7 @@ class Starcraft2Options(PerGameCommonOptions):
     enable_epilogue_missions: EnableEpilogueMissions
     enable_nco_missions: EnableNCOMissions
     enable_race_swap: EnableRaceSwapVariants
+    mission_race_balancing: EnableMissionRaceBalancing
     shuffle_campaigns: ShuffleCampaigns
     shuffle_no_build: ShuffleNoBuild
     starter_unit: StarterUnit

worlds/sc2/regions.py

@@ -1,11 +1,14 @@
 from typing import TYPE_CHECKING, List, Dict, Any, Tuple, Optional
 
 from .locations import LocationData, Location
-from .mission_tables import SC2Mission, SC2Campaign, get_campaign_goal_priority, campaign_final_mission_locations, campaign_alt_final_mission_locations
+from .mission_tables import (
+    SC2Mission, SC2Campaign, MissionFlag, get_campaign_goal_priority,
+    campaign_final_mission_locations, campaign_alt_final_mission_locations
+)
 from .options import (
     get_option_value, ShuffleNoBuild, RequiredTactics, ExtraLocations, ShuffleCampaigns,
     kerrigan_unit_available, TakeOverAIAllies, MissionOrder, get_excluded_missions, get_enabled_campaigns, static_mission_orders,
-    GridTwoStartPositions, KeyMode
+    GridTwoStartPositions, KeyMode, EnableMissionRaceBalancing
 )
 from .mission_order.options import CustomMissionOrder
 from .mission_order.structs import SC2MissionOrder, Difficulty
@@ -31,6 +34,7 @@ def create_mission_order(
     mission_pools = SC2MOGenMissionPools()
     mission_pools.set_exclusions(get_excluded_missions(world), []) # TODO set unexcluded
     adjust_mission_pools(world, mission_pools)
+    setup_mission_pool_balancing(world, mission_pools)
 
     mission_order_type = get_option_value(world, "mission_order")
     if mission_order_type == MissionOrder.option_custom:
@@ -142,6 +146,15 @@ def adjust_mission_pools(world: 'SC2World', pools: SC2MOGenMissionPools):
         # Flashpoint needs just a few items at start but competent comp at the end
         pools.move_mission(SC2Mission.FLASHPOINT, Difficulty.HARD, Difficulty.EASY)
 
+def setup_mission_pool_balancing(world: 'SC2World', pools: SC2MOGenMissionPools):
+    race_mission_balance = get_option_value(world, "mission_race_balancing")
+    flag_ratios: Dict[MissionFlag, int] = {}
+    flag_weights: Dict[MissionFlag, int] = {}
+    if race_mission_balance == EnableMissionRaceBalancing.option_semi_balanced:
+        flag_weights = { MissionFlag.Terran: 1, MissionFlag.Zerg: 1, MissionFlag.Protoss: 1 }
+    elif race_mission_balance == EnableMissionRaceBalancing.option_fully_balanced:
+        flag_ratios = { MissionFlag.Terran: 1, MissionFlag.Zerg: 1, MissionFlag.Protoss: 1 }
+    pools.set_flag_balances(flag_ratios, flag_weights)
 
 def create_regular_mission_order(world: 'SC2World', mission_pools: SC2MOGenMissionPools) -> Dict[str, Dict[str, Any]]:
     mission_order_type = get_option_value(world, "mission_order")

worlds/sc2/test/test_generation.py

@@ -815,4 +815,37 @@ def test_locking_required_items(self):
 
         # These items will be in the pool despite exclusions
         self.assertIn(item_names.KERRIGAN_LEAPING_STRIKE, itempool)
-        self.assertIn(item_names.KERRIGAN_MEND, itempool)
+        self.assertIn(item_names.KERRIGAN_MEND, itempool)
+
+
+    def test_fully_balanced_mission_races(self):
+        """
+        Tests whether fully balanced mission race balancing actually is fully balanced.
+        """
+        campaign_size = 57
+        self.assertEqual(campaign_size % 3, 0, "Chosen test size cannot be perfectly balanced")
+        world_options = {
+            # Reasonably large grid with enough missions to balance races
+            'mission_order': options.MissionOrder.option_grid,
+            'maximum_campaign_size': campaign_size,
+            'enable_wol_missions': True,
+            'enable_prophecy_missions': True,
+            'enable_hots_missions': True,
+            'enable_lotv_prologue_missions': True,
+            'enable_lotv_missions': True,
+            'enable_epilogue_missions': True,
+            'enable_nco_missions': True,
+            'selected_races': options.SelectRaces.option_all,
+            'enable_race_swap': options.EnableRaceSwapVariants.option_shuffle_all,
+            'mission_race_balancing': options.EnableMissionRaceBalancing.option_fully_balanced,
+        }
+
+        self.generate_world(world_options)
+        world_regions = [region.name for region in self.multiworld.regions]
+        world_regions.remove('Menu')
+        missions = [mission_tables.lookup_name_to_mission[region] for region in world_regions]
+        race_flags = [mission_tables.MissionFlag.Terran, mission_tables.MissionFlag.Zerg, mission_tables.MissionFlag.Protoss]
+        race_counts = { flag: sum(flag in mission.flags for mission in missions) for flag in race_flags }
+
+        self.assertEqual(race_counts[mission_tables.MissionFlag.Terran], race_counts[mission_tables.MissionFlag.Zerg])
+        self.assertEqual(race_counts[mission_tables.MissionFlag.Zerg], race_counts[mission_tables.MissionFlag.Protoss])