Filter vision detections to roster defined by vision→cmd mapping

utama_core/config/settings.py

@@ -31,7 +31,7 @@
 BAUD_RATE = 115200
 PORT = "/dev/ttyACM0"
 TIMEOUT = 0.1
-KICKER_COOLDOWN_TIME = 10  # in seconds to prevent kicker from being actuated too frequently
+KICKER_COOLDOWN_TIME = 1  # in seconds to prevent kicker from being actuated too frequently
 KICKER_COOLDOWN_TIMESTEPS = int(KICKER_COOLDOWN_TIME * CONTROL_FREQUENCY)  # in timesteps
 KICKER_PERSIST_TIMESTEPS = 10  # in timesteps to persist the kick command
 

utama_core/data_processing/refiners/position.py

@@ -60,6 +60,8 @@ def __init__(
         full_field_dims: FieldDimensions,
         filtering: bool = True,
         exp_ball: bool = True,
+        allowed_yellow_ids: Optional[frozenset[int]] = None,
+        allowed_blue_ids: Optional[frozenset[int]] = None,
     ):
         # alpha=0 means no change in angle (inf smoothing), alpha=1 means no smoothing
         self.angle_smoother = AngleSmoother(alpha=1)
@@ -80,6 +82,8 @@ def __init__(
         )
 
         self.exp_ball = exp_ball
+        self.allowed_yellow_ids = allowed_yellow_ids
+        self.allowed_blue_ids = allowed_blue_ids
 
         if self.filtering:
             # Instantiate a dedicated Kalman filter for each robot so filtering can be kept independent.
@@ -238,13 +242,22 @@ def _include_vanished_robots(
             game_frame.enemy_robots.keys(),
         )
 
-        # Current vision IDs
-        yellow_present = {r.id for r in vision_data.yellow_robots}
-        blue_present = {r.id for r in vision_data.blue_robots}
+        # Current vision IDs (filtered to roster if provided)
+        yellow_robots = vision_data.yellow_robots
+        blue_robots = vision_data.blue_robots
+        if self.allowed_yellow_ids is not None:
+            yellow_robots = [r for r in yellow_robots if r.id in self.allowed_yellow_ids]
+            yellow_ids_last_frame = yellow_ids_last_frame & self.allowed_yellow_ids
+        if self.allowed_blue_ids is not None:
+            blue_robots = [r for r in blue_robots if r.id in self.allowed_blue_ids]
+            blue_ids_last_frame = blue_ids_last_frame & self.allowed_blue_ids
+
+        yellow_present = {r.id for r in yellow_robots}
+        blue_present = {r.id for r in blue_robots}
 
         # Start with current measurements
-        yellow_vision_dict: dict[int, Optional[VisionRobotData]] = {r.id: r for r in vision_data.yellow_robots}
-        blue_vision_dict: dict[int, Optional[VisionRobotData]] = {r.id: r for r in vision_data.blue_robots}
+        yellow_vision_dict: dict[int, Optional[VisionRobotData]] = {r.id: r for r in yellow_robots}
+        blue_vision_dict: dict[int, Optional[VisionRobotData]] = {r.id: r for r in blue_robots}
 
         # Add None for vanished robots
         for robot_id in yellow_ids_last_frame - yellow_present:
