Google、エージェント向け強化学習ライブラリ「Tunix」を公開
- •Googleは、TPU上でエージェント型強化学習モデルの学習を高速化するオープンソースライブラリ「Tunix」を公開した。
- •Tunixは、非同期ロールアウトとデカップルされたパイプラインを採用し、多段階推論中のアクセラレータのアイドル時間を排除する。
- •本フレームワークはドメイン特化型のプロファイリングとモジュール式のAPIを備え、コード修正なしで外部環境との統合を可能にする。
Googleは、TPUインフラストラクチャ上でのエージェント型ワークフローに向けた強化学習（RL）最適化ライブラリ「Tunix」をリリースした。LLM開発が静的なチャットボットから、コード実行や環境と対話する多段階推論エージェントへと移行する中、ウェブ検索やツール実行といったホスト側の遅延によりAIアクセラレータがアイドル状態になるという学習効率の課題が浮上している。
Tunixは、非同期ロールアウトとバリアフリーなプロデューサー・コンシューマー・パイプラインという2つの主要なアーキテクチャ革新により、このボトルネックを解消する。非同期ロールアウトエンジンは、Pythonのasyncioを活用してTPU推論とホスト側の環境遅延を分離し、個々のエージェントが外部ツールの結果を待機している間もモデル生成を継続させる。同時に、デカップルされたパイプラインが可変長の軌跡を学習器に継続的にストリーミングすることで、バッチ完了を待つ間の同期的な学習ループの停止を防ぐ。
パフォーマンスのオーケストレーションに加え、Tunixはエージェント型RL専用の軽量で継続的なオブザーバビリティ（可観測性）を導入している。一般的なプロファイラーとは異なり、このツールは学習ジョブ全体を通じてドメイン固有のメトリクスを追跡し、開発者がTPUの利用率を可視化してボトルネックを特定できるようにする。また、プロンプトフォーマットや会話履歴を扱うエージェント層と環境層を分離するモジュール式APIを提供しており、コードベースを書き換えることなくSWE-benchやWebArenaなどのオープンソースベンチマークへのプラグアンドプレイ統合を実現する。
TunixはJAXおよびTPUエコシステム向けに最適化されており、主にPyTorchをターゲットとするOpenRLHFやHugging Face TRLなどの既存フレームワークとは一線を画している。XLAコンパイラ最適化とPathways分散学習をネイティブに活用することで、GemmaやQwenファミリーといった推論モデルの学習をサポートする設計だ。現在、本ライブラリはオープンソースプロジェクトとして提供されており、カスタムエージェントや環境実装のための事前構築クラスに加え、一般的なエージェントタスクに向けた詳細なドキュメントとレシピが用意されている。