Add student/teacher distillation algorithms

[zhenkirito123456]

Summary

This is the first core PR split out of #230 (PHP release support), following the review discussion on amazon-far-internal/holosoma-internal#80. It adds the student/teacher distillation algorithms to core and makes them generic. Depth sensing, WBT command changes, inference, and the PHP extension will follow in separate PRs.

• Import the distillation stack from php release support #230 as-is:

  • Distillation (pure DAgger) and DistillationPPO (DAgger blended into PPO on a schedule)
  • their rollout storages and the StudentTeacher* modules
  • TrainingConfig.teacher_checkpoint and the teacher-count/loading logic in train_agent.py
  • registered algo:distillation and algo:distillation-ppo presets
  • a run-id prefix on W&B checkpoints copied into log_dir, so teachers whose checkpoints share a filename don't overwrite each other

  Left out because they are application-specific or belong to later PRs: preprocess_hook, registry_name, wandb_path, W&B latest-checkpoint resolution, depth/terrain presets.

• Reshuffle distillation minibatches every epoch. RolloutStorage.mini_batch_generator reuses one permutation across epochs. The distillation algorithms were tuned against per-epoch reshuffling, as in the reference implementation. This adds reshuffle_each_epoch (default False, so PPO is unchanged) and turns it on for the two distillation storages.

• Make the student's depth input optional:

  • Shared training-loop, checkpoint, teacher-loading, export, and eval code moves into DistillationBase, which removes about 400 duplicated lines between the two algorithms.
  • StudentTeacherModuleConfig.depth_backbone=None builds a proprioception-only student. That student needs no depth observation group and exports a single ONNX policy with PPO's input/output names. The default is unchanged, so existing configs behave the same.
  • The depth observation group and term are configurable (depth_obs_group / depth_obs_term) instead of hardcoded to depth_camera / depth_cam.

Validation

• Unit tests: distillation, config-registry, rollout-storage, and checkpoint tests pass (159 passed). New tests cover:
  • proprioception-only training and single-file export
  • the depth group being required only when a backbone is set
  • configurable depth group and term names
  • per-epoch reshuffling
  • ONNX export checked against act_inference via onnxruntime
• Lint and types: ruff and mypy, run with the pre-commit hook's pinned versions and arguments, are clean on all changed files.
• Bit-for-bit equivalence with php release support #230 (4a0bf04). A seeded CPU harness runs complete learn() loops of both algorithms on a fake env, then saves, resumes and continues training, evaluates, and exports. Losses, learning rates, ppo_coef, weights, optimizer state, checkpoints, eval actions, and ONNX outputs are identical with a depth backbone configured, including the PPO phase with adaptive LR. The same holds for the PHP release's subclasses (multi-teacher routing, masked DAgger loss, PPO gate) on a branch of 4a0bf04 with these commits applied.
• PHP test suite against that branch: same results as on 4a0bf04 (231 passed). The one failure, test_sim2sim_launchers, also fails on 4a0bf04 and is unrelated.
• IsaacSim A/B on 4×L40S per run: one-teacher PHP distillation, 16384 envs, seed 42, 200 iterations, 4a0bf04 vs. 4a0bf04 + these commits.
  • Both finish cleanly with the same 58 logged scalars, checkpoint structure, and exported depth_backbone.onnx + student.onnx.
  • Schedule-determined values (ppo_coef, LR, entropy, noise std, sample count) match exactly.
  • Stochastic metrics agree within run-to-run noise, e.g. final Loss/behavior 0.0457 vs 0.0459 and Train/mean_reward 15.59 vs 15.51.

🤖 Generated with Claude Code

[Mark2000]

Reviewed the distillation / distillation-PPO carve-out, focusing on the non-test changes.

No correctness bugs on the live training path: both reachable depth backbones honour depth_output_dim end to end, a width mismatch would fail loudly at the first forward rather than silently, and the teacher-readiness machinery looks sound — require_loaded_teachers() gates both teacher_act() and learn(), the setup() -> load() -> load_teacher() ordering in train_agent.py is deliberate, and pre-metadata checkpoints fall back to all-False so a legacy resume fails loudly instead of training against unloaded teachers.

Three comments below. The one on RecurrentDepthBackbone is the only one I'd treat as needing action before merge; the other two are a cleanup and a nit.

[zhenkirito123456]

Thanks @Mark2000! All three are addressed in 2396f04. The CPU equivalence check is still bit-identical, and existing checkpoints load unchanged. Would you mind taking another look?

[juan-g-bonilla]

Thanks for splitting this into a reviewable self-contained PR. My main concerns is that some of the classes/logic that is meant to be generic and easily extensible/maintainable is making some specific assumptions that might make it hard to use in the future.

• added some comments from my agent, feel free to disregard if not relevant.

[juan-g-bonilla]

Here you’re building each teacher as a plain MLP (nn.Sequential), but the supplied checkpoint may come from a policy that also includes an encoder.

