update workflow design

Author: kylemontgomery1

rllm/engine/agent_workflow_engine.py

@@ -68,7 +68,7 @@ async def process_task_with_retry(task: dict, uid: str) -> Episode:
             try:
                 for retry_attempt in range(1, self.retry_limit + 1):
                     try:
-                        episode = await workflow(task=task, uid=uid, **kwargs)
+                        episode = await workflow.run_with_termination_handling(task=task, uid=uid, **kwargs)
                         return episode
                     except Exception as e:
                         print(f"Rollout {uid} failed on attempt {retry_attempt}/{self.retry_limit}: {e}")
@@ -177,7 +177,7 @@ def _transform_results_for_verl(self, episodes: list[Episode], task_ids: np.ndar
 
             episode_ids.extend([episode.id] * total_steps)
             is_correct.extend([episode.is_correct] * total_steps)
-            termination_reasons.extend([episode.termination_reason if episode.termination_reason is not None else TerminationReason.ENV_DONE] * total_steps)
+            termination_reasons.extend([episode.termination_reason if episode.termination_reason is not None else TerminationReason.UNKNOWN] * total_steps)
             metrics.extend([episode.metrics] * total_steps)
             repeat_counts.append(total_steps)
 

rllm/workflows/multi_turn_workflow.py

@@ -0,0 +1,46 @@
+from rllm.agents.agent import Episode
+from rllm.workflows.workflow import TerminationEvent, TerminationReason, Workflow
+
+
+class MultiTurnWorkflow(Workflow):
+    def __init__(
+        self,
+        agent_cls,
+        env_cls,
+        agent_args=None,
+        env_args=None,
+        max_steps=5,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        # Initialize mutable defaults
+        agent_args = dict(agent_args) if agent_args is not None else {}
+        env_args = dict(env_args) if env_args is not None else {}
+
+        self.agent = agent_cls(**agent_args)
+        self.register_agent(self.agent)
+        self.env = env_cls(**env_args)
+        self.max_steps = max_steps
+
+    async def run(self, task: dict, uid: str, **kwargs) -> Episode | None:
+        """Execute a multi-step workflow"""
+
+        observation, info = await self.run_in_executor(self.reset, task=task, uid=uid)  # returns observation and info from the environment
+
+        self.agent.update_from_env(observation, 0, False, info)
+
+        for _ in range(1, self.max_steps + 1):
+            response = (await self.get_model_response(self.agent, **kwargs)).text
+            action = self.agent.update_from_model(response)
+
+            next_obs, reward, done, info = await self.run_in_executor(self.env.step, action)
+            self.agent.update_from_env(next_obs, reward, done, info)
+
+            if self._termination_buffer is not None:
+                raise TerminationEvent(self._termination_buffer)
+
+            if done:
+                raise TerminationReason.ENV_DONE
+
+        raise TerminationReason.MAX_TURNS_EXCEEDED

rllm/workflows/single_turn_workflow.py

@@ -0,0 +1,37 @@
+from rllm.agents.agent import Episode
+from rllm.workflows.workflow import TerminationEvent, TerminationReason, Workflow
+
+
+class SingleTurnWorkflow(Workflow):
+    def __init__(
+        self,
+        agent_cls,
+        env_cls,
+        agent_args=None,
+        env_args=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        # Initialize mutable defaults
+        agent_args = dict(agent_args) if agent_args is not None else {}
+        env_args = dict(env_args) if env_args is not None else {}
+
+        self.agent = agent_cls(**agent_args)
+        self.register_agent(self.agent)
+        self.env = env_cls(**env_args)
+
+    async def run(self, task: dict, uid: str, **kwargs) -> Episode | None:
+        observation, info = await self.run_in_executor(self.reset, task=task, uid=uid)  # returns observation and info from the environment
+        self.agent.update_from_env(observation, 0, False, info)
+
+        response = (await self.get_model_response(self.agent, **kwargs)).text
+        action = self.agent.update_from_model(response)
+
+        next_obs, reward, done, info = await self.run_in_executor(self.env.step, action)
+        self.agent.update_from_env(next_obs, reward, done, info)
+
+        if self._termination_buffer is not None:
+            raise TerminationEvent(self._termination_buffer)
+
+        raise TerminationReason.ENV_DONE