Add configs and adapt exporter for RSL-RL distillation

[ClemensSchwarke]

Description

This PR adds configuration classes for Student-Teacher Distillation and adapts the policy exporters to be able to export student policies.

Type of change

• Non-breaking change

Checklist

• I have run the pre-commit checks with ./isaaclab.sh --format
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• I have updated the changelog and the corresponding version in the extension's config/extension.toml file
• I have added my name to the CONTRIBUTORS.md or my name already exists there

[alessandroassirelli98]

Hey! Thanks a lot for this feature — I actually needed something like this, and you already had it covered. Much appreciated!

I had a quick question: how do you think to integrate the teacher-student setup within the broader IsaacLab framework?

At the moment, I’ve defined two separate environments — one for the teacher and one for the student. The only real difference between them is the observation space.
For the agents, I’ve set up two configurations:

• The teacher uses RslRlPpoActorCriticCfg for the policy and RslRlPpoAlgorithmCfg for the algorithm.

• The student uses RslRlDistillationStudentTeacherCfg for the policy and RslRlDistillationAlgorithmCfg for the algorithm.

Does this align with your intended setup?

When using the student though, I always need to initialize the teacher first. So I was thinking of adding a teacher_experiment_name attribute to the RslRlOnPolicyRunnerCfg class. That way, in the training script, we could check if the algorithm is a distillation type and load the corresponding teacher checkpoint accordingly.

This would be different from --resume as resume would try to load the checkpoint for the current experiment which would be just the student

env config:

@configclass
class G1FlatGaitRewardTeacherCfg(G1FlatCfg):
    observations: TeacherObservationsCfg = TeacherObservationsCfg()
    rewards: G1GaitRewards = G1GaitRewards()
    commands: CommandsCfg = CommandsCfg()

    def __post_init__(self):
        # post init of parent
        super().__post_init__()
        self.rewards.feet_air_time = None
        self.rewards.track_lin_vel_xy_exp.params["command_name"] = (
            "gaited_base_velocity"
        )
        self.rewards.track_ang_vel_z_exp.params["command_name"] = "gaited_base_velocity"
        # observation terms (order preserved)

@configclass
class G1FlatGaitRewardStudentCfg(G1FlatGaitRewardTeacherCfg):
    observations: StudentObservationsCfg = StudentObservationsCfg()

    def __post_init__(self):
        # post init of parent
        super().__post_init__()

agent config:

@configclass
class G1FlatStudentGaitRewardPPORunnerCfg(G1RoughPPORunnerCfg):
    def __post_init__(self):
        super().__post_init__()
        self.num_steps_per_env = 64
        self.max_iterations = 300
        self.experiment_name = "g1_flat_gait_student"
        self.teacher_experiment_name = "g1_flat_gait_teacher"
        self.policy = RslRlDistillationStudentTeacherCfg(
            init_noise_std= 0.001,
            student_hidden_dims = [256, 128, 128],
            teacher_hidden_dims = [256, 128, 128],
            activation="elu"
        )
        self.algorithm = RslRlDistillationAlgorithmCfg(
            num_learning_epochs=5,
            learning_rate=1e-03,
            gradient_length=2.
        )

@configclass
class G1FlatTeacherGaitRewardPPORunnerCfg(G1RoughPPORunnerCfg):
    def __post_init__(self):
        super().__post_init__()

        self.max_iterations = 30000
        self.experiment_name = "g1_flat_gait_teacher"
        self.policy.actor_hidden_dims = [256, 128, 128]
        self.policy.critic_hidden_dims = [256, 128, 128]

then in the main script adding something like:

    if isinstance(agent_cfg.policy, RslRlDistillationStudentTeacherCfg):
        teacher_root_path =  os.path.join("logs", "rsl_rl", agent_cfg.teacher_experiment_name)
        teacher_root_path = os.path.abspath(teacher_root_path)
        trained_teacher_path = get_checkpoint_path(teacher_root_path, agent_cfg.load_run_teacher, agent_cfg.load_checkpoint_teacher)

Let me know what you think!

[ClemensSchwarke]

Hey! Thanks a lot for this feature — I actually needed something like this, and you already had it covered. Much appreciated!
...

Does this align with your intended setup?

When using the student though, I always need to initialize the teacher first. So I was thinking of adding a teacher_experiment_name attribute to the RslRlOnPolicyRunnerCfg class. That way, in the training script, we could check if the algorithm is a distillation type and load the corresponding teacher checkpoint accordingly.
...

Let me know what you think!

Hey! Yes, your setup looks good :) However, I don't think you need an attribute for the teacher. Just pass the directory with the load_run flag when starting the training. PS: You also don't need to pass the resume flag for training your student.

[Mayankm96]

@ClemensSchwarke should the parameter by default be 1?

Suggested change

	"""The number of environment steps the gradient flows back."""
	"""The number of rollout steps for gradient propagation.
	This is useful for sequential training of recurrent student network.
	"""

[Mayankm96]

Can you please also update the extension.toml and CHANGELOG?

Thank you! Looks good otherwise.