250 lines
9.7 KiB
Python
250 lines
9.7 KiB
Python
# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION
|
|
# All rights reserved.
|
|
#
|
|
# SPDX-License-Identifier: BSD-3-Clause
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import warnings
|
|
from torch.distributions import Normal
|
|
|
|
from rsl_rl.networks import MLP, EmpiricalNormalization, Memory
|
|
|
|
|
|
class StudentTeacherRecurrent(nn.Module):
|
|
is_recurrent = True
|
|
|
|
def __init__(
|
|
self,
|
|
obs,
|
|
obs_groups,
|
|
num_actions,
|
|
student_obs_normalization=False,
|
|
teacher_obs_normalization=False,
|
|
student_hidden_dims=[256, 256, 256],
|
|
teacher_hidden_dims=[256, 256, 256],
|
|
activation="elu",
|
|
init_noise_std=0.1,
|
|
noise_std_type: str = "scalar",
|
|
rnn_type="lstm",
|
|
rnn_hidden_dim=256,
|
|
rnn_num_layers=1,
|
|
teacher_recurrent=False,
|
|
**kwargs,
|
|
):
|
|
if "rnn_hidden_size" in kwargs:
|
|
warnings.warn(
|
|
"The argument `rnn_hidden_size` is deprecated and will be removed in a future version. "
|
|
"Please use `rnn_hidden_dim` instead.",
|
|
DeprecationWarning,
|
|
)
|
|
if rnn_hidden_dim == 256: # Only override if the new argument is at its default
|
|
rnn_hidden_dim = kwargs.pop("rnn_hidden_size")
|
|
if kwargs:
|
|
print(
|
|
"StudentTeacherRecurrent.__init__ got unexpected arguments, which will be ignored: "
|
|
+ str(kwargs.keys()),
|
|
)
|
|
super().__init__()
|
|
|
|
self.loaded_teacher = False # indicates if teacher has been loaded
|
|
self.teacher_recurrent = teacher_recurrent # indicates if teacher is recurrent too
|
|
|
|
# get the observation dimensions
|
|
self.obs_groups = obs_groups
|
|
num_student_obs = 0
|
|
for obs_group in obs_groups["policy"]:
|
|
assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations."
|
|
num_student_obs += obs[obs_group].shape[-1]
|
|
num_teacher_obs = 0
|
|
for obs_group in obs_groups["teacher"]:
|
|
assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations."
|
|
num_teacher_obs += obs[obs_group].shape[-1]
|
|
|
|
# student
|
|
self.memory_s = Memory(num_student_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim)
|
|
self.student = MLP(rnn_hidden_dim, num_actions, student_hidden_dims, activation)
|
|
|
|
# student observation normalization
|
|
self.student_obs_normalization = student_obs_normalization
|
|
if student_obs_normalization:
|
|
self.student_obs_normalizer = EmpiricalNormalization(num_student_obs)
|
|
else:
|
|
self.student_obs_normalizer = torch.nn.Identity()
|
|
|
|
print(f"Student RNN: {self.memory_s}")
|
|
print(f"Student MLP: {self.student}")
|
|
|
|
# teacher
|
|
if self.teacher_recurrent:
|
|
self.memory_t = Memory(
|
|
num_teacher_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim
|
|
)
|
|
num_teacher_obs = rnn_hidden_dim
|
|
self.teacher = MLP(num_teacher_obs, num_actions, teacher_hidden_dims, activation)
|
|
|
|
# teacher observation normalization
|
|
self.teacher_obs_normalization = teacher_obs_normalization
|
|
if teacher_obs_normalization:
|
|
self.teacher_obs_normalizer = EmpiricalNormalization(num_teacher_obs)
|
|
else:
|
|
self.teacher_obs_normalizer = torch.nn.Identity()
|
|
|
|
if self.teacher_recurrent:
|
|
print(f"Teacher RNN: {self.memory_t}")
|
|
print(f"Teacher MLP: {self.teacher}")
|
|
|
|
# action noise
|
|
self.noise_std_type = noise_std_type
|
|
if self.noise_std_type == "scalar":
|
|
self.std = nn.Parameter(init_noise_std * torch.ones(num_actions))
|
|
elif self.noise_std_type == "log":
|
|
self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions)))
|
|
else:
|
|
raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
|
|
|
|
# action distribution (populated in update_distribution)
|
|
self.distribution = None
|
|
# disable args validation for speedup
|
|
Normal.set_default_validate_args(False)
|
|
|
|
def reset(self, dones=None, hidden_states=None):
|
|
if hidden_states is None:
|
|
hidden_states = (None, None)
|
|
self.memory_s.reset(dones, hidden_states[0])
|
|
if self.teacher_recurrent:
|
|
self.memory_t.reset(dones, hidden_states[1])
|
|
|
|
def forward(self):
|
|
raise NotImplementedError
|
|
|
|
@property
|
|
def action_mean(self):
|
|
return self.distribution.mean
|
|
|
|
@property
|
|
def action_std(self):
|
|
return self.distribution.stddev
|
|
|
|
@property
|
|
def entropy(self):
|
|
return self.distribution.entropy().sum(dim=-1)
|
|
|
|
def update_distribution(self, obs):
|
|
# compute mean
|
|
mean = self.student(obs)
|
|
# compute standard deviation
|
|
if self.noise_std_type == "scalar":
|
|
std = self.std.expand_as(mean)
|
|
elif self.noise_std_type == "log":
|
|
std = torch.exp(self.log_std).expand_as(mean)
|
|
else:
|
|
raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
|
|
# create distribution
|
|
self.distribution = Normal(mean, std)
|
|
|
|
def act(self, obs):
|
|
obs = self.get_student_obs(obs)
|
|
obs = self.student_obs_normalizer(obs)
|
|
out_mem = self.memory_s(obs).squeeze(0)
|
|
self.update_distribution(out_mem)
|
|
return self.distribution.sample()
|
|
|
|
def act_inference(self, obs):
|
|
obs = self.get_student_obs(obs)
|
|
obs = self.student_obs_normalizer(obs)
|
|
out_mem = self.memory_s(obs).squeeze(0)
|
|
return self.student(out_mem)
|
|
|
|
def evaluate(self, obs):
|
|
obs = self.get_teacher_obs(obs)
|
|
obs = self.teacher_obs_normalizer(obs)
|
|
with torch.no_grad():
|
|
if self.teacher_recurrent:
|
|
self.memory_t.eval()
|
|
obs = self.memory_t(obs).squeeze(0)
|
|
return self.teacher(obs)
|
|
|
|
def get_student_obs(self, obs):
|
|
obs_list = []
|
|
for obs_group in self.obs_groups["policy"]:
|
|
obs_list.append(obs[obs_group])
|
|
return torch.cat(obs_list, dim=-1)
|
|
|
|
def get_teacher_obs(self, obs):
|
|
obs_list = []
|
|
for obs_group in self.obs_groups["teacher"]:
|
|
obs_list.append(obs[obs_group])
|
|
return torch.cat(obs_list, dim=-1)
|
|
|
|
def get_hidden_states(self):
|
|
if self.teacher_recurrent:
|
|
return self.memory_s.hidden_states, self.memory_t.hidden_states
|
|
else:
|
|
return self.memory_s.hidden_states, None
|
|
|
|
def detach_hidden_states(self, dones=None):
|
|
self.memory_s.detach_hidden_states(dones)
|
|
if self.teacher_recurrent:
|
|
self.memory_t.detach_hidden_states(dones)
|
|
|
|
def train(self, mode=True):
|
|
super().train(mode)
|
|
# make sure teacher is in eval mode
|
|
self.teacher.eval()
|
|
self.teacher_obs_normalizer.eval()
|
|
|
|
def update_normalization(self, obs):
|
|
if self.student_obs_normalization:
|
|
student_obs = self.get_student_obs(obs)
|
|
self.student_obs_normalizer.update(student_obs)
|
|
|
|
def load_state_dict(self, state_dict, strict=True):
|
|
"""Load the parameters of the student and teacher networks.
|
|
|
|
Args:
|
|
state_dict (dict): State dictionary of the model.
|
|
strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this
|
|
module's state_dict() function.
|
|
|
|
Returns:
|
|
bool: Whether this training resumes a previous training. This flag is used by the `load()` function of
|
|
`OnPolicyRunner` to determine how to load further parameters.
|
|
"""
|
|
|
|
# check if state_dict contains teacher and student or just teacher parameters
|
|
if any("actor" in key for key in state_dict.keys()): # loading parameters from rl training
|
|
# rename keys to match teacher and remove critic parameters
|
|
teacher_state_dict = {}
|
|
teacher_obs_normalizer_state_dict = {}
|
|
for key, value in state_dict.items():
|
|
if "actor." in key:
|
|
teacher_state_dict[key.replace("actor.", "")] = value
|
|
if "actor_obs_normalizer." in key:
|
|
teacher_obs_normalizer_state_dict[key.replace("actor_obs_normalizer.", "")] = value
|
|
self.teacher.load_state_dict(teacher_state_dict, strict=strict)
|
|
self.teacher_obs_normalizer.load_state_dict(teacher_obs_normalizer_state_dict, strict=strict)
|
|
# also load recurrent memory if teacher is recurrent
|
|
if self.teacher_recurrent:
|
|
memory_t_state_dict = {}
|
|
for key, value in state_dict.items():
|
|
if "memory_a." in key:
|
|
memory_t_state_dict[key.replace("memory_a.", "")] = value
|
|
self.memory_t.load_state_dict(memory_t_state_dict, strict=strict)
|
|
# set flag for successfully loading the parameters
|
|
self.loaded_teacher = True
|
|
self.teacher.eval()
|
|
self.teacher_obs_normalizer.eval()
|
|
return False # training does not resume
|
|
elif any("student" in key for key in state_dict.keys()): # loading parameters from distillation training
|
|
super().load_state_dict(state_dict, strict=strict)
|
|
# set flag for successfully loading the parameters
|
|
self.loaded_teacher = True
|
|
self.teacher.eval()
|
|
self.teacher_obs_normalizer.eval()
|
|
return True # training resumes
|
|
else:
|
|
raise ValueError("state_dict does not contain student or teacher parameters")
|