From ddfbb9746cd8437a2849e21a211af01d0970ae38 Mon Sep 17 00:00:00 2001 From: lqzxt <997260326@qq.com> Date: Mon, 6 Jul 2026 17:57:42 +0800 Subject: [PATCH 1/2] feat: adapt Agent-R1 to verl 0.8.0 --- README.md | 2 +- agent_r1/agent_flow/agent_flow.py | 70 +++--- agent_r1/reward_loop/reward_loop.py | 26 +-- agent_r1/trainer/main_agent_ppo.py | 160 ++++--------- agent_r1/trainer/ppo/ray_trainer.py | 208 +++++++---------- agent_r1/workers/actor/__init__.py | 4 +- agent_r1/workers/actor/dp_actor.py | 219 +----------------- agent_r1/workers/critic/__init__.py | 4 +- agent_r1/workers/critic/dp_critic.py | 142 +----------- agent_r1/workers/engine_workers.py | 1 + agent_r1/workers/fsdp_workers.py | 52 +---- docs/getting-started/installation-guide.md | 4 +- docs/zh/getting-started/installation-guide.md | 4 +- examples/gsm8k/run_grpo.sh | 72 ++++++ recipes/alfworld/alfworld_agent_flow.py | 5 +- recipes/hotpotqa/hotpotqa_agent_flow.py | 5 +- .../paper_search/paper_search_agent_flow.py | 5 +- recipes/webshop/webshop_agent_flow.py | 5 +- 18 files changed, 286 insertions(+), 702 deletions(-) create mode 100755 examples/gsm8k/run_grpo.sh diff --git a/README.md b/README.md index fec6a2d..354009a 100644 --- a/README.md +++ b/README.md @@ -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: diff --git a/agent_r1/agent_flow/agent_flow.py b/agent_r1/agent_flow/agent_flow.py index 9281efa..48a4c41 100644 --- a/agent_r1/agent_flow/agent_flow.py +++ b/agent_r1/agent_flow/agent_flow.py @@ -29,11 +29,9 @@ 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.workers.rollout.llm_server import GlobalRequestLoadBalancer, LLMServerClient +from verl.workers.rollout.utils import update_prometheus_config from verl.experimental.agent_loop.utils import resolve_config_path from verl.protocol import DataProto from verl.single_controller.ray.base import RayResourcePool, RayWorkerGroup @@ -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, @@ -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. @@ -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 @@ -501,6 +499,10 @@ 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 @@ -768,15 +770,10 @@ 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 - - client_name = get_random_string(length=6) + import transfer_queue as tq - 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 @@ -784,15 +781,15 @@ 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 + 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): @@ -834,10 +831,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 @@ -890,7 +887,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 = [] @@ -906,7 +912,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: @@ -1014,8 +1020,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.""" diff --git a/agent_r1/reward_loop/reward_loop.py b/agent_r1/reward_loop/reward_loop.py index b5d821f..cca9829 100644 --- a/agent_r1/reward_loop/reward_loop.py +++ b/agent_r1/reward_loop/reward_loop.py @@ -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 @@ -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'" @@ -113,11 +113,11 @@ 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) @@ -125,12 +125,12 @@ async def compute_score(self, input_data) -> dict: 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." @@ -245,8 +245,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 = { diff --git a/agent_r1/trainer/main_agent_ppo.py b/agent_r1/trainer/main_agent_ppo.py index 5e03ab9..3bce8d4 100644 --- a/agent_r1/trainer/main_agent_ppo.py +++ b/agent_r1/trainer/main_agent_ppo.py @@ -23,8 +23,8 @@ from omegaconf import OmegaConf from agent_r1.trainer.ppo.ray_trainer import RayAgentTrainer, need_critic_agent_ppo +from verl.experimental.reward_loop import migrate_legacy_reward_impl from verl.trainer.constants_ppo import get_ppo_ray_runtime_env -from verl.trainer.ppo.reward import load_reward_manager from verl.trainer.ppo.utils import need_reference_policy from verl.utils.config import validate_config from verl.utils.device import auto_set_device, is_cuda_available @@ -39,6 +39,7 @@ def main(config): """ # Automatically set `config.trainer.device = npu` when running on Ascend NPU. auto_set_device(config) + config = migrate_legacy_reward_impl(config) run_ppo_agent(config) @@ -115,74 +116,34 @@ def __init__(self): self.mapping = {} def add_actor_rollout_worker(self, config): - """Add actor rollout worker based on the actor strategy.""" + """Add actor rollout worker using the unified model engine implementation.""" from verl.single_controller.ray import RayWorkerGroup from verl.trainer.ppo.ray_trainer import Role - use_legacy_worker_impl = config.trainer.get("use_legacy_worker_impl", "auto") + from agent_r1.workers.engine_workers import ActorRolloutRefWorker - # use new model engine implementation - if use_legacy_worker_impl == "disable": - from agent_r1.workers.engine_workers import ActorRolloutRefWorker - - actor_rollout_cls = ActorRolloutRefWorker - ray_worker_group_cls = RayWorkerGroup - # NOTE: In new model engine, ref policy and actor rollout are in same ActorRolloutRefWorker, - # while in legacy model engine, ref policy is in a separate ActorRolloutRefWorker. - if config.algorithm.use_kl_in_reward or config.actor_rollout_ref.actor.use_kl_loss: - role = Role.ActorRolloutRef - else: - role = Role.ActorRollout - self.role_worker_mapping[role] = ray.remote(actor_rollout_cls) - self.mapping[role] = "global_pool" - return actor_rollout_cls, ray_worker_group_cls - - # Note: sync mode validation is now handled in RolloutConfig.__post_init__ - # Always use async worker since sync mode is deprecated and rejected - if config.actor_rollout_ref.actor.strategy in {"fsdp", "fsdp2"}: - from agent_r1.workers.fsdp_workers import AsyncActorRolloutRefWorker - - actor_rollout_cls = AsyncActorRolloutRefWorker - ray_worker_group_cls = RayWorkerGroup - - elif config.actor_rollout_ref.actor.strategy == "megatron": - from verl.workers.megatron_workers import AsyncActorRolloutRefWorker - - actor_rollout_cls = AsyncActorRolloutRefWorker - ray_worker_group_cls = RayWorkerGroup + actor_rollout_cls = ActorRolloutRefWorker + ray_worker_group_cls = RayWorkerGroup + lora_rank = config.actor_rollout_ref.model.get("lora", {}).get("rank", 0) + if lora_rank <= 0: + lora_rank = config.actor_rollout_ref.model.get("lora_rank", 0) + ref_in_actor = lora_rank > 0 or config.actor_rollout_ref.model.get("lora_adapter_path") is not None + if need_reference_policy(config) and not ref_in_actor: + role = Role.ActorRolloutRef else: - raise NotImplementedError - - self.role_worker_mapping[Role.ActorRollout] = ray.remote(actor_rollout_cls) - self.mapping[Role.ActorRollout] = "global_pool" + role = Role.ActorRollout + self.role_worker_mapping[role] = ray.remote(actor_rollout_cls) + self.mapping[role] = "global_pool" return actor_rollout_cls, ray_worker_group_cls def add_critic_worker(self, config): - """Add critic worker to role mapping.""" - use_legacy_worker_impl = config.trainer.get("use_legacy_worker_impl", "auto") - if config.critic.strategy in {"fsdp", "fsdp2"}: - if use_legacy_worker_impl in ["auto", "enable"]: - from agent_r1.workers.fsdp_workers import CriticWorker - elif use_legacy_worker_impl == "disable": - # we don't need to specialize critic worker. Just use TrainingWorker - from agent_r1.workers.engine_workers import TrainingWorker - - CriticWorker = TrainingWorker - print("Using new worker implementation") - else: - raise ValueError(f"Invalid use_legacy_worker_impl: {use_legacy_worker_impl}") - - elif config.critic.strategy == "megatron": - # TODO: switch this to TrainingWorker as well - from verl.workers.megatron_workers import CriticWorker - - else: - raise NotImplementedError - + """Add critic worker to role mapping using the unified model engine implementation.""" from verl.trainer.ppo.ray_trainer import Role - self.role_worker_mapping[Role.Critic] = ray.remote(CriticWorker) + from agent_r1.workers.engine_workers import TrainingWorker + + self.role_worker_mapping[Role.Critic] = ray.remote(TrainingWorker) self.mapping[Role.Critic] = "global_pool" def init_resource_pool_mgr(self, config): @@ -192,60 +153,45 @@ def init_resource_pool_mgr(self, config): resource_pool_spec = { global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes, } - # TODO Here you can use the new registration method to support dynamic registration of roles - if config.reward_model.enable_resource_pool: - if config.reward_model.n_gpus_per_node <= 0: - raise ValueError("config.reward_model.n_gpus_per_node must be greater than 0") - if config.reward_model.nnodes <= 0: - raise ValueError("config.reward_model.nnodes must be greater than 0") - - reward_pool = [config.reward_model.n_gpus_per_node] * config.reward_model.nnodes + + if config.reward.reward_model.enable_resource_pool: + if config.reward.reward_model.n_gpus_per_node <= 0: + raise ValueError("config.reward.reward_model.n_gpus_per_node must be greater than 0") + if config.reward.reward_model.nnodes <= 0: + raise ValueError("config.reward.reward_model.nnodes must be greater than 0") + + reward_pool = [config.reward.reward_model.n_gpus_per_node] * config.reward.reward_model.nnodes resource_pool_spec["reward_pool"] = reward_pool + else: + config.reward.reward_model.nnodes = config.trainer.nnodes + config.reward.reward_model.n_gpus_per_node = config.trainer.n_gpus_per_node from verl.trainer.ppo.ray_trainer import ResourcePoolManager resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=self.mapping) return resource_pool_manager - def add_reward_model_worker(self, config): - """Add reward model worker if enabled.""" - from verl.trainer.ppo.ray_trainer import Role + def add_reward_model_resource_pool(self, config): + """Add reward model resource pool if enabled. - if config.reward_model.enable: - use_legacy_worker_impl = config.trainer.get("use_legacy_worker_impl", "auto") - if use_legacy_worker_impl in ["auto", "enable", "disable"]: - if config.reward_model.strategy in {"fsdp", "fsdp2"}: - from verl.workers.fsdp_workers import RewardModelWorker - elif config.reward_model.strategy == "megatron": - from verl.workers.megatron_workers import RewardModelWorker - else: - raise NotImplementedError - # elif use_legacy_worker_impl == "disable": - # from verl.workers.engine_workers import RewardModelWorker - # - # print("Using new worker implementation") - else: - raise ValueError(f"Invalid use_legacy_worker_impl: {use_legacy_worker_impl}") + In v0.8.0, reward model workers are created inside RayPPOTrainer.init_workers() + via RewardLoopManager. We only need to register the resource pool mapping here. + """ + from verl.trainer.ppo.ray_trainer import Role - self.role_worker_mapping[Role.RewardModel] = ray.remote(RewardModelWorker) - if config.reward_model.enable_resource_pool: + if config.reward.reward_model.enable: + if config.reward.reward_model.enable_resource_pool: self.mapping[Role.RewardModel] = "reward_pool" else: self.mapping[Role.RewardModel] = "global_pool" def add_ref_policy_worker(self, config, ref_policy_cls): - """Add reference policy worker if KL loss or KL reward is used.""" - from verl.trainer.ppo.ray_trainer import Role - - # Ref policy has been fused into ActorRolloutRefWorker in new model engine, - # we don't need to add a separate ref policy worker group. - use_legacy_worker_impl = config.trainer.get("use_legacy_worker_impl", "auto") - if use_legacy_worker_impl == "disable": - return + """Ref policy is fused into ActorRolloutRefWorker in the unified model engine. - if config.algorithm.use_kl_in_reward or config.actor_rollout_ref.actor.use_kl_loss: - self.role_worker_mapping[Role.RefPolicy] = ray.remote(ref_policy_cls) - self.mapping[Role.RefPolicy] = "global_pool" + This is now a no-op because the reference policy lives on the same worker group + as the actor/rollout. + """ + return def run(self, config): """Execute the main PPO training workflow. @@ -271,13 +217,7 @@ def run(self, config): actor_rollout_cls, ray_worker_group_cls = self.add_actor_rollout_worker(config) self.add_critic_worker(config) - # We should adopt a multi-source reward function here: - # - for rule-based rm, we directly call a reward score - # - for model-based rm, we call a model - # - for code related prompt, we send to a sandbox if there are test cases - # finally, we combine all the rewards together - # The reward type depends on the tag of the data - self.add_reward_model_worker(config) + self.add_reward_model_resource_pool(config) # Add a reference policy worker if KL loss or KL reward is used. self.add_ref_policy_worker(config, actor_rollout_cls) @@ -285,7 +225,7 @@ def run(self, config): # validate config validate_config( config=config, - use_reference_policy=need_reference_policy(self.role_worker_mapping), + use_reference_policy=need_reference_policy(config), use_critic=need_critic_agent_ppo(config), ) @@ -303,14 +243,6 @@ def run(self, config): # Used for multimodal LLM, could be None processor = hf_processor(local_path, trust_remote_code=trust_remote_code, use_fast=True) - # Load the reward manager for training and validation. - reward_fn = load_reward_manager( - config, tokenizer, num_examine=0, **config.reward_model.get("reward_kwargs", {}) - ) - val_reward_fn = load_reward_manager( - config, tokenizer, num_examine=1, **config.reward_model.get("reward_kwargs", {}) - ) - resource_pool_manager = self.init_resource_pool_mgr(config) from verl.trainer.main_ppo import create_rl_dataset, create_rl_sampler @@ -344,8 +276,6 @@ def run(self, config): role_worker_mapping=self.role_worker_mapping, resource_pool_manager=resource_pool_manager, ray_worker_group_cls=ray_worker_group_cls, - reward_fn=reward_fn, - val_reward_fn=val_reward_fn, train_dataset=train_dataset, val_dataset=val_dataset, collate_fn=collate_fn, diff --git a/agent_r1/trainer/ppo/ray_trainer.py b/agent_r1/trainer/ppo/ray_trainer.py index ec5a2c6..dc7ca55 100644 --- a/agent_r1/trainer/ppo/ray_trainer.py +++ b/agent_r1/trainer/ppo/ray_trainer.py @@ -36,7 +36,6 @@ from agent_r1.trainer.ppo.metric_utils import compute_data_metrics from agent_r1.trainer.ppo.trajectory_batching import prepare_trajectory_mini_batch from verl import DataProto -from verl.experimental.dataset.sampler import AbstractCurriculumSampler from verl.protocol import pad_dataproto_to_divisor from verl.single_controller.ray import RayClassWithInitArgs from verl.single_controller.ray.base import create_colocated_worker_cls @@ -52,7 +51,7 @@ apply_kl_penalty, compute_response_mask, ) -from verl.trainer.ppo.reward import compute_reward_async +from verl.trainer.ppo.reward import extract_reward from verl.trainer.ppo.utils import Role from verl.utils.checkpoint.checkpoint_manager import should_save_ckpt_esi from verl.utils.config import omega_conf_to_dataclass @@ -387,74 +386,66 @@ def _update_actor(self, batch: DataProto) -> DataProto: batch.meta_info["temperature"] = rollout_config.temperature ppo_mini_batch_size = self.config.actor_rollout_ref.actor.ppo_mini_batch_size ppo_mini_batch_size = ppo_mini_batch_size * self.config.actor_rollout_ref.rollout.n - if self.use_legacy_worker_impl == "disable": - from verl.utils import tensordict_utils as tu - from verl.utils.py_functional import rename_dict - from verl.workers.utils.padding import left_right_2_no_padding - - calculate_entropy = self.config.actor_rollout_ref.actor.entropy_coeff != 0.0 - dp_size = self._get_dp_size(self.actor_rollout_wg, "actor") - update_batch = prepare_trajectory_mini_batch(batch, mini_batch_size=ppo_mini_batch_size, dp_size=dp_size) - batch_td = update_batch.to_tensordict() - batch_td = left_right_2_no_padding(batch_td) - ppo_epochs = self.config.actor_rollout_ref.actor.ppo_epochs - seed = self.config.actor_rollout_ref.actor.data_loader_seed - shuffle = self.config.actor_rollout_ref.actor.shuffle - tu.assign_non_tensor( - batch_td, - calculate_entropy=calculate_entropy, - mini_batch_size=ppo_mini_batch_size, - num_mini_batch=update_batch.meta_info["num_mini_batch"], - epochs=ppo_epochs, - seed=seed, - dataloader_kwargs={"shuffle": shuffle}, - ) - actor_output = self.actor_rollout_wg.update_actor(batch_td) - actor_output = tu.get(actor_output, "metrics") - actor_output = rename_dict(actor_output, "actor/") - actor_output["perf/mfu/actor"] = actor_output.pop("actor/mfu") - actor_output = DataProto.from_single_dict(data={}, meta_info={"metrics": actor_output}) - else: - dp_size = self._get_worker_group_dp_size(self.actor_rollout_wg, ("actor",)) - update_batch = prepare_trajectory_mini_batch(batch, mini_batch_size=ppo_mini_batch_size, dp_size=dp_size) - actor_output = self.actor_rollout_wg.update_actor(update_batch) + from verl.utils import tensordict_utils as tu + from verl.utils.py_functional import rename_dict + from verl.workers.utils.padding import left_right_2_no_padding + + calculate_entropy = self.config.actor_rollout_ref.actor.entropy_coeff != 0.0 + dp_size = self._get_dp_size(self.actor_rollout_wg, "actor") + update_batch = prepare_trajectory_mini_batch(batch, mini_batch_size=ppo_mini_batch_size, dp_size=dp_size) + batch_td = update_batch.to_tensordict() + batch_td = left_right_2_no_padding(batch_td) + ppo_epochs = self.config.actor_rollout_ref.actor.ppo_epochs + seed = self.config.actor_rollout_ref.actor.data_loader_seed + shuffle = self.config.actor_rollout_ref.actor.shuffle + tu.assign_non_tensor( + batch_td, + calculate_entropy=calculate_entropy, + mini_batch_size=ppo_mini_batch_size, + num_mini_batch=update_batch.meta_info["num_mini_batch"], + epochs=ppo_epochs, + seed=seed, + dataloader_kwargs={"shuffle": shuffle}, + ) + + actor_output = self.actor_rollout_wg.update_actor(batch_td) + actor_output = tu.get(actor_output, "metrics") + actor_output = rename_dict(actor_output, "actor/") + actor_output["perf/mfu/actor"] = actor_output.pop("actor/mfu") + actor_output = DataProto.from_single_dict(data={}, meta_info={"metrics": actor_output}) return actor_output def _update_critic(self, batch: DataProto) -> DataProto: ppo_mini_batch_size = self.config.critic.ppo_mini_batch_size ppo_mini_batch_size = ppo_mini_batch_size * self.config.actor_rollout_ref.rollout.n - if self.use_legacy_worker_impl == "disable": - from verl.utils import tensordict_utils as tu - from verl.utils.py_functional import rename_dict - from verl.workers.utils.padding import left_right_2_no_padding - - dp_size = self._get_worker_group_dp_size(self.critic_wg, ("train", "critic")) - update_batch = prepare_trajectory_mini_batch(batch, mini_batch_size=ppo_mini_batch_size, dp_size=dp_size) - batch_td = update_batch.to_tensordict() - batch_td = left_right_2_no_padding(batch_td) - ppo_epochs = self.config.critic.ppo_epochs - seed = self.config.critic.data_loader_seed - shuffle = self.config.critic.shuffle - tu.assign_non_tensor( - batch_td, - mini_batch_size=ppo_mini_batch_size, - num_mini_batch=update_batch.meta_info["num_mini_batch"], - epochs=ppo_epochs, - seed=seed, - dataloader_kwargs={"shuffle": shuffle}, - ) - output = self.critic_wg.train_mini_batch(batch_td) - output = output.get() - output = tu.get(output, "metrics") - output = rename_dict(output, "critic/") - output["perf/mfu/critic"] = output.pop("critic/mfu") - output = DataProto.from_single_dict(data={}, meta_info={"metrics": output}) - else: - dp_size = self._get_worker_group_dp_size(self.critic_wg, ("critic",)) - update_batch = prepare_trajectory_mini_batch(batch, mini_batch_size=ppo_mini_batch_size, dp_size=dp_size) - output = self.critic_wg.update_critic(update_batch) + from verl.utils import tensordict_utils as tu + from verl.utils.py_functional import rename_dict + from verl.workers.utils.padding import left_right_2_no_padding + + dp_size = self._get_worker_group_dp_size(self.critic_wg, ("train", "critic")) + update_batch = prepare_trajectory_mini_batch(batch, mini_batch_size=ppo_mini_batch_size, dp_size=dp_size) + batch_td = update_batch.to_tensordict() + batch_td = left_right_2_no_padding(batch_td) + ppo_epochs = self.config.critic.ppo_epochs + seed = self.config.critic.data_loader_seed + shuffle = self.config.critic.shuffle + tu.assign_non_tensor( + batch_td, + mini_batch_size=ppo_mini_batch_size, + num_mini_batch=update_batch.meta_info["num_mini_batch"], + epochs=ppo_epochs, + seed=seed, + dataloader_kwargs={"shuffle": shuffle}, + ) + + output = self.critic_wg.train_mini_batch(batch_td) + output = output.get() + output = tu.get(output, "metrics") + output = rename_dict(output, "critic/") + output["perf/mfu/critic"] = output.pop("critic/mfu") + output = DataProto.from_single_dict(data={}, meta_info={"metrics": output}) return output def _get_worker_group_dp_size(self, worker_group, roles: Sequence[str]) -> int: @@ -588,7 +579,7 @@ def _validate(self): ) # we only do validation on rule-based rm - if self.config.reward_model.enable and test_batch[0].non_tensor_batch["reward_model"]["style"] == "model": + if self.use_rm and test_batch[0].non_tensor_batch["reward_model"]["style"] == "model": return {} sample_uids.extend(test_batch.non_tensor_batch["uid"]) @@ -616,12 +607,8 @@ def _validate(self): test_output_gen_batch.meta_info["validate"] = True # evaluate using reward_function - result = self._compute_or_extract_reward( - test_output_gen_batch, reward_fn=self.val_reward_fn, return_dict=True - ) - reward_tensor = result["reward_tensor"] + reward_tensor, reward_extra_info = extract_reward(test_output_gen_batch) step_scores = reward_tensor.sum(-1).detach().cpu().tolist() - reward_extra_info = result.get("reward_extra_info", {}) step_inputs = self.tokenizer.batch_decode( test_output_gen_batch.batch["input_ids"], skip_special_tokens=True ) @@ -681,7 +668,7 @@ def _validate(self): sample_outputs.extend(batch_traj_outputs) reward_extra_infos_dict["reward"].extend(batch_traj_scores) - if "reward_extra_info" in result: + if batch_traj_extra_info: for key, vals in batch_traj_extra_info.items(): reward_extra_infos_dict[key].extend(make_json_safe(vals)) @@ -760,26 +747,21 @@ def init_workers(self): critic_cfg: CriticConfig = omega_conf_to_dataclass(self.config.critic) - if self.use_legacy_worker_impl == "disable": - # convert critic_cfg into TrainingWorkerConfig - from verl.workers.config.engine import FSDPEngineConfig - from verl.workers.engine_workers import TrainingWorkerConfig - - orig_critic_cfg = critic_cfg - if orig_critic_cfg.strategy == "fsdp": - engine_config: FSDPEngineConfig = orig_critic_cfg.model.fsdp_config - engine_config.infer_max_token_len_per_gpu = critic_cfg.ppo_infer_max_token_len_per_gpu - engine_config.max_token_len_per_gpu = critic_cfg.ppo_max_token_len_per_gpu - else: - raise NotImplementedError(f"Unknown strategy {orig_critic_cfg.strategy=}") - - critic_cfg = TrainingWorkerConfig( - model_type="value_model", - model_config=orig_critic_cfg.model_config, - engine_config=engine_config, - optimizer_config=orig_critic_cfg.optim, - checkpoint_config=orig_critic_cfg.checkpoint, - ) + # convert critic_cfg into TrainingWorkerConfig for the unified model engine worker + from verl.workers.engine_workers import TrainingWorkerConfig + + orig_critic_cfg = critic_cfg + engine_config = orig_critic_cfg.engine + engine_config.infer_max_token_len_per_gpu = critic_cfg.ppo_infer_max_token_len_per_gpu + engine_config.max_token_len_per_gpu = critic_cfg.ppo_max_token_len_per_gpu + + critic_cfg = TrainingWorkerConfig( + model_type="value_model", + model_config=orig_critic_cfg.model, + engine_config=engine_config, + optimizer_config=orig_critic_cfg.optim, + checkpoint_config=orig_critic_cfg.checkpoint, + ) critic_cls = RayClassWithInitArgs(cls=self.role_worker_mapping[Role.Critic], config=critic_cfg) self.resource_pool_to_cls[resource_pool][str(Role.Critic)] = critic_cls @@ -828,17 +810,14 @@ def init_workers(self): if self.use_critic: self.critic_wg = all_wg[str(Role.Critic)] - if self.use_legacy_worker_impl == "disable": - self.critic_wg.reset() - # assign critic loss - from functools import partial + self.critic_wg.reset() + # assign critic loss + from functools import partial - from agent_r1.workers.utils.losses import value_loss + from agent_r1.workers.utils.losses import value_loss - value_loss_ = partial(value_loss, config=orig_critic_cfg) - self.critic_wg.set_loss_fn(value_loss_) - else: - self.critic_wg.init_model() + value_loss_ = partial(value_loss, config=orig_critic_cfg) + self.critic_wg.set_loss_fn(value_loss_) if self.use_reference_policy and not self.ref_in_actor: if str(Role.RefPolicy) in all_wg: @@ -859,10 +838,11 @@ def init_workers(self): # create async rollout manager and request scheduler # Note: mode is always "async" since sync mode is deprecated self.async_rollout_mode = True + self.reward_loop_manager = None from agent_r1.agent_flow import AgentFlowManager - if self.config.reward_model.enable: + if self.use_rm: rm_resource_pool = self.resource_pool_manager.get_resource_pool(Role.RewardModel) else: rm_resource_pool = None @@ -900,7 +880,7 @@ def fit(self): # perform validation before training # currently, we only support validation using the reward_function. - if self.val_reward_fn is not None and self.config.trainer.get("val_before_train", True): + if self.val_dataloader is not None and self.config.trainer.get("val_before_train", True): val_metrics = self._validate() assert val_metrics, f"{val_metrics=}" pprint(f"Initial validation metrics: {val_metrics}") @@ -967,8 +947,6 @@ def fit(self): gen_batch_output.meta_info.pop("timing", None) if AgentAdvantageEstimator(self.config.algorithm.adv_estimator) == AgentAdvantageEstimator.REMAX: - if self.reward_fn is None: - raise ValueError("A reward_fn is required for REMAX advantage estimation.") # TODO: implement REMAX advantage estimation for agent flow. raise NotImplementedError("REMAX advantage estimation is not supported for agent flow.") @@ -1001,15 +979,8 @@ def fit(self): reward_tensor = self.reward_loop_manager.compute_rm_score(batch) batch = batch.union(reward_tensor) - # Compute or extract reward for training - if self.config.reward_model.launch_reward_fn_async: - future_reward = compute_reward_async.remote( - data=batch, config=self.config, tokenizer=self.tokenizer - ) - else: - reward_tensor, reward_extra_infos_dict = self._compute_or_extract_reward( - batch, reward_fn=self.reward_fn, return_dict=False - ) + # extract reward_tensor and reward_extra_infos_dict for training + reward_tensor, reward_extra_infos_dict = extract_reward(batch) # Operating Mode Selection: # - Bypass mode: Sets old_log_probs = rollout_log_probs (2 policies: π_rollout, π_θ) @@ -1066,9 +1037,6 @@ def fit(self): with marked_timer("adv", timing_raw, color="brown"): # we combine with rule-based rm - reward_extra_infos_dict: dict[str, list] - if self.config.reward_model.launch_reward_fn_async: - reward_tensor, reward_extra_infos_dict = ray.get(future_reward) batch.batch["token_level_scores"] = reward_tensor if reward_extra_infos_dict: @@ -1148,7 +1116,7 @@ def fit(self): # validate if ( - self.val_reward_fn is not None + self.val_dataloader is not None and self.config.trainer.test_freq > 0 and (is_last_step or self.global_steps % self.config.trainer.test_freq == 0) ): @@ -1213,7 +1181,7 @@ def fit(self): # Note: mismatch metrics (KL, PPL, etc.) are collected at line 1179 after advantage computation # this is experimental and may be changed/removed in the future in favor of a general-purpose one - if isinstance(self.train_dataloader.sampler, AbstractCurriculumSampler): + if hasattr(self.train_dataloader.sampler, "update"): self.train_dataloader.sampler.update(batch=batch) # TODO: make a canonical logger that supports various backend @@ -1246,11 +1214,9 @@ def fit(self): def _pad_dataproto_to_world_size(self, batch): dp_sizes = [] if self.use_critic and self.critic_wg.world_size != 0: - critic_roles = ("train", "critic") if self.use_legacy_worker_impl == "disable" else ("critic",) - dp_sizes.append(self._get_worker_group_dp_size(self.critic_wg, critic_roles)) + dp_sizes.append(self._get_worker_group_dp_size(self.critic_wg, ("train", "critic"))) if self.use_reference_policy and self.ref_policy_wg.world_size != 0: - ref_roles = ("ref", "actor") if self.use_legacy_worker_impl == "disable" else ("actor", "ref") - dp_sizes.append(self._get_worker_group_dp_size(self.ref_policy_wg, ref_roles)) + dp_sizes.append(self._get_worker_group_dp_size(self.ref_policy_wg, ("ref", "actor"))) if self.hybrid_engine: if self.actor_rollout_wg.world_size != 0: dp_sizes.append(self._get_worker_group_dp_size(self.actor_rollout_wg, ("actor",))) diff --git a/agent_r1/workers/actor/__init__.py b/agent_r1/workers/actor/__init__.py index b46e9dd..091bb98 100644 --- a/agent_r1/workers/actor/__init__.py +++ b/agent_r1/workers/actor/__init__.py @@ -1,3 +1 @@ -from .dp_actor import DataParallelPPOActor - -__all__ = ["DataParallelPPOActor"] +# Removed: verl v0.8.0 no longer has DataParallelPPOActor. diff --git a/agent_r1/workers/actor/dp_actor.py b/agent_r1/workers/actor/dp_actor.py index a1ad6f6..5307d9b 100644 --- a/agent_r1/workers/actor/dp_actor.py +++ b/agent_r1/workers/actor/dp_actor.py @@ -1,217 +1,2 @@ -# Copyright 2025 Agent-R1 Teams -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -""" -Agent-R1 local PPO actor overrides. -""" - -import logging -import os - -import torch - -from agent_r1.trainer.ppo.core_algos import agg_loss, get_policy_loss_fn -from agent_r1.trainer.ppo.trajectory_batching import get_mini_batch_global_info, split_data_proto_by_mini_batch_id -from verl import DataProto -from verl.trainer.ppo.core_algos import kl_penalty -from verl.utils.device import get_device_id -from verl.utils.profiler import GPUMemoryLogger -from verl.utils.py_functional import append_to_dict -from verl.utils.seqlen_balancing import prepare_dynamic_batch -from verl.workers.actor.dp_actor import DataParallelPPOActor as VerlDataParallelPPOActor - -logger = logging.getLogger(__file__) -logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) - - -class DataParallelPPOActor(VerlDataParallelPPOActor): - @GPUMemoryLogger(role="dp actor", logger=logger) - def update_policy(self, data: DataProto): - self.actor_module.train() - - temperature = data.meta_info["temperature"] - - select_keys = [ - "responses", - "response_mask", - "input_ids", - "attention_mask", - "position_ids", - "old_log_probs", - "advantages", - ] - if self.config.use_kl_loss: - select_keys.append("ref_log_prob") - if "rollout_is_weights" in data.batch.keys(): - select_keys.append("rollout_is_weights") - if "rollout_log_probs" in data.batch.keys(): - select_keys.append("rollout_log_probs") - has_planned_mini_batches = "mini_batch_id" in data.batch.keys() - if has_planned_mini_batches: - select_keys.extend( - [ - "mini_batch_id", - "mini_batch_global_size", - "mini_batch_global_token_num", - "mini_batch_global_response_token_num", - "sample_mask", - ] - ) - - has_multi_modal_inputs = "multi_modal_inputs" in data.non_tensor_batch.keys() - non_tensor_select_keys = ["multi_modal_inputs"] if has_multi_modal_inputs else [] - - data = data.select(batch_keys=select_keys, non_tensor_batch_keys=non_tensor_select_keys) - if has_planned_mini_batches: - mini_batches = split_data_proto_by_mini_batch_id(data) - else: - mini_batches = data.split(self.config.ppo_mini_batch_size) - on_policy = len(mini_batches) == 1 and self.config.ppo_epochs == 1 - - metrics = { - "actor/pg_loss": 0.0, - "actor/kl_loss": 0.0, - } - for _ in range(self.config.ppo_epochs): - for mini_batch in mini_batches: - use_global_mini_batch_info = "mini_batch_global_size" in mini_batch.batch.keys() - if use_global_mini_batch_info: - global_info = get_mini_batch_global_info(mini_batch) - if not hasattr(self.config, "global_batch_info") or self.config.global_batch_info is None: - self.config.global_batch_info = {} - dp_size = ( - torch.distributed.get_world_size() // self.ulysses_sequence_parallel_size - if torch.distributed.is_initialized() - else 1 - ) - self.config.global_batch_info.update( - { - "dp_size": dp_size, - "batch_num_tokens": global_info["batch_num_tokens"], - "global_batch_size": global_info["global_batch_size"], - "loss_scale_factor": self.config.loss_scale_factor, - } - ) - elif hasattr(self.config, "global_batch_info"): - self.config.global_batch_info.clear() - - if self.config.use_dynamic_bsz: - max_token_len = self.config.ppo_max_token_len_per_gpu * self.ulysses_sequence_parallel_size - micro_batches, _ = prepare_dynamic_batch(mini_batch, max_token_len=max_token_len) - else: - micro_batches = mini_batch.split(self.config.ppo_micro_batch_size_per_gpu) - - if not has_planned_mini_batches: - micro_batches = [mb for mb in micro_batches if bool(mb.batch["response_mask"].any().item())] - if not micro_batches: - append_to_dict(metrics, {"actor/grad_norm": 0.0}) - continue - - if not self.config.use_dynamic_bsz: - self.gradient_accumulation = len(micro_batches) - - self.actor_optimizer.zero_grad() - - for micro_batch in micro_batches: - micro_batch = micro_batch.to(get_device_id()) - micro_batch_metrics = {} - model_inputs = {**micro_batch.batch, **micro_batch.non_tensor_batch} - response_mask = model_inputs["response_mask"] - advantages = model_inputs["advantages"] - - entropy_coeff = self.config.entropy_coeff - loss_agg_mode = self.config.loss_agg_mode - calculate_entropy = self.config.calculate_entropy or (entropy_coeff != 0) - - if use_global_mini_batch_info: - loss_scale_factor = 1.0 - elif self.config.use_dynamic_bsz: - loss_scale_factor = response_mask.shape[0] / self.config.ppo_mini_batch_size - else: - loss_scale_factor = 1 / self.gradient_accumulation - - entropy, log_prob = self._forward_micro_batch( - model_inputs, temperature=temperature, calculate_entropy=calculate_entropy - ) - - if hasattr(self.config, "use_rollout_log_probs") and self.config.use_rollout_log_probs: - old_log_prob = model_inputs["old_log_probs"] - else: - old_log_prob = log_prob.detach() if on_policy else model_inputs["old_log_probs"] - - loss_mode = self.config.policy_loss.get("loss_mode", "vanilla") - rollout_is_weights = model_inputs.get("rollout_is_weights", None) - policy_loss_fn = get_policy_loss_fn(loss_mode) - pg_loss, pg_metrics = policy_loss_fn( - old_log_prob=old_log_prob, - log_prob=log_prob, - advantages=advantages, - response_mask=response_mask, - loss_agg_mode=loss_agg_mode, - config=self.config, - rollout_is_weights=rollout_is_weights, - ) - micro_batch_metrics.update(pg_metrics) - - rollout_log_prob = model_inputs.get("rollout_log_probs", None) - if loss_mode != "bypass_mode" and rollout_log_prob is not None: - from verl.trainer.ppo.rollout_corr_helper import compute_rollout_corr_metrics_from_logprobs - - rollout_corr_metrics = compute_rollout_corr_metrics_from_logprobs( - log_prob=log_prob, - rollout_log_prob=rollout_log_prob, - response_mask=response_mask, - ) - micro_batch_metrics.update(rollout_corr_metrics) - - policy_loss = pg_loss - if calculate_entropy and entropy is not None: - entropy_agg = agg_loss( - loss_mat=entropy, - loss_mask=response_mask, - loss_agg_mode=loss_agg_mode, - **getattr(self.config, "global_batch_info", {}), - ) - micro_batch_metrics["actor/entropy"] = entropy_agg.detach().item() - if entropy_coeff != 0: - policy_loss -= entropy_agg * entropy_coeff - - if self.config.use_kl_loss: - ref_log_prob = model_inputs["ref_log_prob"] - kld = kl_penalty( - logprob=log_prob, ref_logprob=ref_log_prob, kl_penalty=self.config.kl_loss_type - ) - kl_loss = agg_loss( - loss_mat=kld, - loss_mask=response_mask, - loss_agg_mode=loss_agg_mode, - **getattr(self.config, "global_batch_info", {}), - ) - policy_loss = policy_loss + kl_loss * self.config.kl_loss_coef - metrics["actor/kl_loss"] += kl_loss.detach().item() * loss_scale_factor - micro_batch_metrics["actor/kl_coef"] = self.config.kl_loss_coef - - loss = policy_loss * loss_scale_factor - if self.scaler is not None: - self.scaler.scale(loss).backward() - else: - loss.backward() - - metrics["actor/pg_loss"] += pg_loss.detach().item() * loss_scale_factor - append_to_dict(metrics, micro_batch_metrics) - - grad_norm = self._optimizer_step() - append_to_dict(metrics, {"actor/grad_norm": grad_norm.detach().item()}) - - self.actor_optimizer.zero_grad() - return metrics +# Removed: verl v0.8.0 no longer has verl.workers.actor.dp_actor. +# Custom loss logic is injected via set_loss_fn() in engine_workers.py. diff --git a/agent_r1/workers/critic/__init__.py b/agent_r1/workers/critic/__init__.py index 4a355a7..767e3e4 100644 --- a/agent_r1/workers/critic/__init__.py +++ b/agent_r1/workers/critic/__init__.py @@ -1,3 +1 @@ -from .dp_critic import DataParallelPPOCritic - -__all__ = ["DataParallelPPOCritic"] +# Removed: verl v0.8.0 no longer has DataParallelPPOCritic. diff --git a/agent_r1/workers/critic/dp_critic.py b/agent_r1/workers/critic/dp_critic.py index 991b758..3b00311 100644 --- a/agent_r1/workers/critic/dp_critic.py +++ b/agent_r1/workers/critic/dp_critic.py @@ -1,140 +1,2 @@ -# Copyright 2025 Agent-R1 Teams -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -""" -Agent-R1 local PPO critic overrides. -""" - -import logging -import os - -import torch - -from agent_r1.trainer.ppo.core_algos import compute_value_loss -from agent_r1.trainer.ppo.trajectory_batching import get_mini_batch_global_info, split_data_proto_by_mini_batch_id -from verl import DataProto -from verl.utils.device import get_device_id -from verl.utils.profiler import GPUMemoryLogger -from verl.utils.py_functional import append_to_dict -from verl.utils.seqlen_balancing import prepare_dynamic_batch -from verl.utils.torch_functional import masked_mean -from verl.workers.critic.dp_critic import DataParallelPPOCritic as VerlDataParallelPPOCritic - -logger = logging.getLogger(__file__) -logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) - - -class DataParallelPPOCritic(VerlDataParallelPPOCritic): - @GPUMemoryLogger(role="dp critic", logger=logger) - def update_critic(self, data: DataProto): - self.critic_module.train() - metrics = { - "critic/vf_loss": 0.0, - } - - select_keys = ["input_ids", "responses", "response_mask", "attention_mask", "position_ids", "values", "returns"] - has_planned_mini_batches = "mini_batch_id" in data.batch.keys() - if has_planned_mini_batches: - select_keys.extend( - [ - "mini_batch_id", - "mini_batch_global_size", - "mini_batch_global_token_num", - "mini_batch_global_response_token_num", - "sample_mask", - ] - ) - - has_multi_modal_inputs = "multi_modal_inputs" in data.non_tensor_batch.keys() - non_tensor_select_keys = ["multi_modal_inputs"] if has_multi_modal_inputs else [] - - data = data.select(batch_keys=select_keys, non_tensor_batch_keys=non_tensor_select_keys) - if has_planned_mini_batches: - mini_batches = split_data_proto_by_mini_batch_id(data) - else: - mini_batches = data.split(self.config.ppo_mini_batch_size) - - for _ in range(self.config.ppo_epochs): - for mini_batch in mini_batches: - use_global_mini_batch_info = "mini_batch_global_size" in mini_batch.batch.keys() - global_batch_info = {} - if use_global_mini_batch_info: - global_info = get_mini_batch_global_info(mini_batch) - dp_size = ( - torch.distributed.get_world_size() // self.ulysses_sequence_parallel_size - if torch.distributed.is_initialized() - else 1 - ) - global_batch_info = { - "dp_size": dp_size, - "batch_num_tokens": global_info["batch_num_tokens"], - "global_batch_size": global_info["global_batch_size"], - } - - if self.config.use_dynamic_bsz: - max_token_len = self.config.ppo_max_token_len_per_gpu * self.ulysses_sequence_parallel_size - micro_batches, _ = prepare_dynamic_batch(mini_batch, max_token_len=max_token_len) - else: - micro_batches = mini_batch.split(self.config.ppo_micro_batch_size_per_gpu) - - if not has_planned_mini_batches: - micro_batches = [mb for mb in micro_batches if bool(mb.batch["response_mask"].any().item())] - if not micro_batches: - append_to_dict(metrics, {"critic/grad_norm": 0.0}) - continue - - if not self.config.use_dynamic_bsz: - self.gradient_accumulation = len(micro_batches) - - self.critic_optimizer.zero_grad() - - for micro_batch in micro_batches: - micro_batch = micro_batch.to(get_device_id()) - model_inputs = {**micro_batch.batch, **micro_batch.non_tensor_batch} - response_mask = model_inputs["response_mask"] - values = model_inputs["values"] - returns = model_inputs["returns"] - - vpreds = self._forward_micro_batch(model_inputs) - vf_loss, vf_clipfrac = compute_value_loss( - vpreds=vpreds, - values=values, - returns=returns, - response_mask=response_mask, - cliprange_value=self.config.cliprange_value, - loss_agg_mode=self.config.loss_agg_mode, - **global_batch_info, - ) - if use_global_mini_batch_info: - loss_scale_factor = 1.0 - elif self.config.use_dynamic_bsz: - loss_scale_factor = response_mask.shape[0] / self.config.ppo_mini_batch_size - else: - loss_scale_factor = 1 / self.gradient_accumulation - loss = vf_loss * loss_scale_factor - loss.backward() - - append_to_dict( - metrics, - { - "critic/vf_clipfrac": vf_clipfrac.detach().item(), - "critic/vpred_mean": masked_mean(vpreds, response_mask).detach().item(), - }, - ) - metrics["critic/vf_loss"] += vf_loss.detach().item() * loss_scale_factor - - grad_norm = self._optimizer_step() - append_to_dict(metrics, {"critic/grad_norm": grad_norm.detach().item()}) - - self.critic_optimizer.zero_grad() - return metrics +# Removed: verl v0.8.0 no longer has verl.workers.critic.dp_critic. +# Custom loss logic is injected via set_loss_fn() in engine_workers.py. diff --git a/agent_r1/workers/engine_workers.py b/agent_r1/workers/engine_workers.py index 91580cd..a31e2d6 100644 --- a/agent_r1/workers/engine_workers.py +++ b/agent_r1/workers/engine_workers.py @@ -132,6 +132,7 @@ def train_mini_batch(self, data: TensorDict) -> TensorDict: self.engine.train_mode(disable_auto_offload=disable_auto_offload), Timer(name="train_batch", logger=None), ): + data = data.cpu() output_lst = [] iteration_idx = 0 for epoch in range(epochs): diff --git a/agent_r1/workers/fsdp_workers.py b/agent_r1/workers/fsdp_workers.py index 3dc3df9..56fe7bb 100644 --- a/agent_r1/workers/fsdp_workers.py +++ b/agent_r1/workers/fsdp_workers.py @@ -1,50 +1,2 @@ -# Copyright 2025 Agent-R1 Teams -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -""" -Thin FSDP worker wrappers that swap in Agent-R1 local actor/critic implementations. -""" - -from verl.single_controller.base.decorator import Dispatch, register -from verl.utils.config import omega_conf_to_dataclass -from verl.workers.fsdp_workers import AsyncActorRolloutRefWorker as VerlAsyncActorRolloutRefWorker -from verl.workers.fsdp_workers import CriticWorker as VerlCriticWorker - - -class AsyncActorRolloutRefWorker(VerlAsyncActorRolloutRefWorker): - @register(dispatch_mode=Dispatch.ONE_TO_ALL) - def init_model(self): - super().init_model() - - from agent_r1.workers.actor import DataParallelPPOActor - - if self._is_actor: - actor_cfg = omega_conf_to_dataclass(self.config.actor) - self.actor = DataParallelPPOActor( - config=actor_cfg, actor_module=self.actor_module_fsdp, actor_optimizer=self.actor_optimizer - ) - - if self._is_ref: - self.ref_policy = DataParallelPPOActor(config=self.config.ref, actor_module=self.ref_module_fsdp) - - -class CriticWorker(VerlCriticWorker): - @register(dispatch_mode=Dispatch.ONE_TO_ALL) - def init_model(self): - super().init_model() - - from agent_r1.workers.critic import DataParallelPPOCritic - - self.critic = DataParallelPPOCritic( - config=self.config, critic_module=self.critic_module, critic_optimizer=self.critic_optimizer - ) +# Removed: verl v0.8.0 no longer has verl.workers.fsdp_workers. +# All worker functionality is now in engine_workers.py. diff --git a/docs/getting-started/installation-guide.md b/docs/getting-started/installation-guide.md index baf6ab6..54bfc92 100644 --- a/docs/getting-started/installation-guide.md +++ b/docs/getting-started/installation-guide.md @@ -4,7 +4,7 @@ Agent-R1 uses the same environment setup as `verl`. ## Base Environment -Follow the official [`verl` installation guide](https://verl.readthedocs.io/en/latest/start/install.html), but make sure the environment ends up with `verl==0.7.0`. +Follow the official [`verl` installation guide](https://verl.readthedocs.io/en/latest/start/install.html), but make sure the environment ends up with `verl==0.8.0`. If you want a broader overview of the base training workflow, the [`verl` quickstart](https://verl.readthedocs.io/en/latest/start/quickstart.html) is also useful. @@ -12,7 +12,7 @@ If you want a broader overview of the base training workflow, the [`verl` quicks Once the `verl` environment is working, Agent-R1 should run in the same environment. In practice, that means you can: -- prepare a Python environment with `verl==0.7.0` +- prepare a Python environment with `verl==0.8.0` - clone this repository - run Agent-R1 commands directly from the repository root diff --git a/docs/zh/getting-started/installation-guide.md b/docs/zh/getting-started/installation-guide.md index eb1705a..6ca35c2 100644 --- a/docs/zh/getting-started/installation-guide.md +++ b/docs/zh/getting-started/installation-guide.md @@ -4,7 +4,7 @@ Agent-R1 使用与 `verl` 相同的环境设置。 ## 基础环境 -请参考官方 [`verl` 安装指南](https://verl.readthedocs.io/en/latest/start/install.html)，并确保最终环境中使用 `verl==0.7.0`。 +请参考官方 [`verl` 安装指南](https://verl.readthedocs.io/en/latest/start/install.html)，并确保最终环境中使用 `verl==0.8.0`。 如果你想先了解基础训练流程，也可以参考 [`verl` quickstart](https://verl.readthedocs.io/en/latest/start/quickstart.html)。 @@ -12,7 +12,7 @@ Agent-R1 使用与 `verl` 相同的环境设置。 当 `verl` 环境可以正常工作后，Agent-R1 应该可以在同一个环境中运行。实践中，你需要： -- 准备一个包含 `verl==0.7.0` 的 Python 环境 +- 准备一个包含 `verl==0.8.0` 的 Python 环境 - 克隆本仓库 - 直接在仓库根目录运行 Agent-R1 命令 diff --git a/examples/gsm8k/run_grpo.sh b/examples/gsm8k/run_grpo.sh new file mode 100755 index 0000000..d26f9ae --- /dev/null +++ b/examples/gsm8k/run_grpo.sh @@ -0,0 +1,72 @@ +#!/usr/bin/env bash +set -euo pipefail +set -x + +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" +export VLLM_USE_V1="${VLLM_USE_V1:-1}" +export CUDA_HOME="${CUDA_HOME:-/usr/local/cuda}" +export HF_ENDPOINT=https://hf-mirror.com + +PROJECT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +CONFIG_PATH="$PROJECT_DIR/recipes/gsm8k/base.yaml" + +DATA_DIR="${DATA_DIR:-$HOME/data/gsm8k_tool}" +MODEL_PATH="${MODEL_PATH:-Qwen/Qwen3-4B-Instruct-2507}" +EXP_NAME="${EXP_NAME:-gsm8k_grpo_tool}" +N_GPUS_PER_NODE="${N_GPUS_PER_NODE:-4}" + +python3 -m agent_r1.trainer.main_agent_ppo \ + algorithm.adv_estimator=grpo \ + data.train_files="$DATA_DIR/train.parquet" \ + data.val_files="$DATA_DIR/test.parquet" \ + data.train_batch_size=128 \ + data.max_prompt_length=2048 \ + data.max_response_length=2048 \ + data.filter_overlong_prompts=True \ + data.truncation='error' \ + data.return_raw_chat=True \ + actor_rollout_ref.model.path="$MODEL_PATH" \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.actor.ppo_mini_batch_size=16 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \ + actor_rollout_ref.actor.use_kl_loss=True \ + actor_rollout_ref.actor.kl_loss_coef=0.001 \ + actor_rollout_ref.actor.kl_loss_type=low_var_kl \ + actor_rollout_ref.actor.entropy_coeff=0 \ + actor_rollout_ref.actor.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.tensor_model_parallel_size=2 \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2 \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \ + actor_rollout_ref.rollout.n=5 \ + actor_rollout_ref.rollout.prompt_length=2048 \ + actor_rollout_ref.rollout.response_length=2048 \ + actor_rollout_ref.rollout.max_model_len=4096 \ + actor_rollout_ref.rollout.do_sample=True \ + actor_rollout_ref.rollout.temperature=0.6 \ + actor_rollout_ref.rollout.top_p=0.95 \ + actor_rollout_ref.rollout.top_k=20 \ + actor_rollout_ref.rollout.agent.agent_flow_config_path="$CONFIG_PATH" \ + actor_rollout_ref.rollout.agent.default_agent_flow=gsm8k_tool \ + actor_rollout_ref.rollout.agent.max_steps=5 \ + actor_rollout_ref.rollout.agent.skip_special_tokens=False \ + actor_rollout_ref.rollout.val_kwargs.do_sample=False \ + actor_rollout_ref.rollout.val_kwargs.n=1 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=2 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + algorithm.use_kl_in_reward=False \ + custom_reward_function.path=recipes/gsm8k/reward_fn.py \ + custom_reward_function.name=compute_score \ + trainer.logger='["console"]' \ + trainer.project_name='agent_r1_gsm8k_tool' \ + trainer.experiment_name="$EXP_NAME" \ + trainer.n_gpus_per_node="$N_GPUS_PER_NODE" \ + trainer.nnodes=1 \ + trainer.save_freq=-1 \ + trainer.test_freq=5 \ + trainer.total_epochs=15 \ + trainer.val_before_train=True \ + "$@" \ No newline at end of file diff --git a/recipes/alfworld/alfworld_agent_flow.py b/recipes/alfworld/alfworld_agent_flow.py index 30bc030..5c1ee94 100644 --- a/recipes/alfworld/alfworld_agent_flow.py +++ b/recipes/alfworld/alfworld_agent_flow.py @@ -18,7 +18,8 @@ build_invalid_tool_call_observation, extract_task_text, ) -from verl.experimental.agent_loop.agent_loop import AsyncLLMServerManager, DictConfigWrap +from verl.experimental.agent_loop.agent_loop import DictConfigWrap +from verl.workers.rollout.llm_server import LLMServerClient from verl.experimental.agent_loop.tool_parser import FunctionCall, ToolParser from verl.utils.profiler import simple_timer @@ -62,7 +63,7 @@ class AlfworldAgentFlow(AgentFlowBase): def __init__( self, trainer_config: DictConfigWrap, - server_manager: AsyncLLMServerManager, + server_manager: LLMServerClient, reward_loop_worker: RewardLoopWorker, tokenizer: AutoTokenizer, processor: AutoProcessor, diff --git a/recipes/hotpotqa/hotpotqa_agent_flow.py b/recipes/hotpotqa/hotpotqa_agent_flow.py index 7529893..f594364 100755 --- a/recipes/hotpotqa/hotpotqa_agent_flow.py +++ b/recipes/hotpotqa/hotpotqa_agent_flow.py @@ -35,7 +35,8 @@ HOTPOTQA_TOOL_SCHEMAS, HOTPOTQA_USER_PROMPT, ) -from verl.experimental.agent_loop.agent_loop import AsyncLLMServerManager, DictConfigWrap +from verl.experimental.agent_loop.agent_loop import DictConfigWrap +from verl.workers.rollout.llm_server import LLMServerClient from verl.experimental.agent_loop.tool_parser import FunctionCall, ToolParser from verl.utils.profiler import simple_timer @@ -158,7 +159,7 @@ class HotpotQAAgentFlow(AgentFlowBase): def __init__( self, trainer_config: DictConfigWrap, - server_manager: AsyncLLMServerManager, + server_manager: LLMServerClient, reward_loop_worker: RewardLoopWorker, tokenizer: AutoTokenizer, processor: AutoProcessor, diff --git a/recipes/paper_search/paper_search_agent_flow.py b/recipes/paper_search/paper_search_agent_flow.py index 5d25868..7d521dd 100755 --- a/recipes/paper_search/paper_search_agent_flow.py +++ b/recipes/paper_search/paper_search_agent_flow.py @@ -11,7 +11,8 @@ from recipes.paper_search.prompts import PAPERSEARCH_TOOL_SCHEMAS from recipes.paper_search.runtime import PaperSearchRuntime, PaperSearchRuntimeConfig from recipes.paper_search.utils import recover_tool_calls_from_text -from verl.experimental.agent_loop.agent_loop import AsyncLLMServerManager, DictConfigWrap +from verl.experimental.agent_loop.agent_loop import DictConfigWrap +from verl.workers.rollout.llm_server import LLMServerClient from verl.experimental.agent_loop.tool_parser import ToolParser from verl.utils.profiler import simple_timer @@ -24,7 +25,7 @@ class PaperSearchAgentFlow(AgentFlowBase): def __init__( self, trainer_config: DictConfigWrap, - server_manager: AsyncLLMServerManager, + server_manager: LLMServerClient, reward_loop_worker: RewardLoopWorker, tokenizer: AutoTokenizer, processor: AutoProcessor, diff --git a/recipes/webshop/webshop_agent_flow.py b/recipes/webshop/webshop_agent_flow.py index d3b369d..be759e4 100644 --- a/recipes/webshop/webshop_agent_flow.py +++ b/recipes/webshop/webshop_agent_flow.py @@ -14,7 +14,8 @@ from recipes.webshop.env.client import WebShopEnvClient from recipes.webshop.prompts import WEBSHOP_TOOL_SCHEMAS from recipes.webshop.utils import build_invalid_tool_call_observation, build_webshop_messages -from verl.experimental.agent_loop.agent_loop import AsyncLLMServerManager, DictConfigWrap +from verl.experimental.agent_loop.agent_loop import DictConfigWrap +from verl.workers.rollout.llm_server import LLMServerClient from verl.experimental.agent_loop.tool_parser import FunctionCall, ToolParser from verl.utils.profiler import simple_timer @@ -57,7 +58,7 @@ class WebShopAgentFlow(AgentFlowBase): def __init__( self, trainer_config: DictConfigWrap, - server_manager: AsyncLLMServerManager, + server_manager: LLMServerClient, reward_loop_worker: RewardLoopWorker, tokenizer: AutoTokenizer, processor: AutoProcessor, From 525667adad796544de6d16d230c067ea3a629c72 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 6 Jul 2026 09:59:26 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- agent_r1/agent_flow/agent_flow.py | 9 ++++----- agent_r1/reward_loop/reward_loop.py | 5 ++++- agent_r1/trainer/main_agent_ppo.py | 6 ++---- agent_r1/trainer/ppo/ray_trainer.py | 1 - recipes/alfworld/alfworld_agent_flow.py | 2 +- recipes/hotpotqa/hotpotqa_agent_flow.py | 2 +- recipes/paper_search/paper_search_agent_flow.py | 2 +- recipes/webshop/webshop_agent_flow.py | 2 +- 8 files changed, 14 insertions(+), 15 deletions(-) diff --git a/agent_r1/agent_flow/agent_flow.py b/agent_r1/agent_flow/agent_flow.py index 48a4c41..9d3911d 100644 --- a/agent_r1/agent_flow/agent_flow.py +++ b/agent_r1/agent_flow/agent_flow.py @@ -30,8 +30,6 @@ from agent_r1.reward_loop.reward_loop import RewardLoopWorker from verl.experimental.agent_loop.agent_loop import DictConfigWrap -from verl.workers.rollout.llm_server import GlobalRequestLoadBalancer, LLMServerClient -from verl.workers.rollout.utils import update_prometheus_config from verl.experimental.agent_loop.utils import resolve_config_path from verl.protocol import DataProto from verl.single_controller.ray.base import RayResourcePool, RayWorkerGroup @@ -46,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")) @@ -501,6 +501,7 @@ def __init__( _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: @@ -780,9 +781,7 @@ def create_transferqueue_client( class AgentFlowWorker(AgentFlowWorkerBase): """Agent flow worker takes a batch of messages and run each message in an agent flow.""" - def __init__( - self, config: DictConfig, llm_client: LLMServerClient, 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. diff --git a/agent_r1/reward_loop/reward_loop.py b/agent_r1/reward_loop/reward_loop.py index cca9829..429bbdb 100644 --- a/agent_r1/reward_loop/reward_loop.py +++ b/agent_r1/reward_loop/reward_loop.py @@ -125,7 +125,10 @@ async def compute_score(self, input_data) -> dict: 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.reward.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." diff --git a/agent_r1/trainer/main_agent_ppo.py b/agent_r1/trainer/main_agent_ppo.py index 3bce8d4..5eefeae 100644 --- a/agent_r1/trainer/main_agent_ppo.py +++ b/agent_r1/trainer/main_agent_ppo.py @@ -117,11 +117,10 @@ def __init__(self): def add_actor_rollout_worker(self, config): """Add actor rollout worker using the unified model engine implementation.""" + from agent_r1.workers.engine_workers import ActorRolloutRefWorker from verl.single_controller.ray import RayWorkerGroup from verl.trainer.ppo.ray_trainer import Role - from agent_r1.workers.engine_workers import ActorRolloutRefWorker - actor_rollout_cls = ActorRolloutRefWorker ray_worker_group_cls = RayWorkerGroup @@ -139,9 +138,8 @@ def add_actor_rollout_worker(self, config): def add_critic_worker(self, config): """Add critic worker to role mapping using the unified model engine implementation.""" - from verl.trainer.ppo.ray_trainer import Role - from agent_r1.workers.engine_workers import TrainingWorker + from verl.trainer.ppo.ray_trainer import Role self.role_worker_mapping[Role.Critic] = ray.remote(TrainingWorker) self.mapping[Role.Critic] = "global_pool" diff --git a/agent_r1/trainer/ppo/ray_trainer.py b/agent_r1/trainer/ppo/ray_trainer.py index dc7ca55..d5bc28c 100644 --- a/agent_r1/trainer/ppo/ray_trainer.py +++ b/agent_r1/trainer/ppo/ray_trainer.py @@ -27,7 +27,6 @@ from typing import Optional import numpy as np -import ray import torch from omegaconf import OmegaConf from tqdm import tqdm diff --git a/recipes/alfworld/alfworld_agent_flow.py b/recipes/alfworld/alfworld_agent_flow.py index 5c1ee94..b489ca9 100644 --- a/recipes/alfworld/alfworld_agent_flow.py +++ b/recipes/alfworld/alfworld_agent_flow.py @@ -19,9 +19,9 @@ extract_task_text, ) from verl.experimental.agent_loop.agent_loop import DictConfigWrap -from verl.workers.rollout.llm_server import LLMServerClient from verl.experimental.agent_loop.tool_parser import FunctionCall, ToolParser from verl.utils.profiler import simple_timer +from verl.workers.rollout.llm_server import LLMServerClient logger = logging.getLogger(__file__) logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) diff --git a/recipes/hotpotqa/hotpotqa_agent_flow.py b/recipes/hotpotqa/hotpotqa_agent_flow.py index f594364..4bf2d6a 100755 --- a/recipes/hotpotqa/hotpotqa_agent_flow.py +++ b/recipes/hotpotqa/hotpotqa_agent_flow.py @@ -36,9 +36,9 @@ HOTPOTQA_USER_PROMPT, ) from verl.experimental.agent_loop.agent_loop import DictConfigWrap -from verl.workers.rollout.llm_server import LLMServerClient from verl.experimental.agent_loop.tool_parser import FunctionCall, ToolParser from verl.utils.profiler import simple_timer +from verl.workers.rollout.llm_server import LLMServerClient logger = logging.getLogger(__file__) logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) diff --git a/recipes/paper_search/paper_search_agent_flow.py b/recipes/paper_search/paper_search_agent_flow.py index 7d521dd..e6b2778 100755 --- a/recipes/paper_search/paper_search_agent_flow.py +++ b/recipes/paper_search/paper_search_agent_flow.py @@ -12,9 +12,9 @@ from recipes.paper_search.runtime import PaperSearchRuntime, PaperSearchRuntimeConfig from recipes.paper_search.utils import recover_tool_calls_from_text from verl.experimental.agent_loop.agent_loop import DictConfigWrap -from verl.workers.rollout.llm_server import LLMServerClient from verl.experimental.agent_loop.tool_parser import ToolParser from verl.utils.profiler import simple_timer +from verl.workers.rollout.llm_server import LLMServerClient logger = logging.getLogger(__file__) logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) diff --git a/recipes/webshop/webshop_agent_flow.py b/recipes/webshop/webshop_agent_flow.py index be759e4..dcb29c6 100644 --- a/recipes/webshop/webshop_agent_flow.py +++ b/recipes/webshop/webshop_agent_flow.py @@ -15,9 +15,9 @@ from recipes.webshop.prompts import WEBSHOP_TOOL_SCHEMAS from recipes.webshop.utils import build_invalid_tool_call_observation, build_webshop_messages from verl.experimental.agent_loop.agent_loop import DictConfigWrap -from verl.workers.rollout.llm_server import LLMServerClient from verl.experimental.agent_loop.tool_parser import FunctionCall, ToolParser from verl.utils.profiler import simple_timer +from verl.workers.rollout.llm_server import LLMServerClient logger = logging.getLogger(__file__) logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))