VerlTool: ツール使用を伴うホリスティックなエージェンティック強化学習に向けて

ツールを用いた複雑なタスク解決のためのエージェンティックシステムのトレーニング

成果: 単一の統一されたフレームワークを用いて、イテレーティブに推論し、外部ツールを呼び出し、結果を観察し、多様な問題領域全体で適応するエージェントを構築します。

問題コンテキスト

エージェンティックAIへの既存のアプローチは、ツール使用をドメイン固有のシステムに分断しています。知識QAシステムは検索を異なる方法で処理しますが、コード実行システムとは異なります。SQLエージェントはビジュアル推論システムとは異なるAPIを使用します。研究者が複数ドメインにまたがるタスクを解決するエージェントを求めるとき、または実践者が新しいツール機能を追加する必要があるとき、統合のコストは高くなります。カスタムパイプライン、再実装された調整ロジック、そして繰り返されるインフラストラクチャ投資です。

単一ターンの言語モデルは、ツール使用に必要な連続的な意思決定も自然に処理できません。アクション、そのアクションからの結果の観察、そして新しい情報に基づく次のアクション選択です。強化学習（RL）はこれらの複数ターンの軌跡を最適化できますが、RLフレームワークを多様なツールに対応させるには、軌跡表現、観察トークン化、およびモダリティ全体にわたる報酬アライメントの慎重な検討が必要です。

VerlToolは、検証可能な報酬を伴う強化学習（RLVR）を複数ターンのエージェンティック設定とツール使用に拡張する統一されたモジュール式フレームワークを導入することで、この問題を解決します。

コア概念

VerlToolはRL内のツール使用を、アクションと観察を交互に行う軌跡として扱います。エージェントはアクション（ツール呼び出しと引数）を選択し、ツールが実行して観察を返し、このサイクルはタスクが解決されるまで繰り返されます。重要な洞察は、ツール使用の軌跡内の観察はエージェントの制御外の環境生成事実であり、トレーニング中の政策勾配に影響を与えてはならず、後続の決定のための情報のみを提供するべきであるということです。

フレームワークが「ホリスティック」である理由は、以下を統一するためです：

• 複数のツールタイプ（コード、検索、SQL、ビジョン）を1つのAPI下で統一
• 複数のモダリティ（テキスト、画像、ビデオ）を観察トークンで統一
• 効率性のための非同期実行を伴う複数ターンのRLトレーニング
• VeRLとの上流互換性によるシームレスなメンテナンス

コード実行、検索、SQLのための別々のエージェントを構築する代わりに、チームはこれらのツールの任意の組み合わせを効果的に構成することを学習する単一のエージェントを開発します。

アーキテクチャの概要

VerlToolは、標準化されたAPIを介して通信する2つの主要なサブシステムに分割されます。

• VeRL Workflow（トレーニング側）: 政策トレーニング、報酬計算、モデル更新を調整します。VeRLをサブモジュールとして継承し、上流の改善との連携を保ち、コアRL論理の重複を避けます。

• Tool Server（実行側）: ツール呼び出しの実行を管理します。同期バッチアライメントを強制する代わりに、軌跡ごと単位での非同期ロールアウトを処理し、アイドル待機を排除し、ロールアウト段階で約2倍の高速化を達成します。

モジュール式ツール登録: 各ツールは、アクション解析、環境状態管理、および操作実行のためのメソッドを持つ共通のBaseTool インターフェースを実装します。新しいツールは、トレーニングコードを変更することなく、軽量なPython定義ファイルを作成することで追加されます。このデザインは関心事を分離し、ドメイン専門家がRL知識がなくともツールに貢献できるようにします。

トークン化戦略: 安定性のための重要な詳細：アクション文字列と観察文字列は別々にトークン化され、その後連結されます。これにより、複数ターンロールアウト中のトレーニング不安定性を引き起こす可能性のある境界関連トークンの不一致を防ぎます。フレームワークはすべてのターン全体で一貫したトークンアライメントを保証します。

