Spaces:
Running
Running
| import torch | |
| import torch.nn.functional as F | |
| import numpy | |
| from torch_ac.utils import DictList | |
| # dictionary that defines what head is required for each extra info used for auxiliary supervision | |
| required_heads = {'seen_state': 'binary', | |
| 'see_door': 'binary', | |
| 'see_obj': 'binary', | |
| 'obj_in_instr': 'binary', | |
| 'in_front_of_what': 'multiclass9', # multi class classifier with 9 possible classes | |
| 'visit_proportion': 'continuous01', # continous regressor with outputs in [0, 1] | |
| 'bot_action': 'binary' | |
| } | |
| class ExtraInfoCollector: | |
| ''' | |
| This class, used in rl.algos.base, allows connecting the extra information from the environment, and the | |
| corresponding predictions using the specific heads in the model. It transforms them so that they are easy to use | |
| to evaluate losses | |
| ''' | |
| def __init__(self, aux_info, shape, device): | |
| self.aux_info = aux_info | |
| self.shape = shape | |
| self.device = device | |
| self.collected_info = dict() | |
| self.extra_predictions = dict() | |
| for info in self.aux_info: | |
| self.collected_info[info] = torch.zeros(*shape, device=self.device) | |
| if required_heads[info] == 'binary' or required_heads[info].startswith('continuous'): | |
| # we predict one number only | |
| self.extra_predictions[info] = torch.zeros(*shape, 1, device=self.device) | |
| elif required_heads[info].startswith('multiclass'): | |
| # means that this is a multi-class classification and we need to predict the whole proba distr | |
| n_classes = int(required_heads[info].replace('multiclass', '')) | |
| self.extra_predictions[info] = torch.zeros(*shape, n_classes, device=self.device) | |
| else: | |
| raise ValueError("{} not supported".format(required_heads[info])) | |
| def process(self, env_info): | |
| # env_info is now a tuple of dicts | |
| env_info = [{k: v for k, v in dic.items() if k in self.aux_info} for dic in env_info] | |
| env_info = {k: [env_info[_][k] for _ in range(len(env_info))] for k in env_info[0].keys()} | |
| # env_info is now a dict of lists | |
| return env_info | |
| def fill_dictionaries(self, index, env_info, extra_predictions): | |
| for info in self.aux_info: | |
| dtype = torch.long if required_heads[info].startswith('multiclass') else torch.float | |
| self.collected_info[info][index] = torch.tensor(env_info[info], dtype=dtype, device=self.device) | |
| self.extra_predictions[info][index] = extra_predictions[info] | |
| def end_collection(self, exps): | |
| collected_info = dict() | |
| extra_predictions = dict() | |
| for info in self.aux_info: | |
| # T x P -> P x T -> P * T | |
| collected_info[info] = self.collected_info[info].transpose(0, 1).reshape(-1) | |
| if required_heads[info] == 'binary' or required_heads[info].startswith('continuous'): | |
| # T x P x 1 -> P x T x 1 -> P * T | |
| extra_predictions[info] = self.extra_predictions[info].transpose(0, 1).reshape(-1) | |
| elif type(required_heads[info]) == int: | |
| # T x P x k -> P x T x k -> (P * T) x k | |
| k = required_heads[info] # number of classes | |
| extra_predictions[info] = self.extra_predictions[info].transpose(0, 1).reshape(-1, k) | |
| # convert the dicts to DictLists, and add them to the exps DictList. | |
| exps.collected_info = DictList(collected_info) | |
| exps.extra_predictions = DictList(extra_predictions) | |
| return exps | |
| class SupervisedLossUpdater: | |
| ''' | |
| This class, used by PPO, allows the evaluation of the supervised loss when using extra information from the | |
| environment. It also handles logging accuracies/L2 distances/etc... | |
| ''' | |
| def __init__(self, aux_info, supervised_loss_coef, recurrence, device): | |
| self.aux_info = aux_info | |
| self.supervised_loss_coef = supervised_loss_coef | |
| self.recurrence = recurrence | |
| self.device = device | |
| self.log_supervised_losses = [] | |
| self.log_supervised_accuracies = [] | |
| self.log_supervised_L2_losses = [] | |
| self.log_supervised_prevalences = [] | |
| self.batch_supervised_loss = 0 | |
| self.batch_supervised_accuracy = 0 | |
| self.batch_supervised_L2_loss = 0 | |
| self.batch_supervised_prevalence = 0 | |
| def init_epoch(self): | |
| self.log_supervised_losses = [] | |
| self.log_supervised_accuracies = [] | |
| self.log_supervised_L2_losses = [] | |
| self.log_supervised_prevalences = [] | |
| def init_batch(self): | |
| self.batch_supervised_loss = 0 | |
| self.batch_supervised_accuracy = 0 | |
| self.batch_supervised_L2_loss = 0 | |
| self.batch_supervised_prevalence = 0 | |
| def eval_subbatch(self, extra_predictions, sb): | |
| supervised_loss = torch.tensor(0., device=self.device) | |
| supervised_accuracy = torch.tensor(0., device=self.device) | |
| supervised_L2_loss = torch.tensor(0., device=self.device) | |
| supervised_prevalence = torch.tensor(0., device=self.device) | |
| binary_classification_tasks = 0 | |
| classification_tasks = 0 | |
| regression_tasks = 0 | |
| for pos, info in enumerate(self.aux_info): | |
| coef = self.supervised_loss_coef[pos] | |
| pred = extra_predictions[info] | |
| target = dict.__getitem__(sb.collected_info, info) | |
| if required_heads[info] == 'binary': | |
| binary_classification_tasks += 1 | |
| classification_tasks += 1 | |
| supervised_loss += coef * F.binary_cross_entropy_with_logits(pred.reshape(-1), target) | |
| supervised_accuracy += ((pred.reshape(-1) > 0).float() == target).float().mean() | |
| supervised_prevalence += target.mean() | |
| elif required_heads[info].startswith('continuous'): | |
| regression_tasks += 1 | |
| mse = F.mse_loss(pred.reshape(-1), target) | |
| supervised_loss += coef * mse | |
| supervised_L2_loss += mse | |
| elif required_heads[info].startswith('multiclass'): | |
| classification_tasks += 1 | |
| supervised_accuracy += (pred.argmax(1).float() == target).float().mean() | |
| supervised_loss += coef * F.cross_entropy(pred, target.long()) | |
| else: | |
| raise ValueError("{} not supported".format(required_heads[info])) | |
| if binary_classification_tasks > 0: | |
| supervised_prevalence /= binary_classification_tasks | |
| else: | |
| supervised_prevalence = torch.tensor(-1) | |
| if classification_tasks > 0: | |
| supervised_accuracy /= classification_tasks | |
| else: | |
| supervised_accuracy = torch.tensor(-1) | |
| if regression_tasks > 0: | |
| supervised_L2_loss /= regression_tasks | |
| else: | |
| supervised_L2_loss = torch.tensor(-1) | |
| self.batch_supervised_loss += supervised_loss.item() | |
| self.batch_supervised_accuracy += supervised_accuracy.item() | |
| self.batch_supervised_L2_loss += supervised_L2_loss.item() | |
| self.batch_supervised_prevalence += supervised_prevalence.item() | |
| return supervised_loss | |
| def update_batch_values(self): | |
| self.batch_supervised_loss /= self.recurrence | |
| self.batch_supervised_accuracy /= self.recurrence | |
| self.batch_supervised_L2_loss /= self.recurrence | |
| self.batch_supervised_prevalence /= self.recurrence | |
| def update_epoch_logs(self): | |
| self.log_supervised_losses.append(self.batch_supervised_loss) | |
| self.log_supervised_accuracies.append(self.batch_supervised_accuracy) | |
| self.log_supervised_L2_losses.append(self.batch_supervised_L2_loss) | |
| self.log_supervised_prevalences.append(self.batch_supervised_prevalence) | |
| def end_training(self, logs): | |
| logs["supervised_loss"] = numpy.mean(self.log_supervised_losses) | |
| logs["supervised_accuracy"] = numpy.mean(self.log_supervised_accuracies) | |
| logs["supervised_L2_loss"] = numpy.mean(self.log_supervised_L2_losses) | |
| logs["supervised_prevalence"] = numpy.mean(self.log_supervised_prevalences) | |
| return logs | |