Capturing token IDs during agentic interactions for better reinforcement learning

Summary

Reinforcement learning for agentic tasks requires updating model parameters against the precise token sequences generated during rollouts. Standard agent harnesses frequently modify conversation transcripts through message compaction and tool formatting, causing retokenization drift that degrades policy-gradient training. Turnstile addresses this issue as a Rust-based proxy situated between agent harnesses and inference backends such as SGLang. By intercepting OpenAI-compatible Chat Completions traffic, the proxy records exact token IDs, log probabilities, loss masks, weight versions, and mixture-of-experts routing traces at the moment of generation. Validation across text-based coding and multimodal computer-use agents demonstrated steady reinforcement learning improvements without requiring modifications to agent harnesses.

Takeaways

1. 1

   Text-based transcripts from agent harnesses often suffer from retokenization and chat template drift, generating token sequences that diverge from what the model actually processed.

2. 2

   Capturing rollout state at the proxy layer preserves exact token IDs, per-token log probabilities, and loss masks without requiring modifications to existing agent harnesses.

3. 3

   Turnstile records mixture-of-experts routing traces to split mismatched trajectories and stores processed image features directly to maintain consistency in multimodal RL training.