政策最適化における観察マスキング: 観察はトレーニング中のモデルに対して環境生成かつ非ポリシーであるため、フレームワークは勾配計算中に観察トークンをマスクします。アクショントークンと最終報酬のみが政策更新に寄与します。これにより、モデルが古い観察データとの見かけ上の相関を学習することを防ぎ、RL安定性を保ちます。

実装

1. VerlToolモジュール式アーキテクチャの設定

VeRLを継承し、ツールサーバーを別個のコンポーネントとして構造化することから始めます。これにより、トレーニング論理がツール実行インフラストラクチャから分離されたままになります。

# verltool/config.py
# Configuration for dual-component architecture

from typing import Dict, List, Optional
from dataclasses import dataclass, field
from verl.config import BaseConfig

@dataclass
class ToolServerConfig(BaseConfig):
    """Configuration for asynchronous tool execution server."""
    host: str = "localhost"
    port: int = 8888
    max_workers: int = 32  # Async worker threads for tool calls
    timeout_seconds: int = 300
    enable_async_rollout: bool = True  # Critical for 2× speedup

@dataclass
class VerlToolConfig(BaseConfig):
    """Top-level VerlTool configuration."""
    verl_config: Dict = field(default_factory=dict)  # Inherited VeRL settings
    tool_server: ToolServerConfig = field(default_factory=ToolServerConfig)
    modalities: List[str] = field(default_factory=lambda: ["text"])  # text, image, video
    observation_masking_enabled: bool = True  # Mask off-policy observations
    tokenization_mode: str = "separate"  # Separate action/observation tokens

2. 統一されたツールインターフェースの実装

すべてのツールが継承する基本的な抽象化を定義します。これにより、動的な登録と多様なツールの一貫した処理が可能になります。

# verltool/tools/base.py
# Unified tool interface for consistent behavior across domains

from abc import ABC, abstractmethod
from typing import Any, Dict, Optional, Union
from dataclasses import dataclass

@dataclass
class ToolAction:
    """Structured representation of a tool invocation."""
    tool_name: str
    arguments: Dict[str, Any]
    timestamp: Optional[float] = None

@dataclass
class ToolObservation:
    """Result returned by tool execution."""
    content: Union[str, bytes]
    modality: str = "text"  # text, image, video
    token_count: Optional[int] = None
    error: Optional[str] = None

class BaseTool(ABC):
    """Base class for all tools in VerlTool ecosystem."""

    def __init__(self, name: str, modalities: List[str] = None):
        self.name = name
        self.modalities = modalities or ["text"]
        self.state = {}  # Maintains context across turns

    @abstractmethod
    def parse_action(self, action_string: str) -> ToolAction:
        """Convert raw action string to structured ToolAction."""
        pass

    @abstractmethod
    def execute(self, action: ToolAction) -> ToolObservation:
        """Execute the tool and return observation."""
        pass

    def update_state(self, observation: ToolObservation):
        """Update internal state based on execution result."""
        self.state['last_result'] = observation

    def reset(self):
        """Clear state for new trajectory."""
        self.state = {}

3. ツールの動的登録

トレーニング論理にハードコードすることなく、ツールを発見および管理するレジストリを作成します。

# verltool/tools/registry.py
# Dynamic tool registration system for extensibility

from typing import Dict, Type, Optional
from verltool.tools.base import BaseTool
import importlib
import os

