Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ The main loop is:

## Getting Started

Agent-R1 uses the same environment setup as [verl](https://verl.readthedocs.io/en/latest/start/install.html), and the current version requires `verl==0.7.0`. You only need to clone this repository; there is no separate Agent-R1 installation step.
Agent-R1 uses the same environment setup as [verl](https://verl.readthedocs.io/en/latest/start/install.html), and the current version requires `verl==0.8.0`. You only need to clone this repository; there is no separate Agent-R1 installation step.

The recommended path is:

Expand Down
73 changes: 44 additions & 29 deletions agent_r1/agent_flow/agent_flow.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,11 +29,7 @@
from transformers import AutoProcessor, AutoTokenizer

from agent_r1.reward_loop.reward_loop import RewardLoopWorker
from verl.experimental.agent_loop.agent_loop import (
AsyncLLMServerManager,
DictConfigWrap,
)
from verl.experimental.agent_loop.prometheus_utils import update_prometheus_config
from verl.experimental.agent_loop.agent_loop import DictConfigWrap
from verl.experimental.agent_loop.utils import resolve_config_path
from verl.protocol import DataProto
from verl.single_controller.ray.base import RayResourcePool, RayWorkerGroup
Expand All @@ -48,7 +44,9 @@
rollout_trace_attr,
)
from verl.utils.transferqueue_utils import tqbridge
from verl.workers.rollout.llm_server import GlobalRequestLoadBalancer, LLMServerClient
from verl.workers.rollout.replica import get_rollout_replica_class
from verl.workers.rollout.utils import update_prometheus_config

logger = logging.getLogger(__file__)
logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
Expand Down Expand Up @@ -133,7 +131,7 @@ class AgentFlowBase(ABC):
def __init__(
self,
trainer_config: DictConfigWrap,
server_manager: AsyncLLMServerManager,
server_manager: LLMServerClient,
reward_loop_worker: RewardLoopWorker,
tokenizer: AutoTokenizer,
processor: AutoProcessor,
Expand All @@ -145,7 +143,7 @@ def __init__(

Args:
trainer_config (DictConfigWrap): trainer config.
server_manager (AsyncLLMServerManager): OpenAI compatible LLM server manager.
server_manager (LLMServerClient): LLM server client for generation requests.
reward_loop_worker (RewardLoopWorker): Reward loop worker.
tokenizer (AutoTokenizer): Tokenizer for tokenize messages.
processor (AutoProcessor): Processor for process messages.
Expand Down Expand Up @@ -471,20 +469,20 @@ class AgentFlowWorkerBase:
def __init__(
self,
config: DictConfig,
server_handles: list[ray.actor.ActorHandle],
llm_client: LLMServerClient,
reward_router_address: str = None,
):
"""Initialize agent flow manager.

Args:
config (DictConfig): YAML config.
server_handles (List[ray.actor.ActorHandle]): OpenAI compatible LLM server actor handles.
llm_client (LLMServerClient): LLM server client for generation requests.
"""
self.config = config

# for recipe to change
if not hasattr(self, "server_manager"):
self.server_manager = AsyncLLMServerManager(config, server_handles)
self.server_manager = llm_client

self.dataset_cls = get_dataset_class(config.data)
self.reward_router_address = reward_router_address
Expand All @@ -501,6 +499,11 @@ def __init__(
agent_flow_configs = OmegaConf.load(resolved_path)
for agent_flow_config in agent_flow_configs:
_agent_flow_registry[agent_flow_config.name] = agent_flow_config
if "_target_" in agent_flow_config:
import importlib

module_path = agent_flow_config["_target_"].rsplit(".", 1)[0]
importlib.import_module(module_path)
if self.config.actor_rollout_ref.model.get("custom_chat_template", None) is not None:
if self.processor is not None:
self.processor.chat_template = self.config.actor_rollout_ref.model.custom_chat_template
Expand Down Expand Up @@ -768,31 +771,24 @@ def create_transferqueue_client(
self,
):
"""Create a client for data system (TransferQueue)."""
from verl.single_controller.ray.base import get_random_string
from verl.utils.transferqueue_utils import create_transferqueue_client
import transfer_queue as tq

client_name = get_random_string(length=6)

self.tq_client = create_transferqueue_client(
client_id=f"AgentLoopWorker_{client_name}",
config=self.config.transfer_queue,
)
tq.init()
self.tq_client = tq.get_client()


@ray.remote
class AgentFlowWorker(AgentFlowWorkerBase):
"""Agent flow worker takes a batch of messages and run each message in an agent flow."""

def __init__(
self, config: DictConfig, server_handles: list[ray.actor.ActorHandle], reward_router_address: str = None
):
def __init__(self, config: DictConfig, llm_client: LLMServerClient, reward_router_address: str = None):
"""Initialize agent flow manager.
Args:
config (DictConfig): YAML config.
server_handles (List[ray.actor.ActorHandle]): OpenAI compatible LLM server actor handles.
llm_client (LLMServerClient): LLM server client for generation requests.
reward_router_address (str): reward router address.
"""
super().__init__(config, server_handles, reward_router_address)
super().__init__(config, llm_client, reward_router_address)


async def get_trajectory_info(step, index, validate):
Expand Down Expand Up @@ -834,10 +830,10 @@ def __init__(
self.worker_group = worker_group
self.reward_model_manager = None
self.reward_router_address = None
if self.config.reward_model.enable:
if self.config.reward.reward_model.enable:
from verl.experimental.reward_loop import RewardModelManager

self.reward_model_manager = RewardModelManager(config.reward_model, rm_resource_pool)
self.reward_model_manager = RewardModelManager(config.reward.reward_model, rm_resource_pool)
self.reward_router_address = self.reward_model_manager.get_router_address()

# for recipe to change
Expand Down Expand Up @@ -890,7 +886,16 @@ def _initialize_llm_servers(self):
if rollout_config.prometheus.enable:
if rollout_config.disable_log_stats:
raise ValueError("PROMETHEUS needs disable_log_stats==False, but it is currently True.")
update_prometheus_config(rollout_config.prometheus, self.server_addresses)
update_prometheus_config(rollout_config.prometheus, self.server_addresses, rollout_config.name)

# Create global load balancer and LLM server client
self.global_load_balancer = GlobalRequestLoadBalancer.remote(
servers=dict(zip(self.server_addresses, self.server_handles, strict=True)),
)
self.llm_client = LLMServerClient(
config=self.config,
load_balancer_handle=self.global_load_balancer,
)

def _init_agent_flow_workers(self):
self.agent_flow_workers = []
Expand All @@ -906,7 +911,7 @@ def _init_agent_flow_workers(self):
scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(
node_id=node_id, soft=True
),
).remote(self.config, self.server_handles, self.reward_router_address)
).remote(self.config, self.llm_client, self.reward_router_address)
)

def generate_sequences(self, prompts: DataProto) -> DataProto:
Expand Down Expand Up @@ -1014,8 +1019,18 @@ def _performance_metrics(
return timing

def wake_up(self):
"""Wake up all rollout replica instances."""
self._run_all([replica.wake_up() for replica in self.rollout_replicas])
"""Sync training weights to vLLM and wake up rollout replicas.

In verl 0.8.0 the vLLMHttpServer.wake_up() only restores engine
state from sleep — it no longer calls the training workers to push
model weights. We call update_weights(mode="naive") through the
worker group dispatch so the FSDP parameters are sent to the vLLM
engine via the ServerAdapter (matching the verl 0.7.0 behaviour).
"""
if self.worker_group:
ray.get(self.worker_group.update_weights(mode="naive"))
else:
self._run_all([replica.wake_up() for replica in self.rollout_replicas])

def sleep(self):
"""Sleep all rollout replica instances."""
Expand Down
29 changes: 16 additions & 13 deletions agent_r1/reward_loop/reward_loop.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
import ray
from omegaconf import DictConfig

from verl.experimental.reward_loop.reward_loop import get_reward_manager_cls
from verl.experimental.reward_loop.reward_manager import get_reward_manager_cls
from verl.protocol import DataProto
from verl.trainer.ppo.reward import get_custom_reward_fn
from verl.utils import hf_tokenizer
Expand Down Expand Up @@ -59,24 +59,24 @@ def _init_reward_fn(self):
input_tokenizer_local_path = copy_to_local(self.config.actor_rollout_ref.model.path)
self.input_tokenizer = hf_tokenizer(input_tokenizer_local_path, trust_remote_code=True)
self.reward_model_tokenizer = None
if self.config.reward_model.enable:
reward_model_tokenizer_local_path = copy_to_local(self.config.reward_model.model.path)
if self.config.reward.reward_model.enable:
reward_model_tokenizer_local_path = copy_to_local(self.config.reward.reward_model.model.path)
self.reward_model_tokenizer = hf_tokenizer(reward_model_tokenizer_local_path, trust_remote_code=True)
self.reward_fn = get_custom_reward_fn(self.config)

# Load reward loop manager class
# Support both registry and importlib loading methods
reward_loop_source = self.config.reward_model.get("reward_loop_source", "register")
reward_loop_source = self.config.reward.reward_model.get("reward_loop_source", "register")

if reward_loop_source == "register":
# Load from registry (default behavior)
reward_manager_cls = get_reward_manager_cls(self.config.reward_model.reward_manager)
reward_manager_cls = get_reward_manager_cls(self.config.reward.reward_manager.name)
elif reward_loop_source == "importlib":
# Load from external module using importlib
from verl.utils.import_utils import load_extern_object

reward_loop_module_path = self.config.reward_model.get("reward_loop_module_path", None)
reward_loop_class_name = self.config.reward_model.get("reward_loop_class_name", None)
reward_loop_module_path = self.config.reward.reward_model.get("reward_loop_module_path", None)
reward_loop_class_name = self.config.reward.reward_model.get("reward_loop_class_name", None)

assert reward_loop_module_path is not None, (
"reward_loop_module_path must be set when reward_loop_source='importlib'"
Expand Down Expand Up @@ -113,24 +113,27 @@ async def compute_score(self, input_data) -> dict:
if isinstance(input_data, DataProto):
data = input_data
assert len(data) == 1, "RewardLoopWorker only support single data item"
if self.config.custom_reward_function.path is not None:
if self.config.reward.custom_reward_function.path is not None:
# directly use user-customized reward function
return await self.reward_loop.run_single(data)
else:
if self.config.reward_model.enable:
if self.config.reward.reward_model.enable:
# we assume the rm is disrm
# genrm must set custom_reward_function
return await self.compute_score_disrm(data)
else:
return await self.reward_loop.run_single(data)

# For now, only support DataProto-less inputs for DisRM.
if getattr(self.config, "custom_reward_function", None) is not None and self.config.custom_reward_function.path:
if (
getattr(self.config, "custom_reward_function", None) is not None
and self.config.reward.custom_reward_function.path
):
raise NotImplementedError(
"RewardLoopWorker currently supports non-DataProto inputs only in the DisRM path. "
"When custom_reward_function is configured, you must pass DataProto."
)
if not self.config.reward_model.enable:
if not self.config.reward.reward_model.enable:
raise NotImplementedError(
"RewardLoopWorker only supports non-DataProto inputs in DisRM mode. "
"Set reward_model.enable=True (and do not use custom_reward_function), or pass DataProto."
Expand Down Expand Up @@ -245,8 +248,8 @@ async def compute_score_disrm(self, input_data) -> dict:
"Supported: DataProto | str | messages(list[dict])"
)

engine_name = self.config.reward_model.rollout.name
model_name = self.config.reward_model.model.path
engine_name = self.config.reward.reward_model.rollout.name
model_name = self.config.reward.reward_model.model.path
if engine_name == "vllm":
# TODO (dyy): the "activation" has been changed to "use_activation" in vllm 0.11.2
payloads = {
Expand Down
Loading