load_teacher_state_dict renames the checkpoint keys and filters out parameters that aren’t present in the constructed MLP before calling load_state_dict. Consequently, even strict=True cannot detect those discarded parameters.

If the remaining expected keys and tensor shapes match, loading succeeds and marks the teacher as ready, despite omitting its encoder and associated input processing. This can silently change the teacher’s behavior. We should either reconstruct the complete teacher policy or explicitly reject unsupported learned components.

[juan-g-bonilla]

of course my preference would be to support a wider array of teacher policies

[juan-g-bonilla]

I dont know how I feel with all the prefix/key massaging here because it makes checkpoint compatibility depend on internal naming conventions. Those names can change during a refactor without changing the model's behaviour, and matching parameter names and shapes doesn't really establish that we're reconstructing the same computation.

Here we're already assuming an MLP, then adapting several possible checkpoint layouts to it. I worry that supporting more teacher implementations will mean accumulating special cases in this general class, without an explicit definition of the source teacher.

Could we require the teacher's model type and construction config, either saved with the checkpoint or supplied by the caller? We could then construct that policy and load its weights. PyTorch documents constructing a known model before loading its state dict, and Hugging Face's auto classes normally select the model through its configuration. That's the separation I'd expect here. https://docs.pytorch.org/tutorials/beginner/saving_loading_models.html https://huggingface.co/docs/transformers/en/model_doc/auto

Known older checkpoint formats could still have small, explicit conversion functions outside this class. Each would handle a documented naming change and reject ambiguous mappings. I'd be more comfortable with that than having the general loader accumulate assumptions about the internals of different teacher implementations.

[juan-g-bonilla]

Using an all-zero action to represent “no teacher applies” seems ambiguous, since a valid teacher can also return all zeros. A downstream consumer checking for this sentinel would incorrectly discard those valid responses. A consumer that doesn't handle it would instead treat the missing response as a real zero-action target.

Could we return a separate validity mask alongside the actions and keep it with each sample through to the loss? That would let consumers distinguish a valid zero action from a missing teacher response without inferring it from the action values.

[juan-g-bonilla]

I'm not sure these depth-specific assumptions belong in the shared distillation base class. For example, _extract_obs selects one configured sensor input and exposes it as depth_obs, and get_inference_policy forwards only actor_obs and depth_obs. A policy that also needs rgb_obs would lose that input here, even if the caller supplied it.

The depth export contract also assumes separate depth_backbone and student modules, with the student consuming [actor_obs | depth_latent]. These assumptions describe a particular policy architecture. Supporting another sensor input or a different network composition would therefore require overriding parts of the algorithm as well as implementing the policy.

Could we let the policy declare its required inputs and define its export layout, with the depth-specific handling kept in the depth policy or an adapter? The shared base could then collect and forward those inputs through that interface. I'd prefer that over adding more sensor-specific branches to the base as we support more applications.

[juan-g-bonilla]

I think this duplicates checkpoint-format handling already in StudentTeacher.load_teacher_state_dict. It also preserves only the actor weights and observation normalizer, so any additional model config or preprocessing state would be discarded before the policy loader sees it.

Could we pass the complete checkpoint through and keep format-specific conversion in a single teacher loader?

[zhenkirito123456]

Thanks Juan, these are really helpful. Fixed some, for the remaining, here's what I'd propose before changing code:

• Teacher loading (Release Holosoma Humanoid Policies on Hugging Face #2/When using unitree_sdk2_python for inference, I suspect the G1 might not need a network cable. #3/SMPL-X Format Support for Motion Retargeting #4/Any plan to support multi-motion tracking for WBT? #7): build each teacher from its own checkpoint config instead of fitting keys into a fixed MLP. holosoma PPO checkpoints already embed experiment_config, so the teacher actor can be reconstructed through PPO's own module construction and loaded strictly; nothing gets silently dropped. A small TeacherPolicy interface (from_checkpoint, act) with a PPO implementation first leaves room for other teacher types. Older formats (legacy teacher.* keys, pre-normalizer) get explicit converters that reject ambiguous mappings. The base and the policy then share a single loader.
• Validity mask (README training examples fail with GLXBadFBConfig error in headless environments #5): teacher_act returns (actions, valid), valid is stored per sample, and the DAgger loss masks with it. No more zero-action sentinel. I'll go ahead with this one.
• Policy-declared inputs/export (Retargeting for new robot #6): agreed on the direction. The policy declares the observation groups and terms it consumes and owns its export (Any plans for benchmarking FastSAC and building a standalone version? #10 already moved export into export_modules()); the base only collects and forwards. The depth-specific pieces move into the depth policy, and the deployed two-file ONNX contract stays the same. Since amazon-far/holosoma#232 is already large, would you be OK doing this in a follow-up PR right after it, or would you rather have it here?

Does that match what you had in mind?