feat: spatial reasoning tasks

Author: MagdalenaKotynia

src/rai_bench/rai_bench/examples/images/image_1.jpg

@@ -0,0 +1,100 @@
+# Copyright (C) 2025 Robotec.AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List, Sequence
+
+from rai_bench.tool_calling_agent_bench.agent_tasks_interfaces import (
+    SpatialReasoningAgentTask,
+)
+from rai_bench.tool_calling_agent_bench.spatial_reasoning_tasks import (
+    BoolImageTask,
+    BoolImageTaskInput,
+)
+
+inputs: List[BoolImageTaskInput] = [
+    BoolImageTaskInput(
+        question="Is the door on the left from the desk?",
+        images_paths=["src/rai_bench/rai_bench/examples/images/image_1.jpg"],
+        expected_response=True,
+    ),
+    BoolImageTaskInput(
+        question="Is the door open?",
+        images_paths=["src/rai_bench/rai_bench/examples/images/image_1.jpg"],
+        expected_response=False,
+    ),
+    BoolImageTaskInput(
+        question="Is someone in the room?",
+        images_paths=["src/rai_bench/rai_bench/examples/images/image_1.jpg"],
+        expected_response=False,
+    ),
+    BoolImageTaskInput(
+        question="Is the light on in the room?",
+        images_paths=["src/rai_bench/rai_bench/examples/images/image_2.jpg"],
+        expected_response=True,
+    ),
+    BoolImageTaskInput(
+        question="Do you see the plant?",
+        images_paths=["src/rai_bench/rai_bench/examples/images/image_2.jpg"],
+        expected_response=True,
+    ),
+    BoolImageTaskInput(
+        question="Do you see the plant?",
+        images_paths=["src/rai_bench/rai_bench/examples/images/image_3.jpg"],
+        expected_response=False,
+    ),
+    BoolImageTaskInput(
+        question="Are there any pictures on the wall?",
+        images_paths=["src/rai_bench/rai_bench/examples/images/image_3.jpg"],
+        expected_response=True,
+    ),
+    BoolImageTaskInput(
+        question="Are there 3 pictures on the wall?",
+        images_paths=["src/rai_bench/rai_bench/examples/images/image_4.jpg"],
+        expected_response=True,
+    ),
+    BoolImageTaskInput(
+        question="Are there 4 pictures on the wall?",
+        images_paths=["src/rai_bench/rai_bench/examples/images/image_4.jpg"],
+        expected_response=False,
+    ),
+    BoolImageTaskInput(
+        question="Is there a rack on the left from the sofa?",
+        images_paths=["src/rai_bench/rai_bench/examples/images/image_4.jpg"],
+        expected_response=False,
+    ),
+    BoolImageTaskInput(
+        question="Is there a plant behind the rack?",
+        images_paths=["src/rai_bench/rai_bench/examples/images/image_5.jpg"],
+        expected_response=True,
+    ),
+    BoolImageTaskInput(
+        question="Is there a plant on the right from the window?",
+        images_paths=["src/rai_bench/rai_bench/examples/images/image_6.jpg"],
+        expected_response=False,
+    ),
+    BoolImageTaskInput(
+        question="Is there a pillow on the armchain?",
+        images_paths=["src/rai_bench/rai_bench/examples/images/image_7.jpg"],
+        expected_response=True,
+    ),
+    BoolImageTaskInput(
+        question="Is there a red pillow on the armchair?",
+        images_paths=["src/rai_bench/rai_bench/examples/images/image_7.jpg"],
+        expected_response=False,
+    ),
+]
+
+tasks: Sequence[SpatialReasoningAgentTask] = [
+    BoolImageTask(task_input=input_item) for input_item in inputs
+]

src/rai_bench/rai_bench/examples/images/image_2.jpg

@@ -0,0 +1,74 @@
+# Copyright (C) 2025 Robotec.AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from datetime import datetime
+from pathlib import Path
+
+from rai.agents.conversational_agent import create_conversational_agent
+from rai.utils.model_initialization import (
+    get_llm_model,
+    get_llm_model_config_and_vendor,
+)
+
+from rai_bench.examples.spatial_reasoning_tasks import tasks
+from rai_bench.tool_calling_agent_bench.agent_bench import ToolCallingAgentBenchmark
+
+if __name__ == "__main__":
+    current_test_name = Path(__file__).stem
+
+    now = datetime.now()
+    experiment_dir = (
+        Path("src/rai_bench/rai_bench/experiments")
+        / current_test_name
+        / now.strftime("%Y-%m-%d_%H-%M-%S")
+    )
+    experiment_dir.mkdir(parents=True, exist_ok=True)
+    log_filename = experiment_dir / "benchmark.log"
+    results_filename = experiment_dir / "results.csv"
+
+    file_handler = logging.FileHandler(log_filename)
+    file_handler.setLevel(logging.DEBUG)
+    formatter = logging.Formatter(
+        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    )
+    file_handler.setFormatter(formatter)
+
+    bench_logger = logging.getLogger("Benchmark logger")
+    bench_logger.setLevel(logging.INFO)
+    bench_logger.addHandler(file_handler)
+
+    agent_logger = logging.getLogger("Agent logger")
+    agent_logger.setLevel(logging.INFO)
+    agent_logger.addHandler(file_handler)
+
+    for task in tasks:
+        task.logger = bench_logger
+
+    benchmark = ToolCallingAgentBenchmark(
+        tasks=tasks, logger=bench_logger, results_filename=results_filename
+    )
+
+    model_type = "simple_model"
+    model_config = get_llm_model_config_and_vendor(model_type=model_type)[0]
+    model_name = getattr(model_config, model_type)
+
+    for task in tasks:
+        agent = create_conversational_agent(
+            llm=get_llm_model(model_type=model_type),
+            tools=task.expected_tools,
+            system_prompt=task.get_system_prompt(),
+            logger=agent_logger,
+        )
+        benchmark.run_next(agent=agent, model_name=model_name)

src/rai_bench/rai_bench/examples/images/image_3.jpg

@@ -28,6 +28,7 @@
 from rai.messages.multimodal import HumanMultimodalMessage
 
 from rai_bench.tool_calling_agent_bench.agent_tasks_interfaces import (
+    SpatialReasoningAgentTask,
     ToolCallingAgentTask,
 )
 from rai_bench.tool_calling_agent_bench.scores_tracing import ScoreTracingHandler
@@ -160,10 +161,26 @@ def run_next(self, agent: CompiledStateGraph, model_name: str) -> None:
 
             ts = time.perf_counter()
             try:
-                response = agent.invoke(
-                    {"messages": [HumanMultimodalMessage(content=task.get_prompt())]},
-                    config=config,
-                )
+                if isinstance(task, SpatialReasoningAgentTask):
+                    response = agent.invoke(
+                        {
+                            "messages": [
+                                HumanMultimodalMessage(
+                                    content=task.get_prompt(), images=task.get_images()
+                                )
+                            ]
+                        },
+                        config=config,
+                    )
+                else:
+                    response = agent.invoke(
+                        {
+                            "messages": [
+                                HumanMultimodalMessage(content=task.get_prompt())
+                            ]
+                        },
+                        config=config,
+                    )
                 task.verify_tool_calls(response=response)
             except GraphRecursionError as e:
                 task.log_error(msg=f"Graph Recursion Error: {e}")