added storing results, bug with get_object_pose

Author: jmatejcz

src/rai_bench/rai_bench/benchmark_model.py

@@ -187,6 +187,8 @@ def run_next(self, agent):
             self._logger.info(
                 f"RUNNING SCENARIO NUMBER {i+1}, TASK: {scenario.task.get_prompt()}"
             )
+            initial_result = scenario.task.calculate_result(self.simulation_bridge)
+            self._logger.info(f"RESULT OF THE INITIAL SETUP: {initial_result}")
             ts = time.perf_counter()
             for state in agent.stream(
                 {"messages": [HumanMessage(content=scenario.task.get_prompt())]}
@@ -208,15 +210,26 @@ def run_next(self, agent):
                     raise ValueError(f"Unexpected type of message: {type(msg)}")
 
                 self._logger.info(f"AI Message: {msg}")
-                # TODO (jm) figure out how to get number of tool calls
             te = time.perf_counter()
 
             result = scenario.task.calculate_result(self.simulation_bridge)
 
             total_time = te - ts
             self._logger.info(f"TASK SCORE: {result}, TOTAL TIME: {total_time:.3f}")
 
-            # self.results.append{""})
+            self.results.append(
+                {
+                    "task": scenario.task.get_prompt(),
+                    "initial_score": initial_result,
+                    "final_score": result,
+                    "total_time": f"{total_time:.3f}",
+                    # TODO (jm) figure out how to get number of tool calls
+                    "tool_calls": None,
+                }
+            )
 
         except StopIteration:
             print("No more scenarios left to run.")
+
+    def get_results(self) -> list[dict]:
+        return self.results

src/rai_bench/rai_bench/main.py

@@ -46,8 +46,6 @@
     node = connector.node
     node.declare_parameter("conversion_ratio", 1.0)
 
-    o3de = O3DEngineArmManipulationBridge(connector)
-
     # define model
     llm = get_llm_model(model_type="complex_model", streaming=True)
 
@@ -154,6 +152,7 @@
     # custom request to arm
     base_arm_pose = PoseModel(translation=Translation(x=0.3, y=0.0, z=0.5))
 
+    o3de = O3DEngineArmManipulationBridge(connector, logger=agent_logger)
     # define benchamrk
     benchmark = Benchmark(
         simulation_bridge=o3de,

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

@@ -36,6 +36,8 @@ def calculate_result(self, simulation_bridge: SimulationBridge) -> float:
         )
         num_initial_carrots = len(initial_carrots)
 
+        self.logger.info(initial_carrots)
+        self.logger.info(final_carrots)
         if num_initial_carrots != len(final_carrots):
             raise EntitiesMismatchException(
                 "Number of initially spawned entities does not match number of entities present at the end."
@@ -51,6 +53,8 @@ def calculate_result(self, simulation_bridge: SimulationBridge) -> float:
                         # NOTE the specific coords that refer to for example
                         # middle of the table can differ across simulations,
                         # take that into consideration
+                        self.logger.info(initial_y)
+                        self.logger.info(final_y)
                         if (
                             initial_y <= 0.0
                         ):  # Carrot started in the incorrect place (right side)

src/rai_sim/rai_sim/o3de/o3de_bridge.py

@@ -175,6 +175,7 @@ def _despawn_entity(self, entity: SpawnedEntity):
             )
 
     def get_object_pose(self, entity: SpawnedEntity) -> PoseModel:
+        self.logger.info(f"GET OBJECT POSE: {entity}")
         object_frame = entity.name + "/"
         ros2_pose = do_transform_pose(
             Pose(), self.connector.get_transform(object_frame + "odom", object_frame)
@@ -193,6 +194,7 @@ def get_scene_state(self) -> SceneState:
         entities: list[SpawnedEntity] = []
         for entity in self.spawned_entities:
             current_pose = self.get_object_pose(entity)
+            self.logger.info(f"AFTER GET OBJECT POSE: {current_pose}")
             entities.append(
                 SpawnedEntity(
                     id=entity.id,