Technical Summary: Teacher-Student Representational Alignment for Reinforcement Learning-Driven Imitation Learning

Problem Statement

Imitation Learning (IL) is frequently employed to overcome the sample complexity and "curse of dimensionality" associated with Reinforcement Learning (RL) in high-dimensional observation spaces, such as those found in robotics. A common paradigm involves training an RL-based "teacher" on low-dimensional, analytical state representations and then distilling this behavior into a "student" policy that operates on high-dimensional observations (e.g., images).

However, a significant challenge arises when the teacher and student are trained in isolation. The teacher often relies on "privileged information" (state variables available to the teacher but not inferable from the student's observations, such as exact goal coordinates or hidden object locations). This creates an irreducible imitation gap: the teacher learns a policy that is optimal for its privileged view but impossible for the student to imitate because the necessary information is missing from the student's input.

Traditionally, this gap is addressed by adding an RL fine-tuning phase to the student after IL. However, this reintroduces the high sample complexity and training instability of RL. Alternatively, some methods modify the teacher's reward function or action space to force imitability, but these approaches often require heavy feature engineering, introduce difficult-to-tune hyperparameters, or degrade the teacher's task performance.

Methodology

The authors propose a novel, task-agnostic algorithm that learns a shared embedding space to bridge the imitation gap by construction. Instead of modifying rewards or fine-tuning the student post-hoc, the method ensures the teacher learns a policy that is inherently imitable by the student.

Core Architecture

The approach utilizes a multi-view representation learning framework where the teacher and student observe the same underlying environment state $s_t$ through different modalities ($o_T$ and $o_S$). These observations are modeled as containing:

1. Common components ($c_t$): Information shared by both agents (e.g., relative obstacle positions).
2. Private components ($p_T, p_S$): Agent-specific information (e.g., exact goal coordinates for the teacher).

The goal is to learn an encoder that maps both $o_T$ and $o_S$ into a latent space containing only $c_t$, effectively hiding private information from the teacher during its RL training.

Training Phases

The system operates in alternating phases involving two modules:

1. Contrastive Representation Learning (Embedding Module):

   • Objective: Learn a mapping from teacher and student observations to a common latent space.
   • Mechanism: Uses InfoNCE contrastive loss. Paired observations $(o_T, o_S)$ from the same timestep are treated as positive pairs, while unpaired observations in the batch serve as negatives.
   • Alignment Loss ($L_{alignment}$): Explicitly maximizes the similarity between paired embeddings to ensure the policy network sees consistent representations for both agents.
   • Stability Loss ($L_{stability}$): Prevents large deviations in the latent space during training by penalizing shifts in the raw logits of the policy network relative to previous states. This is crucial because RL policies are sensitive to small embedding shifts.
   • Gradient Isolation: Crucially, gradients from the RL policy (Phase 2) are not backpropagated through the encoder networks. This prevents the policy from "cheating" by learning to extract private information from the teacher's raw observations to maximize rewards.

2. Reinforcement Learning (Policy Module):

   • The teacher policy ($\pi_T$) is trained using Proximal Policy Optimization (PPO) to maximize the task reward.
   • However, the policy operates only on the learned latent embeddings, not the raw observations.
   • Because the embeddings lack private information, the teacher is forced to maximize the reward by exploiting only the shared information available to the student.

Key Distinctions from Prior Work

• Unlike SITT [8], which penalizes the teacher for visiting non-imitable states via a KL-divergence term in the reward, this method does not modify the reward function or introduce trade-off coefficients ($\alpha$).
• Unlike methods requiring domain knowledge to select contrastive pairs, this approach is fully task-agnostic and relies on the temporal alignment of teacher-student trajectories.

Experimental Results

The method was evaluated in two simulation environments designed to highlight the imitation gap:

1. CollectHealth: A 3D navigation task where the teacher knows exact box locations, while the student has a limited Field of View (FOV) and must reorient to find targets.
2. TunnelVision: A grid-world navigation task where the teacher sees all surrounding cells, while the student sees only three cells in its forward direction.

Performance Metrics

The study compared the proposed method against:

• Behavior Cloning (BC): Standard IL using a teacher trained in isolation.
• SITT: The state-of-the-art baseline that penalizes non-imitable teacher behaviors.

Key Findings:

• Imitation Gap Reduction: The proposed method achieved a near-zero imitation gap ($\Delta \approx 0.00$ to $0.01$) in both environments, significantly outperforming BC (gap $\approx 0.53-0.81$) and SITT (gap $\approx 0.19-0.43$).
• Student Performance: The distilled student policies achieved success rates of 1.00 (TunnelVision) and 0.98 (CollectHealth), matching or exceeding the teacher's performance without requiring RL fine-tuning.
• Teacher Performance: Unlike SITT, which suffered a drop in teacher performance due to reward modification, the proposed method maintained high teacher performance (1.00 and 0.99) because the original reward function remained untouched.
• Ablation Studies:
  • Removing the alignment loss caused a noticeable drop in student performance, particularly in CollectHealth, confirming the need for explicit embedding similarity.
  • Removing the stability loss caused a significant drop in TunnelVision (a high-sensitivity environment), demonstrating that stable embeddings are critical for preventing policy collapse in sensitive settings.

Significance and Contributions

The paper claims the following contributions:

1. Novel Task-Agnostic Method: It introduces a plug-and-play approach to bridge the imitation gap by learning a shared latent representation, eliminating the need for reward function modifications or complex feature engineering.
2. Construction of Imitability: By restricting the teacher to operate on a shared embedding space via contrastive learning, the method ensures the teacher learns behaviors that are inherently reproducible by the student, avoiding the irreducible gap caused by privileged information.
3. Superior Empirical Performance: The approach consistently outperforms strong baselines (BC and SITT) in terms of student success rates and imitation gap reduction while preserving the teacher's optimal task performance.
4. Stability Mechanisms: The introduction of explicit alignment and stability losses addresses the sensitivity of RL policies to representation shifts, enabling robust training without post-hoc fine-tuning.

The authors conclude that this method offers a robust solution for distilling high-dimensional policies from privileged teachers, making it suitable for integration into existing IL frameworks without altering the fundamental reward structure of the task.