🤖 machine learning

Test-Time Graph Search for Goal-Conditioned Reinforcement Learning

This paper introduces Test-Time Graph Search (TTGS), a lightweight, training-free planning wrapper that leverages the inherent geometric structure of existing offline goal-conditioned RL policies to dramatically improve success rates on long-horizon tasks without requiring additional supervision or parameter updates.

Original authors: Evgenii Opryshko, Junwei Quan, Claas Voelcker, Yilun Du, Igor Gilitschenski

Published 2026-05-26✓ Author reviewed ⓘ

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Evgenii Opryshko, Junwei Quan, Claas Voelcker, Yilun Du, Igor Gilitschenski

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have trained a very smart robot to navigate a maze. You've shown it thousands of different paths, and it has learned how to move from point A to point B when those points are close together. However, when you ask it to cross a massive, complex maze from one side to the other, it gets confused. It tries to take a giant leap, misses its mark, gets stuck in a corner, or runs out of time. This is a common problem in robotics and AI: short-term planning works well, but long-term planning often fails.

This paper introduces a clever, "plug-and-play" solution called Test-Time Graph Search (TTGS). It doesn't require retraining the robot or teaching it new skills. Instead, it gives the robot a "map" and a "guide" right before it starts moving.

Here is how it works, using simple analogies:

1. The Problem: The "Giant Leap" Trap

Think of your trained robot as a hiker who knows the terrain perfectly for the next 10 steps. If you tell them to walk 100 steps to a specific tree, they might try to sprint the whole way. Because they can't see that far ahead clearly, they might trip over a rock or walk into a dead end. In the paper's terms, the robot's "value function" (its internal guess of how good a move is) gets noisy and unreliable over long distances.

2. The Solution: The "Relay Race" Strategy

Instead of asking the robot to run the whole marathon in one go, TTGS breaks the journey into a series of short, manageable dashes. It turns the robot's journey into a relay race.

The Map (The Graph): The system looks at the massive library of practice runs (the offline dataset) the robot has already done. It picks out key "waypoints" from these old runs and connects them like dots on a map.
The Guide (The Shortest Path): When you give the robot a new goal, the system uses a classic math algorithm (Dijkstra's algorithm) to find the shortest, safest path between the start and the finish using only the dots from the old practice runs.
The Handoffs (Subgoals): The robot doesn't look at the final destination yet. It only looks at the next "waypoint" on the map. Once it reaches that, it gets a new instruction to go to the next waypoint. It keeps doing this until it reaches the goal.

3. The Secret Sauce: The "Soft Penalty"

There is a catch: sometimes the "map" might suggest a shortcut that looks short but is actually dangerous (like a bridge that looks solid but is actually broken). The paper's authors noticed that the robot's internal "guessing" about distance can be wrong.

To fix this, they added a soft penalty. Imagine the map has a rule: "If a path looks too long or risky, we don't delete it, but we put a huge 'tax' on it." The robot's planner will still see the risky path, but it will prefer a slightly longer, safer route made of small, reliable steps. This prevents the robot from trying to jump over gaps it can't actually cross, while still keeping the map connected.

4. Why It's Special

No Retraining: You don't need to teach the robot anything new. You just take the robot you already built, give it this "map wrapper," and it works better immediately.
Works with "Frozen" Policies: The robot's brain is "frozen" (it can't learn new things during the test), but this method helps it use what it already knows more effectively.
It Knows When to Stop: If the map doesn't have enough "waypoints" to bridge the gap between the start and the goal (like trying to cross a canyon with no stepping stones), the system is smart enough to say, "I can't plan this safely," and just lets the robot try its best on its own. It doesn't force a bad plan.

The Results

The researchers tested this on a benchmark called OGBench, which includes complex mazes for robots like ants and humanoids.

Before: On the hardest mazes, the robots often failed completely (0% success rate).
After: With TTGS, success rates jumped to over 90% in many cases.
Comparison: This performance matched or beat much more complex methods that required extra training, expensive computer models, or online practice, all while taking less than a second to plan.

Summary

Think of TTGS as giving a skilled but short-sighted hiker a GPS that only shows them the next few safe steps, based on a map of where other hikers have successfully walked before. It turns a terrifying, long-distance journey into a series of easy, confident steps, allowing the robot to solve problems it previously couldn't touch.

Technical Summary: Test-Time Graph Search for Goal-Conditioned Reinforcement Learning

Problem Statement

Offline goal-conditioned reinforcement learning (GCRL) aims to train agents to reach user-specified goals using precollected datasets, avoiding the costs and safety risks of online interaction. While effective for moderate-horizon tasks, standard offline GCRL policies often fail on long-horizon tasks. In complex environments (e.g., large mazes), errors in value estimation accumulate, leading to unreliable policies that cannot stitch together behaviors from disparate parts of the dataset.

Prior solutions to this problem have introduced significant complexity, such as computationally intensive generative planners, training distributional value function ensembles, utilizing specialized geometric representations, or relying on additional data beyond the offline dataset. These approaches often operate under the assumption that standard goal-conditioned value functions are inherently too noisy or unreliable to support direct long-horizon planning.

Methodology: Test-Time Graph Search (TTGS)

