Section 17.5: Teacher-student and privileged-information distillation

"The teacher knew where the ball was. The student learned to catch it anyway."

A Patient Distillation Engineer

Imagine handing a student pilot a map that shows every hidden crosswind in the room. The pilot aces the test. Then you take the map away the instant they touch real controls. That is the bargain privileged-information distillation strikes. Get it wrong, and the policy that looked flawless in simulation falls over on its first step on concrete. For this to hold up, the training artifact must version four things together: simulator fidelity, Proximal Policy Optimization (PPO) rollout semantics, reward terms, and the reset distribution.

This section develops the distillation contract for fast simulator-trained policies. The teacher is trained or evaluated with privileged state such as terrain heights, contact impulses, object poses, or exact base velocity; the student learns to imitate useful teacher actions from deployable signals such as proprioception, commands, history, and exteroception. "Teacher-student" names the two roles in this pattern, and "privileged-information distillation" names what passes between them: the teacher's privileged-state advantage is distilled into action targets the student can learn from its own deployable observations. The rest of this section builds, audits, and evaluates exactly that contract. Figure 17.5A captures the core asymmetry: the teacher reads simulator secrets that simply do not exist on real hardware, so the student must learn the same policy from what the robot can actually sense.

The key question is practical: which information is legal at deployment, and how do we prove that the student evaluation uses only that legal interface? Figure 17.5B traces the full pipeline that answers this question, from privileged teacher inputs through the valid-step mask to the exported student, and shows where the teacher-to-deployment path is deliberately blocked.

Theory

A distillation objective that reduces to a single mean-squared term is $$\mathcal{L}_{\text{distill}} = \frac{1}{B}\sum_{i=1}^{B}\| \pi_s(o_i^{\text{deploy}}) - \pi_t(o_i^{\text{priv}}) \|_2^2,$$ where \(\pi_t\) is the privileged teacher, \(\pi_s\) is the student, \(o_i^{\text{priv}}\) includes simulator-only information, and \(o_i^{\text{deploy}}\) contains only signals available on the robot.

The formula is easy to write and easy to misuse. The batch must cover the same command, terrain, contact, and disturbance distribution used for evaluation, and the loss should be reported alongside closed-loop student performance, not as a standalone success metric. To see why this matters at scale: in representative locomotion benchmarks (as of 2024), a student trained without any privileged teacher on rough terrain may need 50 million simulator steps to reach 60% success on a held-out panel; a student distilled from a teacher that observed exact terrain heights reaches the same threshold in under 5 million steps, a roughly 10x compression bought entirely by the quality of the supervision signal.

Worked Example

The 10x supervision speedup the theory promises only materializes if the imitation targets are clean, which is exactly what the valid-step mask guarantees, so the smallest worked example worth building is the masked loss itself.

Code Fragment 17.5.1 computes a tiny distillation loss. The mask keeps failed or invalid teacher steps from becoming imitation targets, which matters when the teacher is still imperfect.

# Compute a masked action-distillation loss for teacher-student RL.
# The student matches teacher actions only on valid rollout steps.
import numpy as np

teacher_actions = np.array([
    [0.20, -0.10, 0.05],
    [0.35, -0.08, 0.02],
    [1.50, 0.90, -1.20],
])
student_actions = np.array([
    [0.18, -0.12, 0.04],
    [0.30, -0.10, 0.01],
    [0.10, 0.05, -0.02],
])
valid_teacher_step = np.array([1.0, 1.0, 0.0])

squared_error = ((student_actions - teacher_actions) ** 2).mean(axis=1)
masked_loss = (squared_error * valid_teacher_step).sum() / valid_teacher_step.sum()

print(f"per-step error: {squared_error.round(4)}")
print(f"masked distillation loss: {masked_loss:.4f}")

Expected output: the trace should show both per-step error and the masked loss. A distillation run that reports only average imitation loss can hide whether failed teacher states were included as targets.

