186 lines
6.4 KiB
Python
186 lines
6.4 KiB
Python
# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION
|
|
# All rights reserved.
|
|
#
|
|
# SPDX-License-Identifier: BSD-3-Clause
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from rsl_rl.modules import StudentTeacher, StudentTeacherRecurrent
|
|
from rsl_rl.storage import RolloutStorage
|
|
from rsl_rl.utils import resolve_optimizer
|
|
|
|
|
|
class Distillation:
|
|
"""Distillation algorithm for training a student model to mimic a teacher model."""
|
|
|
|
policy: StudentTeacher | StudentTeacherRecurrent
|
|
"""The student teacher model."""
|
|
|
|
def __init__(
|
|
self,
|
|
policy,
|
|
num_learning_epochs=1,
|
|
gradient_length=15,
|
|
learning_rate=1e-3,
|
|
max_grad_norm=None,
|
|
loss_type="mse",
|
|
optimizer="adam",
|
|
device="cpu",
|
|
# Distributed training parameters
|
|
multi_gpu_cfg: dict | None = None,
|
|
):
|
|
# device-related parameters
|
|
self.device = device
|
|
self.is_multi_gpu = multi_gpu_cfg is not None
|
|
# Multi-GPU parameters
|
|
if multi_gpu_cfg is not None:
|
|
self.gpu_global_rank = multi_gpu_cfg["global_rank"]
|
|
self.gpu_world_size = multi_gpu_cfg["world_size"]
|
|
else:
|
|
self.gpu_global_rank = 0
|
|
self.gpu_world_size = 1
|
|
|
|
# distillation components
|
|
self.policy = policy
|
|
self.policy.to(self.device)
|
|
self.storage = None # initialized later
|
|
|
|
# initialize the optimizer
|
|
self.optimizer = resolve_optimizer(optimizer)(self.policy.parameters(), lr=learning_rate)
|
|
|
|
# initialize the transition
|
|
self.transition = RolloutStorage.Transition()
|
|
self.last_hidden_states = None
|
|
|
|
# distillation parameters
|
|
self.num_learning_epochs = num_learning_epochs
|
|
self.gradient_length = gradient_length
|
|
self.learning_rate = learning_rate
|
|
self.max_grad_norm = max_grad_norm
|
|
|
|
# initialize the loss function
|
|
loss_fn_dict = {
|
|
"mse": nn.functional.mse_loss,
|
|
"huber": nn.functional.huber_loss,
|
|
}
|
|
if loss_type in loss_fn_dict:
|
|
self.loss_fn = loss_fn_dict[loss_type]
|
|
else:
|
|
raise ValueError(f"Unknown loss type: {loss_type}. Supported types are: {list(loss_fn_dict.keys())}")
|
|
|
|
self.num_updates = 0
|
|
|
|
def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape):
|
|
# create rollout storage
|
|
self.storage = RolloutStorage(
|
|
training_type,
|
|
num_envs,
|
|
num_transitions_per_env,
|
|
obs,
|
|
actions_shape,
|
|
self.device,
|
|
)
|
|
|
|
def act(self, obs):
|
|
# compute the actions
|
|
self.transition.actions = self.policy.act(obs).detach()
|
|
self.transition.privileged_actions = self.policy.evaluate(obs).detach()
|
|
# record the observations
|
|
self.transition.observations = obs
|
|
return self.transition.actions
|
|
|
|
def process_env_step(self, obs, rewards, dones, extras):
|
|
# update the normalizers
|
|
self.policy.update_normalization(obs)
|
|
|
|
# record the rewards and dones
|
|
self.transition.rewards = rewards
|
|
self.transition.dones = dones
|
|
# record the transition
|
|
self.storage.add_transitions(self.transition)
|
|
self.transition.clear()
|
|
self.policy.reset(dones)
|
|
|
|
def update(self):
|
|
self.num_updates += 1
|
|
mean_behavior_loss = 0
|
|
loss = 0
|
|
cnt = 0
|
|
|
|
for epoch in range(self.num_learning_epochs):
|
|
self.policy.reset(hidden_states=self.last_hidden_states)
|
|
self.policy.detach_hidden_states()
|
|
for obs, _, privileged_actions, dones in self.storage.generator():
|
|
|
|
# inference the student for gradient computation
|
|
actions = self.policy.act_inference(obs)
|
|
|
|
# behavior cloning loss
|
|
behavior_loss = self.loss_fn(actions, privileged_actions)
|
|
|
|
# total loss
|
|
loss = loss + behavior_loss
|
|
mean_behavior_loss += behavior_loss.item()
|
|
cnt += 1
|
|
|
|
# gradient step
|
|
if cnt % self.gradient_length == 0:
|
|
self.optimizer.zero_grad()
|
|
loss.backward()
|
|
if self.is_multi_gpu:
|
|
self.reduce_parameters()
|
|
if self.max_grad_norm:
|
|
nn.utils.clip_grad_norm_(self.policy.student.parameters(), self.max_grad_norm)
|
|
self.optimizer.step()
|
|
self.policy.detach_hidden_states()
|
|
loss = 0
|
|
|
|
# reset dones
|
|
self.policy.reset(dones.view(-1))
|
|
self.policy.detach_hidden_states(dones.view(-1))
|
|
|
|
mean_behavior_loss /= cnt
|
|
self.storage.clear()
|
|
self.last_hidden_states = self.policy.get_hidden_states()
|
|
self.policy.detach_hidden_states()
|
|
|
|
# construct the loss dictionary
|
|
loss_dict = {"behavior": mean_behavior_loss}
|
|
|
|
return loss_dict
|
|
|
|
"""
|
|
Helper functions
|
|
"""
|
|
|
|
def broadcast_parameters(self):
|
|
"""Broadcast model parameters to all GPUs."""
|
|
# obtain the model parameters on current GPU
|
|
model_params = [self.policy.state_dict()]
|
|
# broadcast the model parameters
|
|
torch.distributed.broadcast_object_list(model_params, src=0)
|
|
# load the model parameters on all GPUs from source GPU
|
|
self.policy.load_state_dict(model_params[0])
|
|
|
|
def reduce_parameters(self):
|
|
"""Collect gradients from all GPUs and average them.
|
|
|
|
This function is called after the backward pass to synchronize the gradients across all GPUs.
|
|
"""
|
|
# Create a tensor to store the gradients
|
|
grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None]
|
|
all_grads = torch.cat(grads)
|
|
# Average the gradients across all GPUs
|
|
torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM)
|
|
all_grads /= self.gpu_world_size
|
|
# Update the gradients for all parameters with the reduced gradients
|
|
offset = 0
|
|
for param in self.policy.parameters():
|
|
if param.grad is not None:
|
|
numel = param.numel()
|
|
# copy data back from shared buffer
|
|
param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data))
|
|
# update the offset for the next parameter
|
|
offset += numel
|