On Learning Informative Trajectory Embeddings for Imitation, Classification and Regression

Abstract: In real-world sequential decision making tasks like autonomous driving, robotics, and healthcare, learning from observed state-action trajectories is critical for tasks like imitation, classification, and clustering. For example, self-driving cars must replicate human driving behaviors, while robots and healthcare systems benefit from modeling decision sequences, whether or not they come from expert data. Existing trajectory encoding methods often focus on specific tasks or rely on reward signals, limiting their ability to generalize across domains and tasks. Inspired by the success of embedding models like CLIP and BERT in static domains, we propose a novel method for embedding state-action trajectories into a latent space that captures the skills and competencies in the dynamic underlying decision-making processes. This method operates without the need for reward labels, enabling better generalization across diverse domains and tasks. Our contributions are threefold: (1) We introduce a trajectory embedding approach that captures multiple abilities from state-action data. (2) The learned embeddings exhibit strong representational power across downstream tasks, including imitation, classification, clustering, and regression. (3) The embeddings demonstrate unique properties, such as controlling agent behaviors in IQ-Learn and an additive structure in the latent space. Experimental results confirm that our method outperforms traditional approaches, offering more flexible and powerful trajectory representations for various applications. Our code is available at https://github.com/Erasmo1015/vte.

• The paper introduces VTE, an unsupervised embedding approach that extracts latent policy ability from state-action trajectories without relying on reward signals.
• It combines hierarchical state space models with a transformer-based VAE to enable robust imitation, clustering, classification, and regression across varying ability levels.
• Empirical results demonstrate that VTE embeddings form clear ability-level clusters and support disentangled control, enhancing performance in sequential decision-making tasks.

Informative Trajectory Embeddings: Unsupervised Representation of Policy Ability for Sequential Decision Making

Problem Context and Motivation

Learning from state-action trajectories is a fundamental problem in sequential decision making, underpinning tasks such as imitation learning, classification, regression, and clustering in domains including autonomous driving, robotics, and healthcare. Existing trajectory encoding approaches typically depend on explicit reward or goal labels and often target specific downstream applications, thereby limiting their generality and adaptability. The ubiquity and versatility of unsupervised representation learning models in static domains (e.g., CLIP, BERT) motivate the exploration of analogous techniques for dynamic sequential environments. This paper proposes an unsupervised trajectory embedding method that extracts and encodes latent ability of policies solely from state-action trajectories, operating independently of reward or goal signals.

Figure 1: Illustration of the VTE Framework, showing skill extraction and sequential encoding to produce a latent trajectory embedding.

Methodological Contributions

The proposed approach, Variational Trajectory Encoding (VTE), is structured as a two-stage pipeline:

1. Skill Extraction via Hierarchical State Space Models (HSSM): Leveraging the LOVE framework, temporal abstractions (skills) and their boundaries are inferred from trajectories by maximizing sequence likelihood while minimizing description length. Rather than relying on sampled skill assignments, VTE utilizes the logit representations of skill and boundary posteriors, preserving fine-grained information across the trajectory.
2. Sequential Latent Encoding via Transformer-VAE: The sequence of skill and boundary logits is individually projected through MLPs, concatenated per timestep, and then processed by a transformer. The transformer's pooled output is mapped to the mean and variance parameters of a Gaussian distribution, forming the trajectory's latent embedding as per the VAE paradigm. The decoding mechanism reconstructs actions conditioned on states and the trajectory embedding, optimizing the standard VAE evidence lower bound (ELBO).

This embedding captures both the identity and temporal order of skills, resulting in an ability-conditioned latent vector that characterizes the generating policy. Unlike reward- or goal-conditioned methods (e.g., GCPC [zeng2024goal]), VTE operates entirely without external supervision or reward labels.

Empirical Results

Experiments are conducted on MuJoCo environments (Hopper, Walker2D, Half-Cheetah), with datasets generated from policies at three ability levels (low, medium, expert). VTE embeddings are evaluated on multiple downstream tasks:

• Clustering (tSNE/PCA): The latent embeddings separate trajectories by ability level, forming well-defined clusters.

  Figure 2: tSNE analysis reveals ability-level separation in trajectory embedding space.

• Ability-Conditioned Imitation: IQ-Learn is extended to use trajectory embeddings as conditioning variables, recovering diverse policies corresponding to each ability level. VTE outperforms relevant baselines, including reward-utilizing methods, in achieving return fidelity analogous to upper-bound known-ability policies.

  Figure 3: Evaluation curves for Hopper show returns matching target ability levels via VTE embedding conditioning.

• Classification and Regression: MLP classifiers trained on VTE embeddings achieve perfect accuracy in ability prediction. Regression of returns yields low relative L2 error, surpassing competitive methods.
• Behavioral Disentanglement: Perturbing individual embedding dimensions induces coherent and interpretable behavioral modifications in agents, demonstrating disentanglement across latent variables.

  Figure 4: Behavior modulation in Walker2D and Hopper achieved by varying single dimensions of the trajectory embedding.

• Embedding Distribution Analysis: Wasserstein distance metrics between distributions of embeddings for different ability levels further validate the structural discriminability of VTE.

  Figure 5: Heatmap of inter- and intra-ability Wasserstein distances, revealing distinct separation among ability-conditioned embeddings.

• Ablation Studies: Naive mean pooling of skill embeddings fails to capture trajectory-level distinctions or produce meaningful imitation; the sequential transformer encoder is critical for structural fidelity and task performance.

Theoretical and Practical Implications

The VTE approach advances unsupervised trajectory representation learning in dynamic settings by:

• Reward and Label Independence: It eschews reliance on explicit reward or task indicators, enabling deployment in domains where such supervision is unavailable or unreliable.
• Generalization Across Policy Mixtures: The embedding robustly distinguishes and reconstructs policies across the full suboptimal-optimal spectrum.
• Enabling Disentangled Control: Its additive latent structure supports modular behavior manipulation and interpretability.
• Supporting Multiple Downstream Tasks: The same embedding facilitates imitation, clustering, classification, and regression, unifying sequential decision task pipelines.

Speculation on Future Developments

The trajectory embedding paradigm established here opens avenues toward:

• Transfer Learning Across Heterogeneous Domains: Unsupervised embeddings can potentially bridge task or environment gaps, offering policy reuse and adaptation.
• Meta-Learning and Continual Learning: Trajectory representations may enhance policy generalization, fast adaptation, and robustness in nonstationary settings.
• Hierarchical Policy Structure Discovery: Extending the framework to decompose complex behaviors into latent sequences could yield hierarchical or compositional skill representations.
• Causal Inference in Sequential Decision Processes: Embedding disentanglement may inform causal relationships among policy factors and outcomes.

Conclusion

This work introduces a robust unsupervised embedding methodology for state-action trajectories, grounded in hierarchical skill abstraction and sequential encoding. The resulting latent vectors demonstrate discriminative representational capacity, disentangled behavior modulation, and strong performance across imitation, classification, and regression without reward labels. The approach offers theoretical and practical advancements for representation learning in sequential decision tasks, and is poised to enable broad applications in multi-policy environments, adaptive robotics, and real-world agent modeling.