GoogleがTPUのアイドル時間を排除するエージェントRL学習ライブラリ「Tunix」を新規公開 — 非同期パイプラインでAIアクセラレータの無駄な待ち時間をなくす

7月21日、Haoyu GaoらGoogle Software Engineerチームが「Scaling Agentic RL: High-Throughput Agentic Training with Tunix- Google Developers Blog」と題した記事を公開した。この記事では、Googleが新規公開したポストトレーニングライブラリ「Tunix」を用いてTPUアイドル時間を排除し、エージェント強化学習を高スループットでスケールさせる手法について詳しく紹介されている。

エージェントRLのボトルネックはTPUの「待ち時間」だった

LLMのアライメント手法は、静的なチャットボット向け学習から、ツール呼び出しや複数ステップの意思決定を伴う「エージェント型ワークフロー」へと急速にシフトしている。この変化がインフラ側に新たな課題をもたらしている。

エージェントがコードを実行したり、Web検索を待ったりしている間、TPUは何もせずアイドル状態になる。高価なAIアクセラレータの稼働率がそのまま落ちる。これがエージェントRL特有のボトルネックだ。

Tunixはこの問題を2つのメカニズムで正面から解決する新規ライブラリとして公開された。元記事ではPerfettoトレースによる可視化を用いてTPU飢餓の解消を実証しており、非同期パイプライン導入後にトレーナーTPUの待ち時間がほぼゼロになる様子がタイムライン上で示されている。

核心：TPUをアイドルにさせない2つの仕組み

1. 非同期ロールアウト（Asynchronous Rollouts）

従来の同期型ロールアウトには2つの問題がある。

実行バブル：環境の初期化やステップ応答を同期的に待つ間、アクセラレータが止まる
ストラグラー効果：バッチ内で最も遅いトラジェクトリの完了を全体が待つ「長尾問題」

TunixはPythonのasyncioを活用した非同期トラジェクトリコレクターエンジンでこれを解消する。あるエージェントがツール実行で待機している間、推論エンジンは即座に別のアクティブなトラジェクトリのトークン生成に切り替える。vLLM-TPUやSGLang-Jaxとネイティブに統合しており、ノンブロッキングなサンプリングを実現している。

2. デカップルドパイプライニング（Barrier-Free Pipelining）

非同期ロールアウトだけでは不十分だ。可変長・長尾のロールアウトを、厳密に同期動作するトレーニングループに橋渡しする部分にも詰まりが生じる。ナイーブな実装では、バッチ全体のトラジェクトリが揃うまでトレーナーTPUが待ちぼうけになる。

Tunixはロールアウトとトレーニングをプロデューサー・コンシューマーパイプラインとして完全に分離する。

プロデューサー：非同期ロールアウトオーケストレーターが完了したトラジェクトリを高スループットキューに流し続ける
コンシューマー：AgenticRLLearnerがキューから消費し、GRPOのようにプロンプトごとに複数の推論パスを必要とするアルゴリズム向けに、トラジェクトリを動的にグループ化する

トラジェクトリグループが揃った瞬間にポストプロセス・スコアリングしてトレーナーへストリーミングするため、同期型トレーナーへの供給が途切れない。

プラグアンドプレイな環境抽象化

エージェントRL開発における別の摩擦点は、アルゴリズムと環境ループの密結合だ。SWE-benchやWebArenaといった外部ベンチマークを新たに追加しようとすると、大規模なコード書き直しが発生する。

TunixはConversationAgentBaseとBaseTaskEnvを継承するだけで、任意のエージェントや環境をトレーニングワークフローに組み込める設計になっている。GymnasiumのようなOSSライブラリとの接続も数十行で済む。

import gymnasium as gym
from tunix.rl.agentic.agentic_grpo_learner import GRPOLearner
from tunix.rl.agentic.environments.base_environment import BaseTaskEnv, EnvStepResult

class MyEnv(BaseTaskEnv):
    def _initial_observation(self):
        # 環境の初期化と最初の観測値を返す
        self.env = gym.make("your_chosen_env")
        observation, info = self.env.reset(seed=42)
        return observation

    def _step_impl(self, action):
        # エージェントのアクションを環境に渡し、次の観測・報酬・終了フラグを受け取る
        action = self.env.action_space.sample()
        obs, reward, done, info = self.env.step(action)
        return EnvStepResult(obs, reward, done, info)

    def close(self):
        # 環境リソースの解放
        self.env.close()

# カスタム環境クラスをGRPOLearnerに渡すだけで訓練ループに組み込める
learner = GRPOLearner(env_class=MyEnv, ...)

エージェント側も同様で、update_from_modelにカスタムロジックを記述するだけ。モデルがQwenであろうとLlamaであろうとGemmaであろうと、エージェントコードはモデルの詳細を知る必要がない。

ブラックボックスを排除する軽量プロファイリング

XProfのような標準プロファイラはオペレーターレベルの詳細なトレースを取れるが、オーバーヘッドが高く、常時計測は現実的でない。非同期エージェント訓練において「どのフェーズがボトルネックか」を俯瞰するには粒度が細かすぎる。

Tunixが導入したのは、RL固有の高レベルメトリクスを常時・軽量に計測し続ける仕組みだ。ロールアウト・トレーニング・重み同期の各フェーズをPerfettoトレースとして可視化し、CPUスレッドとTPUデバイスの実行タイムラインを一目で把握できる。

これにより次のことが可能になる。

TPU飢餓の特定：Pythonのツール呼び出しがどこで非同期パイプラインをブロックしているかを可視化
パイプライン整合性の検証：トレーナーがロールアウト生成を待っていないか、重み同期が遅延を引き起こしていないかを確認
設定のチューニング：スレッドプールとTPUアイドル時間を対比させてロールアウト並列数を最適化、HBM制約に合わせてマイクロバッチサイズを調整

他フレームワークとの比較

※以下の比較は元記事の記述に基づくが、OpenRLHFやTRLとの明示的な比較が元記事に含まれているかは確認が取れていない部分がある。編集部の補足として読んでいただきたい。（編集部注）

vs. OpenRLHF / veRL：Ray + vLLMベースで進化してきたが、PyTorchエコシステムが前提。TunixはJAX/TPUエコシステムにネイティブに対応し、XLAコンパイラ最適化とPathwaysマルチホスト分散トレーニングをそのまま活用できる。
vs. Hugging Face TRL：TRLは広く使われているが、エージェント型マルチターン訓練に特化した非同期パイプラインはTunixの設計範囲に入る。

JAX/TPU環境でエージェントRLを動かすチームにとって、今まで選択肢が事実上なかった領域をTunixが埋める格好だ。

詳細はScaling Agentic RL: High-Throughput Agentic Training with Tunix- Google Developers Blogを参照していただきたい。