Technical Summary: ConTraIRL – Factorized Contrastive Abstractions for Transferable IRL

1. Problem Statement

Inverse Reinforcement Learning (IRL) aims to recover the reward function underlying expert behavior, enabling agents to learn complex skills without manual reward design. A critical challenge in this domain is compositional generalization: the ability to transfer a learned reward to target environments where the context is an unseen pairing of previously observed environmental dynamics ($d$) and task goals ($g$).

Standard IRL methods typically model the reward as a monolithic function of the state. This approach entangles dynamics-dependent and goal-dependent features. Consequently, when an agent is deployed in a target environment where $d$ and $g$ are individually familiar but their specific combination $(d_{tar}, g_{tar})$ was never observed during training, the learned reward is queried under an out-of-distribution context. This leads to unreliable generalization and poor reward recovery.

The paper addresses this by studying a setting where:

• Training data consists of multiple environments where each factor ($d$ and $g$) appears across several contexts.
• A subset of factor pairings is held out as target contexts.
• In target contexts, only a fixed subset of states from a single expert trajectory is available (few-shot supervision), rather than full demonstrations.

2. Methodology: ConTraIRL

The authors propose ConTraIRL (Contrastive Abstractions for Transferable IRL), a framework designed to enable compositional reward transfer by learning decoupled latent representations of dynamics and goals.

2.1 Factorized Contrastive Abstractions

The core of the framework is a dual-encoder architecture that maps observations into two separate latent manifolds:

1. Dynamics Space ($Z_d$): A latent representation invariant to the goal $g$.
2. Goal Space ($Z_g$): A latent representation invariant to the dynamics $d$.

This factorization ensures that the recovered reward does not encode spurious correlations induced by the joint training distribution. To achieve this, the authors enforce Factor Orthogonality:

• $z_d = \Phi_d(s, d)$ must be invariant to $g$.
• $z_g = \Phi_g(s, g)$ must be invariant to $d$.

2.2 Temporal Phase Alignment

The authors argue that factorization alone is insufficient for sequential control tasks, as states along a trajectory represent different stages of execution. ConTraIRL introduces Temporal Phase Alignment, requiring that states corresponding to similar relative progress within a behavior map to nearby latent coordinates across trajectories.

• States with the same dynamics and similar phase remain close in $Z_d$, regardless of the goal.
• States with the same goal and similar phase remain close in $Z_g$, regardless of the dynamics.
• Phase $\psi(s)$ is defined as the normalized timestep $t/T$.

2.3 Reward Recovery via Latent Similarity

The recovered reward function $R(s, d, g)$ is defined as the average similarity between the current state's latent representation and phase-conditioned expert abstractions:
$$R(s, d, g) = \frac{1}{2} \left( \text{sim}(\Phi_d(s, d), \bar{z}_d(\psi(s))) + \text{sim}(\Phi_g(s, g), \bar{z}_g(\psi(s))) \right)$$
Here, $\bar{z}_d(\cdot)$ and $\bar{z}_g(\cdot)$ are centroids of expert states grouped by their relative temporal phase. This formulation converts an out-of-distribution pairing into an in-distribution recombination of factors.

2.4 Training Objectives

The framework employs a dual contrastive objective:

1. Expert Structure Learning ($L_E$): A margin-based contrastive loss applied to expert pairs to structure the latent manifolds. Pairs sharing the same factor and similar phase are treated as positives; others are negatives. This shapes the expert manifolds independently of the learner.
2. Expert-Learner Contrastive Calibration ($L_L$): A noise-contrastive objective that calibrates the reward by distinguishing between expert states ($x_E$) and learner states ($x_L$). This effectively minimizes the Jensen–Shannon divergence between expert and learner state occupancies.

The total objective is $L_{total} = L_E - L_L$.

2.5 Implementation Details

• Context-Modulated Encoder: Instead of simple concatenation, the framework uses a hypernetwork-inspired input layer where the context descriptor ($d$ or $g$) generates specific weight matrices and bias vectors to transform the state. This dynamically influences feature extraction, encouraging the isolation of factor-relevant structure early in the network.
• Expert Abstraction Buffer: A buffer maintains phase-conditioned expert centroids, which are refreshed after each encoder update to ensure consistency with the evolving latent manifold.

3. Key Contributions

• Factorized Representation: The proposal of a dual-encoder architecture that explicitly disentangles dynamics and goals into orthogonal latent spaces, preventing the entanglement of features that plagues monolithic IRL models.
• Temporal Phase Alignment: The integration of relative temporal phase into the latent alignment criterion, providing dense guidance for reward recovery in sequential tasks rather than relying on sparse state matching.
• Compositional Transfer with Few-Shot Supervision: A framework capable of recovering rewards for unseen dynamics-goal pairings using only a partial subset of expert states from the target environment.
• Expert Structure Learning: An explicit objective to shape expert manifolds independently of the learner distribution, providing a stationary reference for reward recovery.

4. Experimental Results

The authors evaluated ConTraIRL on four MuJoCo continuous control benchmarks (Ant, HalfCheetah, Walker, Swimmer) with compositional variations in dynamics and goals.

• Performance on Unseen Pairings: ConTraIRL consistently outperformed baselines (TraIRL, C-AIRL, and SFM) in target contexts. For example, on the Ant environment, ConTraIRL achieved a normalized return of 0.93 ± 0.02, compared to 0.88 for SFM and 0.57 for TraIRL.
• Robustness to Data Scarcity: In ablation studies where the number of target expert states was reduced (from 20% to 5%), baseline methods suffered catastrophic performance drops (e.g., TraIRL dropped to 0.05), while ConTraIRL maintained significantly higher performance (0.27 at 5% states), demonstrating effective few-shot transfer.
• Ablation Analysis:
  • Removing factorization ("w/o F") caused a severe drop in performance, confirming the necessity of disentangled representations.
  • Removing temporal phase alignment ("w/o TA") reduced performance, indicating that phase alignment provides essential dense guidance beyond simple factorization.
• Robustness to Label Noise: ConTraIRL demonstrated greater robustness to noisy factor labels compared to baselines, attributed to the reliance on relative similarity constraints rather than direct supervised prediction.
• Latent Space Visualization: UMAP visualizations confirmed that the learned abstractions successfully separated samples by their corresponding factor (dynamics or goal) while remaining invariant to the other, and that temporal phase was preserved as a smooth gradient within clusters.

5. Significance and Claims

The paper claims that ConTraIRL addresses a fundamental limitation in transfer IRL: the inability to generalize to unseen combinations of known environmental factors. By factorizing the context into orthogonal dynamics and goal components and aligning them temporally, the framework enables compositional generalization.

The authors emphasize that their approach allows for reliable reward recovery in target environments where the specific dynamics-goal pairing was absent during training, provided that the individual factors were observed. The learned reward is shown to track true task performance closely, serving as a reliable proxy for policy optimization even when ground-truth rewards are unavailable.

The work is modestly scoped to compositional transfer across known contextual factors with limited target supervision. The authors explicitly note that the current formulation assumes access to factor labels and comparable temporal phases, and they identify extending the approach to settings with latent factors, variable-length behaviors, and fully zero-shot transfer as important directions for future work.