Technical Summary: Zero-shot Imitation Learning by Latent Topology Mapping (ZALT)

Problem Definition

The paper addresses the challenge of zero-shot behavior composition in long-horizon, goal-conditioned imitation learning. In this setting, an agent is provided with a fixed, offline dataset of expert demonstrations that cover only a subset of possible start–goal tasks. The objective is for the agent to solve unseen start–goal combinations without additional environment interaction or policy optimization.

The core difficulty lies in the accumulation of errors over long trajectories. Existing methods, such as behavioral cloning or model-based imitation, often fail in this regime because small policy, model, or distribution-shift errors compound over hundreds of primitive action steps when the agent attempts to compose behaviors outside the demonstrated start–goal distribution. The paper posits that while a fixed dataset may lack complete trajectories for new tasks, it likely contains the necessary local behavioral segments; the challenge is learning to recombine these segments reliably.

Methodology: ZALT

The authors propose Zero-shot Agents from Latent Topologies (ZALT), a method that converts a fixed set of demonstrations into a composable, high-level structure for planning. The approach consists of four main stages:

1. Latent Hub Discovery

ZALT first constructs a behavior topology by identifying "hub states" within the demonstration data. These hubs act as reusable junctions where trajectories converge or diverge.

• Latent Space: To handle high-dimensional observations, ZALT uses a low-level encoder–dynamics–decoder model to map states into a dynamics-aligned latent space. This space is trained to predict future states based on actions, ensuring that behaviorally similar states (rather than just visually similar ones) are clustered together.
• Hub Identification: Hubs are defined as latent clusters where:
  • Convergences: Multiple distinct prior trajectories enter the same cluster.
  • Divergences: A single cluster leads to multiple distinct future trajectories.
  • Starts/Terminations: Initial and terminal states are also treated as hubs.
• Significance: By identifying these hubs, ZALT compresses long primitive-action horizons into shorter sequences of abstract transitions, isolating errors to the links between hubs rather than the entire trajectory.

2. Learning a High-Level Hub Dynamics Model

Once the topology (graph of hubs and edges) is established, ZALT learns a high-level dynamics model ($M^{high}$) that predicts the next hub given the sequence of hubs visited so far ($h_{0:m}$).

• Context Awareness: Unlike a simple graph traversal, this model uses a GRU to capture path-dependent information. For example, reaching a hallway hub might lead to different outcomes depending on whether the agent previously collected a specific key.
• Masking: The model is trained with topology masking; transitions not supported by any demonstration edge are assigned zero probability, preventing the planner from attempting unsupported behaviors.

3. Learning Hub-to-Hub Policies

For each edge in the topology (transition from hub $h_i$ to $h_j$), ZALT trains a specific low-level policy ($\pi_{h_i}$) conditioned on the target hub embedding.

• Behavior Cloning: These policies are trained via behavior cloning on the primitive action segments from demonstrations that connect the specific source and target hubs.
• Robustness: To handle timing variations and boundary detection errors, training includes stochastic boundary perturbation (truncating or extending segments) and observation noise.

4. Zero-Shot Inference

At test time, given a new start–goal pair:

1. Mapping: The start observation is encoded and matched to a start hub ($h_s$); the goal is matched to a set of terminal success hubs ($H_g$).
2. Search: A search algorithm (using a bottleneck transition cost) explores the behavior topology to find the most plausible route from $h_s$ to any $h_g \in H_g$. The search relies on the high-level dynamics model to score path plausibility based on history.
3. Execution: The resulting abstract plan (a sequence of hubs) is executed by chaining the corresponding learned hub-to-hub policies.

Key Contributions

• Latent Topology Mapping: The introduction of a method to explicitly identify and utilize latent hub states (convergences/divergences) as reusable junctions for behavior composition.
• Horizon Compression: A mechanism that reduces long-horizon planning from hundreds of primitive decisions to a short sequence of abstract hub transitions, significantly reducing the opportunity for error accumulation.
• Path-Dependent Dynamics: A high-level model that conditions future predictions on the sequence of visited hubs, rather than just the current state, allowing the agent to reason about context (e.g., inventory state) implicitly encoded in the path history.
• Zero-Shot Composition: A framework that enables solving tasks with start–goal pairs never seen in training, provided the necessary local behavioral segments exist in the demonstration dataset.

Experimental Results

The authors evaluate ZALT in a complex 3D maze environment (MiniWorld/CARL-style) requiring ordered object manipulation (keys, locked doors, diamonds, and a barrel). The task involves navigating to specific keys, unlocking doors, collecting items, and depositing them in a specific order.

• Dataset: 18 successful demonstrations and 120 failed demonstrations covering only 18 of 36 possible start–goal pairs.
• Zero-Shot Performance: ZALT achieves a 55% success rate on the 18 unseen start–goal pairs.
• Baseline Comparison:
  • DQN+BC: 6% success (the strongest baseline).
  • CQL, IQL, Option-Critic: 0% success.
  • Seen Tasks: ZALT achieves 72.2% success on seen tasks, compared to 11% for the strongest baseline, indicating the environment is difficult even for standard imitation learning.
• Efficiency: Successful ZALT plans compress roughly 244–281 primitive timesteps into about 33 hub-level transitions (average edge length ~7.9 actions).
• Ablation Studies:
  • Removing the high-level dynamics model (replacing it with BFS) results in 0% success, proving that raw topology connectivity is insufficient without path-dependent prediction.
  • Removing the GRU memory from the low-level encoder results in 0% success, proving that hub discovery requires behavioral context (e.g., collected keys) beyond immediate visual observation.
• Online Adaptation: After 500 episodes of online interaction, baselines failed to recover performance, suggesting the difficulty stems from the long horizon and sparse rewards rather than a lack of practice.

Significance and Claims

The paper claims that ZALT demonstrates that explicit behavior topology is a critical component for zero-shot generalization in long-horizon imitation learning. The authors argue that while standard methods struggle when the full start–goal route is out-of-distribution, ZALT succeeds by leveraging the "local behavioral pieces" present in the data.

The significance lies in the ability to solve complex, compositional tasks without requiring a complete set of demonstrations for every possible task variation. By compressing the planning horizon and explicitly modeling the connectivity of reusable behavioral segments, ZALT overcomes the error accumulation problem that plagues long-horizon policies. The authors conclude that this approach offers a promising direction for training agents in interactive environments where collecting expert trajectories for every possible task combination is infeasible.

The paper remains modest regarding its scope, noting that the method assumes the existence of reusable behavioral connectivity in the demonstrations and that hub quality is dependent on the learned latent representation. Future work is suggested to explore the forces driving hub formation and to adapt the framework to online or hybrid learning settings.