Adding pub_coordination_mini scenario that is very small and fast to …

…run, but still has a meaningful challenge to it

PiperOrigin-RevId: 672882785
Change-Id: I27f4d3aced79229c85feb1abcd9480e0777fdc6e

examples/modular/environment/pub_coordination.py

@@ -64,13 +64,9 @@
 
 TIME_INCREMENT_BETWEEN_SCENES = datetime.timedelta(hours=24)
 
-NUM_MAIN_PLAYERS = 4
-NUM_SUPPORTING_PLAYERS = 1
-
 USE_CONVERSATION_GM = True
 
 FIRST_NAMES = player_names.FIRST_NAMES
-NUM_GAMES = 3
 
 
 @dataclasses.dataclass
@@ -89,6 +85,9 @@ class WorldConfig:
       dataclasses.field(default_factory=dict)
   )
   formative_memory_prompts: Mapping[str, Sequence[str]] | None = None
+  num_main_players: int = 4
+  num_supporting_players: int = 1
+  num_games: int = 3
 
 
 def get_shared_memories_and_context(
@@ -188,13 +187,16 @@ def configure_players(sampled_settings: Any) -> tuple[
     main_player_configs: configs for the main characters
     supporting_player_configs: configs for the supporting characters
   """
-  names = sampled_settings.people[: NUM_MAIN_PLAYERS + NUM_SUPPORTING_PLAYERS]
+  names = sampled_settings.people[
+      : sampled_settings.num_main_players
+      + sampled_settings.num_supporting_players
+  ]
   all_players = ', '.join(names)
   player_configs = []
 
   num_pubs = len(sampled_settings.venues)
 
-  for i in range(NUM_MAIN_PLAYERS):
+  for i in range(sampled_settings.num_main_players):
     name = names[i]
     favorite_pub = sampled_settings.venues[i % num_pubs]
     gender = sampled_settings.person_data[name]['gender']
@@ -209,8 +211,8 @@ def configure_players(sampled_settings: Any) -> tuple[
     )
     player_configs.append(config)
 
-  for i in range(NUM_SUPPORTING_PLAYERS):
-    name = names[NUM_MAIN_PLAYERS + i]
+  for i in range(sampled_settings.num_supporting_players):
+    name = names[sampled_settings.num_main_players + i]
     gender = sampled_settings.person_data[name]['gender']
     favorite_pub = sampled_settings.venues[0]
     config = configure_player(
@@ -448,7 +450,7 @@ def configure_scenes(
 
   secondary_environments = []
 
-  for i in range(NUM_GAMES):
+  for i in range(sampled_settings.num_games):
     closed_pub = None
     if random.random() < pub_closed_probability:
       closed_pub = random.choice(sampled_settings.venues)
@@ -595,6 +597,9 @@ def __init__(
       time_and_place_module: str | None = None,
       pub_closed_probability: float = 0.0,
       use_relational_matrix: bool = False,
+      num_games: int = 3,
+      num_main_players: int = 4,
+      num_supporting_players: int = 1,
   ):
     """Initialize the simulation object.
 
@@ -619,6 +624,9 @@ def __init__(
       pub_closed_probability: the probability that a pub is closed. Zero by
         default.
       use_relational_matrix: whether to use relational matrix or not.
+      num_games: the number of games to play.
+      num_main_players: the number of main players.
+      num_supporting_players: the number of supporting players.
     """
     # Support for these parameters will be added in a future addition coming
     # very imminently.
@@ -647,6 +655,10 @@ def __init__(
         )
     )
 
+    sampled_settings.num_main_players = num_main_players
+    sampled_settings.num_supporting_players = num_supporting_players
+    sampled_settings.num_games = num_games
+
     start_time = datetime.datetime(
         year=time_and_place_params.YEAR,
         month=time_and_place_params.MONTH,

examples/modular/environment/pub_coordination_mini.py

@@ -0,0 +1,74 @@
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A Concordia Environment Configuration."""
+
+from collections.abc import Callable, Sequence
+import types
+
+from examples.modular.environment import pub_coordination
+from concordia.factory.agent import basic_agent
+from concordia.language_model import language_model
+from concordia.utils import measurements as measurements_lib
+import numpy as np
+
+
+class Simulation(pub_coordination.Simulation):
+  """Simulation with pub closures."""
+
+  def __init__(
+      self,
+      model: language_model.LanguageModel,
+      embedder: Callable[[str], np.ndarray],
+      measurements: measurements_lib.Measurements,
+      agent_module: types.ModuleType = basic_agent,
+      resident_visitor_modules: Sequence[types.ModuleType] | None = None,
+      supporting_agent_module: types.ModuleType | None = None,
+      time_and_place_module: str | None = None,
+  ):
+    """Initialize the simulation object.
+
+    The launch script assumes this API object has a run() method.
+
+    Args:
+      model: the language model to use.
+      embedder: the sentence transformer to use.
+      measurements: the measurements object to use.
+      agent_module: the agent module to use for all main characters.
+      resident_visitor_modules: optionally, use different modules for majority
+        and minority parts of the focal population.
+      supporting_agent_module: agent module to use for all supporting players. A
+        supporting player is a non-player character with a persistent memory
+        that plays a specific role in defining the task environment. Their role
+        is not incidental but rather is critcal to making the task what it is.
+        Supporting players are not necessarily interchangeable with focal or
+        background players.
+      time_and_place_module: optionally, specify a module containing settings
+        that create a sense of setting in a specific time and place. If not
+        specified, a random module will be chosen from the default options.
+    """
+
+    super().__init__(
+        model=model,
+        embedder=embedder,
+        measurements=measurements,
+        agent_module=agent_module,
+        resident_visitor_modules=resident_visitor_modules,
+        supporting_agent_module=supporting_agent_module,
+        time_and_place_module=time_and_place_module,
+        pub_closed_probability=0.0,
+        num_games=1,
+        num_main_players=2,
+        num_supporting_players=0,
+    )

examples/modular/scenario/scenarios.py

@@ -99,8 +99,10 @@ class ScenarioConfig:
         supporting_agent_module='basic_puppet_agent',
     ),
     reality_show=SubstrateConfig(
-        description=('players are contestants on a reality show featuring '
-                     'social dilemma games alternating with conversation'),
+        description=(
+            'players are contestants on a reality show featuring '
+            'social dilemma games alternating with conversation'
+        ),
         environment='reality_show',
         supporting_agent_module=None,
     ),
@@ -142,6 +144,17 @@ class ScenarioConfig:
             'role playing',
         ),
     ),
