Agent that follows a pre-defined action sequence
| 92 | |
| 93 | |
| 94 | class TeacherForcingAgent(Agent): |
| 95 | """Agent that follows a pre-defined action sequence""" |
| 96 | |
| 97 | def __init__(self) -> None: |
| 98 | super().__init__() |
| 99 | |
| 100 | @beartype |
| 101 | def set_action_set_tag(self, tag: str) -> None: |
| 102 | self.action_set_tag = tag |
| 103 | |
| 104 | @beartype |
| 105 | def set_actions(self, action_seq: str | list[str]) -> None: |
| 106 | if isinstance(action_seq, str): |
| 107 | action_strs = action_seq.strip().split("\n") |
| 108 | else: |
| 109 | action_strs = action_seq |
| 110 | action_strs = [a.strip() for a in action_strs] |
| 111 | |
| 112 | actions = [] |
| 113 | for a_str in action_strs: |
| 114 | try: |
| 115 | if self.action_set_tag == "playwright": |
| 116 | cur_action = create_playwright_action(a_str) |
| 117 | elif self.action_set_tag == "id_accessibility_tree": |
| 118 | cur_action = create_id_based_action(a_str) |
| 119 | else: |
| 120 | raise ValueError( |
| 121 | f"Unknown action type {self.action_set_tag}" |
| 122 | ) |
| 123 | except ActionParsingError as e: |
| 124 | cur_action = create_none_action() |
| 125 | |
| 126 | cur_action["raw_prediction"] = a_str |
| 127 | actions.append(cur_action) |
| 128 | |
| 129 | self.actions: list[Action] = actions |
| 130 | |
| 131 | @beartype |
| 132 | def next_action( |
| 133 | self, trajectory: Trajectory, intent: str, meta_data: Any |
| 134 | ) -> Action: |
| 135 | """Predict the next action given the observation""" |
| 136 | return self.actions.pop(0) |
| 137 | |
| 138 | @beartype |
| 139 | def reset( |
| 140 | self, |
| 141 | test_config_file: str, |
| 142 | ) -> None: |
| 143 | with open(test_config_file) as f: |
| 144 | ref_actions = json.load(f)["reference_action_sequence"] |
| 145 | tag = ref_actions["action_set_tag"] |
| 146 | action_seq = ref_actions["action_sequence"] |
| 147 | self.set_action_set_tag(tag) |
| 148 | self.set_actions(action_seq) |
| 149 | |
| 150 | |
| 151 | class PromptAgent(Agent): |