[rollout] feat: pass agent_data to tool calling (verl-project#4469)

Author: wuxibin89

recipe/fully_async_policy/agent_loop/agent_loop.py

@@ -27,8 +27,8 @@
     AgentLoopOutput,
     AgentLoopWorkerBase,
     AsyncLLMServerManager,
+    DictConfigWrap,
     _agent_loop_registry,
-    _DummyConfig,
     get_trajectory_info,
 )
 from verl.experimental.agent_loop.prometheus_utils import update_prometheus_config
@@ -182,7 +182,7 @@ async def _partial_run_agent_loop(
                 agent_loop_config = _agent_loop_registry[agent_name]
                 agent_loop = hydra.utils.instantiate(
                     config=agent_loop_config,
-                    trainer_config=_DummyConfig(config=self.config),
+                    trainer_config=DictConfigWrap(config=self.config),
                     server_manager=self.server_manager,
                     tokenizer=self.tokenizer,
                     processor=self.processor,

verl/experimental/agent_loop/agent_loop.py

@@ -175,21 +175,20 @@ class _InternalAgentLoopOutput(AgentLoopOutput):
     """Extra fields for dynamic addition."""
 
 
-# make hydra.utils.instantiate happy
-class _DummyConfig:
-    def __init__(self, config: DictConfig) -> None:
+class DictConfigWrap:
+    """Wrapper for DictConfig to avoid hydra.utils.instantiate recursive resolve."""
+
+    def __init__(self, config: DictConfig):
         self.config = config
 
 
 class AgentLoopBase(ABC):
     """An agent loop takes an input message, chat with OpenAI compatible LLM server and interact with various
     environments."""
 
-    _class_initialized = False
-
     def __init__(
         self,
-        trainer_config: _DummyConfig,
+        trainer_config: DictConfigWrap,
         server_manager: AsyncLLMServerManager,
         tokenizer: AutoTokenizer,
         processor: AutoProcessor,
@@ -198,32 +197,17 @@ def __init__(
         """Initialize agent loop, each sample will have its own loop instance.
 
         Args:
-            trainer_config (_DummyConfig): trainer config.
+            trainer_config (DictConfigWrap): trainer config.
             server_manager (AsyncLLMServerManager): OpenAI compatible LLM server manager.
             tokenizer (AutoTokenizer): Tokenizer for tokenize messages.
             processor (AutoProcessor): Processor for process messages.
         """
-        self.init_class(config=trainer_config.config, tokenizer=tokenizer, processor=processor, **kwargs)
         self.config = trainer_config.config
         self.server_manager = server_manager
         self.tokenizer = tokenizer
         self.processor = processor
         self.loop = asyncio.get_running_loop()
 
-    @classmethod
-    def init_class(cls, config: DictConfig, tokenizer: AutoTokenizer, processor: AutoProcessor, **kwargs):
-        """This is used to do heavy initialization work that should shared across all instances. It's only called once.
-
-        Args:
-            config (DictConfig): trainer config.
-            tokenizer (AutoTokenizer): Tokenizer for tokenize messages.
-            processor (AutoProcessor): Processor for process multi_modal data.
-            **kwargs: extra kwargs from config file passed in by `hydra.utils.instantiate`.
-        """
-        if cls._class_initialized:
-            return
-        cls._class_initialized = True
-
     @abstractmethod
     async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutput:
         """Run agent loop to interact with LLM server and environment.
@@ -420,7 +404,7 @@ async def _run_agent_loop(
             agent_loop_config = _agent_loop_registry[agent_name]
             agent_loop = hydra.utils.instantiate(
                 config=agent_loop_config,
-                trainer_config=_DummyConfig(config=self.config),
+                trainer_config=DictConfigWrap(config=self.config),
                 server_manager=self.server_manager,
                 tokenizer=self.tokenizer,
                 processor=self.processor,

verl/experimental/agent_loop/tool_agent_loop.py

@@ -20,7 +20,15 @@
 from typing import Any, Optional
 from uuid import uuid4
 
-from verl.experimental.agent_loop.agent_loop import AgentLoopBase, AgentLoopOutput, register
+from transformers import AutoProcessor, AutoTokenizer
+
+from verl.experimental.agent_loop.agent_loop import (
+    AgentLoopBase,
+    AgentLoopOutput,
+    AsyncLLMServerManager,
+    DictConfigWrap,
+    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.interactions.base import BaseInteraction
@@ -44,7 +52,8 @@ class AgentState(Enum):
 
 
 class AgentData:
-    """Encapsulates all state variables for the agent loop."""
+    """Encapsulates all state variables for the agent loop. AgentData is passed to tool calling in case that
+    tool may need to access full history state. User can store any tool session data in `extra_fields`."""
 
     def __init__(
         self,
@@ -77,44 +86,49 @@ def __init__(
         # Temporary state for tool calls
         self.tool_calls: list[FunctionCall] = []
 
-        # Extra fields for dynamic addition
+        # Extra fields for dynamic addition, e.g., tool session data
         self.extra_fields: dict[str, Any] = {}
 
 
 @register("tool_agent")
 class ToolAgentLoop(AgentLoopBase):
-    @classmethod
-    def init_class(cls, config, tokenizer, processor, **kwargs):
-        if cls._class_initialized:
-            return
-        cls._class_initialized = True
-        print("Performing class-level ToolAgentLoop initialization")
+    def __init__(
+        self,
+        trainer_config: DictConfigWrap,
+        server_manager: AsyncLLMServerManager,
+        tokenizer: AutoTokenizer,
+        processor: AutoProcessor,
+        **kwargs,
+    ):
+        super().__init__(trainer_config, server_manager, tokenizer, processor, **kwargs)
+        config = trainer_config.config
 
         # Initialize tools from config file
-        cls.tokenizer = tokenizer
-        cls.processor = processor
-        cls.max_user_turns = config.actor_rollout_ref.rollout.multi_turn.max_user_turns
-        cls.max_assistant_turns = config.actor_rollout_ref.rollout.multi_turn.max_assistant_turns
-        cls.max_parallel_calls = config.actor_rollout_ref.rollout.multi_turn.max_parallel_calls
-        cls.max_tool_response_length = config.actor_rollout_ref.rollout.multi_turn.max_tool_response_length
-        cls.tool_response_truncate_side = config.actor_rollout_ref.rollout.multi_turn.tool_response_truncate_side
+        self.max_user_turns = config.actor_rollout_ref.rollout.multi_turn.max_user_turns
+        self.max_assistant_turns = config.actor_rollout_ref.rollout.multi_turn.max_assistant_turns
+        self.max_parallel_calls = config.actor_rollout_ref.rollout.multi_turn.max_parallel_calls
+        self.max_tool_response_length = config.actor_rollout_ref.rollout.multi_turn.max_tool_response_length
+        self.tool_response_truncate_side = config.actor_rollout_ref.rollout.multi_turn.tool_response_truncate_side
         tool_config_path = config.actor_rollout_ref.rollout.multi_turn.tool_config_path
         tool_list = initialize_tools_from_config(tool_config_path) if tool_config_path else []
-        cls.tools = {tool.name: tool for tool in tool_list}
-        cls.tool_schemas = [tool.tool_schema.model_dump(exclude_unset=True, exclude_none=True) for tool in tool_list]
-        cls.tool_parser = ToolParser.get_tool_parser(config.actor_rollout_ref.rollout.multi_turn.format, cls.tokenizer)
-        cls.tool_parser_name = config.actor_rollout_ref.rollout.multi_turn.format
-        print(f"Initialized tools: {cls.tools}")
+        self.tools = {tool.name: tool for tool in tool_list}
+        self.tool_schemas = [tool.tool_schema.model_dump(exclude_unset=True, exclude_none=True) for tool in tool_list]
+        self.tool_parser = ToolParser.get_tool_parser(
+            config.actor_rollout_ref.rollout.multi_turn.format, self.tokenizer
+        )
+        self.tool_parser_name = config.actor_rollout_ref.rollout.multi_turn.format
 
-        cls.apply_chat_template_kwargs = config.data.get("apply_chat_template_kwargs", {})
-        cls.prompt_length = config.actor_rollout_ref.rollout.prompt_length
-        cls.response_length = config.actor_rollout_ref.rollout.response_length
-        cls.system_prompt = initialize_system_prompt(cls.tokenizer, **cls.apply_chat_template_kwargs)
+        self.apply_chat_template_kwargs = config.data.get("apply_chat_template_kwargs", {})
+        self.prompt_length = config.actor_rollout_ref.rollout.prompt_length
+        self.response_length = config.actor_rollout_ref.rollout.response_length
+        self.system_prompt = initialize_system_prompt(self.tokenizer, **self.apply_chat_template_kwargs)
 
         # Initialize interactions from config file
-        cls.interaction_config_file = config.actor_rollout_ref.rollout.multi_turn.interaction_config_path
-        if cls.interaction_config_file:
-            cls.interaction_map: dict[str, BaseInteraction] = cls._initialize_interactions(cls.interaction_config_file)
+        self.interaction_config_file = config.actor_rollout_ref.rollout.multi_turn.interaction_config_path
+        if self.interaction_config_file:
+            self.interaction_map: dict[str, BaseInteraction] = self._initialize_interactions(
+                self.interaction_config_file
+            )
 
     @rollout_trace_op
     async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutput:
@@ -271,7 +285,7 @@ async def _handle_processing_tools_state(self, agent_data: AgentData) -> AgentSt
         tasks = []
         tool_call_names = []
         for tool_call in agent_data.tool_calls[: self.max_parallel_calls]:
-            tasks.append(self._call_tool(tool_call, agent_data.tools_kwargs))
+            tasks.append(self._call_tool(tool_call, agent_data.tools_kwargs, agent_data))
             tool_call_names.append(tool_call.name)
 
         with simple_timer("tool_calls", agent_data.metrics):
@@ -434,7 +448,7 @@ async def _handle_interacting_state(self, agent_data: AgentData) -> AgentState:
             return AgentState.GENERATING
 
     async def _call_tool(
-        self, tool_call: FunctionCall, tools_kwargs: dict[str, Any]
+        self, tool_call: FunctionCall, tools_kwargs: dict[str, Any], agent_data: AgentData
     ) -> tuple[ToolResponse, float, dict]:
         """Call tool and return tool response."""
         tool, instance_id = None, None
@@ -445,7 +459,9 @@ async def _call_tool(
             tool = self.tools[tool_name]
             kwargs = tools_kwargs.get(tool_name, {})
             instance_id, _ = await tool.create(create_kwargs=kwargs.get("create_kwargs", {}))
-            tool_execution_response, tool_reward, res = await tool.execute(instance_id, tool_args)
+            tool_execution_response, tool_reward, res = await tool.execute(
+                instance_id, tool_args, agent_data=agent_data
+            )
         except Exception as e:
             logger.warning(f"Error when executing tool: {e}")
             return (
@@ -481,8 +497,7 @@ async def _call_tool(
 
         return ToolResponse(**tool_response_kwargs), tool_reward, res
 
-    @classmethod
-    def _initialize_interactions(cls, interaction_config_file):
+    def _initialize_interactions(self, interaction_config_file):
         """Initialize interactions from configuration.
         Returns:
             dict[str, BaseInteraction]: A dictionary mapping interaction names to interaction instances.
@@ -491,5 +506,4 @@ def _initialize_interactions(cls, interaction_config_file):
             return {}
 
         interaction_map = initialize_interactions_from_config(interaction_config_file)
-        logger.info(f"Initialize interactions from configuration: interaction_map: {list(interaction_map.keys())}")
         return interaction_map

verl/workers/rollout/replica.py

@@ -18,7 +18,7 @@
 from enum import Enum
 from typing import Any, Callable, Optional
 
-from omegaconf import DictConfig, OmegaConf
+from omegaconf import DictConfig
 from pydantic import BaseModel
 from ray.actor import ActorHandle
 
@@ -90,18 +90,7 @@ def __init__(
     ) -> None:
         self.replica_rank = replica_rank
         self.config = omega_conf_to_dataclass(config)
-        # TODO: make lora config irrelevant to the model engine choice
-        # Convert megatron lora config to HFModelConfig
-        # If model_config is not an OmegaConf object, convert it first
-        if OmegaConf.is_config(model_config):
-            model_config_dict = OmegaConf.to_container(model_config)
-            model_config_dict.pop("lora", None)
-
-            self.model_config: HFModelConfig = omega_conf_to_dataclass(
-                OmegaConf.create(model_config_dict), dataclass_type=HFModelConfig
-            )
-        else:
-            self.model_config: HFModelConfig = model_config
+        self.model_config: HFModelConfig = model_config
 
         self.world_size = (
             self.config.tensor_model_parallel_size