🤖 machine learning

Preserve Support, Not Correspondence: Dynamic Routing for Offline Reinforcement Learning

DROL (Dynamic Routing for Offline Reinforcement Learning) improves one-step offline RL actors by using a top-1 dynamic routing mechanism that assigns dataset actions to the nearest of $K$ latent-conditioned candidates, allowing the model to perform local policy improvements without the conflicting gradients caused by traditional pointwise extraction.

Original authors: Zhancun Mu, Guangyu Zhao, Yiwu Zhong, Chi Zhang

Published 2026-04-27

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Zhancun Mu, Guangyu Zhao, Yiwu Zhong, Chi Zhang

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 or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Problem: The "Stubborn Student" Dilemma

Imagine you are a student trying to learn how to navigate a complex, crowded city. You have a Teacher (a super-intelligent GPS) who can tell you exactly which path to take. However, you want to be a "One-Step" Student: you don't want to have to consult the GPS every single second; you want to be able to look at a street corner and just know where to go instantly.

In current AI (Offline Reinforcement Learning), students learn by trying to mimic the Teacher. But there is a catch: the Teacher often gives a specific instruction for a specific moment (e.g., "At this exact intersection, turn left").

The problem is that the student becomes too obsessed with the specific instruction. If the Teacher says "Turn left," and the student realizes that "Turning slightly left" would actually be much faster and safer, the student often refuses to change. They feel "tethered" to the Teacher's exact words. They end up compromising—taking a mediocre middle path—instead of finding the best local route. They are preserving the correspondence (the exact instruction) rather than the support (the fact that the road is actually drivable).

The Solution: DROL (The "Dynamic Roommate" Strategy)

The researchers proposed a new method called DROL. Instead of one student trying to memorize every specific instruction, DROL gives the student a team of "Candidates" (think of them as a group of roommates) during training.

1. The Roommate Assignment (Dynamic Routing)

Instead of one student trying to be everything to everyone, DROL says: "For every street corner in the city, we will send out $K$ different roommates. Each roommate will try to 'claim' a different part of the neighborhood."

When a piece of data comes in (e.g., "At this corner, the driver turned right"), DROL doesn't force a specific student to memorize it. Instead, it looks at the roommates and asks: "Which roommate is currently standing closest to that right turn?"

The roommate who is closest is the "Winner." Only that winner gets the instruction to "Learn how to turn right."

2. The "Hand-off" (Responsibility Transfer)

This is the magic part. Because the roommates are constantly moving and learning, they can swap responsibilities.

Imagine Roommate A is currently responsible for "The North Side" of the park. As Roommate A learns to be a better navigator, they might realize, "Hey, I'm actually getting really good at the fast highway route!" As Roommate A moves toward the highway, Roommate B might step in and say, "I'll take over the North Side of the park; I'm closer to it now."

In old methods, the "ownership" of a path was stuck to a specific latent code (a specific student). In DROL, ownership is fluid. This allows one student to move toward a "better" action (higher Q-value) without leaving a "hole" in the map, because another student can immediately step in to cover the old spot.

Why This Matters (The "TL;DR")

Old Way (Pointwise Correspondence): The student is a perfectionist who is afraid to deviate from the Teacher's exact words, even if the Teacher's path is slightly inefficient.
DROL Way (Dynamic Routing): The student is a coordinated team. They focus on covering the map (preserving support) rather than mimicking the words (preserving correspondence).

The Result:
The AI becomes much better at complex, "multimodal" tasks (situations where there are many different "right" ways to do something, like navigating a maze). It stays fast and cheap to use (one-step inference), but it's much smarter because it learned to specialize its "roommates" to cover the whole territory efficiently.

Technical Summary: Preserve Support, Not Correspondence: Dynamic Routing for Offline Reinforcement Learning (DROL)

1. Problem Statement

In Offline Reinforcement Learning (RL), a major challenge is improving a policy (the actor) to maximize expected return while ensuring the policy stays within the "support" of the fixed dataset (avoiding distribution shift). This is particularly difficult in multimodal environments (e.g., OGBench), where a single state may require several different, valid actions.

