feat: Let AgentFlow use is_validation and preserve TerminationReason in metrics (#523)

* add is_validation to AgentConfig and let TerminationReason flow to metrics

* clean up tests

Author: JasonWei05

rllm/experimental/engine/agent_flow_engine.py

@@ -231,6 +231,7 @@ async def _run_single(self, task: dict, uid: str, is_validation: bool = False) -
             base_url=session_url,
             model=self.model,
             session_uid=uid,
+            is_validation=is_validation,
         )
 
         # 3. Run agent flow (prefers arun if available, else run in executor)
@@ -276,7 +277,8 @@ async def _run_single(self, task: dict, uid: str, is_validation: bool = False) -
         for signal in eval_output.signals:
             enriched.metrics[signal.name] = signal.value
 
-        enriched.termination_reason = TerminationReason.ENV_DONE
+        if enriched.termination_reason is None:
+            enriched.termination_reason = TerminationReason.ENV_DONE
         return enriched
 
     def _enrich_episode(

rllm/experimental/unified_trainer.py

@@ -418,6 +418,12 @@ async def _train_batch_async(self, batch: Any, trainer_state: TrainerState) -> N
             return
 
         workflow_metrics, termination_counts = self._collect_workflow_metrics_from_episodes(trainer_state.episodes)
+        for key, value in workflow_metrics.items():
+            trainer_state.metrics[f"batch/{key}"] = np.mean(value)
+
+        total_counts = max(sum(termination_counts.values()), 1)
+        for r in TerminationReason:
+            trainer_state.metrics[f"batch/termination_reason/{r.value}"] = termination_counts[r.value] / total_counts
 
         # stage 2: transform episodes to trajectory groups (sync)
         trajectory_groups, transform_metrics = transform_episodes_to_trajectory_groups(trainer_state.episodes, self.transform_config, self.cf_config, traj_grouping_hook=self.traj_grouping_hook)
@@ -461,13 +467,6 @@ async def _train_batch_async(self, batch: Any, trainer_state: TrainerState) -> N
                 show_workflow_metadata=True,
             )
 
-        for key, value in workflow_metrics.items():
-            trainer_state.metrics[f"batch/{key}"] = np.mean(value)
-
-        total_counts = max(sum(termination_counts.values()), 1)
-        for r in TerminationReason:
-            trainer_state.metrics[f"batch/termination_reason/{r.value}"] = termination_counts[r.value] / total_counts
-
     # =========================================================================
     # Fully-asynchronous training pipeline
     # =========================================================================
@@ -809,14 +808,15 @@ def shutdown(self):
     # =========================================================================
     # Helper functions
     # =========================================================================
-    def _collect_workflow_metrics_from_episodes(self, episodes: list[Episode]) -> tuple[dict, Counter]:
+    @staticmethod
+    def _collect_workflow_metrics_from_episodes(episodes: list[Episode]) -> tuple[dict, Counter]:
         workflow_metrics = defaultdict(list)
         termination_counts = Counter()
         for episode in episodes:
             for k, v in episode.metrics.items():
                 workflow_metrics[k].append(v)
-            if episode.termination_reason is not None:
-                termination_counts[episode.termination_reason.value] += 1
+            reason = episode.termination_reason or TerminationReason.UNKNOWN
+            termination_counts[getattr(reason, "value", reason)] += 1
         # reduce the metrics to a scalar value, with error handling
         reduced_workflow_metrics = {}
         for k, v in workflow_metrics.items():

rllm/types.py

@@ -144,6 +144,7 @@ class AgentConfig:
     model: str
     session_uid: str
     metadata: dict = field(default_factory=dict)
+    is_validation: bool = False
 
 
 @runtime_checkable

tests/engine/test_agent_flow_engine.py

@@ -0,0 +1,65 @@
+import asyncio
+
+from rllm.agents.agent import Episode, Trajectory
+from rllm.eval.types import EvalOutput
+from rllm.experimental.engine.agent_flow_engine import AgentFlowEngine
+from rllm.workflows.workflow import TerminationReason
+
+
+class _Agent:
+    def __init__(self):
+        self.config = None
+
+    async def arun(self, task, config):
+        self.config = config
+        return Episode(
+            id=task.id,
+            termination_reason=TerminationReason.ERROR,
+            trajectories=[Trajectory(name="solver")],
+        )
+
+
+class _Evaluator:
+    def evaluate(self, task, episode):
+        return EvalOutput(reward=0.0, is_correct=False)
+
+
+class _Gateway:
+    def __init__(self):
+        self.created = None
+        self.deleted = None
+
+    async def acreate_session(self, session_id, is_validation=False):
+        self.created = (session_id, is_validation)
+
+    def get_session_url(self, session_id):
+        return f"http://gateway/{session_id}"
+
+    async def aget_traces(self, session_id):
+        return []
+
+    async def adelete_session(self, session_id):
+        self.deleted = session_id
+
+
+def test_run_single_passes_validation_flag_and_preserves_termination_reason():
+    agent = _Agent()
+    gateway = _Gateway()
+    engine = AgentFlowEngine(
+        agent_flow=agent,
+        evaluator=_Evaluator(),
+        gateway=gateway,
+        model="test-model",
+        n_parallel_tasks=1,
+    )
+
+    try:
+        episode = asyncio.run(engine._run_single({"question": "q"}, "task:0", is_validation=True))
+    finally:
+        engine.shutdown()
+
+    assert gateway.created == ("task:0", True)
+    assert gateway.deleted == "task:0"
+    assert agent.config.is_validation is True
+    assert agent.config.session_uid == "task:0"
+    assert episode.termination_reason == TerminationReason.ERROR

tests/eval/test_eval_types.py

@@ -188,6 +188,7 @@ def test_signals(self):
 def test_agent_config_defaults():
     config = AgentConfig(base_url="http://localhost:8000", model="test-model", session_uid="s1")
     assert config.metadata == {}
+    assert config.is_validation is False
 
 
 # ---------------------------------------------------------------------------

tests/unified_trainer/test_termination_metrics.py

@@ -0,0 +1,16 @@
+from rllm.agents.agent import Episode
+from rllm.experimental.unified_trainer import UnifiedTrainer
+from rllm.workflows.workflow import TerminationReason
+
+
+def test_collect_workflow_metrics_counts_unknown_termination_reason():
+    episodes = [
+        Episode(id="task:0", termination_reason=None, metrics={"custom": 1.0}),
+        Episode(id="task:1", termination_reason=TerminationReason.ERROR, metrics={"custom": 3.0}),
+    ]
+
+    workflow_metrics, termination_counts = UnifiedTrainer._collect_workflow_metrics_from_episodes(episodes)
+
+    assert workflow_metrics["custom"] == 2.0
+    assert termination_counts[TerminationReason.UNKNOWN.value] == 1
+    assert termination_counts[TerminationReason.ERROR.value] == 1