add terminal agent with openai-agents

Signed-off-by: CormickKneey <[email protected]>

Author: CormickKneey

areal/dataset/__init__.py

@@ -10,7 +10,7 @@
     from transformers.processing_utils import ProcessorMixin
     from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
 
-VALID_DATASETS = ["gsm8k", "clevr_count_70k", "geometry3k", "hh-rlhf", "torl_data"]
+VALID_DATASETS = ["gsm8k", "clevr_count_70k", "geometry3k", "hh-rlhf", "torl_data", "terminal_bench"]
 
 logger = logging.getLogger("Dataset")
 
@@ -105,6 +105,16 @@ def _get_custom_dataset(
             max_length=max_length,
             **kwargs,
         )
+    elif "terminal_bench" in path and type == "rl":
+        from .terminal_bench import get_terminal_bench_rl_dataset
+
+        return get_terminal_bench_rl_dataset(
+            path=path,
+            split=split,
+            tokenizer=tokenizer,
+            max_length=max_length,
+            **kwargs,
+        )        
     else:
         raise ValueError(
             f"Dataset {path} with split {split} and training type {type} is not supported. "

areal/dataset/terminal_bench.py

@@ -0,0 +1,58 @@
+from typing import TYPE_CHECKING
+
+from datasets import load_dataset
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedTokenizerFast
+
+
+def get_terminal_bench_rl_dataset(
+    path: str,
+    split: str,
+    tokenizer: "PreTrainedTokenizerFast",
+    max_length: int | None = None,
+):
+    """Load terminal-bench dataset for RL training.
+
+    The dataset should be in parquet format with the following columns:
+    - prompt: The formatted prompt for the task
+    - task_name: Name of the task
+    - instruction: Raw instruction text
+    - extra_info: JSON string containing task metadata
+    """
+    # Load from parquet file
+    dataset = load_dataset("parquet", data_files={split: path}, split=split)
+
+    # The dataset already has the right format from the converter:
+    # - prompt: contains the formatted conversation
+    # - task_name, instruction, extra_info: metadata fields
+
+    # For RL training, we need to extract messages from the prompt or extra_info
+    def process(sample):
+        # The prompt is already formatted, but we need to extract the instruction
+        # to create a messages structure for the workflow
+        instruction = sample.get("instruction", "")
+        task_name = sample.get("task_name", "")
+
+        # Return data in the format expected by the workflow
+        return {
+            "instruction": instruction,
+            "task_name": task_name,
+            "extra_info": sample.get("extra_info", ""),
+            "data_source": sample.get("data_source", "terminal_bench"),
+        }
+
+    dataset = dataset.map(process)
+
+    # Filter out sequences longer than max_length if specified
+    if max_length is not None:
+
+        def filter_length(samples):
+            # Tokenize instructions in batches for efficiency
+            instructions = samples["instruction"]
+            tokens_list = tokenizer(instructions, add_special_tokens=False)["input_ids"]
+            return [len(tokens) <= max_length for tokens in tokens_list]
+
+        dataset = dataset.filter(filter_length, batched=True)
+
+    return dataset

areal/experimental/openai/client.py

@@ -165,7 +165,7 @@ async def create(
             input_ids=prompt_token_ids,
             gconfig=gconfig,
             rid=str(uuid.uuid4()),
-            metadata=metadata if metadata is not NOT_GIVEN else {},
+            metadata=metadata if not is_omitted(metadata) else {},
             tokenizer=self.tokenizer,
         )
 
@@ -276,6 +276,27 @@ async def create(
         if input is NOT_GIVEN or input is None:
             raise ValueError("input is required for Responses.create")
 
+        def _convert_tool_output_format(item: dict) -> dict:
+            """Convert custom tool output format to standard chat template format.
+
+            Converts from: {'call_id': ..., 'output': ..., 'type': 'function_call_output'}
+            To: {'role': 'tool', 'content': ..., 'tool_call_id': ...}
+            """
+            if (
+                item
+                and item.get("output")
+                and item.get("type") == "function_call_output"
+            ):
+                converted = {
+                    "role": "tool",
+                    "content": item["output"],
+                }
+                # Add tool_call_id if present
+                if "call_id" in item:
+                    converted["tool_call_id"] = item["call_id"]
+                return converted
+            return item
+
         def _build_messages_list(item: ResponseInputItemParam) -> list[dict]:
             messages_list = []
             if "content" in item:
@@ -286,13 +307,17 @@ def _build_messages_list(item: ResponseInputItemParam) -> list[dict]:
                 elif isinstance(item["content"], Iterable):
                     for content in item["content"]:
                         if isinstance(content, dict):
-                            messages_list.append(deepcopy(content))
+                            # Convert tool output format if needed
+                            converted = _convert_tool_output_format(content)
+                            messages_list.append(deepcopy(converted))
                         else:
                             raise ValueError("Unsupported content format")
                 else:
                     raise ValueError("Unsupported input item format")
             else:
-                messages_list.append(deepcopy(item))
+                # Convert tool output format if needed
+                converted = _convert_tool_output_format(item)
+                messages_list.append(deepcopy(converted))
             return messages_list
 
         if isinstance(input, str):
@@ -335,7 +360,7 @@ def _build_messages_list(item: ResponseInputItemParam) -> list[dict]:
         temp = 1.0 if temperature is NOT_GIVEN else (temperature or 0.0)
         top_p_val = 1.0 if top_p is NOT_GIVEN else (top_p or 1.0)
         max_new_tokens = 512
-        if max_output_tokens is not NOT_GIVEN and max_output_tokens is not None:
+        if not is_omitted(max_output_tokens):
             max_new_tokens = max_output_tokens
 
         stop = kwargs.get("stop", None)
@@ -359,7 +384,7 @@ def _build_messages_list(item: ResponseInputItemParam) -> list[dict]:
             input_ids=prompt_token_ids,
             gconfig=gconfig,
             rid=str(uuid.uuid4()),
-            metadata=metadata if metadata is not NOT_GIVEN else {},
+            metadata=metadata if not is_omitted(metadata) else {},
             tokenizer=self.tokenizer,
         )
 
@@ -420,14 +445,14 @@ def _build_messages_list(item: ResponseInputItemParam) -> list[dict]:
             created_at=current_time,
             error=None,
             incomplete_details=None,
-            instructions=None if instructions is NOT_GIVEN else instructions,
-            metadata=None if metadata is NOT_GIVEN else metadata,
+            instructions=None if is_omitted(instructions) else instructions,
+            metadata=None if is_omitted(metadata) else metadata,
             model="None",
             object="response",
             output=resp_output,
             parallel_tool_calls=False,
             temperature=temp,
-            tool_choice=tool_choice if tool_choice is not NOT_GIVEN else "none",
+            tool_choice=tool_choice if not is_omitted(tool_choice) else "none",
             tools=tools,
             top_p=top_p_val,
             background=None,
@@ -751,3 +776,13 @@ def export_responses(
             "export_responses is deprecated. Please use export_interactions instead."
         )
         return self.export_interactions(style)
+
+
+def is_omitted(value) -> bool:
+    """Check if a value is NOT_GIVEN or Omit type."""
+    if value is NOT_GIVEN or value is None:
+        return True
+    # Check by class name to handle both NotGiven and Omit
+    if hasattr(value, "__class__"):
+        return value.__class__.__name__ in ("NotGiven", "Omit")
+    return False

assets/qwen3_4b_thinking_terminal.png

@@ -0,0 +1,190 @@
+import asyncio
+import logging
+import os
+
+from agents import Agent as OpenAIAgent
+from agents import ModelSettings, OpenAIProvider, RunConfig, SQLiteSession
+from agents import Runner as OpenAIRunner
+from terminal.env import TerminalEnv
+from terminal.prompt import SYSTEM_PROMPT
+from transformers import PreTrainedTokenizerFast
+
+from areal.api.cli_args import GenerationHyperparameters
+from areal.api.workflow_api import RolloutWorkflow
+from areal.experimental.openai import ArealOpenAI
+from areal.utils import stats_tracker
+
+logger = logging.getLogger(__name__)
+
+
+class TerminalAgent:
+    def __init__(
+        self,
+        tokenizer: PreTrainedTokenizerFast,
+        max_tokens_per_turn: int = 1024,
+        max_turns: int = 8,
+        max_total_tokens: int = 32768,
+        dump_dir: str | None = None,
+        rollout_stat_scope: str = "rollout",
+    ):
+        self.tokenizer = tokenizer
+        self.max_tokens_per_turn = max_tokens_per_turn
+        self.max_turns = max_turns
+        self.max_total_tokens = max_total_tokens
+        self.dump_dir = dump_dir
+        self.rollout_stat_scope = rollout_stat_scope
+
+    async def run_agent(self, data, client: ArealOpenAI):
+        """Run the agent workflow for terminal task execution."""
+        run_config = RunConfig(
+            model_provider=OpenAIProvider(
+                openai_client=client,
+                use_responses=True,
+            ),
+            tracing_disabled=True,
+            model_settings=ModelSettings(
+                temperature=1.0,
+                extra_args={"max_completion_tokens": self.max_tokens_per_turn},
+                tool_choice="auto",
+                store=True,
+            ),
+        )
+
+        async with TerminalEnv(
+            task_name=data["task_name"],
+            dump_dir=self.dump_dir,
+            rollout_stat_scope=self.rollout_stat_scope,
+        ) as env:
+            # Create agent workflow with terminal tools
+            agent = OpenAIAgent(
+                name="Terminal Task Agent",
+                instructions=SYSTEM_PROMPT,
+                tools=env.get_tools(),
+            )
+            session = SQLiteSession("terminal")
+            content = data["instruction"]
+
+            max_attempts = self.max_turns
+            reward = 0
+            tracker = stats_tracker.get(self.rollout_stat_scope)
+
+            with tracker.record_timing("run_agent_total"):
+                error_count = 0.0
+                attempts_used = 0.0
+                for attempt in range(max_attempts):
+                    attempts_used = float(attempt + 1)
+                    try:
+                        with tracker.record_timing("openai_runner_run"):
+                            result = await OpenAIRunner.run(
+                                agent,
+                                input=content,
+                                session=session,
+                                run_config=run_config,
+                                max_turns=30,
+                            )
+                    except Exception as e:
+                        logger.error(f"Error running agent: {e}")
+                        error_count += 1.0
+                        break
+
+                    with tracker.record_timing("env_validate_reward"):
+                        reward = env.reward()
+                    if reward is None:
+                        reward = 0.0
+                    tracker.scalar(
+                        attempt_reward=reward,
+                        attempt_index=float(attempt),
+                        input_chars=float(len(content) if content else 0.0),
+                        output_chars=float(
+                            len(getattr(result, "final_output", "") or "")
+                        ),
+                    )
+
+                    if isinstance(reward, float) and reward >= 0.99:
+                        tracker.scalar(success=1.0)
+                        break
+
+                    if attempt < max_attempts - 1:
+                        content = f"""The previous attempt didn't complete the task successfully.
+                        Please try a different approach.
+                        Original task: {data["instruction"]}
+
+                        Previous attempt result: {result.final_output}
+
+                        Please analyze what went wrong and try again with a corrected approach."""
+                    else:
+                        content = f"""This is your final attempt. Please be extremely careful.
+                        Original task: {data["instruction"]}
+
+                        Previous attempts: {result.final_output}
+
+                        Please provide a final, carefully executed solution."""
+                    tracker.scalar(success=0.0)
+
+            tracker.scalar(
+                final_reward=reward, attempts_used=attempts_used, errors=error_count
+            )
+
+            client.set_final_reward(reward)
+
+            return reward
+
+
+class TerminalAgentWorkflow(RolloutWorkflow):
+    def __init__(
+        self,
+        gconfig: GenerationHyperparameters,
+        tokenizer: PreTrainedTokenizerFast,
+        dump_dir: str | None = None,
+        rollout_stat_scope: str = "rollout",
+        n_trajs: int = 1,
+        max_tokens: int = 32768,
+        max_turns: int = 8,
+    ):
+        self.gconfig = gconfig
+        self.gconfig.n_samples = 1
+        self.tokenizer = tokenizer
+        self.dump_dir = dump_dir
+        self.max_tokens = max_tokens
+        self.rollout_stat_scope = rollout_stat_scope
+        if self.dump_dir is not None and not os.path.exists(self.dump_dir):
+            os.makedirs(self.dump_dir, exist_ok=True)
+
+        # Search hyper-parameters
+        self.n_trajs = n_trajs
+        self.agent = TerminalAgent(
+            tokenizer=self.tokenizer,
+            max_tokens_per_turn=self.gconfig.max_new_tokens,
+            max_turns=max_turns,
+            max_total_tokens=max_tokens,
+            dump_dir=self.dump_dir,
+            rollout_stat_scope=self.rollout_stat_scope,
+        )
+
+    async def arun_episode(self, engine, data):
+        clients = [
+            ArealOpenAI(
+                engine=engine, tokenizer=self.tokenizer, tool_call_parser="qwen25"
+            )
+            for _ in range(self.n_trajs)
+        ]
+
+        # Collect trajectories
+        rewards = await asyncio.gather(
+            *[
+                self.agent.run_agent(
+                    data=data,
+                    client=clients[i],
+                )
+                for i in range(self.n_trajs)
+            ]
+        )
+        for reward in rewards:
+            stats_tracker.get(self.rollout_stat_scope).scalar(reward=reward)
+
+        interactions_with_reward = {}
+        for client in clients:
+            client.apply_reward_discount(turn_discount=0.9)
+            interactions = client.export_interactions(style="individual")
+            interactions_with_reward.update(interactions)
+        return interactions_with_reward

examples/__init__.py

@@ -19,7 +19,7 @@ cluster:
     type: nfs
     nfs_record_root: /tmp/areal/name_resolve
 
-allocation_mode: sglang.d4p1t1+d4p1t1
+allocation_mode: sglang.d4p1t1+d1p1t1c4
 
 rollout:
   experiment_name: ${experiment_name}