Current state-of-the-art "one-step" actors (which are computationally efficient for inference) often use a pointwise extraction method. In this setup, a strong iterative "teacher" model provides a specific target action for a specific latent code. The student actor is then trained to simultaneously move toward higher Q-values and stay close to that specific teacher endpoint.

The core issue: If the direction of higher Q-value and the direction of the teacher's endpoint disagree, the loss forces a "compromise" on that specific sample. This prevents the actor from moving toward a better local action even if that action is well-supported by the data, because the student is "tethered" to a specific, potentially sub-optimal, teacher-latent correspondence.

2. Methodology: Dynamic Routing (DROL)

The authors propose DROL (Dynamic Routing for Offline RL), which shifts the focus from preserving pointwise correspondence (latent $\to$ teacher) to preserving local action support (the set of valid actions).

Key components of the DROL mechanism:

Candidate Set Generation: Instead of training on a single output, for each state $s$ , the actor samples $K$ candidate actions $\{\hat{a}_1, \dots, \hat{a}_K\}$ from a bounded latent prior.
Top-1 Dynamic Routing: For every dataset action $(s, a)$ in a training batch, DROL identifies the "winner"—the candidate $\hat{a}_k$ that is closest to the dataset action $a$ in Euclidean space:
$k^*(s, a) = \arg \min_{k \in [K]} \|\hat{a}_k - a\|_2^2$
Winner-Only Update: Only the routed winner receives the gradient update. The loss function for the actor is:
$\mathcal{L}_{\text{actor}}(\theta) = \mathbb{E}_{(s,a), z_{1:K}} \left[ \|\hat{a}_{k^*(s,a)} - a\|_2^2 - \alpha Q_\phi(s, \hat{a}_{k^*(s,a)}) \right]$
This combines Behavior Cloning (BC) (to anchor the winner to the data) and Critic Guidance (to push the winner toward higher Q-values).
Responsibility Transfer: Because routing is recomputed at every gradient step based on the current geometry of the candidates, the "ownership" of a data region can shift. If one candidate moves toward a higher Q-value, another candidate can "take over" the responsibility of covering the original data region.

3. Key Contributions

Conceptual Shift: Identifies that pointwise latent-to-teacher correspondence is an unnecessary constraint. The goal should be preserving the local support of the action space.
DROL Algorithm: Introduces a routed candidate-set actor that enables "responsibility transfer," allowing for local policy improvement without losing data support.
Theoretical Insights:
- Proves that winner-only routing avoids "mode collapse" in multimodal distributions.
- Demonstrates that routing replaces a "persistent tether" (which biases the policy toward sub-optimal teacher endpoints) with a "transferable anchor."
- Shows that increasing the routing budget $K$ improves the probability of covering all local modes (coverage).

4. Experimental Results

The authors evaluated DROL on OGBench (a multimodal benchmark) and D4RL.

Performance: DROL is highly competitive with much heavier, multi-step iterative policies. On OGBench, the default DROL(16) configuration matches or outperforms the one-step FQL baseline in 6 out of 10 task groups.
Efficiency: Crucially, DROL maintains the one-step inference advantage. While it uses $K$ candidates during training to build a better structure, at test time, it only requires a single latent draw, making it as fast as standard one-step actors.
Scaling: Results show that increasing the routing budget $K$ generally improves performance on complex, highly multimodal tasks (like antmaze-giant), as it provides a finer "quantization" of the supported action space.
Mechanism Validation: Visualizations and metrics confirm that the candidate set does not collapse; instead, candidates specialize into distinct local modes.

5. Significance

DROL represents a significant step forward for efficient, high-performance offline RL. It bridges the gap between the high performance of slow, iterative generative models (like diffusion or flow-based teachers) and the deployment efficiency of fast, one-step actors. By treating the actor as a dynamic set of candidates rather than a rigid mapping, DROL allows policies to be both "ambitious" (following the critic) and "safe" (staying within data support).