Technical Summary: Hitting Time Isomorphism for Multi-Stage Planning with Foundation Policies

Problem Statement

The paper addresses the challenge of offline goal-conditioned reinforcement learning (RL) for foundation policy learning. The objective is to learn a reusable, goal-agnostic policy from a static dataset of trajectories that can generalize to an open set of goals in a zero-shot manner.

Existing approaches for offline state representation learning fall into three categories, each with critical limitations for multi-stage planning:

1. Successor Measure Methods: While unsupervised, they often fail to satisfy the triangle inequality, preventing the composition of distances across intermediate sub-goals.
2. Reachability Distance Methods: These often impose metric or quasimetric structures but rely on goal-conditioned supervision for policy training, violating the goal-agnostic requirement for foundation policies.
3. Hitting Time Methods: These satisfy the triangle inequality and support unsupervised training but typically rely on symmetric Euclidean distances. This fails to capture the directed, irreversible temporal geometry of real-world environments, leading to suboptimal planning where routes pass through subgoals unreachable in the required direction.

The authors identify the missing primitive as directional temporal geometry measured in discrete transition steps. They propose expected hitting time as the fundamental quantity to recover, aiming to learn a representation where hitting times are linear functionals of latent displacements, preserving both direction and scale.

Methodology: Isomorphic Embedding Learning (IEL)

The core contribution is a new operator-theoretic framework and the Isomorphic Embedding Learning (IEL) algorithm.

1. Theoretical Framework: Hitting Time Isomorphism

The authors formalize the environment as a Controlled Markov Process (CMP) within a Hilbert space. They define an $\epsilon$-Sufficient CMP where transition dynamics are approximately linear in a latent feature space $\phi$.

• Linear Representation of Hitting Times: Under the condition that the latent dynamics close linearly and a specific adjoint range condition is met, the expected hitting time $V^\pi(x, g)$ from state $x$ to goal $g$ admits an exact linear representation:
  $$V^\pi(x, g) = \langle \phi(g) - \phi(x), \omega^\pi_g \rangle_H$$
  Here, $\omega^\pi_g$ is a task identifier vector in the Hilbert space. This proves that hitting-time functionals are equivalent to latent displacements up to a bounded linear isomorphism.
• Identifiability: The paper proves that any representation recovering this hitting-time geometry is uniquely identifiable up to a bounded linear change of coordinates (Theorem 2.3). This ensures the learned geometry reflects the environment's true directed temporal structure rather than an arbitrary coordinate system.
• Error Bounds: Theoretical analysis shows that global hitting-time error is bounded by the one-step latent transition error amplified by the environment's transient spectral radius (Theorem 3.3). Finite-sample guarantees decompose error into approximation, statistical complexity, coverage, and trajectory-label mismatch terms (Theorem 3.5).

2. The IEL Algorithm

IEL operationalizes this theory through four components:

1. Task Identifier Learning: A contrastive learning approach (InfoNCE loss) learns a mapping $\omega_\psi(g)$ that transforms a target state into a task identifier. This establishes the directional axis in the latent space, enabling asymmetric hitting-time calculations.
2. Hitting Time Regression: The algorithm samples intermediate points from offline trajectories and performs expectile regression to predict hitting-time progress. This biases the geometry toward shorter, in-distribution distances.
3. Temporal Difference (TD) Learning: To handle the scarcity of long-horizon data, IEL bridges the hitting-time formula with HILP-style objectives. It introduces a directional bias via a transformation:
   $$V^\pi(x, g) \approx \|\phi(g) - \phi(x)\| \exp(\beta(1 - \cos(\xi_{x \to g, g})))$$
   where $\xi$ is the angle between the embedding difference and the task identifier. This penalizes angular misalignment, acting as a soft directional constraint absent in symmetric methods.
4. Graph-Based Multi-Stage Planning: At test time, IEL constructs a directed graph over a coreset of latent states. Edge costs are defined by the learned asymmetric hitting-time score. Navigation is performed via Single Source Shortest Path (SSSP) search on the reversed graph, converting local temporal predictions into global zero-shot navigation plans.

Key Contributions

• Theoretical: Proves the existence and identifiability of a displacement-linear representation for hitting times in sufficient CMPs, establishing a rigorous link between operator theory and hitting-time functionals.
• Algorithmic: Introduces IEL, the first foundation policy learning algorithm that simultaneously satisfies unsupervised policy training, directional distance, and the triangle inequality.
• Practical: Demonstrates that incorporating directional asymmetry into the embedding geometry significantly improves multi-stage planning capabilities compared to symmetric baselines.

Experimental Results

The authors evaluate IEL on six offline goal-conditioned RL benchmarks from the D4RL dataset (AntMaze and Kitchen suites).

• Baselines: Comparisons are made against HILP (the primary baseline for foundation policies), GC-IQL, and GC-CQL.
• Performance:
  • IEL with asymmetric graph planning achieves state-of-the-art results, significantly outperforming HILP and recursive midpoint planning in environments like antmaze-ultra-diverse (79.0 vs 39.8) and antmaze-large-diverse (71.8 vs 68.0).
  • Even the symmetric variant of IEL (using symmetric distances but the same graph planning) often outperforms HILP's recursive midpoint strategy, validating the efficacy of the graph-based planning mechanism itself.
  • The asymmetric variant provides further gains in four of the six environments, confirming the necessity of directional geometry for complex navigation.
• Efficiency: IEL's training cost is comparable to HILP. The asymmetric graph planning runs at speeds comparable to recursive midpoint planning, while the symmetric variant is approximately three times faster due to graph reusability.

Significance and Claims

The paper claims that IEL provides the first framework to satisfy the three critical capabilities required for foundation policies in multi-stage planning: unsupervised training, directional distance, and the triangle inequality.

The authors argue that by replacing symmetric distance heuristics with a directed temporal geometry grounded in operator theory, agents can navigate complex, irreversible environments with provable efficiency. They posit that this work provides a mathematical foundation for transitioning Vision-Language-Action (VLA) research from reliance on imitation learning toward reward-based policy learning, potentially enabling more autonomous and trustworthy general-purpose robots. The framework is presented as a rigorous alternative to distance-based models that lack global consistency guarantees.