Section 17.5: Teacher-student and privileged-information distillation

A Careful Control Loop

For Teacher-student and privileged-information distillation, GPU RL depends on simulator fidelity, PPO rollout semantics, reward terms, and reset distribution being versioned in the same training artifact.

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.

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?

Theory

A simple distillation objective 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.

Worked Example

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.

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.

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.

The graduate-level habit 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.

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.