Test-Time Graph Search for Goal-Conditioned Reinforcement Learning
Abstract
Offline goal-conditioned reinforcement learning (GCRL) often struggles with long-horizon tasks, where errors in value estimation accumulate and produce unreliable policies. It is typically assumed that effective long-term planning is infeasible without specialized training. In contrast, our work demonstrates that existing GCRL policies can complete long-horizon tasks when combined with a lightweight, training-free planning wrapper. We find that standard goal-conditioned value functions encode locally consistent geometric structure sufficient for planning. Our approach, Test-Time Graph Search (TTGS), constructs a graph over the offline dataset and employs an adaptive subgoal selection strategy. To address unreliable value estimates during shortest-path search, we propose a novel mechanism that softly penalizes long-distance transitions. Our method incurs negligible computational overhead and requires no additional supervision or parameter updates. On the OGBench benchmark, TTGS significantly boosts success rates across multiple base learners and tasks, with primary gains on challenging long-horizon locomotion tasks where some success rates are improved from near-zero to over 90\%, often matching or outperforming methods that require complex auxiliary training. Code and videos can be found at https://ktolnos.github.io/ttgs.
Lay Summary
Teaching an AI agent to perform a long sequence of actions is much harder than teaching it short ones: small mistakes accumulate, and the agent gets lost. Existing fixes add generative planners or extra neural networks, but they are heavy and require redoing the training pipeline. We asked whether an agent that already takes reliable short steps could handle long journeys without any retraining. We built Test-Time Graph Search (TTGS), a lightweight tool that wraps around an already-trained agent. It treats observations from the training data as waypoints on a map and connects them using the agent's own sense of how close two situations are. A shortest-path search picks a chain of nearby waypoints that guides the agent one short step at a time. Because the agent's sense of distance is trustworthy for nearby places but noisy for far ones, we bias the search toward paths made of many small hops rather than a few long, uncertain ones. On a standard benchmark, TTGS lifted success rates on the hardest navigation tasks from near zero to over 90%, while adding less than a second of computation. This lets practitioners unlock long-horizon performance from agents they have already trained, without extra data or training.