class ToolRegistry:
    """Registry for tool discovery and management."""

    def __init__(self, tool_dir: str = "verltool/tools"):
        self.tools: Dict[str, BaseTool] = {}
        self.tool_dir = tool_dir

    def register(self, tool_class: Type[BaseTool]):
        """Explicitly register a tool class."""
        instance = tool_class()
        self.tools[instance.name] = instance
        return self

    def discover_from_directory(self):
        """Auto-discover tools from tool_dir by importing modules."""
        for filename in os.listdir(self.tool_dir):
            if filename.endswith('_tool.py') and not filename.startswith('_'):
                module_name = filename[:-3]
                try:
                    module = importlib.import_module(f'verltool.tools.{module_name}')
                    # Assume each module exports TOOL_CLASS
                    if hasattr(module, 'TOOL_CLASS'):
                        self.register(module.TOOL_CLASS)
                except Exception as e:
                    print(f"Warning: Could not load tool {module_name}: {e}")
        return self

    def get_tool(self, name: str) -> Optional[BaseTool]:
        """Retrieve a tool by name."""
        return self.tools.get(name)

    def list_tools(self) -> Dict[str, BaseTool]:
        """Return all registered tools."""
        return self.tools.copy()

4. コード実行ツールの実装

推論およびソフトウェアエンジニアリングタスクの一般的な要件であるPythonコード実行のための具体的なツール実装を示します。

# verltool/tools/code_execution_tool.py
# Code execution tool with sandboxed environment support

from verltool.tools.base import BaseTool, ToolAction, ToolObservation
from typing import Dict, Any
import subprocess
import tempfile
import os

class CodeExecutionTool(BaseTool):
    """Execute Python code and capture output."""

    def __init__(self, timeout_seconds: int = 30):
        super().__init__(
            name="code_execution",
            modalities=["text"]
        )
        self.timeout = timeout_seconds
        self.execution_history = []

    def parse_action(self, action_string: str) -> ToolAction:
        """Extract Python code from action string."""
        # Expected format: <code>python_code_here</code>
        start = action_string.find('<code>')
        end = action_string.find('</code>')

        if start == -1 or end == -1:
            return ToolAction(
                tool_name=self.name,
                arguments={"code": action_string}
            )

        code = action_string[start + 6:end].strip()
        return ToolAction(
            tool_name=self.name,
            arguments={"code": code}
        )

    def execute(self, action: ToolAction) -> ToolObservation:
        """Run code in subprocess and capture output."""
        code = action.arguments.get("code", "")

        if not code:
            return ToolObservation(
                content="Error: No code provided",
                modality="text",
                error="empty_code"
            )

        try:
            with tempfile.NamedTemporaryFile(
                mode='w',
                suffix='.py',
                delete=False
            ) as f:
                f.write(code)
                temp_file = f.name

            result = subprocess.run(
                ['python', temp_file],
                capture_output=True,
                text=True,
                timeout=self.timeout
            )

            output = result.stdout
            if result.stderr:
                output += "\nStderr:\n" + result.stderr

            self.execution_history.append({
                "code": code,
                "output": output,
                "returncode": result.returncode
            })

            return ToolObservation(
                content=output or "(No output)",
                modality="text",
                error=None if result.returncode == 0 else "execution_error"
            )

        except subprocess.TimeoutExpired:
            return ToolObservation(
                content=f"Error: Code execution timeout (>{self.timeout}s)",
                modality="text",
                error="timeout"
            )
        except Exception as e:
            return ToolObservation(
                content=f"Error: {str(e)}",
                modality="text",
                error="execution_failed"
            )
        finally:
            if os.path.exists(temp_file):
                os.remove(temp_file)

TOOL_CLASS = CodeExecutionTool

5. 観察マスキング付きの複数ターン軌跡処理の構築

アクションと観察を交互に行う軌跡構築を実装し、非ポリシー観察の適切なマスキングを行います。

# verltool/training/trajectory.py
# Multi-turn trajectory representation with observation masking

from typing import List, Optional, Dict, Any
from dataclasses import dataclass, field
import torch

