[Question] How can I load Pretrained Teacher model in Rsl-Rl distillation?

[pumfish]

Question

Hi, I'm learning to use the RSL-RL Distillation in IsaacLab
I first check source/isaaclab_rl/isaaclab_rl/rsl_rl/rl_cfg.py, and find that

IsaacLab/source/isaaclab_rl/isaaclab_rl/rsl_rl/rl_cfg.py

Lines 150 to 154 in 28ed0bb

	policy: RslRlPpoActorCriticCfg | RslRlDistillationStudentTeacherCfg = MISSING
	"""The policy configuration."""
	algorithm: RslRlPpoAlgorithmCfg | RslRlDistillationAlgorithmCfg = MISSING
	"""The algorithm configuration."""

Then I checked source/isaaclab_rl/isaaclab_rl/rsl_rl/distillation_cfg.py, but I noticed that RslRlDistillationStudentTeacherCfg doesn't include parameters for loading pre-trained teacher model weights - it only contains the hidden layer lists for both teacher and student networks.

IsaacLab/source/isaaclab_rl/isaaclab_rl/rsl_rl/distillation_cfg.py

Lines 31 to 35 in 28ed0bb

	student_hidden_dims: list[int] = MISSING
	"""The hidden dimensions of the student network."""
	teacher_hidden_dims: list[int] = MISSING
	"""The hidden dimensions of the teacher network."""

How should I load the pre-trained teacher network? Or, how does IsaacLab's wrapped policy configuration set up the Distillation in RSL-RL?
https://github.com/leggedrobotics/rsl_rl/blob/750e84566d91877a8bbabc7971578be24429bca8/rsl_rl/algorithms/distillation.py#L16-L20

[RandomOakForest]

Thank you for posting this. I'll move this post to our Discussions for follow up.

[ClemensSchwarke]

Hi @pumfish!
You shouldn't need to modify anything, rsl_rl detects if you load a teacher model or continue a distillation training (i.e. load a teacher and a student model) automatically. After training the teacher network, just make sure you create the StudentTeacher class with the correct teacher_hidden_dims. All you need to do then is starting the distillation training with the path to the teacher model, i.e., pass --load_run=<your teacher> as a command line argument. FYI, if you want to continue a previous distillation training, just pass the student-teacher model instead of the teacher and everything should be loaded correctly.
Hope this helps :)

[miromoemen]

Hi @ClemensSchwarke @RandomOakForest ,
I have a question related to this topic. I want to load a teacher model that has different (more) observations than the student model. I'm unable to find any documentation for such task.
Thanks.

[ClemensSchwarke]

Hi @miromoemen,
You just have to create an observation group named "teacher" that contains your teacher observations. The "policy" observation group should contain the student observations.

[miromoemen]

@ClemensSchwarke Many thanks for the prompt response!