refactor: funtions naming, typing and rotation of arm change

Author: jmatejcz

src/rai_bench/rai_bench/benchmark_model.py

@@ -15,7 +15,7 @@
 import logging
 import time
 from abc import ABC, abstractmethod
-from typing import List, TypeVar, Union
+from typing import Any, Dict, List, TypeVar, Union
 
 from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
 from rclpy.impl.rcutils_logger import RcutilsLogger
@@ -59,9 +59,8 @@ def get_prompt(self) -> str:
         pass
 
     @abstractmethod
-    def validate_scene(self, simulation_config: SimulationConfig) -> bool:
-        """Task should be able to verify if given scene is suitable for specific task
-        for example: GrabCarrotTask should verify if there is any carrots in the scene
+    def validate_config(self, simulation_config: SimulationConfig) -> bool:
+        """Task should be able to verify if given config is suitable for specific task
 
         Args:
             simulation_config (SimulationConfig): initial scene setup
@@ -136,7 +135,7 @@ class Scenario:
     """Single instances are run separatly by benchmark"""
 
     def __init__(self, task: Task, simulation_config: SimulationConfig) -> None:
-        if not task.validate_scene(simulation_config):
+        if not task.validate_config(simulation_config):
             raise ValueError("This scene is invalid for this task.")
         self.task = task
         self.simulation_config = simulation_config
@@ -149,13 +148,13 @@ class Benchmark:
 
     def __init__(
         self,
-        simulation_bridge: SimulationBridge,
+        simulation_bridge: SimulationBridge[SimulationConfig],
         scenarios: list[Scenario],
         logger: loggers_type | None = None,
     ) -> None:
         self.simulation_bridge = simulation_bridge
         self.scenarios = enumerate(iter(scenarios))
-        self.results = []
+        self.results: List[Dict[str, Any]] = []
         if logger:
             self._logger = logger
         else:
@@ -164,7 +163,7 @@ def __init__(
     @classmethod
     def create_scenarios(
         cls, tasks: List[Task], simulation_configs: List[SimulationConfig]
-    ):
+    ) -> list[Any]:
         scenarios = []
         for task in tasks:
             for sim_conf in simulation_configs:
@@ -235,7 +234,7 @@ def run_next(self, agent):
                     "initial_score": initial_result,
                     "final_score": result,
                     "total_time": f"{total_time:.3f}",
-                    "numer_of_tool_calls": tool_calls_num,
+                    "number_of_tool_calls": tool_calls_num,
                 }
             )
 

src/rai_bench/rai_bench/main.py

@@ -35,7 +35,7 @@
     O3DExROS2SimulationConfig,
     PoseModel,
 )
-from rai_sim.simulation_bridge import Translation
+from rai_sim.simulation_bridge import Rotation, Translation
 
 if __name__ == "__main__":
     rclpy.init()
@@ -147,7 +147,10 @@
     )
 
     # custom request to arm
-    base_arm_pose = PoseModel(translation=Translation(x=0.5, y=0.1, z=0.3))
+    base_arm_pose = PoseModel(
+        translation=Translation(x=0.5, y=0.1, z=0.3),
+        rotation=Rotation(x=1.0, y=0.0, z=0.0, w=0.0),
+    )
 
     o3de = O3DEngineArmManipulationBridge(connector, logger=agent_logger)
     # define benchamrk

src/rai_bench/rai_bench/o3de_test_bench/tasks/grab_carrot_task.py

@@ -26,7 +26,7 @@ class GrabCarrotTask(Task):
     def get_prompt(self) -> str:
         return "Manipulate objects, so that all carrots to the left side of the table (positive y)"
 
-    def validate_scene(self, simulation_config: SimulationConfig) -> bool:
+    def validate_config(self, simulation_config: SimulationConfig) -> bool:
         for ent in simulation_config.entities:
             if ent.prefab_name == "carrot":
                 return True

src/rai_bench/rai_bench/o3de_test_bench/tasks/place_cubes_task.py

@@ -23,7 +23,7 @@ class PlaceCubesTask(Task):
     def get_prompt(self) -> str:
         return "Manipulate objects, so that all cubes are adjacent to at least one cube"
 
-    def validate_scene(self, simulation_config: SimulationConfig) -> bool:
+    def validate_config(self, simulation_config: SimulationConfig) -> bool:
         cube_types = ["red_cube", "blue_cube", "yellow_cube"]
         cubes_num = 0
         for ent in simulation_config.entities:

src/rai_sim/rai_sim/o3de/o3de_bridge.py

@@ -214,11 +214,11 @@ def _is_robotic_stack_ready(
             services = self.connector.node.get_service_names_and_types()
             topics_names = [tp[0] for tp in topics]
             service_names = [srv[0] for srv in services]
-            self.logger.debug(
+            self.logger.info(
                 f"required services: {simulation_config.required_services}"
             )
-            self.logger.debug(f"required topics: {simulation_config.required_topics}")
-            self.logger.debug(f"required actions: {simulation_config.required_actions}")
+            self.logger.info(f"required topics: {simulation_config.required_topics}")
+            self.logger.info(f"required actions: {simulation_config.required_actions}")
             # NOTE actions will be listed in services and topics
             if (
                 all(srv in service_names for srv in simulation_config.required_services)
@@ -227,7 +227,7 @@ def _is_robotic_stack_ready(
                     ac in service_names for ac in simulation_config.required_actions
                 )
             ):
-                self.logger.debug("All required services are available.")
+                self.logger.info("All required services are available.")
                 return True
 
             time.sleep(5)
@@ -380,11 +380,6 @@ def move_arm(
             request.target_pose.pose.orientation.y = pose.rotation.y
             request.target_pose.pose.orientation.z = pose.rotation.z
             request.target_pose.pose.orientation.w = pose.rotation.w
-        else:
-            request.target_pose.pose.orientation.x = 1.0
-            request.target_pose.pose.orientation.y = 0.0
-            request.target_pose.pose.orientation.z = 0.0
-            request.target_pose.pose.orientation.w = 0.0
 
         client = self.connector.node.create_client(
             ManipulatorMoveTo,