@dataclass
class MultiTurnTrajectory:
    """Represents a complete multi-turn agent trajectory."""

    actions: List[str] = field(default_factory=list)  # Agent actions/tool calls
    observations: List[str] = field(default_factory=list)  # Environment observations
    action_tokens: List[torch.Tensor] = field(default_factory=list)
    observation_tokens: List[torch.Tensor] = field(default_factory=list)
    observation_masks: List[bool] = field(default_factory=list)  # True = mask out
    reward: float = 0.0
    episode_done: bool = False

    def add_turn(
        self,
        action: str,
        observation: str,
        action_token_ids: torch.Tensor,
        observation_token_ids: torch.Tensor,
        mask_observation: bool = True
    ):
        """Add a single action-observation pair to trajectory."""
        self.actions.append(action)
        self.observations.append(observation)
        self.action_tokens.append(action_token_ids)
        self.observation_tokens.append(observation_token_ids)
        # Observation tokens are off-policy and should be masked
        self.observation_masks.append(mask_observation)

    def get_concatenated_tokens(self) -> torch.Tensor:
        """Interleave action and observation tokens."""
        sequence = []
        for i in range(len(self.actions)):
            sequence.append(self.action_tokens[i])
            if i < len(self.observation_tokens):
                sequence.append(self.observation_tokens[i])
        return torch.cat(sequence, dim=0)

    def get_loss_mask(self) -> torch.Tensor:
        """Create mask: True where loss should be computed, False otherwise."""
        mask = []
        for i in range(len(self.actions)):
            # Action tokens contribute to loss
            mask.append(torch.ones_like(self.action_tokens[i], dtype=torch.bool))
            # Observation tokens are masked (do not contribute to loss)
            if i < len(self.observation_masks):
                is_masked = self.observation_masks[i]
                obs_mask = torch.zeros_like(
                    self.observation_tokens[i],
                    dtype=torch.bool
                ) if is_masked else torch.ones_like(
                    self.observation_tokens[i],
                    dtype=torch.bool
                )
                mask.append(obs_mask)
        return torch.cat(mask, dim=0)

class TrajectoryBuilder:
    """Construct trajectories from agent rollouts."""

    def __init__(self, tokenizer):
        self.tokenizer = tokenizer

    def build_from_rollout(
        self,
        actions: List[str],
        observations: List[str],
        reward: float,
        mask_all_observations: bool = True
    ) -> MultiTurnTrajectory:
        """Convert rollout data into a trajectory with proper tokenization."""
        trajectory = MultiTurnTrajectory(reward=reward)

        for i, action in enumerate(actions):
            # Tokenize separately to avoid boundary issues
            action_tokens = self.tokenizer.encode(action)
            action_tensor = torch.tensor(action_tokens)

            if i < len(observations):
                obs = observations[i]
                obs_tokens = self.tokenizer.encode(obs)
                obs_tensor = torch.tensor(obs_tokens)

                trajectory.add_turn(
                    action=action,
                    observation=obs,
                    action_token_ids=action_tensor,
                    observation_token_ids=obs_tensor,
                    mask_observation=mask_all_observations
                )

        return trajectory

6. 2倍の高速化のための非同期ツールサーバーの実装

ロールアウトがツール完了の待機でブロックされないよう、ツール実行を非同期にデプロイします。

# verltool/server/async_tool_server.py
# Asynchronous tool execution with trajectory-level batching

import asyncio
from typing import List, Dict, Any, Callable
from concurrent.futures import ThreadPoolExecutor
import queue
from verltool.tools.registry import ToolRegistry
from verltool.tools.base import ToolAction, ToolObservation

