Tunix는 Google이 공개한 JAX 네이티브 강화학습(reinforcement learning) 라이브러리로, 에이전트가 도구와 환경을 오가며 긴 trajectory를 만드는 훈련 병목을 줄이는 데 초점을 둔다. 특히 TPU에서 rollout 생성과 학습을 분리해 가속기를 놀리지 않는 구조를 강조한다.
해결하려는 병목
에이전틱 RL은 일반적인 단일 응답 학습보다 불규칙하다. 한 prompt에서 몇 번의 도구 호출로 끝날 수도 있고, 어떤 작업은 긴 탐색과 검증 루프를 거친다. 이 때문에 rollout 길이가 제각각이고, 동기식 배치가 완성될 때까지 trainer가 대기하는 시간이 생긴다.
Tunix는 이 문제를 두 축으로 푼다.
| 병목 | Tunix 접근 |
|---|---|
| rollout 생성 지연 | vLLM-TPU, SGLang-Jax 같은 비동기 inference engine과 연동 |
| trainer idle | rollout producer와 training consumer를 큐 기반 pipeline으로 분리 |
| variable-length trajectory | 완료된 trajectory를 동적으로 그룹화하고 후처리 |
| 커스텀 환경 연결 | OSS environment와 reward function을 꽂을 수 있는 구조 |
구조
prompts
-> async rollout orchestrator
-> completed trajectory queue
-> AgenticRLLearner
-> scoring / grouping / GRPO-style update
-> TPU trainer비동기 rollout orchestrator는 완료된 trajectory를 계속 큐에 넣고, AgenticRLLearner는 필요한 그룹이 완성되는 즉시 학습 단계로 넘긴다. GRPO처럼 prompt당 여러 reasoning path를 비교하는 알고리즘에서는 trajectory가 완성되는 순서대로 동적으로 묶는 방식이 중요하다.
어디에 맞는가
- 도구 사용 에이전트 훈련: 검색, 코드 실행, 브라우저 조작처럼 rollout 길이가 불규칙한 작업
- TPU 기반 연구팀: JAX 생태계와 TPU 가속을 유지하면서 에이전틱 RL 처리량을 높이고 싶을 때
- RLHF/RLAIF 확장 실험: 단일 응답 선호학습을 넘어 환경 상호작용과 reward를 포함하려는 경우
- 벤치마크 학습 파이프라인: SWE, web, math, tool-use benchmark에서 trajectory 병목을 줄이려는 경우
주의할 점
Tunix는 일반 앱 개발자가 바로 쓰는 에이전트 프레임워크라기보다 훈련 인프라에 가깝다. JAX, TPU, rollout engine, reward 설계, 데이터 큐잉까지 이해해야 장점이 드러난다. 또한 throughput 최적화가 곧바로 더 좋은 에이전트 행동을 보장하지는 않는다. reward 품질과 환경 설계가 여전히 핵심이다.
관련 문서
- single-rollout-asynchronous-optimization — 에이전틱 RL을 비동기로 안정화하는 학습 기법
- agent-harness — 에이전트 실행 하네스 설계
- agentgym2 — 현실적 노이즈와 불확실성을 반영한 에이전트 벤치마크
참고 자료
- Scaling Agentic RL: High-Throughput Agentic Training with Tunix — Google Developers Blog (2026-07-22)