@@ -330,6 +343,13 @@ def _combine_both_teams_game_vision_positions(
             old_yellow_robots = game_frame.enemy_robots.copy()
             old_blue_robots = game_frame.friendly_robots.copy()
 
+        if self.allowed_yellow_ids is not None:
+            yellow_vision_robots = [r for r in yellow_vision_robots if r.id in self.allowed_yellow_ids]
+            old_yellow_robots = {k: v for k, v in old_yellow_robots.items() if k in self.allowed_yellow_ids}
+        if self.allowed_blue_ids is not None:
+            blue_vision_robots = [r for r in blue_vision_robots if r.id in self.allowed_blue_ids]
+            old_blue_robots = {k: v for k, v in old_blue_robots.items() if k in self.allowed_blue_ids}
+
         new_yellow_robots = self._combine_single_team_positions(
             old_yellow_robots,
             yellow_vision_robots,

utama_core/run/game_gater.py

@@ -1,5 +1,5 @@
 import time
-from typing import Deque, List, Optional, Tuple
+from typing import Deque, Dict, List, Optional, Tuple
 
 from utama_core.data_processing.refiners import PositionRefiner
 from utama_core.entities.data.raw_vision import RawVisionData
@@ -8,6 +8,29 @@
 
 
 class GameGater:
+    @staticmethod
+    def _validate_mapping_covers_game_frame(mapping: Dict[int, int], observed_ids: set, team_label: str) -> None:
+        """Verify that mapping keys exactly match observed vision IDs in the first game frame.
+
+        Raises ValueError if any robot visible in the frame has no mapping entry,
+        or if the mapping names IDs that were not observed.
+        """
+        if not mapping:
+            return
+        missing = observed_ids - mapping.keys()
+        extra = mapping.keys() - observed_ids
+        if missing or extra:
+            parts = []
+            if missing:
+                parts.append(f"missing entries for observed IDs {sorted(missing)}")
+            if extra:
+                parts.append(f"extra entries for unseen IDs {sorted(extra)}")
+            raise ValueError(
+                f"vision_to_cmd_mapping for {team_label} team does not match observed vision IDs: "
+                + "; ".join(parts)
+                + f". Observed IDs: {sorted(observed_ids)}."
+            )
+
     @staticmethod
     def wait_until_game_valid(
         my_team_is_yellow: bool,
@@ -20,6 +43,8 @@ def wait_until_game_valid(
         is_pvp: bool,
         rsim_env: SSLBaseEnv = None,
         wait_before_warn: float = 3.0,
+        my_vision_to_cmd_mapping: Optional[Dict[int, int]] = None,
+        opp_vision_to_cmd_mapping: Optional[Dict[int, int]] = None,
     ) -> Tuple[GameFrame, Optional[GameFrame]]:
         """
         Waits until the game frame has the expected number of robots and a ball.
@@ -31,8 +56,14 @@ def wait_until_game_valid(
 
         def print_current_vision(game_frame: GameFrame):
             print("Waiting for valid game frame...")
-            print(f"Friendly robots: {len(game_frame.friendly_robots)}/{exp_friendly}")
-            print(f"Enemy robots: {len(game_frame.enemy_robots)}/{exp_enemy}")
+            friendly_ids = sorted(game_frame.friendly_robots.keys())
+            enemy_ids = sorted(game_frame.enemy_robots.keys())
+            print(f"Friendly robots: {len(game_frame.friendly_robots)}/{exp_friendly}  observed IDs: {friendly_ids}")
+            if my_vision_to_cmd_mapping is not None:
+                print(f"  expected IDs (from mapping): {sorted(my_vision_to_cmd_mapping.keys())}")
+            print(f"Enemy robots: {len(game_frame.enemy_robots)}/{exp_enemy}  observed IDs: {enemy_ids}")
+            if opp_vision_to_cmd_mapping is not None:
+                print(f"  expected IDs (from mapping): {sorted(opp_vision_to_cmd_mapping.keys())}")
             print(f"Ball present: {game_frame.ball is not None} (exp: {exp_ball})\n")
 
         def _add_frame(my_game_frame: GameFrame, opp_game_frame: GameFrame) -> Tuple[GameFrame, Optional[GameFrame]]:
@@ -88,4 +119,20 @@ def _add_frame(my_game_frame: GameFrame, opp_game_frame: GameFrame) -> Tuple[Gam
         if len(my_game_frame.enemy_robots) > exp_enemy:
             raise ValueError(f"Too many enemy robots: {len(my_game_frame.enemy_robots)} > {exp_enemy}")
 
+        # Validate that the vision→cmd mapping keys exactly match the observed robot IDs.
+        # Done here (after first valid frame) because actual vision IDs are non-contiguous
+        # in some deployments and are not known until a real frame arrives.
+        if my_vision_to_cmd_mapping is not None:
+            GameGater._validate_mapping_covers_game_frame(
+                my_vision_to_cmd_mapping,
+                set(my_game_frame.friendly_robots.keys()),
+                "friendly",
+            )
+        if opp_vision_to_cmd_mapping is not None and opp_game_frame is not None:
+            GameGater._validate_mapping_covers_game_frame(
+                opp_vision_to_cmd_mapping,
+                set(opp_game_frame.friendly_robots.keys()),
+                "opponent",
+            )
+
         return my_game_frame, opp_game_frame