Improve state formation environment

PiperOrigin-RevId: 686542040
Change-Id: I7f3df2b2c70d65bf169b570bc92bb8b71758b3a2

examples/modular/environment/modules/pre_state_villages.py

@@ -1913,10 +1913,20 @@ def sample_parameters(
   config.sample_event_treaty_in_effect = lambda: rng.choice(
       _TREATY_IN_EFFECT_DESCRIPTIONS)
 
-  config.villager_how_things_are_constant = (
+  config.villager_how_things_are_constant = config_dict.ConfigDict()
+  config.villager_how_things_are_constant.village_a = (
       'Everyone in {name}\'s family has always been a '
       'farmer. As long as anyone can remember, they have farmed their '
-      'ancestral lands near {village_name}.'
+      'ancestral lands near {village_name}. {name} hates the idea of spending '
+      'time on any activity other than farming or leisure.'
+  )
+  config.villager_how_things_are_constant.village_b = (
+      '{name}\'s family values strength and believes in the importance of '
+      'training for war. They have always felt that they would like to train '
+      'more but the village\'s lack of food security has always made it hard. '
+      'If one spends too much time training then they compromise farming and '
+      'risk starvation. But if freedom from the threat of famine could be '
+      'achieved then {name} would gladly spend more time training for war.'
   )
   config.negotiation_phase_extra_premise = (
       'Agriculture is critical to both villages. But, the more time spent '

examples/modular/environment/state_formation.py

@@ -894,9 +894,17 @@ def __init__(
           state=(f'{player_config.name} lives in {village_name} and never '
                  'travels elsewhere.'),
       )
-      ancestral_job = agent_components.constant.Constant(
+      if village_name == config.village_a_name:
+        how_things_are_string = (
+            config.villager_how_things_are_constant.village_a)
+      elif village_name == config.village_b_name:
+        how_things_are_string = (
+            config.villager_how_things_are_constant.village_b)
+      else:
+        raise ValueError(f'Unknown village name: {village_name}')
+      how_things_are = agent_components.constant.Constant(
           pre_act_key='\nHow things are',
-          state=config.villager_how_things_are_constant.format(
+          state=how_things_are_string.format(
               name=player_config.name,
               village_name=village_name,
           ),
@@ -909,7 +917,7 @@ def __init__(
       )
       supporter_extra_components[player_config.name] = {
           'Home village': which_village,
-          'How things are': ancestral_job,
+          'How things are': how_things_are,
           'Guiding principle of good conversation': conversation_style,
       }