robot_activity_model内の配列をなるべく減らすために，dataclassを使うように変更 (#504)

Author: uchikun2493

consai_game/consai_game/world_model/robot_activity_model.py

@@ -25,7 +25,7 @@
 
 from consai_msgs.msg import MotionCommand
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import Enum, auto
 
 
@@ -45,7 +45,96 @@ class OurRobotsArrived:
     arrived: bool = False
 
 
+@dataclass
+class RobotInfo:
+    """単一のロボット情報を保持するデータクラス."""
+
+    # ロボットID
+    robot_id: int = 0
+
+    # 目標位置までの距離
+    desired_distance: float = float("inf")
+    # ボールまでの距離
+    ball_distance: float = float("inf")
+    # プレースメント位置までの距離
+    placement_distance: float = float("inf")
+
+    # 目標位置に到着しているかのフラグ
+    arrived: bool = False
+
+
+@dataclass
+class RobotsInfo:
+    """自ロボットの情報を保持するデータクラス."""
+
+    robots: dict[int, RobotInfo] = field(default_factory=dict)
+
+    def clear(self):
+        """全ロボット情報を初期化して空にするメソッド."""
+        self.robots.clear()
+
+    def visible_ids(self) -> list[int]:
+        """可視ロボットのIDリストを返すメソッド."""
+        return list(self.robots.keys())
+
+    def arrived_ids(self) -> list[int]:
+        """目標位置に到達したロボットのIDリストを返すメソッド."""
+        return [r.robot_id for r in self.robots.values() if r.arrived]
+
+    def all_arrived(self) -> bool:
+        """全ロボットが目標位置に到達しているかを返すメソッド."""
+        return all(r.arrived for r in self.robots.values())
+
+    def get(self, robot_id: int) -> RobotInfo:
+        """指定したロボットIDのRobotInfoを返す。存在しない場合はKeyErrorメソッド."""
+        return self.robots[robot_id]
+
+    def __getitem__(self, robot_id: int) -> RobotInfo:
+        """辞書のようにロボットIDでRobotInfoへアクセスできるようにするメソッド."""
+        return self.robots[robot_id]
+
+    def __setitem__(self, robot_id: int, value: RobotInfo):
+        """辞書のようにロボットIDでRobotInfoを設定できるようにするメソッド."""
+        self.robots[robot_id] = value
+
+    def __contains__(self, robot_id: int) -> bool:
+        """ロボットIDが含まれているか判定するメソッド."""
+        return robot_id in self.robots
+
+    def __len__(self):
+        """可視ロボット数を返すメソッド."""
+        return len(self.robots)
+
+    def keys(self):
+        """可視ロボットのID一覧を返すメソッド."""
+        return self.robots.keys()
+
+    def values(self):
+        """可視ロボットのRobotInfo一覧を返すメソッド."""
+        return self.robots.values()
+
+    def items(self):
+        """可視ロボットの(ID, RobotInfo)タプル一覧を返すメソッド."""
+        return self.robots.items()
+
+
+@dataclass
+class OurRobotsInfo:
+    """自ロボットの情報を保持するデータクラス."""
+
+    our_robots: RobotsInfo = field(default_factory=RobotsInfo)
+
+
+@dataclass
+class TheirRobotsInfo:
+    """敵ロボットの情報を保持するデータクラス."""
+
+    their_robots: RobotsInfo = field(default_factory=RobotsInfo)
+
+
 class ProhibitedKickRobotSearchState(Enum):
+    """キック禁止ロボット探索状態を表す列挙型."""
+
     BEFORE_SEARCH = 0
     SHOULD_FIRST_SEARCH = auto()
     SHOULD_SECOND_SEARCH = auto()
@@ -60,18 +149,12 @@ class RobotActivityModel:
     INVALID_ROBOT_ID = -1
 
     def __init__(self):
-        """ロボットの可視状態と順序リストの初期化関数."""
-        self.ordered_our_visible_robots: list[int] = []
-        self.ordered_their_visible_robots: list[int] = []
-        self.our_visible_robots: list[int] = []
-        self.their_visible_robots: list[int] = []
-        self.our_robots_by_ball_distance: list[int] = []
-        self.their_robots_by_ball_distance: list[int] = []
-        self.our_robots_by_placement_distance: list[int] = []
+        """RobotActivityModelの初期化."""
+        self.our_visible_robots = RobotsInfo()
+        self.their_visible_robots = RobotsInfo()
         self.our_ball_receive_score: list[ReceiveScore] = []
-        self.our_robots_arrived_list: list[OurRobotsArrived] = []
-        self.our_prohibited_kick_robot_id: int = self.INVALID_ROBOT_ID  # 直近のキック禁止ロボットID
-        self.prohibited_kick_robot_candidate_id: int = self.INVALID_ROBOT_ID  # キック禁止ロボット候補ID
+        self.our_prohibited_kick_robot_id: int = self.INVALID_ROBOT_ID
+        self.prohibited_kick_robot_candidate_id: int = self.INVALID_ROBOT_ID
         self.prohibited_kick_robot_search_state = ProhibitedKickRobotSearchState.BEFORE_SEARCH
         self.number_of_their_robots_in_our_area: int = 0
 
@@ -83,67 +166,76 @@ def update(
         game_config: GameConfigModel,
         referee: RefereeModel,
     ):
-        """ロボットの可視状態を更新し, 順序づけされたIDリストを更新する関数."""
-        self.our_visible_robots = [robot.robot_id for robot in robots.our_robots.values() if robot.is_visible]
-        self.their_visible_robots = [robot.robot_id for robot in robots.their_robots.values() if robot.is_visible]
+        """ロボットの可視状態を更新し, 距離や状態をセットする関数."""
 
-        self.ordered_our_visible_robots = self.ordered_merge(
-            self.ordered_our_visible_robots,
-            self.our_visible_robots,
-        )
-        self.ordered_their_visible_robots = self.ordered_merge(
-            self.ordered_their_visible_robots,
-            self.their_visible_robots,
-        )
+        # 可視ロボットをRobotInfoで管理
+        self.our_visible_robots.clear()
+        for robot in robots.our_robots.values():
+            if robot.is_visible:
+                info = RobotInfo(
+                    robot_id=robot.robot_id,
+                    ball_distance=tools.get_distance(robot.pos, ball.pos),
+                    placement_distance=tools.get_distance(robot.pos, referee.placement_pos),
+                )
+                self.our_visible_robots[robot.robot_id] = info
 
-        # ボールに近い順にリストを作る
+        self.their_visible_robots.clear()
+        for robot in robots.their_robots.values():
+            if robot.is_visible:
+                info = RobotInfo(
+                    robot_id=robot.robot_id,
+                    ball_distance=tools.get_distance(robot.pos, ball.pos),
+                    placement_distance=tools.get_distance(robot.pos, referee.placement_pos),
+                )
+                self.their_visible_robots[robot.robot_id] = info
+
+        # ボールに近い順
         self.our_robots_by_ball_distance = [
-            robot_id
-            for robot_id, _ in self.robot_ball_distances(
-                robots.our_visible_robots,
-                ball,
-            )
+            r.robot_id for r in sorted(self.our_visible_robots.values(), key=lambda x: x.ball_distance)
         ]
         self.their_robots_by_ball_distance = [
-            robot_id
-            for robot_id, _ in self.robot_ball_distances(
-                robots.their_visible_robots,
-                ball,
-            )
+            r.robot_id for r in sorted(self.their_visible_robots.values(), key=lambda x: x.ball_distance)
         ]
-
-        # プレースメント位置に近い順にリストを作る
+        # プレースメント位置に近い順
         self.our_robots_by_placement_distance = [
-            robot_id
-            for robot_id, _ in self.robot_placement_distances(
-                robots.our_visible_robots,
-                referee,
-            )
+            r.robot_id for r in sorted(self.our_visible_robots.values(), key=lambda x: x.placement_distance)
         ]
 
-        # ボールを受け取れるスコアを計算する
+        # ボール受け取りスコア
         self.our_ball_receive_score = self.calc_ball_receive_score_list(
-            robots=robots.our_visible_robots,
+            robots={rid: robots.our_robots[rid] for rid in self.our_visible_robots.visible_ids()},
             ball=ball,
             ball_activity=ball_activity,
             game_config=game_config,
         )
 
-        # 自陣にいる相手ロボットの台数を計算する
-        self.number_of_their_robots_in_our_area = self.count_their_robots(robots.their_visible_robots)
+        # 自陣にいる相手ロボット数
+        self.number_of_their_robots_in_our_area = sum(
+            1 for r in self.their_visible_robots.values() if robots.their_robots[r.robot_id].pos.x < 0
+        )
 
-        # ダブルタッチ防止のために、キック禁止ロボット情報を更新する
         self.update_prohibited_kick_robot(ball_activity, referee)
 
-    def ordered_merge(self, prev_list: list[int], new_list: list[int]) -> list[int]:
-        """過去の順序を保ちながら, 新しいリストでマージする関数."""
-        # new_listに存在するものを残す
-        output_list = [r for r in prev_list if r in new_list]
+    def update_our_robots_arrived(self, robots: dict[int, Robot], commands: list[MotionCommand]):
+        """各ロボットが目標位置に到達したかをRobotInfoにセット"""
+        for command in commands:
+            if command.robot_id not in robots:
+                continue
+            robot = robots[command.robot_id]
+            dist = tools.get_distance(robot.pos, command.desired_pose)
+            if command.robot_id in self.our_visible_robots:
+                self.our_visible_robots[command.robot_id].arrived = dist < self.DIST_ROBOT_TO_DESIRED_THRESHOLD
 
-        # 新しい要素を追加する
-        output_list.extend([r for r in new_list if r not in output_list])
+    @property
+    def our_robots_arrived(self) -> bool:
+        """すべての自ロボットが目標位置に到達したか"""
+        return self.our_visible_robots.all_arrived()
 
-        return output_list
+    def our_robot_arrived(self, robot_id: int) -> bool:
+        """指定したロボットが目標位置に到達したか"""
+        if robot_id in self.our_visible_robots:
+            return self.our_visible_robots[robot_id].arrived
+        return False
 
     def count_their_robots(self, robots: dict[int, Robot]) -> int:
         """自陣にいる相手ロボットの数を返す関数."""
@@ -231,33 +323,6 @@ def calc_intercept_time(self, robot: Robot, ball: BallModel, game_config: GameCo
             return intercept_time
         return float("inf")
 
-    def update_our_robots_arrived(self, our_visible_robots: dict[int, Robot], commands: list[MotionCommand]) -> bool:
-        """各ロボットが目標位置に到達したか判定する関数."""
-
-        # 初期化
-        self.our_robots_arrived_list = []
-        # エラー処理
-        if len(our_visible_robots) == 0 or len(commands) == 0:
-            return
-
-        # 更新
-        for command in commands:
-            if command.robot_id not in our_visible_robots.keys():
-                continue
-
-            robot = our_visible_robots[command.robot_id]
-            robot_pos = robot.pos
-            desired_pose = command.desired_pose
-            # ロボットと目標位置の距離を計算
-            dist_robot_to_desired = tools.get_distance(robot_pos, desired_pose)
-            # 目標位置に到達したか判定結果をリストに追加
-            self.our_robots_arrived_list.append(
-                OurRobotsArrived(
-                    robot_id=robot.robot_id,
-                    arrived=dist_robot_to_desired < self.DIST_ROBOT_TO_DESIRED_THRESHOLD,
-                )
-            )
-
     def update_prohibited_kick_robot(self, ball_activity: BallActivityModel, referee: RefereeModel):
         """レフェリー信号を見て、ダブルタッチをしてはいけないロボットを更新する関数."""
         # ストップゲームで初期化する
@@ -300,15 +365,3 @@ def update_prohibited_kick_robot(self, ball_activity: BallActivityModel, referee
                     # 違うロボットがボールを蹴ったら探索終了
                     self.prohibited_kick_robot_search_state = ProhibitedKickRobotSearchState.SEARCH_COMPLETED
                     self.our_prohibited_kick_robot_id = self.INVALID_ROBOT_ID
-
-    @property
-    def our_robots_arrived(self) -> bool:
-        """すべての自ロボットが目標位置に到達したかを返す関数."""
-        return all([robot.arrived for robot in self.our_robots_arrived_list])
-
-    def our_robot_arrived(self, robot_id: int) -> bool:
-        """指定したロボットが目標位置に到達したかを返す関数."""
-        for robot in self.our_robots_arrived_list:
-            if robot.robot_id == robot_id:
-                return robot.arrived
-        return False