Map enemy robots as well

Author: nycrat

src/software/gameplay_tests/field_test_runner.py

@@ -85,13 +85,22 @@ def set_world_state(self, world_state: WorldState):
 
     @override
     def set_tactics(self, blue_tactics={}, yellow_tactics={}):
-        """Override AI tactics, remapping the friendly team's simulated robot ids
-        to the robot ids that are actually available on the field.
+        """Override AI tactics, remapping each team's simulated robot ids to the
+        robot ids that are actually available on the field.
         """
         if self.is_yellow_friendly:
-            yellow_tactics = self._map_tactics_to_field_ids(yellow_tactics)
+            blue_sim_to_field_id = self.enemy_sim_to_field_id
+            yellow_sim_to_field_id = self.friendly_sim_to_field_id
         else:
-            blue_tactics = self._map_tactics_to_field_ids(blue_tactics)
+            blue_sim_to_field_id = self.friendly_sim_to_field_id
+            yellow_sim_to_field_id = self.enemy_sim_to_field_id
+
+        blue_tactics = self._map_tactics_to_field_ids(
+            blue_tactics, blue_sim_to_field_id, "blue"
+        )
+        yellow_tactics = self._map_tactics_to_field_ids(
+            yellow_tactics, yellow_sim_to_field_id, "yellow"
+        )
 
         super().set_tactics(blue_tactics=blue_tactics, yellow_tactics=yellow_tactics)
 
@@ -205,16 +214,21 @@ def _survey_field_robots(self):
         # survey field for available robot ids
         survey_start_time = time.time()
         self.friendly_robot_ids_field = []
+        self.enemy_robot_ids_field = []
         while time.time() - survey_start_time < WORLD_BUFFER_TIMEOUT:
             try:
                 world = self.world_buffer.get(block=True, timeout=0.1)
                 self.initial_world = world
                 self.friendly_robot_ids_field = [
                     robot.id for robot in world.friendly_team.team_robots
                 ]
+                self.enemy_robot_ids_field = [
+                    robot.id for robot in world.enemy_team.team_robots
+                ]
 
                 if len(self.friendly_robot_ids_field) > 0:
                     logger.info(f"friendly team ids {self.friendly_robot_ids_field}")
+                    logger.info(f"enemy team ids {self.enemy_robot_ids_field}")
                     break
             except queue.Empty:
                 continue
@@ -225,20 +239,41 @@ def _survey_field_robots(self):
         # Simulated tests create robots with contiguous ids (0, 1, 2, ...), but
         # the robots physically on the field may have arbitrary ids and there may
         # be fewer of them. Map each simulated id to an available field id, in
-        # ascending order, so tactics written against simulated ids are sent to
-        # the robots that are actually present. e.g. field robots [4, 5] map
-        # simulated id 0 -> 4 and 1 -> 5; field robots [0, 1, 5] map 0 -> 0,
-        # 1 -> 1 and 2 -> 5.
-        self.sim_to_field_robot_id = {
-            sim_id: field_id
-            for sim_id, field_id in enumerate(sorted(self.friendly_robot_ids_field))
-        }
-        # Inverse mapping, used to relabel the world's field ids back to the
+        # ascending order and per team, so tactics and validations written against
+        # simulated ids refer to the robots that are actually present. e.g. field
+        # robots [4, 5] map simulated id 0 -> 4 and 1 -> 5; field robots [0, 1, 5]
+        # map 0 -> 0, 1 -> 1 and 2 -> 5.
+        self.friendly_sim_to_field_id = self._build_sim_to_field_mapping(
+            self.friendly_robot_ids_field
+        )
+        self.enemy_sim_to_field_id = self._build_sim_to_field_mapping(
+            self.enemy_robot_ids_field
+        )
+        # Inverse mappings, used to relabel the world's field ids back to the
         # simulated ids that validations are written against.
-        self.field_to_sim_robot_id = {
-            field_id: sim_id for sim_id, field_id in self.sim_to_field_robot_id.items()
+        self.friendly_field_to_sim_id = {
+            field_id: sim_id
+            for sim_id, field_id in self.friendly_sim_to_field_id.items()
+        }
+        self.enemy_field_to_sim_id = {
+            field_id: sim_id for sim_id, field_id in self.enemy_sim_to_field_id.items()
         }
-        logger.info(f"simulated id to field id mapping {self.sim_to_field_robot_id}")
+        logger.info(
+            f"friendly simulated id to field id mapping {self.friendly_sim_to_field_id}"
+        )
+        logger.info(
+            f"enemy simulated id to field id mapping {self.enemy_sim_to_field_id}"
+        )
+
+    @staticmethod
+    def _build_sim_to_field_mapping(field_ids):
+        """Map contiguous simulated robot ids (0, 1, 2, ...) onto the available
+        field ids for a team, in ascending order.
+
+        :param field_ids: the robot ids present on the field for a team
+        :return: dict of simulated_robot_id -> field_robot_id
+        """
+        return {sim_id: field_id for sim_id, field_id in enumerate(sorted(field_ids))}
 
     def _create_move_tactics(self, robot_states):
         """Create a MoveTactic for each robot to drive it to its world state.
@@ -259,58 +294,89 @@ def _create_move_tactics(self, robot_states):
             for robot_id, robot_state in robot_states.items()
         }
 
