Multistep Quasimetric Learning for Scalable Goal-conditioned Reinforcement Learning

1. Problem Statement

The paper addresses the challenge of Goal-Conditioned Reinforcement Learning (GCRL), specifically the difficulty of reasoning over long horizons in offline settings.

• The Core Issue: Modern offline GCRL methods struggle to generalize to tasks significantly longer than those seen in the training data.
• The Trade-off:
  • Temporal Difference (TD) Learning: Provides optimality guarantees (can recover the optimal $Q^*$) but relies on local, single-step updates. This leads to compounding errors over long horizons and poor horizon generalization.
  • Monte Carlo (MC) Learning: Uses global updates (multi-step returns) which work well in practice for long horizons but typically only recover the behavior policy's value function ($Q_\beta$) rather than the optimal one, lacking theoretical optimality guarantees.
• The Gap: Existing methods often separate these approaches or rely on hierarchical structures that are complex to implement. There is a need for a method that integrates the global value propagation of MC with the optimality guarantees of TD within a unified, scalable architecture.

2. Methodology: Multistep Quasimetric Estimation (MQE)

The authors propose MQE, an offline GCRL method that learns a quasimetric distance between states (and state-action pairs) and goals using multistep returns.

Key Components:

A. Quasimetric Architecture
Instead of standard Q-networks, MQE parameterizes the goal-conditioned value function as a distance metric.

• Definition: The Q-value is defined as $Q_g(s, a) = V_g(g)e^{-d((s,a), g)}$, where $d$ is a learned distance.
• Metric Residual Network (MRN): The distance function $d$ is implemented using an MRN architecture. This enforces the Triangle Inequality ($d(x, z) \leq d(x, y) + d(y, z)$) and non-negativity, ensuring the learned distance has a valid geometric structure suitable for planning.

B. Multistep Returns with Waypoints
To bridge the gap between TD and MC, MQE introduces a multistep backup mechanism:

• Waypoint Sampling ($s^w_t$): Instead of bootstrapping only from the immediate next state $s'$, the algorithm samples intermediate "waypoints" from the future trajectory. These are sampled using a mixture of Bernoulli and Geometric distributions (Eq. 8).
• Multistep Objective ($T_\beta$): The loss function regresses the distance from the current state to the goal against the discounted distance from a sampled waypoint to the goal:
  $$e^{-d((s,a),g)} \leftarrow \mathbb{E}[\gamma^{k'} \cdot e^{-d(s^w_t, g)}]$$
  This allows value information to propagate globally across the trajectory in a single update step, mimicking MC benefits while maintaining the structure of a Bellman-like update.

C. Action Invariance Constraint
To ensure the learned distance represents a valid value function (where the value of a state is independent of the action taken to reach it if the state is the same), MQE enforces Action Invariance:

• Constraint: $d(s, (s, a)) = 0$ for all actions $a$.
• Implementation: A specific loss term ($L_I$) is added to penalize deviations where the distance between a state and its state-action pair is non-zero. This prevents the network from collapsing to trivial solutions and stabilizes training.

D. Policy Extraction
The policy is extracted by minimizing the learned distance between the current state and the goal, regularized by behavior cloning (DDPG + BC):
$$\min_\pi \mathbb{E}[d((s, \pi(s, g)), g) - \alpha \log \pi(a|s, g)]$$

3. Key Contributions

1. First Integration of Multistep Returns with Quasimetrics: MQE is the first method to effectively combine multistep Monte Carlo returns with quasimetric distance architectures, enabling both local and global value propagation without explicit hierarchy.
2. Horizon Generalization: The method demonstrates superior ability to generalize to horizons significantly longer (up to 1000% longer) than those present in the training data.
3. Real-World "Stitching": MQE enables compositional generalization (stitching) in real-world robotics. It can combine short, independent skills learned from an unlabeled dataset to solve complex, multi-stage tasks (e.g., opening a drawer then placing an object inside) that were never seen as a single trajectory in the data.
4. Theoretical Guarantee: The authors provide a proof (Theorem 1) showing that in a tabular setting, MQE performs policy improvement over the behavior policy, recovering a policy with $V^\pi \geq V^{\pi_\beta}$.

4. Experimental Results

Simulated Benchmarks (OGBench)

• Tasks: Evaluated on 95 tasks including locomotion (AntMaze, HumanoidMaze) and manipulation (Cube, Puzzle) with both state-based and pixel-based observations.
• Long-Horizon Performance: MQE significantly outperformed state-of-the-art baselines (including TMD, CRL, HIQL, and GCIQL).
  • On the "Colossal" AntMaze (50% larger than existing benchmarks, requiring ~4000 steps), MQE was the only method to achieve non-zero success rates.
  • In "stitch" datasets (where training trajectories are short, e.g., 4m, but testing requires traversing 40m), MQE showed robust generalization, whereas other methods failed completely.
• Visual Observations: MQE successfully handled high-dimensional visual inputs (64x64x3), demonstrating scalability.

Real-World Robotics (BridgeData)

• Setup: A 6-DoF WidowX250 robot performing manipulation tasks.
• Compositional Tasks: The agent was tested on sequences of up to 4 "Pick and Place" tasks and a "Open Drawer + Place" task. The dataset contained no examples of these specific long-horizon compositions.
• Results:
  • Quadruple Pick and Place: MQE achieved a 13.3% success rate (2/15 trials), while all other baselines (GCBC, GCIQL, TRA) achieved 0%.
  • Open Drawer & Place: MQE achieved 33.3% success (5/15), significantly outperforming baselines which failed to chain the dependencies.
  • Task Progress: MQE maintained high task progress as the number of required sequential steps increased, whereas other methods' performance degraded rapidly.

5. Significance and Impact

• Solving the Long-Horizon Bottleneck: MQE addresses a fundamental limitation in offline RL: the inability to scale to long-horizon tasks without explicit hierarchical planning or massive amounts of long-horizon data.
• Data Efficiency: By leveraging the geometric structure of quasimetrics and multistep returns, MQE can learn complex behaviors from unlabeled, short-horizon offline datasets.
• Real-World Applicability: The successful deployment on a real robot for "stitching" tasks demonstrates that these theoretical advances translate to practical robotic manipulation, moving beyond simulation.
• Simplicity: Unlike hierarchical RL methods that require complex sub-policy training, MQE is an end-to-end method using a single critic and policy, making it easier to implement and deploy.

In summary, MQE represents a significant step forward in making offline goal-conditioned RL scalable, robust, and capable of compositional reasoning in both simulated and real-world environments.