Repository: isi-vista/adam

Issue:
In SituationTemplates allow `before_action_relations` and `after_action_relations`

Comments:

Patch:
diff --git a/adam/situation/templates/phase1_templates.py b/adam/situation/templates/phase1_templates.py
index 93532153..afa6dc4f 100644
--- a/adam/situation/templates/phase1_templates.py
+++ b/adam/situation/templates/phase1_templates.py
@@ -208,6 +208,34 @@ class Phase1SituationTemplate(SituationTemplate):
     A set of `TemplateObjectVariables` s which are the focus of the speaker. 
     Defaults to all semantic role fillers of situation actions.
     """
+    before_action_relations: ImmutableSet[Relation[TemplateObjectVariable]] = attrib(
+        converter=_to_immutableset, kw_only=True, default=immutableset()
+    )
+    """
+    The relations which hold in this `SituationTemplate`,
+    before, but not necessarily after, any actions which occur.
+    
+    It is not necessary to state every relationship which holds in a situation.
+    Rather this should contain the salient relationships
+    which should be expressed in the linguistic description.
+    
+    Do not specify those relations here which are *implied* by any actions which occur.
+    Those are handled automatically. 
+    """
+    after_action_relations: ImmutableSet[Relation[TemplateObjectVariable]] = attrib(
+        converter=_to_immutableset, kw_only=True, default=immutableset()
+    )
+    """
+    The relations which hold in this `SituationTemplate`,
+    after, but not necessarily before, any actions which occur.
+
+    It is not necessary to state every relationship which holds in a situation.
+    Rather this should contain the salient relationships
+    which should be expressed in the linguistic description.
+
+    Do not specify those relations here which are *implied* by any actions which occur.
+    Those are handled automatically. 
+    """
 
     def __attrs_post_init__(self) -> None:
         check_arg(
@@ -527,6 +555,14 @@ class _Phase1SituationTemplateGenerator(
                 )
                 for action in template.actions
             ],
+            before_action_relations=[
+                relation.copy_remapping_objects(object_var_to_instantiations)
+                for relation in template.before_action_relations
+            ],
+            after_action_relations=[
+                relation.copy_remapping_objects(object_var_to_instantiations)
+                for relation in template.after_action_relations
+            ],
             syntax_hints=template.syntax_hints,
             axis_info=self._compute_axis_info(object_var_to_instantiations),
             gazed_objects=immutableset(


Test patch:
diff --git a/tests/situation/templates/phase1_template_test.py b/tests/situation/templates/phase1_template_test.py
index 51796f87..d9975d6d 100644
--- a/tests/situation/templates/phase1_template_test.py
+++ b/tests/situation/templates/phase1_template_test.py
@@ -23,9 +23,18 @@ from adam.ontology.phase1_ontology import (
     near,
     GAILA_PHASE_1_ONTOLOGY,
     MOM,
+    GROUND,
+    BOX,
+    ROLL,
+    AGENT,
+    ROLL_SURFACE_AUXILIARY,
+    on,
+    far,
 )
 from adam.ontology.structural_schema import ObjectStructuralSchema
 from adam.random_utils import RandomChooser
+from adam.relation import flatten_relations
+from adam.situation import Action
 from adam.situation.templates.phase1_templates import (
     Phase1SituationTemplate,
     all_possible,
@@ -190,3 +199,34 @@ def test_learner_as_default_addressee():
 
     assert situation_with_addressee[0].axis_info
     assert situation_with_addressee[0].axis_info.addressee
+
+
+def test_before_after_relations_asserted():
+    ball = object_variable("ball", root_node=BALL)
+    box = object_variable("box", root_node=BOX)
+    ground = object_variable("ground", root_node=GROUND)
+
+    template_action = Phase1SituationTemplate(
+        "Before/After Relation",
+        salient_object_variables=[ball, box],
+        background_object_variables=[ground],
+        actions=[
+            Action(
+                ROLL,
+                argument_roles_to_fillers=[(AGENT, ball)],
+                auxiliary_variable_bindings=[(ROLL_SURFACE_AUXILIARY, ground)],
+            )
+        ],
+        before_action_relations=flatten_relations([on(ball, box)]),
+        after_action_relations=flatten_relations([far(ball, box)]),
+    )
+
+    situation_with_relations = sampled(
+        template_action,
+        ontology=GAILA_PHASE_1_ONTOLOGY,
+        chooser=RandomChooser.for_seed(0),
+        max_to_sample=1,
+    )
+
+    assert situation_with_relations[0].before_action_relations
+    assert situation_with_relations[0].after_action_relations