[rollout, trainer] feat: extend agent loop for custom implementations

verl/experimental/agent_loop/tool_agent_loop.py

@@ -30,12 +30,17 @@
     register,
 )
 from verl.experimental.agent_loop.tool_parser import FunctionCall, ToolParser
-from verl.experimental.agent_loop.utils import add_generation_prompt_for_gpt_oss, format_gpt_oss_tool_response_manually
+from verl.experimental.agent_loop.utils import build_gpt_oss_tool_response_text
 from verl.interactions.base import BaseInteraction
 from verl.interactions.utils.interaction_registry import initialize_interactions_from_config
 from verl.tools.schemas import ToolResponse
 from verl.tools.utils.tool_registry import initialize_tools_from_config
-from verl.utils.chat_template import initialize_system_prompt
+from verl.utils.chat_template import (
+    apply_chat_template_with_processor,
+    apply_chat_template_with_tokenizer,
+    initialize_system_prompt,
+    tokenize_with_processor,
+)
 from verl.utils.profiler import simple_timer
 from verl.utils.rollout_trace import rollout_trace_op
 
@@ -203,25 +208,28 @@ async def _handle_pending_state(self, agent_data: AgentData, sampling_params: di
         if self.processor is not None:
             raw_prompt = await self.loop.run_in_executor(
                 None,
-                lambda: self.processor.apply_chat_template(
+                lambda: apply_chat_template_with_processor(
+                    self.processor,
                     agent_data.messages,
                     tools=self.tool_schemas,
                     add_generation_prompt=True,
-                    tokenize=False,
-                    **self.apply_chat_template_kwargs,
+                    apply_chat_template_kwargs=self.apply_chat_template_kwargs,
                 ),
             )
-            model_inputs = self.processor(text=[raw_prompt], images=agent_data.image_data, return_tensors="pt")
-            agent_data.prompt_ids = model_inputs.pop("input_ids").squeeze(0).tolist()
+            agent_data.prompt_ids = tokenize_with_processor(
+                self.processor,
+                raw_prompt=raw_prompt,
+                image_data=agent_data.image_data,
+            )
         else:
             agent_data.prompt_ids = await self.loop.run_in_executor(
                 None,
-                lambda: self.tokenizer.apply_chat_template(
+                lambda: apply_chat_template_with_tokenizer(
+                    self.tokenizer,
                     agent_data.messages,
                     tools=self.tool_schemas,
                     add_generation_prompt=True,
-                    tokenize=True,
-                    **self.apply_chat_template_kwargs,
+                    apply_chat_template_kwargs=self.apply_chat_template_kwargs,
                 ),
             )
         return AgentState.GENERATING
@@ -346,35 +354,35 @@ async def _handle_processing_tools_state(self, agent_data: AgentData) -> AgentSt
         if self.processor is not None:
             raw_tool_response = await self.loop.run_in_executor(
                 None,
-                lambda: self.processor.apply_chat_template(
+                lambda: apply_chat_template_with_processor(
+                    self.processor,
                     add_messages,
                     add_generation_prompt=True,
-                    tokenize=False,
-                    **self.apply_chat_template_kwargs,
+                    apply_chat_template_kwargs=self.apply_chat_template_kwargs,
                 ),
             )
             # Use only the new images from this turn for processing tool responses
-            current_images = new_images_this_turn if new_images_this_turn else None  # Using local variable
-            model_inputs = self.processor(text=[raw_tool_response], images=current_images, return_tensors="pt")
-            response_ids = model_inputs.pop("input_ids").squeeze(0).tolist()
+            current_images = new_images_this_turn if new_images_this_turn else None
+            response_ids = tokenize_with_processor(
+                self.processor,
+                raw_prompt=raw_tool_response,
+                image_data=current_images,
+            )
         else:
             if self.tool_parser_name == "gpt-oss":
                 logger.info("manually format tool responses for gpt-oss")
-                # Format tool responses manually
-                tool_response_texts = []
-                for i, tool_msg in enumerate(add_messages):
-                    actual_tool_name = tool_call_names[i]
-                    formatted = format_gpt_oss_tool_response_manually(tool_msg["content"], actual_tool_name)
-                    tool_response_texts.append(formatted)
-
-                tool_response_text = add_generation_prompt_for_gpt_oss("".join(tool_response_texts))
+                tool_response_text = build_gpt_oss_tool_response_text(add_messages, tool_call_names)
                 response_ids = await self.loop.run_in_executor(
                     None, lambda: self.tokenizer.encode(tool_response_text, add_special_tokens=False)
                 )
             else:
                 response_ids = await self.loop.run_in_executor(
                     None,
-                    lambda: self.tokenizer.apply_chat_template(add_messages, add_generation_prompt=True, tokenize=True),
+                    lambda: apply_chat_template_with_tokenizer(
+                        self.tokenizer,
+                        add_messages,
+                        add_generation_prompt=True,
+                    ),
                 )
                 response_ids = response_ids[len(self.system_prompt) :]
         if len(agent_data.response_mask) + len(response_ids) >= self.response_length:
@@ -414,23 +422,26 @@ async def _handle_interacting_state(self, agent_data: AgentData) -> AgentState:
         if reward is not None:
             agent_data.turn_scores.append(reward)
 
-        # Update prompt with user responses (similar to _handle_processing_tools_state)
+        # Update prompt with user responses
         if self.processor is not None:
             raw_user_response = await self.loop.run_in_executor(
                 None,
-                lambda: self.processor.apply_chat_template(
+                lambda: apply_chat_template_with_processor(
+                    self.processor,
                     add_messages,
                     add_generation_prompt=True,
-                    tokenize=False,
-                    **self.apply_chat_template_kwargs,
+                    apply_chat_template_kwargs=self.apply_chat_template_kwargs,
                 ),
             )
-            model_inputs = self.processor(text=[raw_user_response], images=None, return_tensors="pt")
-            response_ids = model_inputs.pop("input_ids").squeeze(0).tolist()
+            response_ids = tokenize_with_processor(self.processor, raw_prompt=raw_user_response, image_data=None)
         else:
             response_ids = await self.loop.run_in_executor(
                 None,
-                lambda: self.tokenizer.apply_chat_template(add_messages, add_generation_prompt=True, tokenize=True),
+                lambda: apply_chat_template_with_tokenizer(
+                    self.tokenizer,
+                    add_messages,
+                    add_generation_prompt=True,
+                ),
             )
         response_ids = response_ids[len(self.system_prompt) :]
 

verl/experimental/agent_loop/utils.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+from typing import Any
 
 
 def resolve_config_path(config_path: str) -> str:
@@ -95,3 +96,13 @@ def add_generation_prompt_for_gpt_oss(message_content: str) -> str:
         Message content string with generation prompt
     """
     return message_content + "<|start|>assistant"
+
+
+def build_gpt_oss_tool_response_text(messages: list[dict[str, Any]], tool_call_names: list[str]) -> str:
+    """Build gpt-oss tool response text (manual formatting + generation prompt)."""
+    tool_response_texts: list[str] = []
+    for i, tool_msg in enumerate(messages):
+        actual_tool_name = tool_call_names[i]
+        formatted = format_gpt_oss_tool_response_manually(tool_msg["content"], actual_tool_name)
+        tool_response_texts.append(formatted)
+    return add_generation_prompt_for_gpt_oss("".join(tool_response_texts))

verl/trainer/ppo/ray_trainer.py

@@ -18,6 +18,7 @@
 This trainer supports model-agonistic model initialization with huggingface
 """
 
+import importlib
 import json
 import os
 import uuid
@@ -847,7 +848,14 @@ def init_workers(self):
         # create async rollout manager and request scheduler
         self.async_rollout_mode = False
         if self.config.actor_rollout_ref.rollout.mode == "async":
-            from verl.experimental.agent_loop import AgentLoopManager
+            # Support custom AgentLoopManager via config
+            manager_class = self.config.actor_rollout_ref.rollout.get("agent", {}).get("agent_loop_manager_class")
+            if manager_class:
+                module_path, class_name = manager_class.rsplit(".", 1)
+                module = importlib.import_module(module_path)
+                AgentLoopManager = getattr(module, class_name)
+            else:
+                from verl.experimental.agent_loop import AgentLoopManager
 
             self.async_rollout_mode = True
             if self.config.reward_model.enable and self.config.reward_model.enable_resource_pool:

verl/utils/chat_template.py

@@ -1,6 +1,7 @@
 # Copyright 2025 Bytedance Ltd. and/or its affiliates
 import logging
 import os
+from typing import Any
 
 from jinja2 import TemplateError
 
@@ -40,3 +41,84 @@ def extract_system_prompt_and_generation(tokenizer):
     generate_prompt = token3[len(token1) :]
 
     return system_prompt, generate_prompt
+
+
+# =============================================================================
+# Tokenization utility functions
+# These functions are extracted from ToolAgentLoop for reusability in external
+# projects like remote rollout servers.
+# =============================================================================
+
+
+def apply_chat_template_with_processor(
+    processor,
+    messages: list[dict[str, Any]],
+    *,
+    tools: list[dict[str, Any]] | None = None,
+    add_generation_prompt: bool = True,
+    apply_chat_template_kwargs: dict[str, Any] | None = None,
+) -> str:
+    """Apply chat template via processor (tokenize=False).
+
+    NOTE: To preserve ToolAgentLoop's original behavior, callers should decide
+    whether this runs in an executor or on the event loop thread.
+    """
+    if apply_chat_template_kwargs is None:
+        apply_chat_template_kwargs = {}
+    if tools is None:
+        return processor.apply_chat_template(
+            messages,
+            add_generation_prompt=add_generation_prompt,
+            tokenize=False,
+            **apply_chat_template_kwargs,
+        )
+    return processor.apply_chat_template(
+        messages,
+        tools=tools,
+        add_generation_prompt=add_generation_prompt,
+        tokenize=False,
+        **apply_chat_template_kwargs,
+    )
+
+
+def apply_chat_template_with_tokenizer(
+    tokenizer,
+    messages: list[dict[str, Any]],
+    *,
+    tools: list[dict[str, Any]] | None = None,
+    add_generation_prompt: bool = True,
+    apply_chat_template_kwargs: dict[str, Any] | None = None,
+) -> list[int]:
+    """Apply chat template via tokenizer (tokenize=True).
+
+    IMPORTANT: Some call sites intentionally do NOT pass any extra kwargs
+    (e.g. incremental tool/user message tokenization). When no kwargs are needed,
+    simply omit the apply_chat_template_kwargs parameter.
+    """
+    if apply_chat_template_kwargs is None:
+        apply_chat_template_kwargs = {}
+    if tools is None:
+        return tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=add_generation_prompt,
+            tokenize=True,
+            **apply_chat_template_kwargs,
+        )
+    return tokenizer.apply_chat_template(
+        messages,
+        tools=tools,
+        add_generation_prompt=add_generation_prompt,
+        tokenize=True,
+        **apply_chat_template_kwargs,
+    )
+
+
+def tokenize_with_processor(
+    processor,
+    *,
+    raw_prompt: str,
+    image_data: Any | None,
+) -> list[int]:
+    """Tokenize a pre-rendered prompt string via processor."""
+    model_inputs = processor(text=[raw_prompt], images=image_data, return_tensors="pt")
+    return model_inputs.pop("input_ids").squeeze(0).tolist()

verl/workers/config/rollout.py

@@ -71,6 +71,7 @@ class AgentLoopConfig(BaseConfig):
     default_agent_loop: str = "single_turn_agent"
     agent_loop_config_path: Optional[str] = None
     custom_async_server: CustomAsyncServerConfig = field(default_factory=CustomAsyncServerConfig)
+    agent_loop_manager_class: Optional[str] = None
 
 
 @dataclass
@@ -179,6 +180,9 @@ class RolloutConfig(BaseConfig):
     # Use Prometheus to collect and monitor rollout statistics
     prometheus: PrometheusConfig = field(default_factory=PrometheusConfig)
 
+    # Extension point for custom configurations
+    custom: Optional[dict] = None
+
     update_weights_bucket_megabytes: int = 512
 
     skip_rollout: bool = False