Teacher-Student Training

Training has two stages. First, a privileged teacher that observes the true wrench is trained with PPO. Then a student that only sees arm proprioception is trained with PPO, an intent estimator, and a KL term that pulls its action distribution towards the frozen teacher's.

Tasks

Register two tasks on the same transport environment: a teacher task whose actor observes the teacher observation, and a student task whose actor observes the student observation. Observation sizes, reward weights and the simulation settings are on the Transport Environment page.

# legged_gym/envs/paint/paint_config.py
class PaintCfg(LeggedRobotCfg):  # student task
    class env(LeggedRobotCfg.env):
        num_actions = 9
        num_observations = ...  # TODO: student history size
        num_privileged_obs = ...  # TODO: critic history size
        episode_length_s = 10
        teacher_mode = False

    class terrain(LeggedRobotCfg.terrain):
        ...  # TODO: rough terrain with curriculum

    class control(LeggedRobotCfg.control):
        ...  # TODO: arm PD gains, arm action scale 0.5, decimation 4

    class rewards(LeggedRobotCfg.rewards):
        class scales:
            ...  # TODO: weights of Table I


class PaintTeacherCfg(PaintCfg):  # teacher task
    class env(PaintCfg.env):
        num_observations = ...  # TODO: teacher history size
        teacher_mode = True


class PaintTeacherCfgPPO(LeggedRobotCfgPPO):
    ...  # TODO: PPO settings below, 15k iterations


class PaintCfgPPO(PaintTeacherCfgPPO):
    ...  # TODO: PaintRunner, lambda_KL, teacher checkpoint path, 20k iterations


# legged_gym/envs/__init__.py
task_registry.register("paint_teacher", PaintEnv, PaintTeacherCfg(), PaintTeacherCfgPPO())
task_registry.register("paint", PaintEnv, PaintCfg(), PaintCfgPPO())

# legged_gym/utils/task_registry.py
# TODO: in make_alg_runner, create the runner class named by train_cfg.runner_class_name

PPO settings

Both stages use the legged_gym PPO defaults with smaller networks:

Stage 1: privileged teacher

The teacher sees the true wrench, so the unmodified rsl_rl PPO trains it. It learns an admittance-like mapping from the wrench to base and arm motion.

python legged_gym/scripts/train.py --task=paint_teacher --headless

Stage 2: student with intent estimator

At each step the intent estimator maps the student's observation history o_t to an estimated normalized wrench ŵ, and the student actor acts on [o_t, ŵ]. The estimate is detached before it enters the actor, so policy gradients do not flow into the estimator.

Intent estimator

An MLP with the same architecture as the actor (128-128-128, ELU). Its input is the student observation history (144 values for H = 8) and its output the estimated wrench β^est = [F̂^xy, Θ̂^yaw], expressed normalized as c(t).

Losses

The student maximizes the expected return while staying close to the frozen teacher, which sees the teacher observation of the same transition (Eq. (8) and (9) of the paper):

In practice J_KL is added to the PPO loss with weight λ_KL = 10⁻⁶. Both policies are diagonal Gaussians, so the KL divergence has a closed form:

The estimator is trained with the regression loss of Eq. (7), the squared error between the estimated and the true force and torque. It has its own Adam optimizer (learning rate 1e−3, not adapted) and is updated on the same mini-batches.

Changes to rsl_rl

• Algorithm: extend PPO with the frozen teacher, the intent estimator and its optimizer, the KL term, and the estimator regression.
• Storage: also store the teacher observation of every transition, so the teacher can be queried and the true wrench read on each mini-batch.
• Runner: build the student actor-critic (input 144 + 3), the teacher (loaded from the stage-1 checkpoint) and the estimator, pass the teacher observation from the environment to the algorithm, and save the estimator with the policy.

# rsl_rl/modules/intent_estimator.py
class IntentEstimator(nn.Module):
    def __init__(self, num_obs, num_out=3, hidden_dims=(128, 128, 128)):
        super().__init__()
        # TODO: MLP num_obs -> hidden_dims -> num_out with ELU

    def forward(self, obs):
        # TODO: return the estimated normalized wrench
        ...


# rsl_rl/algorithms/paint_ppo.py
class PaintPPO(PPO):
    def __init__(self, actor_critic, teacher, estimator, distill_coef, **kwargs):
        super().__init__(actor_critic, **kwargs)
        # TODO: keep the frozen teacher, the estimator and a separate Adam optimizer for it

    def act(self, obs, teacher_obs, critic_obs):
        # TODO: w_hat = estimator(obs), detached; act on [obs, w_hat] as in PPO.act
        # TODO: also store teacher_obs in the transition
        ...

    def update(self):
        # TODO: per mini-batch, as in PPO.update, plus:
        #   - student distribution on [obs, w_hat.detach()],
        #     teacher distribution on teacher_obs
        #   - add lambda_KL * KL(student || teacher) to the PPO loss
        #   - estimator step: MSE between w_hat and the true wrench (from teacher_obs)
        ...


# rsl_rl/storage/rollout_storage.py
# TODO: add a teacher_observations buffer and return it from mini_batch_generator


# rsl_rl/runners/paint_runner.py
class PaintRunner(OnPolicyRunner):
    def __init__(self, env, train_cfg, log_dir=None, device="cpu"):
        # TODO: student ActorCritic (input = student obs + 3), teacher ActorCritic
        #       loaded from the stage-1 checkpoint, IntentEstimator, and PaintPPO
        ...

    def learn(self, num_learning_iterations, init_at_random_ep_len=False):
        # TODO: as OnPolicyRunner.learn, passing extras["teacher_obs"] to alg.act
        ...

    # TODO: save and load the estimator together with the policy

python legged_gym/scripts/train.py --task=paint --headless

What to expect

• The base velocity converges to the wrench reference, and the mean episode length approaches the full 500 steps.
• The estimator loss is measured on the normalized wrench. Multiply by 40 N and 10 N·m to read it in physical units.
• For reference, our student with H = 8 reaches these errors in simulation for 2 / 4 / 8 kg payloads (Table II of the paper):