class AsyncToolServer:
    """Non-blocking tool execution server for efficient rollouts."""

    def __init__(self, tool_registry: ToolRegistry, max_workers: int = 32):
        self.registry = tool_registry
        self.executor = ThreadPoolExecutor(max_workers=max_workers)
        self.pending_tasks: Dict[str, asyncio.Future] = {}
        self.task_counter = 0

    def execute_tool_async(
        self,
        tool_name: str,
        action_string: str
    ) -> str:
        """Submit tool execution without blocking, return task ID."""
        tool = self.registry.get_tool(tool_name)
        if not tool:
            return None

        task_id = f"task_{self.task_counter}"
        self.task_counter += 1

        def _run():
            action = tool.parse_action(action_string)
            observation = tool.execute(action)
            tool.update_state(observation)
            return observation

        future = asyncio.get_event_loop().run_in_executor(
            self.executor,
            _run
        )
        self.pending_tasks[task_id] = future
        return task_id

    async def wait_for_result(self, task_id: str, timeout: int = 300) -> ToolObservation:
        """Await tool completion and retrieve result."""
        if task_id not in self.pending_tasks:
            raise ValueError(f"Unknown task: {task_id}")

        try:
            result = await asyncio.wait_for(
                self.pending_tasks[task_id],
                timeout=timeout
            )
            del self.pending_tasks[task_id]
            return result
        except asyncio.TimeoutError:
            return ToolObservation(
                content=f"Timeout waiting for {task_id}",
                modality="text",
                error="timeout"
            )

    async def process_trajectory_batch(
        self,
        trajectories: List[Dict[str, Any]]
    ) -> List[List[ToolObservation]]:
        """Execute all tool calls in a batch of trajectories concurrently."""
        all_observations = []

        for traj in trajectories:
            tool_calls = traj.get("tool_calls", [])
            tasks = [
                self.execute_tool_async(call["tool"], call["action"])
                for call in tool_calls
            ]

            observations = []
            for task_id in tasks:
                result = await self.wait_for_result(task_id)
                observations.append(result)

            all_observations.append(observations)

        return all_observations

7. 報酬とロス計算の設定

ツール使用軌跡の報酬の計算方法と、ロスが観察マスキングをどのように尊重するかを定義します。

# verltool/training/reward.py
# Reward computation for multi-turn agentic trajectories

from typing import Optional, Dict, Any
import torch
import torch.nn as nn

class VerlToolRewardComputer:
    """Compute rewards from verifiable task outcomes."""

    def __init__(self, reward_fn: callable, use_per_step_rewards: bool = False):
        """
        reward_fn: Function taking (final_output, expected_output) -> float
        use_per_step_rewards: If True, grant intermediate rewards for progress
        """
        self.reward_fn = reward_fn
        self.use_per_step_rewards = use_per_step_rewards

    def compute(
        self,
        final_output: str,
        expected_output: Optional[str] = None,
        intermediate_outputs: Optional[Dict[int, str]] = None
    ) -> float:
        """
        Compute scalar reward for trajectory.
        Follows RLVR principle: reward only depends on verifiable task outcome.
        """
        if expected_output is None:
            # Fallback: treat any valid output as success
            return 1.0 if final_output else 0.0

        base_reward = self.reward_fn(final_output, expected_output)

        if self.use_per_step_rewards and intermediate_outputs:
            # Optional: add small bonuses for intermediate progress
            bonus = 0.0
            for step, output in intermediate_outputs.items():
                if output and step < len(intermediate_outputs) - 1:
                    bonus += 0.05  # Small step reward
            return min(base_reward + bonus, 1.0)

        return base_reward

class PolicyLoss(nn.Module):
    """Compute policy gradient loss with observation masking."""

    def __init__(self, model: nn.Module):
        super().__init__()
        self.model = model

    def forward(
        self,
        trajectory,  # MultiTurnTrajectory instance
        logits: torch.Tensor
    ) -> torch.Tensor:
        """
        Compute policy loss, masking observation tokens.
        logits: Model output logits for entire sequence
        """
        loss_mask = trajectory.get_loss_mask()  # Bool tensor

        # Shift logits for language modeling objective
        shift_logits = logits[:-1]
        target_tokens = trajectory.get_concatenated_tokens()[1:]

        # Cross-entropy loss only on unmasked (action) tokens
        ce_loss = torch.nn.functional.cross_entropy(
            shift_logits.view(-1, logits.shape[-1]),
            target_tokens.view(-1),
            reduction='none'
        )

        # Apply mask to loss
        ce_loss = ce_loss * loss_mask[1:].float()

        # Weight by trajectory reward
        trajectory_reward = max(trajectory.reward, 0.0)
        weighted_loss = ce_loss.mean() * trajectory_reward

        return weighted_loss

