feat: unify pydantic models for ros2 types (#545)

Author: boczekbartek

src/rai_bench/rai_bench/manipulation_o3de/benchmark.py

@@ -22,11 +22,7 @@
 from rclpy.impl.rcutils_logger import RcutilsLogger
 
 from rai_bench.manipulation_o3de.interfaces import Task
-from rai_sim.simulation_bridge import (
-    Entity,
-    SimulationBridge,
-    SimulationConfigT,
-)
+from rai_sim.simulation_bridge import Entity, SimulationBridge, SimulationConfigT
 
 loggers_type = Union[RcutilsLogger, logging.Logger]
 EntityT = TypeVar("EntityT", bound=Entity)

src/rai_bench/rai_bench/manipulation_o3de/interfaces.py

@@ -17,14 +17,15 @@
 from collections import defaultdict
 from typing import Dict, List, Set, Tuple, TypeVar, Union
 
+from rai.types import Pose
 from rclpy.impl.rcutils_logger import RcutilsLogger
 
 from rai_sim.simulation_bridge import (
     Entity,
-    Pose,
     SimulationBridge,
     SimulationConfig,
     SimulationConfigT,
+    SpawnedEntity,
 )
 
 loggers_type = Union[RcutilsLogger, logging.Logger]
@@ -114,9 +115,9 @@ def filter_entities_by_object_type(
     def euclidean_distance(self, pos1: Pose, pos2: Pose) -> float:
         """Calculate euclidean distance between 2 positions"""
         return (
-            (pos1.translation.x - pos2.translation.x) ** 2
-            + (pos1.translation.y - pos2.translation.y) ** 2
-            + (pos1.translation.z - pos2.translation.z) ** 2
+            (pos1.position.x - pos2.position.x) ** 2
+            + (pos1.position.y - pos2.position.y) ** 2
+            + (pos1.position.z - pos2.position.z) ** 2
         ) ** 0.5
 
     def is_adjacent(self, pos1: Pose, pos2: Pose, threshold_distance: float):
@@ -217,7 +218,9 @@ def build_neighbourhood_list(
                 other
                 for other in entities
                 if entity != other
-                and self.is_adjacent(entity.pose, other.pose, threshold_distance)
+                and self.is_adjacent(
+                    entity.pose.pose, other.pose.pose, threshold_distance
+                )
             ]
         return neighbourhood_graph
 
@@ -335,15 +338,16 @@ def group_entities_along_z_axis(
         """
 
         entities = sorted(
-            entities, key=lambda ent: (ent.pose.translation.x, ent.pose.translation.y)
+            entities,
+            key=lambda ent: (ent.pose.pose.position.x, ent.pose.pose.position.y),
         )
 
         groups: List[List[EntityT]] = []
         for entity in entities:
             placed = False
             for group in groups:
-                dx = group[0].pose.translation.x - entity.pose.translation.x
-                dy = group[0].pose.translation.y - entity.pose.translation.y
+                dx = group[0].pose.pose.position.x - entity.pose.pose.position.x
+                dy = group[0].pose.pose.position.y - entity.pose.pose.position.y
                 if math.sqrt(dx * dx + dy * dy) <= margin:
                     group.append(entity)
                     placed = True
@@ -424,12 +428,14 @@ def validate_config(self, simulation_config: SimulationConfig) -> bool:
             return False
 
     @abstractmethod
-    def calculate_correct(self, entities: List[EntityT]) -> Tuple[int, int]:
+    def calculate_correct(
+        self, entities: List[Entity] | List[SpawnedEntity]
+    ) -> Tuple[int, int]:
         """Method to calculate how many objects are placed correctly
 
         Parameters
         ----------
-        entities : List[EntityT]
+        entities : List[Entity]
             list of ALL entities present in the simulaiton scene
 
         Returns

src/rai_bench/rai_bench/manipulation_o3de/tasks/move_object_to_left_task.py

@@ -73,7 +73,7 @@ def calculate_correct(self, entities: List[Entity]) -> Tuple[int, int]:
             entities=entities, object_types=self.obj_types
         )
         correct = sum(
-            1 for ent in selected_type_objects if ent.pose.translation.y > 0.0
+            1 for ent in selected_type_objects if ent.pose.pose.position.y > 0.0
         )
         incorrect: int = len(selected_type_objects) - correct
         return correct, incorrect

src/rai_bench/rai_bench/manipulation_o3de/tasks/place_at_coord_task.py

@@ -90,8 +90,8 @@ def calculate_correct(self, entities: List[Entity]) -> Tuple[int, int]:
         correct = 0
 
         for ent in target_objects:
-            dx = ent.pose.translation.x - self.target_position[0]
-            dy = ent.pose.translation.y - self.target_position[1]
+            dx = ent.pose.pose.position.x - self.target_position[0]
+            dy = ent.pose.pose.position.y - self.target_position[1]
             distance = math.sqrt(dx**2 + dy**2)
             if distance <= self.allowable_displacement:
                 correct = 1  # Only one correct placement is needed.

src/rai_bench/rai_bench/manipulation_o3de/tasks/place_cubes_task.py

@@ -89,8 +89,8 @@ def calculate_correct(self, entities: List[Entity]) -> Tuple[int, int]:
             1
             for ent in entities
             if self.is_adjacent_to_any(
-                ent.pose,
-                [e.pose for e in entities if e != ent],
+                ent.pose.pose,
+                [e.pose.pose for e in entities if e != ent],
                 self.threshold_distance,
             )
         )

src/rai_bench/rai_bench/manipulation_o3de/tasks/rotate_object_task.py

@@ -16,12 +16,13 @@
 import math
 from typing import List, Tuple, Union
 
+from rai.types import Quaternion
 from rclpy.impl.rcutils_logger import RcutilsLogger
 
 from rai_bench.manipulation_o3de.interfaces import (
     ManipulationTask,
 )
-from rai_sim.simulation_bridge import Entity, Rotation, SimulationConfig
+from rai_sim.simulation_bridge import Entity, SimulationConfig
 
 loggers_type = Union[RcutilsLogger, logging.Logger]
 
@@ -30,7 +31,7 @@ class RotateObjectTask(ManipulationTask):
     def __init__(
         self,
         obj_types: List[str],
-        target_quaternion: Rotation,
+        target_quaternion: Quaternion,
         logger: loggers_type | None = None,
     ):
         # NOTE (jmatejcz) for now manipulaiton tool does not support passing rotation
@@ -103,14 +104,14 @@ def calculate_correct(
         incorrect = 0
         for entity in entities:
             if entity.prefab_name in self.obj_types:
-                if not entity.pose.rotation:
+                if not entity.pose.pose.orientation:
                     ValueError("Entity has no rotation defined.")
                 else:
                     dot = (
-                        entity.pose.rotation.x * self.target_quaternion.x
-                        + entity.pose.rotation.y * self.target_quaternion.y
-                        + entity.pose.rotation.z * self.target_quaternion.z
-                        + entity.pose.rotation.w * self.target_quaternion.w
+                        entity.pose.pose.orientation.x * self.target_quaternion.x
+                        + entity.pose.pose.orientation.y * self.target_quaternion.y
+                        + entity.pose.pose.orientation.z * self.target_quaternion.z
+                        + entity.pose.pose.orientation.w * self.target_quaternion.w
                     )
                     # Account for the double cover: q and -q represent the same rotation.
                     dot = abs(dot)

src/rai_bench/rai_bench/tool_calling_agent/messages/actions.py

@@ -14,66 +14,66 @@
 
 from typing import Any, Dict, Optional
 
-from rai_bench.tool_calling_agent.messages.base import PoseStamped, Ros2BaseModel, Time
+from rai.types import PoseStamped, ROS2BaseModel, Time
 
 
-class TaskGoal(Ros2BaseModel):
+class TaskGoal(ROS2BaseModel):
     task: Optional[str] = ""
     description: Optional[str] = ""
     priority: Optional[str] = ""
 
 
-class TaskResult(Ros2BaseModel):
+class TaskResult(ROS2BaseModel):
     success: Optional[bool] = False
     report: Optional[str] = ""
 
 
-class TaskFeedback(Ros2BaseModel):
+class TaskFeedback(ROS2BaseModel):
     current_status: Optional[str] = ""
 
 
-class LoadMapRequest(Ros2BaseModel):
+class LoadMapRequest(ROS2BaseModel):
     filename: Optional[str] = ""
 
 
-class LoadMapResponse(Ros2BaseModel):
+class LoadMapResponse(ROS2BaseModel):
     success: Optional[bool] = False
 
 
-class NavigateToPoseGoal(Ros2BaseModel):
+class NavigateToPoseGoal(ROS2BaseModel):
     pose: Optional[PoseStamped] = None
     behavior_tree: Optional[str] = None
 
 
-class ActionResult(Ros2BaseModel):
+class ActionResult(ROS2BaseModel):
     result: Optional[Dict[str, Any]] = None
 
 
-class NavigateToPoseFeedback(Ros2BaseModel):
+class NavigateToPoseFeedback(ROS2BaseModel):
     current_pose: Optional[PoseStamped] = None
     navigation_time: Optional[Time] = None
     estimated_time_remaining: Optional[Time] = None
     number_of_recoveries: Optional[int] = None
     distance_remaining: Optional[float] = None
 
 
-class NavigateToPoseAction(Ros2BaseModel):
+class NavigateToPoseAction(ROS2BaseModel):
     goal: Optional[NavigateToPoseGoal] = None
     result: Optional[ActionResult] = None
     feedback: Optional[NavigateToPoseFeedback] = None
 
 
-class SpinGoal(Ros2BaseModel):
+class SpinGoal(ROS2BaseModel):
     target_yaw: Optional[float] = None
     time_allowance: Optional[Time] = None
 
 
-class SpinFeedback(Ros2BaseModel):
+class SpinFeedback(ROS2BaseModel):
     angle_turned: Optional[float] = None
     remaining_yaw: Optional[float] = None
 
 
-class SpinAction(Ros2BaseModel):
+class SpinAction(ROS2BaseModel):
     goal: Optional[SpinGoal] = None
     result: Optional[ActionResult] = None
     feedback: Optional[SpinFeedback] = None

src/rai_bench/rai_bench/tool_calling_agent/messages/base.py

@@ -12,93 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Optional
+from rai.types import ROS2BaseModel, Time
 
-from pydantic import BaseModel, ConfigDict
 
-
-class Ros2BaseModel(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-
-class Time(Ros2BaseModel):
-    sec: Optional[int] = 0
-    nanosec: Optional[int] = 0
-
-
-class Header(Ros2BaseModel):
-    stamp: Optional[Time] = Time()
-    frame_id: Optional[str] = ""
-
-
-class RegionOfInterest(Ros2BaseModel):
-    x_offset: Optional[int] = 0
-    y_offset: Optional[int] = 0
-    height: Optional[int] = 0
-    width: Optional[int] = 0
-    do_rectify: Optional[bool] = False
-
-
-class Position(Ros2BaseModel):
-    x: Optional[float] = 0.0
-    y: Optional[float] = 0.0
-    z: Optional[float] = 0.0
-
-
-class Orientation(Ros2BaseModel):
-    x: Optional[float] = 0.0
-    y: Optional[float] = 0.0
-    z: Optional[float] = 0.0
-    w: Optional[float] = 1.0
-
-
-class Pose(Ros2BaseModel):
-    position: Optional[Position] = Position()
-    orientation: Optional[Orientation] = Orientation()
-
-
-class PoseStamped(Ros2BaseModel):
-    header: Optional[Header] = Header()
-    pose: Optional[Pose] = Pose()
-
-
-class Clock(Ros2BaseModel):
-    clock: Optional[Time] = Time()
-
-
-class ObjectHypothesis(Ros2BaseModel):
-    class_id: Optional[str] = ""
-    score: Optional[float] = 0.0
-
-
-class PoseWithCovariance(Ros2BaseModel):
-    pose: Optional[Pose] = Pose()
-    covariance: Optional[List[float]] = [0.0] * 36
-
-
-class ObjectHypothesisWithPose(Ros2BaseModel):
-    hypothesis: Optional[ObjectHypothesis] = ObjectHypothesis()
-    pose: Optional[PoseWithCovariance] = PoseWithCovariance()
-
-
-class Point2D(Ros2BaseModel):
-    x: Optional[float] = 0.0
-    y: Optional[float] = 0.0
-
-
-class Pose2D(Ros2BaseModel):
-    position: Optional[Point2D] = Point2D()
-    theta: Optional[float] = 0.0
-
-
-class BoundingBox2D(Ros2BaseModel):
-    center: Optional[Pose2D] = Pose2D()
-    size_x: Optional[float] = 0.0
-    size_y: Optional[float] = 0.0
-
-
-class Detection2D(Ros2BaseModel):
-    header: Optional[Header] = Header()
-    results: Optional[List[ObjectHypothesisWithPose]] = []
-    bbox: Optional[BoundingBox2D] = BoundingBox2D()
-    id: Optional[str] = ""
+class Clock(ROS2BaseModel):
+    _prefix: str = "rosgraph_msgs/msg"
+    clock: Time = Time()