Add benchmark for documented prompts

isaaclab_arena/agentic_environment_generation/spec_inference.py

@@ -33,6 +33,11 @@ class SpecInference:
     def __init__(self, inference_backend: InferenceBackend):
         self._inference_backend = inference_backend
         self._schema = build_strict_schema(ArenaEnvGraphSpec)
+        self.last_infer_stats: dict[str, Any] = {
+            "num_calls": 0,
+            "num_retries": 0,
+            "retry_errors": [],
+        }
 
     def infer(
         self,
@@ -64,7 +69,10 @@ def infer(
         )
         user_message = base_user_message
         data: dict[str, Any] = {}
+        retry_errors: list[str] = []
+        num_calls = 0
         for call_index in range(MAX_SPEC_INFERENCE_CALLS):
+            num_calls += 1
             data = self._inference_backend.run_json(
                 StructuredOutputRequest(
                     schema_name="ArenaEnvGraphSpec",
@@ -86,16 +94,28 @@ def infer(
                 spec = None
                 validation_traces = format_validation_error(exc)
             if spec is not None and not validation_traces:
+                self.last_infer_stats = {
+                    "num_calls": num_calls,
+                    "num_retries": len(retry_errors),
+                    "retry_errors": list(retry_errors),
+                }
                 return spec, data
             if call_index + 1 < MAX_SPEC_INFERENCE_CALLS:
+                error_text = "; ".join(validation_traces)
+                retry_errors.append(error_text)
                 print(
                     f"[generate_spec] critic retry {call_index + 1}/{MAX_SPEC_INFERENCE_CALLS - 1} "
-                    f"after validation failed: {'; '.join(validation_traces)}",
+                    f"after validation failed: {error_text}",
                     flush=True,
                 )
                 user_message = self._critic_user_message(base_user_message, data, validation_traces)
                 continue
             traces.extend(validation_traces)
+        self.last_infer_stats = {
+            "num_calls": num_calls,
+            "num_retries": len(retry_errors),
+            "retry_errors": list(retry_errors),
+        }
         return None, data
 
     @staticmethod

isaaclab_arena/tests/test_spec_inference.py

@@ -105,6 +105,11 @@ def test_infer_retries_after_api_error_then_succeeds(spec_inference):
     assert isinstance(spec, ArenaEnvGraphSpec)
     assert spec.background.registry_name == "maple_table_robolab"
     assert client.chat.completions.create.call_count == 2
+    assert inference.last_infer_stats == {
+        "num_calls": 1,
+        "num_retries": 0,
+        "retry_errors": [],
+    }
 
 
 def test_infer_returns_none_with_validation_traces_on_invalid_spec(spec_inference):
@@ -119,6 +124,10 @@ def test_infer_returns_none_with_validation_traces_on_invalid_spec(spec_inferenc
     assert traces
     assert any("registry_name" in line for line in traces)
     assert client.chat.completions.create.call_count == 3
+    assert inference.last_infer_stats["num_calls"] == 3
+    assert inference.last_infer_stats["num_retries"] == 2
+    assert len(inference.last_infer_stats["retry_errors"]) == 2
+    assert all("registry_name" in error for error in inference.last_infer_stats["retry_errors"])
 
 
 def test_infer_feeds_validation_errors_back_and_recovers(spec_inference):
@@ -139,6 +148,10 @@ def test_infer_feeds_validation_errors_back_and_recovers(spec_inference):
     assert "CRITIC FEEDBACK" in critic_message
     assert "not_a_real_asset" in critic_message
     assert "registry_name" in critic_message
+    assert inference.last_infer_stats["num_calls"] == 2
+    assert inference.last_infer_stats["num_retries"] == 1
+    assert len(inference.last_infer_stats["retry_errors"]) == 1
+    assert "registry_name" in inference.last_infer_stats["retry_errors"][0]
 
 
 def test_infer_feeds_catalog_validation_errors_back_and_recovers(spec_inference):