ObsSpec speeds up reinforcement learning by guessing tool outputs early
ObsSpec: Accelerating Agentic Reinforcement Learning with Observation Speculation
Machine Learning
Summary
Running agents that use tools in reinforcement learning can be slow because they wait for each tool's results before deciding the next step. The authors created ObsSpec, which guesses what the tools will output ahead of time, so the agent can start thinking about the next step sooner. If the guess matches the actual output exactly, the work is saved, making the process faster without losing accuracy. This approach reduced delays and worked well with existing techniques to speed up learning.
What this means in practice
- •For reinforcement learning engineers: Accelerate training of tool-using reinforcement learning agents by overlapping prediction and execution phases.
- •For game ai developers: Speed up decision-making cycles in game AI agents that rely on tools by preemptively guessing tool effects.
Authors
Kyle Montgomery, Wang Zhang, Chenguang Wang
Abstract
Tool-using agents alternate between model generation and tool execution, waiting for each tool call to complete before starting the next generation. This sequential execution is costly during agentic reinforcement learning (RL), which requires large batches of multi-turn rollouts. We introduce ObsSpec, an observation speculation pipeline for agentic RL that uses a small world model to predict tool outputs while the environment executes, allowing the policy to begin its next generation early. ObsSpec retains speculative work only when the predicted and actual observations match exactly, making rollout collection lossless, while co-training keeps the world model up to date as the policy learns. On terminal-use tasks, predictions from a 0.6B world model exactly match actual observations in 66.4% of attempts and arrive roughly 16 times faster. Using these predictions, ObsSpec reduces agent latency by 9.3% and 14.2% for 8B and 32B policies. Across controlled experiments, we characterize the regimes in which ObsSpec reduces agent latency under shared policy resources. Finally, because ObsSpec overlaps policy generation with tool execution, it is architecturally compatible with existing techniques for accelerating RL. Code is available at https://github.com/kylemontgomery1/verl-recipe/tree/obsspec/obsspec.