-    def _map_tactics_to_field_ids(self, tactics):
+    def _map_tactics_to_field_ids(self, tactics, sim_to_field_id, team_name):
         """Remap a tactics dict keyed by simulated robot ids to the robot ids
         available on the field. Tactics for simulated ids with no available
         field robot are dropped with a warning.
 
         :param tactics: None or dict of simulated_robot_id -> tactic
+        :param sim_to_field_id: the team's simulated-id -> field-id mapping
+        :param team_name: the team colour, used in warnings
         :return: None if tactics is None, else dict of field_robot_id -> tactic
         """
         if tactics is None:
             return None
 
         mapped_tactics = {}
         for sim_id, tactic in tactics.items():
-            if sim_id not in self.sim_to_field_robot_id:
+            if sim_id not in sim_to_field_id:
                 logger.warning(
-                    f"no field robot available for simulated robot id {sim_id}; "
-                    "skipping its tactic"
+                    f"no {team_name} field robot available for simulated robot id "
+                    f"{sim_id}; skipping its tactic"
                 )
                 continue
-            mapped_tactics[self.sim_to_field_robot_id[sim_id]] = tactic
+            mapped_tactics[sim_to_field_id[sim_id]] = tactic
 
         return mapped_tactics
 
+    @staticmethod
+    def _remap_robot_state_ids(robot_states, sim_to_field_id):
+        """Remap a world-state robot map keyed by simulated ids to field ids,
+        dropping any simulated id with no available field robot.
+
+        :param robot_states: map of simulated_robot_id -> RobotState
+        :param sim_to_field_id: the team's simulated-id -> field-id mapping
+        :return: dict of field_robot_id -> RobotState
+        """
+        return {
+            sim_to_field_id[sim_id]: robot_state
+            for sim_id, robot_state in robot_states.items()
+            if sim_id in sim_to_field_id
+        }
+
     def _relabel_world_to_sim_ids(self, world: World) -> World:
-        """Return a copy of the world with friendly robot ids translated from
+        """Return a copy of the world with both teams' robot ids translated from
         field ids back to the simulated ids the test uses.
 
         Validations match robots by the contiguous simulated ids (0, 1, 2, ...),
         but the world reports the robots' actual field ids. Relabeling the world
         here lets every id-based validation find the correct robot without each
-        validation needing to know about the mapping. Friendly robots with no
-        mapping (e.g. extras not part of the test) keep their field id.
+        validation needing to know about the mapping. Robots with no mapping
+        (e.g. extras not part of the test) keep their field id.
 
         :param world: The World proto reported by the field full system
-        :return: a copy of the world with friendly robot ids in simulated-id space
+        :return: a copy of the world with robot ids in simulated-id space
         """
         relabeled_world = World()
         relabeled_world.CopyFrom(world)
-        for robot in relabeled_world.friendly_team.team_robots:
-            if robot.id in self.field_to_sim_robot_id:
-                robot.id = self.field_to_sim_robot_id[robot.id]
+        self._relabel_team_to_sim_ids(
+            relabeled_world.friendly_team, self.friendly_field_to_sim_id
+        )
+        self._relabel_team_to_sim_ids(
+            relabeled_world.enemy_team, self.enemy_field_to_sim_id
+        )
         return relabeled_world
 
+    @staticmethod
+    def _relabel_team_to_sim_ids(team, field_to_sim_id):
+        """Relabel a team's robot ids in place from field ids to simulated ids.
+
+        :param team: a Team proto whose robots' ids to relabel
+        :param field_to_sim_id: the team's field-id -> simulated-id mapping
+        """
+        for robot in team.team_robots:
+            if robot.id in field_to_sim_id:
+                robot.id = field_to_sim_id[robot.id]
+
     def _wait_for_robots_at_world_state(self, world_state: WorldState):
         """Block until every robot in the world state is within
         ROBOT_SETUP_TOLERANCE_M of its target position.
 
         Robot positions are read from the friendly full system's World in the
         same fashion as _survey_field_robots. The friendly team maps to the
         friendly-coloured robots in the world state, and the enemy team to the
-        other colour. The friendly targets are remapped from simulated ids to
-        the field ids they were assigned to in set_tactics.
+        other colour. Each team's targets are remapped from simulated ids to the
+        field ids they were assigned to in set_tactics.
 
         :param world_state: The WorldState proto with the desired robot states
         :raises Exception: if the robots do not reach their targets within
@@ -323,13 +389,14 @@ def _wait_for_robots_at_world_state(self, world_state: WorldState):
             friendly_targets = world_state.blue_robots
             enemy_targets = world_state.yellow_robots
 
-        # Remap the friendly targets onto the field ids they were commanded to,
+        # Remap each team's targets onto the field ids they were commanded to,
         # dropping any simulated id that has no available field robot
-        friendly_targets = {
-            self.sim_to_field_robot_id[sim_id]: robot_state
-            for sim_id, robot_state in friendly_targets.items()
-            if sim_id in self.sim_to_field_robot_id
-        }
+        friendly_targets = self._remap_robot_state_ids(
+            friendly_targets, self.friendly_sim_to_field_id
+        )
+        enemy_targets = self._remap_robot_state_ids(
+            enemy_targets, self.enemy_sim_to_field_id
+        )
 
         logger.info("waiting for robots to reach their target positions")
         wait_start_time = time.time()