Matt Strong

Matt Strong

CS PhD at Stanford University

← Paper Digest

RAIL: Representation-Aligned Imitation Learning for Student-Compatible Teacher Policies

Meraj Mammadov, Pedro Zuidberg Dos Martires, Johannes Andreas Stork

November 2026

RAIL policy in TunnelVision, turning to reveal obstacles on the way to the green goal.

TL;DR Train the privileged RL teacher on a latent space that the student's observations can also reach, so the student runs the same policy without distillation. In four simulated tasks the teacher-student gap nearly vanishes (success 1.00, 1.00 and 0.98) and the student beats RL from raw observations on the two pixel-based tasks.

What they do

Teacher-student RL trains a teacher on privileged state and distills it into a student that sees only deployment observations. When the teacher acts on information the student cannot recover, no student can match it, and the imitation gap is irreducible. RAIL removes the gap during teacher training: the teacher policy may only use a representation that student observations also produce, and the student then runs the same policy.

How it works

  • Shared latent space. Teacher and student encoders feed a shared encoder. Paired observations of the same state are positives in a symmetric InfoNCE loss, which keeps shared information and suppresses what only one side sees.
  • Alignment and stability losses. A cosine term pulls paired embeddings together so the policy transfers directly; a stability term penalizes changes in policy outputs caused by encoder updates.
  • Gradient blocking. Each PPO iteration first updates the encoders, then freezes them and updates the policy on teacher embeddings. RL gradients never reach the encoders, so teacher-private information does not leak in.
  • Deployment. Swap the teacher encoder for the student encoder; no distillation or student RL.

Results

  • Four simulated tasks (two grid worlds, MiniWorld navigation from first-person images, a HalfCheetah variant from half-masked images), five seeds each. Student success is 1.00 on Corners and TunnelVision and 0.98 on CollectHealth, with gaps of 0.00, 0.00 and 0.01. The best prior baseline, SITT, reaches 0.37, 0.74 and 0.49.
  • Against PPO on student observations: 0.98 vs. 0.78 on CollectHealth and a return of 2510 vs. 1824 on QuarterCheetah.
  • Letting RL gradients into the encoder drops student success to 0.28 on Corners and 0.18 on TunnelVision. A probe on the teacher embedding predicts the hidden target corner at 25.4%, chance level.
  • With frozen encoders, a new task’s policy trained on teacher observations alone transfers zero-shot and matches RAIL trained from scratch.

Limitations

Simulation only, on small navigation and locomotion tasks. RAIL needs aligned teacher-student observations from the same states, comes with no formal guarantee that the learned policy is realizable from student observations, and its zero-shot transfer was tested only within the same environments and observation spaces.

Videos

Behavior-cloned student in TunnelVision, for comparison.
RAIL policy in Corners, searching for the hidden target corner.

Figures

From Figure 2: Left: teacher and student observations are generated from the same underlying state via different obse...
From Figure 2: Left: teacher and student observations are generated from the same underlying state via different observation mappings. Middle: teacher and student encoders, followed by a shared encoder, map both observations into a common latent space. Right: a shared policy operates on these representations to produce action distributions. Blue arrows indicate the losses used to train the encoders.
From Figure 1: Observation mismatch in TunnelVision: Teacher (blue) vs. student (orange).
From Figure 1: Observation mismatch in TunnelVision: Teacher (blue) vs. student (orange).
From Figure 3: Student evaluation curves for CollectHealth and QuarterCheetah. CollectHealth reports success rate, wh...
From Figure 3: Student evaluation curves for CollectHealth and QuarterCheetah. CollectHealth reports success rate, while QuarterCheetah reports episodic return. BC and DAgger are shown as horizontal lines corresponding to their final student performance.
Abstract

Reinforcement learning (RL) from raw sensory inputs such as images or onboard sensors can be challenging due to high-dimensional observations and sparse rewards. A common strategy is to train a teacher policy with access to privileged state information and then distill it into a student policy that acts from raw inputs alone. However, when the teacher relies on information unavailable to the student, exact imitation may be impossible, creating an irreducible imitation gap. Existing approaches typically mitigate this mismatch through careful reward shaping or additional reinforcement learning on the student policy after the teacher has been trained. Instead, we introduce Representation-Aligned Imitation Learning (RAIL), which addresses the mismatch during teacher training itself. RAIL learns a latent representation shared across teacher and student observations using contrastive learning, and trains the teacher policy directly in this space. Because the teacher policy is constrained to operate on this shared space, it learns behaviors that are reproducible from student observations. This substantially reduces the imitation gap while preserving task performance. Across multiple environments, RAIL outperforms strong baselines without reward modification or post-hoc student fine-tuning, and surpasses direct reinforcement learning from raw observations. The learned representations further enable zero-shot transfer to new tasks through teacher-only training. Project site: https://merajmammadov.com/rail/.

Summaries are written automatically from each paper's text. Figures and videos come from the paper and its project page; all credit goes to the authors.