The valid-step mask matters in embodied AI because a real robot has no undo. Suppose the student imitates a teacher action taken during a fall or a joint-limit violation. It then trains toward a behavior that ends in hardware damage or an unrecoverable configuration. Including unmasked fall steps can double the closed-loop episodes the student needs to reach deployment-level stability. In representative rough-terrain training runs, roughly 8% of rollout steps are typically fall or reset events, and each one injects a target that points directly away from balance. Filtering those steps is not optional hygiene. It is the difference between a student that recovers from perturbations and one that has learned, implicitly, that falling forward is acceptable.

Mechanically, the mask sets a per-step weight to zero for any rollout step tagged as a fall, timeout, or saturated-action event. The distillation loss is then a weighted average over surviving steps only, so the gradient update never sees the bad teacher moment. The denominator uses the sum of valid weights rather than total steps, keeping the loss scale stable as the fraction of masked steps varies across terrain difficulties.

Practical Recipe

1. List privileged fields and deployable fields before training the teacher.
2. Train or select a teacher whose advantage comes from privileged state, not from evaluation leakage.
3. Collect teacher actions on a diverse rollout panel and mark invalid teacher steps.
4. Train the student from deployable observations using action matching, feature matching, or DAgger-style relabeling when needed.
5. Evaluate only the student on held-out seeds with privileged actor inputs disabled.

A student that imitates a teacher's actions without ever facing the teacher's hardest moments is not a trained policy; it is a transcript of good weather.

The idea in this section becomes useful when privilege is treated as a controlled variable. A teacher may use more information, but every extra field must be named, justified, and blocked from the exported actor. Otherwise distillation becomes information leakage with a nicer name.

Three claims to evaluate separately

The disciplined approach is to evaluate three claims separately. The teacher claim says privileged state improves expert behavior. The imitation claim says the student matches useful teacher actions on valid states. The deployment claim says the student still succeeds when privileged actor inputs are removed.

The imitation claim can pass while the deployment claim fails, and knowing why saves wasted iterations. A student with low offline action-matching loss still fails closed-loop because small errors compound: each slightly wrong action shifts the robot into a state the teacher never visited, so the next target lies outside the training distribution. DAgger-style interactive relabeling breaks this loop by generating new states from the student's own rollouts, having the teacher label them, and turning recovery actions into explicit targets. Bootstrap with offline distillation; close the gap with interactive relabeling.

Each of those tools earns its place only if the schema diff in the last row can actually run, and that audit is only possible when the field partitions are written down explicitly, so a robust implementation starts with an observation schema manifest. The manifest makes leakage visible by listing teacher-only, student, and critic-only fields separately.

1. Freeze the deployable observation schema before collecting teacher data.
2. Store teacher-only fields and critic-only fields as explicit lists.
3. Save the mask rule for excluding failed teacher steps.
4. Report imitation loss and closed-loop student metrics from the same held-out panel.
5. Export a student checkpoint with a wrapper that rejects privileged actor inputs.

# Record the observation schema for privileged-information distillation.
# The exported actor must accept only fields listed under student_obs.
from dataclasses import dataclass, asdict

@dataclass
class DistillationSchema:
    teacher_only: tuple[str, ...]
    student_obs: tuple[str, ...]
    critic_only: tuple[str, ...]
    mask_rule: str
    eval_wrapper: str

    def as_row(self) -> dict[str, object]:
        return asdict(self)

schema = DistillationSchema(
    teacher_only=("exact_terrain_heights", "contact_impulses"),
    student_obs=("joint_pos", "joint_vel", "command", "history"),
    critic_only=("base_velocity",),
    mask_rule="exclude falls, timeouts, and saturated actions",
    eval_wrapper="student_actor_only",
)
print(schema.as_row())

When a distilled student fails, decide whether the issue is teacher quality, state coverage, observation insufficiency, or leakage removal. A student that imitates well offline but fails closed-loop usually needs recovery-state data, history, or an interactive relabeling loop, not another blind epoch over the same expert states. This divergence between offline imitation loss and live success is called the deployment credibility gap, and closing it is typically the central engineering challenge in distillation pipelines of this kind.