The authors propose Test-Time Graph Search (TTGS), a lightweight, training-free planning wrapper designed to augment frozen GCRL policies. TTGS operates entirely at test time, requiring no additional supervision, parameter updates, or online interaction.

Core Components

Graph Construction:
- A subset of states is sampled from the offline dataset to serve as graph vertices.
- A distance signal is used to assign edge weights between these vertices. This signal can be derived from:
  - Value-derived distances: Mapping the pretrained goal-conditioned value function $V(s, g)$ to an estimated step count $\hat{d}(s, g)$ . The paper provides closed-form mappings for common reward conventions (e.g., sparse terminal rewards or per-step penalties).
  - Domain-specific distances: Using geometric metrics like Euclidean distance between body positions when available.
- Soft Penalty Mechanism: To address the issue of value functions overestimating connectivity (creating "wormholes" across gaps), TTGS employs a superlinear penalty function $p(\cdot)$ for transitions exceeding a trust region threshold $\tau$ . Edges with distances $D_{ij} < \tau$ retain their raw distance, while those with $D_{ij} \ge \tau$ are assigned a weight $p(D_{ij})$ (e.g., $x \cdot 1000^{x/\tau}$ ). This preserves local metric structure while steering the planner toward chains of short, trustworthy hops, preventing the graph from fragmenting (as a hard threshold would) or allowing infeasible long jumps.
Test-Time Planning:
- Shortest-Path Search: Given a start state $s_0$ and a goal $g$ , the algorithm locates the nearest vertices in the graph and computes a shortest path using Dijkstra's algorithm. This guide path is computed once per episode.
- Adaptive Subgoal Selection: During execution, the agent does not follow the path blindly. Instead, it selects a subgoal from the precomputed path based on a step budget $T$ . The agent identifies the farthest waypoint on the path that is within the budget $T$ from the current state. If no such waypoint exists, it selects the next immediate waypoint to ensure forward progress.
- Policy Execution: The frozen base policy $\pi(a | s, \tilde{g})$ is invoked with the current state and the selected subgoal $\tilde{g}$ .

Diagnostic Tool

The paper introduces a largest hop ratio ( $\rho$ ) as a rollout-free diagnostic to predict whether TTGS will be beneficial for a new task. $\rho$ is the ratio of the maximum edge cost along the planned path to the total estimated distance. A low $\rho$ indicates sufficient intermediate state coverage (favorable for TTGS), while a high $\rho$ suggests the dataset lacks bridging states, in which case TTGS defaults to the base policy behavior.

Key Contributions

Re-evaluation of Value Functions: The authors demonstrate that standard offline GCRL value functions, previously considered too unreliable for search, encode sufficient local geometric structure for robust long-horizon planning when processed with a novel soft-penalty mechanism.
Training-Free Framework: TTGS provides a method to improve long-horizon performance by reusing existing policies and value functions without modifying the training pipeline, requiring no additional supervision or online interaction.
Performance with Minimal Overhead: Empirical results show TTGS significantly boosts success rates for state-of-the-art offline GCRL agents (HIQL, QRL, GCIQL, SAW, OTA) with negligible computational overhead (planning time < 1 second per episode).

Experimental Results

The method was evaluated on the OGBench benchmark, covering locomotion (pointmaze, antmaze, humanoidmaze) and manipulation tasks.

Locomotion Tasks: TTGS achieved substantial gains on long-horizon "stitch" tasks where base policies typically fail.
- On pointmaze-giant-stitch-v0, success rates for HIQL improved from 0.0% to 80.9%, and for GCIQL from 0.0% to 98.0%.
- On humanoidmaze-giant-stitch-v0, HIQL improved from 4.4% to 78.1%.
- These improvements often matched or outperformed complex methods requiring auxiliary training or generative models.
Manipulation Tasks: On tasks where evaluation goals lie outside the data manifold (e.g., scene-play-v0), TTGS provided minimal gains (0–2 percentage points) but did not degrade performance, correctly defaulting to the base policy when intermediate states were missing.
Ablations:
- Soft vs. Hard Penalties: Removing the soft penalty (using raw distances) or using a hard threshold (removing long edges) significantly degraded performance on stitching tasks due to graph disconnections or infeasible paths.
- Subgoal Selection: The adaptive strategy of selecting the farthest reachable waypoint outperformed a simpler strategy of always picking the immediate next waypoint.
- Distance Sources: TTGS worked effectively with both value-derived distances and simple Euclidean body-position distances, demonstrating flexibility.

Significance and Claims

The paper claims that the additional complexity often assumed necessary for long-horizon planning is frequently unnecessary. By leveraging the latent geometric structure already present in standard value functions and applying a lightweight graph search at test time, practitioners can significantly extend the capabilities of existing offline GCRL agents.

The authors emphasize that TTGS is a conservative, safe planning wrapper: it relies solely on retrieved dataset states, avoiding the synthesis of physically impossible trajectories common in generative approaches. When the offline data lacks the necessary intermediate states to bridge the start and goal, TTGS gracefully defaults to the base policy's behavior rather than producing unreliable plans. The work suggests that value-based GCRL agents often leave useful geometric structure unused, and TTGS provides a recipe to recover this structure for long-horizon planning without retraining.