+    pub_coordination_mini=ScenarioConfig(
+        description=(
+            'a mini scenario with one focal and one rational visitor and no'
+            ' supporting agents. Intended for fast testing.'
+        ),
+        substrate_config=SUBSTRATE_CONFIGS['pub_coordination'],
+        background_agent_module='rational_agent',
+        time_and_place_module='pub_coordination_london',
+        focal_is_resident=True,
+        tags=('coordination', 'persuasion'),
+    ),
     pub_coordination_0=ScenarioConfig(
         description=(
             'resident population of focal agents in a pub coordination scenario'
@@ -215,7 +228,8 @@ class ScenarioConfig:
         substrate_config=SUBSTRATE_CONFIGS['reality_show'],
         background_agent_module='paranoid_agent',
         time_and_place_module=(
-            'early_2000s_american_reality_show__stag_hunt_4_players'),
+            'early_2000s_american_reality_show__stag_hunt_4_players'
+        ),
         focal_is_resident=False,
         tags=(
             'discouraging antisocial behavior',
@@ -255,7 +269,8 @@ def build_simulation(
   if substrate_config.supporting_agent_module is not None:
     supporting_agent_module = importlib.import_module(
         f'{support_agent_base_module}.'
-        f'{substrate_config.supporting_agent_module}')
+        f'{substrate_config.supporting_agent_module}'
+    )
   else:
     supporting_agent_module = None
   runnable_simulation = simulation.Simulation(