実践的ガイダンス

ハイパーパラメータ設定:

パラメータ	ドメイン	推奨値	根拠
max_turns	数学/SQL	8–12	過度なシーケンスなしで十分なツール呼び出しを可能にします
max_turns	コード実行	5–8	より短い：限定的なデバッグイテレーションが必要
max_turns	ウェブ検索	4–6	検索クエリは迅速な結果を返します
observation_masking	すべて	True	非ポリシー観察の勾配へのリークを防ぎます
async_workers	すべて	32–64	並行処理のバランスを取ります。32以上で2倍の高速化
tokenization_mode	すべて	"separate"	ターン全体のトークン境界の不一致を回避します
timeout_seconds	コード/検索	30–60	暴走するツール呼び出しを防ぎます。ドメイン別に調整してください
learning_rate	すべて	1e-5 から 5e-5	ツール使用RLは慎重なチューニングが必要

VerlToolを使用するべき場合:

• コード、検索、SQL、ビジョンツールが共存する必要がある複数ドメインエージェント
• 複数のチームが拡張する共有RL インフラストラクチャを持つ組織
• 結果が検証可能な複数ターン推論が必要なタスク（報酬は計算可能）
• 高速なロールアウトスループットが必要なシナリオ（非同期実行は2倍の改善を提供）
• 観察トークンが豊富で、マスキングが勾配のノイズを防ぐ環境

VerlToolを使用すべきでない場合:

• ドメイン固有の最適化が重要な単一タスクシステム（特化されたエージェントがパフォーマンスを上回ります）
• ツール観察が報酬と政策を直接決定するシナリオ（観察マスキングは信号を破棄します）
• 厳密なレイテンシバジェットを持つリアルタイムシステム（非同期デザインはキューイングオーバーヘッドを追加します）
• 非検証可能な報酬の環境（RLVR拡張はグラウンドトゥルースのタスク結果を想定）
• RLの専門知識がないチームがトレーニングインフラストラクチャを管理（慎重な報酬エンジニアリングと収束チューニングが必要）

回避すべき落とし穴:

1. ツールタイムアウト構成の無視: ツールごとに現実的なタイムアウトを設定します。コード実行は60秒必要かもしれません。検索はより迅速にタイムアウトすべきです。タイムアウトが長すぎるとトレーニングボトルネックが生じます。短すぎると不正な失敗が発生します。

2. 不適切な報酬アライメント: 報酬関数が実際のタスク目標を反映していることを確認します。弱い報酬信号はRLの高い分散につながり、収束が不十分になります。検証関数が評価メトリックと一致していることを確認します。

3. ツール状態のリセット忘却: 各軌跡は新しく始まるべきです。前のエピソードからの古い状態は観察を破損させます。軌跡間で常に tool.reset() を呼び出してください。

4. 観察マスキング戦略の混合: 観察マスキングを有効にしたら、トレーニング全体で一貫して適用します。矛盾したマスキングは勾配を不安定化させ、突然のパフォーマンス低下を引き起こします。

5. 非同期サーバーのバッチサイズ不足: 高速化は、ワーカーを飽和させるのに十分な同時実行軌跡を想定します。小さなバッチサイズ（< 8）は非同期利点を無効にします。非同期ロールアウトを有効にするときはバッチサイズ ≥ 16 を使用します。

6. ツールレジストリディスカバリのスキップ: ツールをトレーニングループにハードコードするとモジュール性が損なわれます。レジストリの discover_from_directory() メソッドを使用して、真のプラグアンドプレイツール統合を可能にします。

リファレンス

論文: VerlTool