diff --git a/align_system/algorithms/pipeline_adm.py b/align_system/algorithms/pipeline_adm.py index 95b56a46..7813948c 100644 --- a/align_system/algorithms/pipeline_adm.py +++ b/align_system/algorithms/pipeline_adm.py @@ -11,6 +11,11 @@ class PipelineADM(ActionBasedADM): def __init__(self, steps: list[ADMComponent]): self.steps = steps + def reset_history(self): + for step in self.steps: + if hasattr(step, 'reset_history'): + step.reset_history() + def choose_action(self, scenario_state, available_actions, diff --git a/align_system/algorithms/world_state_adm_component.py b/align_system/algorithms/world_state_adm_component.py new file mode 100644 index 00000000..069dbca2 --- /dev/null +++ b/align_system/algorithms/world_state_adm_component.py @@ -0,0 +1,206 @@ +from copy import deepcopy + +from align_system.algorithms.abstracts import ADMComponent +from align_system.utils import logging + +log = logging.getLogger(__name__) + + +def _get_field(obj, field, default=None): + if isinstance(obj, dict): + return obj.get(field, default) + return getattr(obj, field, default) + + +def _set_field(obj, field, value): + if isinstance(obj, dict): + obj[field] = value + else: + setattr(obj, field, value) + + +def _get_items(obj): + if isinstance(obj, dict): + return list(obj.items()) + elif hasattr(obj, 'to_dict'): + return list(obj.to_dict().items()) + elif hasattr(obj, '__dict__'): + return list(obj.__dict__.items()) + return [] + + +def _merge_dict_or_obj(existing_obj, new_obj, preserve_if_new_none=False): + """ + Recursively merges new_obj into existing_obj. + Supports both dicts and objects with properties/attributes. + """ + if new_obj is None: + return existing_obj + + if existing_obj is None: + return new_obj + + for key, new_val in _get_items(new_obj): + if new_val is None and preserve_if_new_none: + continue + + existing_val = _get_field(existing_obj, key) + + if (isinstance(new_val, dict) or hasattr(new_val, 'to_dict')) and not isinstance(new_val, list): + if existing_val is not None: + merged_sub = _merge_dict_or_obj(existing_val, new_val, preserve_if_new_none=preserve_if_new_none) + _set_field(existing_obj, key, merged_sub) + else: + _set_field(existing_obj, key, new_val) + else: + _set_field(existing_obj, key, new_val) + + return existing_obj + + +class WorldStateTrackerADMComponent(ADMComponent): + """ + ADM Component that maintains cumulative world state across scenes/probes in a scenario. + + Handles recursive merging of character details (e.g. `vitals`). Preserves detailed character + information acquired when `nearby == True` and vitals assessed during interaction, preventing + less-detailed descriptions from overwriting them when moving away (`nearby == False`). Dynamic + attributes (e.g. `tag`, `visited`, `unseen`, top-level `supplies`) continue to update normally. + """ + def __init__(self, + output_scenario_state_key='scenario_state', + output_original_state_key='original_scenario_state', + output_stale_info_key='world_state_stale_info'): + self.output_scenario_state_key = output_scenario_state_key + self.output_original_state_key = output_original_state_key + self.output_stale_info_key = output_stale_info_key + + self.current_scenario_id = None + self.world_state = None + self.stale_characters = {} + + def reset_history(self): + log.info("[bold]*Resetting WorldStateTrackerADMComponent history*[/bold]", + extra={"markup": True}) + self.current_scenario_id = None + self.world_state = None + self.stale_characters = {} + + def run_returns(self): + return (self.output_scenario_state_key, + self.output_original_state_key, + self.output_stale_info_key) + + def run(self, scenario_state, scenario_id=None): + # Auto-reset if scenario_id changes unexpectedly without an explicit reset_history call + if scenario_id is not None and scenario_id != self.current_scenario_id: + if self.current_scenario_id is not None: + log.info(f"Scenario ID changed from {self.current_scenario_id} to {scenario_id}, resetting world state") + self.reset_history() + self.current_scenario_id = scenario_id + + original_scenario_state = scenario_state + # Deepcopy the incoming scenario_state up front to ensure complete reference independence + new_state = deepcopy(scenario_state) + + if self.world_state is None: + # First scene of the scenario: initialize world state as deep copy of incoming state + self.world_state = new_state + self.stale_characters = {} + else: + # Merge incoming state into cumulative world_state + self._merge_scenario_state(new_state) + + stale_info = { + 'stale_characters': deepcopy(self.stale_characters) + } + + return self.world_state, original_scenario_state, stale_info + + def _merge_scenario_state(self, new_state): + # 1. Update top-level fields (unstructured narrative, elapsed_time, meta_info, supplies, etc.) + for key, val in _get_items(new_state): + if key != 'characters': + _set_field(self.world_state, key, val) + + # 2. Merge characters + existing_chars = _get_field(self.world_state, 'characters') + if existing_chars is None: + existing_chars = [] + _set_field(self.world_state, 'characters', existing_chars) + + existing_char_map = {} + for char in existing_chars: + cid = _get_field(char, 'id') or _get_field(char, 'name') + if cid: + existing_char_map[cid] = char + + new_chars = _get_field(new_state, 'characters') or [] + + for new_char in new_chars: + cid = _get_field(new_char, 'id') or _get_field(new_char, 'name') + if not cid: + existing_chars.append(new_char) + continue + + if cid not in existing_char_map: + existing_chars.append(new_char) + existing_char_map[cid] = new_char + else: + existing_char = existing_char_map[cid] + self._merge_character(existing_char, new_char, cid) + + def _merge_character(self, existing_char, new_char, cid): + old_nearby = bool(_get_field(existing_char, 'nearby', False)) + new_nearby = bool(_get_field(new_char, 'nearby', False)) + stale_fields_for_char = [] + + if old_nearby and not new_nearby: + # Character was previously nearby (rich detail) but is now not nearby + old_unstructured = _get_field(existing_char, 'unstructured') + old_vitals = deepcopy(_get_field(existing_char, 'vitals')) + + # Update all fields from new_char + for key, val in _get_items(new_char): + if key not in {'unstructured', 'vitals'}: + _set_field(existing_char, key, val) + + # Preserve richer unstructured description + if old_unstructured is not None: + _set_field(existing_char, 'unstructured', old_unstructured) + stale_fields_for_char.append('unstructured') + + # Merge/preserve vitals non-destructively + new_vitals = _get_field(new_char, 'vitals') + if old_vitals is not None: + if new_vitals is None: + _set_field(existing_char, 'vitals', old_vitals) + else: + merged_vitals = _merge_dict_or_obj(old_vitals, new_vitals, preserve_if_new_none=True) + _set_field(existing_char, 'vitals', merged_vitals) + stale_fields_for_char.append('vitals') + elif new_vitals is not None: + _set_field(existing_char, 'vitals', new_vitals) + + if stale_fields_for_char: + self.stale_characters[cid] = { + 'stale_fields': stale_fields_for_char, + 'nearby': False + } + log.debug(f"Preserving detailed info ({stale_fields_for_char}) for non-nearby character '{cid}'") + else: + # Character is now nearby (or was never nearby) -> recursive update + old_vitals = _get_field(existing_char, 'vitals') + + for key, val in _get_items(new_char): + if key == 'vitals': + if old_vitals is not None and val is not None: + merged_vitals = _merge_dict_or_obj(old_vitals, val) + _set_field(existing_char, key, merged_vitals) + elif val is not None: + _set_field(existing_char, key, val) + else: + _set_field(existing_char, key, val) + + if cid in self.stale_characters: + del self.stale_characters[cid] diff --git a/align_system/configs/adm_component/misc/world_state_tracker.yaml b/align_system/configs/adm_component/misc/world_state_tracker.yaml new file mode 100644 index 00000000..d538a008 --- /dev/null +++ b/align_system/configs/adm_component/misc/world_state_tracker.yaml @@ -0,0 +1 @@ +_target_: align_system.algorithms.world_state_adm_component.WorldStateTrackerADMComponent diff --git a/tests/test_world_state_adm_component.py b/tests/test_world_state_adm_component.py new file mode 100644 index 00000000..6c675af4 --- /dev/null +++ b/tests/test_world_state_adm_component.py @@ -0,0 +1,229 @@ +import pytest +from copy import deepcopy + +from align_system.algorithms.world_state_adm_component import WorldStateTrackerADMComponent +from align_system.algorithms.pipeline_adm import PipelineADM +from align_system.algorithms.random_adm_component import RandomChoiceADMComponent + + +def test_world_state_tracker_initialization_and_update(): + component = WorldStateTrackerADMComponent() + + scene_1 = { + "unstructured": "Scene 1: Initial area.", + "elapsed_time": 0, + "meta_info": {"scene_id": "scene_1"}, + "supplies": [ + {"type": "Tourniquet", "quantity": 999}, + {"type": "Pressure bandage", "quantity": 999} + ], + "characters": [ + { + "id": "Patient 1", + "name": "Patient 1", + "unstructured": "Local civilian with moderate bleeding from their thigh", + "nearby": True, + "unseen": False, + "visited": False, + "tag": None, + "vitals": { + "avpu": "ALERT", + "breathing": "NORMAL", + "heart_rate": "FAST" + } + }, + { + "id": "Patient 2", + "name": "Patient 2", + "unstructured": "Military member of your unit with a broken left wrist, no bleeding", + "nearby": True, + "unseen": False, + "visited": False, + "tag": None, + "vitals": { + "avpu": "ALERT", + "breathing": "NORMAL", + "heart_rate": "FAST" + } + }, + { + "id": "Patient 6", + "name": "Patient 6", + "unstructured": "Civilian likely responsible for explosion", + "nearby": False, + "unseen": False, + "visited": False, + "tag": None + } + ] + } + + # First probe call + world_state, original_state, stale_info = component.run(scene_1, scenario_id="scenario_1") + + assert world_state["unstructured"] == "Scene 1: Initial area." + assert len(world_state["characters"]) == 3 + assert world_state["characters"][0]["vitals"]["avpu"] == "ALERT" + assert original_state == scene_1 + assert stale_info["stale_characters"] == {} + + # Second probe call: + # - Move away from Patient 1 (nearby becomes False, vitals omitted) + # - Patient 6 becomes nearby with vitals assessed + # - Tourniquet used (quantity decreases from 999 to 998) + scene_2 = { + "unstructured": "Scene 2: Moved to patient 6 area.", + "elapsed_time": 10, + "meta_info": {"scene_id": "scene_2"}, + "supplies": [ + {"type": "Tourniquet", "quantity": 998}, + {"type": "Pressure bandage", "quantity": 999} + ], + "characters": [ + { + "id": "Patient 1", + "name": "Patient 1", + "unstructured": "Patient 1 far away.", + "nearby": False, + "unseen": False, + "visited": True, # dynamic field updated + "tag": "RED" + # vitals omitted in scene_2 + }, + { + "id": "Patient 2", + "name": "Patient 2", + "unstructured": "Military member of your unit with a broken left wrist, no bleeding", + "nearby": True, + "unseen": False, + "visited": True, + "tag": None, + "vitals": { + "avpu": "ALERT", + "breathing": "NORMAL", + "heart_rate": "FAST" + } + }, + { + "id": "Patient 6", + "name": "Patient 6", + "unstructured": "Civilian likely responsible for the explosion with moderate bleeding from their torso; possible internal bleeding", + "nearby": True, + "unseen": False, + "visited": True, + "tag": None, + "vitals": { + "avpu": "VOICE", + "breathing": "FAST", + "heart_rate": "FAST" + } + } + ] + } + + world_state_2, original_state_2, stale_info_2 = component.run(scene_2, scenario_id="scenario_1") + + # Patient 1 should RETAIN detailed unstructured description AND vitals from scene 1 + p1 = next(c for c in world_state_2["characters"] if c["id"] == "Patient 1") + assert p1["unstructured"] == "Local civilian with moderate bleeding from their thigh" + assert p1["vitals"]["avpu"] == "ALERT" + assert p1["vitals"]["breathing"] == "NORMAL" + assert p1["visited"] is True + assert p1["tag"] == "RED" + assert p1["nearby"] is False + + # Patient 6 was nearby=False in scene 1, nearby=True in scene 2 -> Should update with new detailed description and vitals + p6 = next(c for c in world_state_2["characters"] if c["id"] == "Patient 6") + assert p6["unstructured"] == "Civilian likely responsible for the explosion with moderate bleeding from their torso; possible internal bleeding" + assert p6["vitals"]["avpu"] == "VOICE" + assert p6["nearby"] is True + + # Supplies quantity should be updated to 998 + t_supply = next(s for s in world_state_2["supplies"] if s["type"] == "Tourniquet") + assert t_supply["quantity"] == 998 + + # Stale info should track Patient 1 as having stale fields + assert "Patient 1" in stale_info_2["stale_characters"] + assert "unstructured" in stale_info_2["stale_characters"]["Patient 1"]["stale_fields"] + assert "vitals" in stale_info_2["stale_characters"]["Patient 1"]["stale_fields"] + + +def test_vitals_recursive_accumulation(): + component = WorldStateTrackerADMComponent() + + scene_1 = { + "characters": [ + { + "id": "P1", + "nearby": True, + "vitals": {"avpu": "ALERT"} + } + ] + } + + # Step 1: Initial vitals (avpu) + ws1, _, _ = component.run(scene_1, scenario_id="s1") + assert ws1["characters"][0]["vitals"] == {"avpu": "ALERT"} + + # Step 2: Assessment adds spo2 + scene_2 = { + "characters": [ + { + "id": "P1", + "nearby": True, + "vitals": {"spo2": "98%"} + } + ] + } + ws2, _, _ = component.run(scene_2, scenario_id="s1") + # Both avpu and spo2 should be present + assert ws2["characters"][0]["vitals"]["avpu"] == "ALERT" + assert ws2["characters"][0]["vitals"]["spo2"] == "98%" + + +def test_world_state_tracker_reset_history_and_scenario_change(): + component = WorldStateTrackerADMComponent() + + scene_s1 = { + "unstructured": "Scenario 1 Scene", + "characters": [{"id": "P1", "unstructured": "Details S1", "nearby": True}] + } + + component.run(scene_s1, scenario_id="scenario_1") + assert component.world_state is not None + + # Reset history + component.reset_history() + assert component.world_state is None + assert component.current_scenario_id is None + + # Auto reset on scenario ID change + component.run(scene_s1, scenario_id="scenario_1") + assert component.world_state["unstructured"] == "Scenario 1 Scene" + + scene_s2 = { + "unstructured": "Scenario 2 Scene", + "characters": [{"id": "P10", "unstructured": "Details S2", "nearby": True}] + } + + world_state_s2, _, _ = component.run(scene_s2, scenario_id="scenario_2") + assert world_state_s2["unstructured"] == "Scenario 2 Scene" + assert len(world_state_s2["characters"]) == 1 + assert world_state_s2["characters"][0]["id"] == "P10" + + +def test_pipeline_adm_reset_history_propagation(): + tracker = WorldStateTrackerADMComponent() + random_comp = RandomChoiceADMComponent() + + pipeline = PipelineADM(steps=[tracker, random_comp]) + + scene = {"unstructured": "Test", "characters": []} + tracker.run(scene, scenario_id="test_scen") + + assert tracker.world_state is not None + + # Call reset_history on PipelineADM + pipeline.reset_history() + + assert tracker.world_state is None