From 08eb8988120d5db677258803280494020264f6f0 Mon Sep 17 00:00:00 2001 From: Arec Date: Thu, 22 Jul 2021 12:06:42 -0700 Subject: [PATCH 1/2] adding dtypes and fixing matrix indexing --- src/data_utils.py | 26 ++++++++++++++------------ 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/src/data_utils.py b/src/data_utils.py index 84f6ba3..2644bdb 100644 --- a/src/data_utils.py +++ b/src/data_utils.py @@ -14,7 +14,7 @@ class InteractionDatasetMultiAgent(Dataset): class InteractionDatasetSingleAgent(Dataset): """Class to load states and actions for individual agents.""" - def __init__(self, output_dir='expert_data', loc:int = 0, tracks:list = [0]): + def __init__(self, output_dir='expert_data', loc:int = 0, tracks:list = [0], dtype=torch.float32): """ Args: output_dir (string): Directory with all the images. @@ -24,6 +24,7 @@ class InteractionDatasetSingleAgent(Dataset): self.output_dir = output_dir self.loc = loc self.tracks = tracks + self.dtype = dtype self._load_dataset() def _load_dataset(self): @@ -43,24 +44,25 @@ class InteractionDatasetSingleAgent(Dataset): for t in range(T): nni = ~torch.isnan(observations[t]['state'][:,0]) max_nv = max(max_nv,nni.count_nonzero()) - self.raw_data['state'].append(observations[t]['state'][nni].float()) - self.raw_data['relative_state'].append(observations[t]['relative_state'][nni.nonzero(),nni.nonzero()].float()) - self.raw_data['action'].append(actions[t][nni].float()) - self.raw_data['path_x'].append(observations[t]['paths'][0][nni].float()) - self.raw_data['path_y'].append(observations[t]['paths'][1][nni].float()) + self.raw_data['state'].append(observations[t]['state'][nni]) + self.raw_data['relative_state'].append(observations[t]['relative_state'].index_select(0, + nni.nonzero()[:,0]).index_select(1, nni.nonzero()[:,0])) + self.raw_data['action'].append(actions[t][nni]) + self.raw_data['path_x'].append(observations[t]['paths'][0][nni]) + self.raw_data['path_y'].append(observations[t]['paths'][1][nni]) # cat lists - self.raw_data['state'] = torch.cat(self.raw_data['state']) - self.raw_data['action'] = torch.cat(self.raw_data['action']) - self.raw_data['path_x'] = torch.cat(self.raw_data['path_x']) - self.raw_data['path_y'] = torch.cat(self.raw_data['path_y']) + self.raw_data['state'] = torch.cat(self.raw_data['state']).type(self.dtype) + self.raw_data['action'] = torch.cat(self.raw_data['action']).type(self.dtype) + self.raw_data['path_x'] = torch.cat(self.raw_data['path_x']).type(self.dtype) + self.raw_data['path_y'] = torch.cat(self.raw_data['path_y']).type(self.dtype) # pad second dimension of relative state for i in range(len(self.raw_data['relative_state'])): nv1, nv2, d = self.raw_data['relative_state'][i].shape - pad = torch.zeros(nv1, max_nv-nv2, d) * np.nan + pad = torch.zeros(nv1, max_nv-nv2, d, dtype=self.dtype) * np.nan self.raw_data['relative_state'][i] = torch.cat((self.raw_data['relative_state'][i], pad), dim=1) - self.raw_data['relative_state'] = torch.cat(self.raw_data['relative_state']) + self.raw_data['relative_state'] = torch.cat(self.raw_data['relative_state']).type(self.dtype) # mandate equal length assert len(self.raw_data['state']) == len(self.raw_data['relative_state']) \ From 563a2cfbd4266a95fe1c4acca19ef679be738ade Mon Sep 17 00:00:00 2001 From: Arec Date: Thu, 22 Jul 2021 12:37:36 -0700 Subject: [PATCH 2/2] making a differentiable transform for use for pytorch, making sure the fitting function treats nans properly while fitting. next issue: forward pass is returning nans --- src/bc/bc.py | 16 +++++++------- src/util/transform.py | 51 ++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 58 insertions(+), 9 deletions(-) diff --git a/src/bc/bc.py b/src/bc/bc.py index 207f528..4c0e120 100644 --- a/src/bc/bc.py +++ b/src/bc/bc.py @@ -1,11 +1,11 @@ import torch import torch.nn as nn -from torch.utils.data import DataLoader, RandomSampler +from torch.utils.data import DataLoader import pickle #from torch.utils.tensorboard import SummaryWriter from src.policies import DeepSetsPolicy -from src.util.transform import SciKitMinMaxScaler +from src.util.transform import MinMaxScaler import json5 class BehaviorCloningPolicy(): @@ -101,11 +101,11 @@ def generate_transforms(dataset): dataset (Dataset): dataset of demo observations and actions """ transforms = { - 'action': SciKitMinMaxScaler(), - 'state': SciKitMinMaxScaler(), - 'relative_state': SciKitMinMaxScaler(reduce_dim=2), - 'path_x': SciKitMinMaxScaler(reduce_dim=2), - 'path_y': SciKitMinMaxScaler(reduce_dim=2), + 'action': MinMaxScaler(), + 'state': MinMaxScaler(), + 'relative_state': MinMaxScaler(reduce_dim=2), + 'path_x': MinMaxScaler(reduce_dim=2), + 'path_y': MinMaxScaler(reduce_dim=2), } for key in transforms.keys(): transforms[key].fit(dataset[:][key]) @@ -151,7 +151,7 @@ def train(train_dataset, cv_dataset, policy, filestr, **kwargs): # compute loss and step optimizer optimizer.zero_grad() - loss.backwards() + loss.backward() optimizer.step() epoch_loss += loss.item() / len(train_dataset) diff --git a/src/util/transform.py b/src/util/transform.py index acc55de..0b244e9 100644 --- a/src/util/transform.py +++ b/src/util/transform.py @@ -1,6 +1,6 @@ import torch from torch import nn - +import numpy as np from sklearn import preprocessing class Transform(nn.Module): @@ -39,6 +39,55 @@ class Transform(nn.Module): def forward(self, X): return self.transform(X) +class MinMaxScaler(Transform): + """ + Scale tensor so each feature is in [0, 1] + """ + def __init__(self, reduce_dim:int=None): + """ + Initialize SciKitTransform + Args: + reduce_dim (int): dimension to start calculating featues from + e.g. with reduce_dim=2, (A, B, C, D, E) will be reshaped to (A*B, C*D*E) + """ + self.reduce_dim = reduce_dim + super(MinMaxScaler, self).__init__() + + def fit(self, X): + nd = X.ndim + if self.reduce_dim: + self.nfeatures = int(torch.tensor(X.shape[self.reduce_dim:]).prod()) + else: + assert nd==2, 'Invalid ndim' + self.nfeatures = X.shape[1] + + X = X.reshape((-1,self.nfeatures)) + nans = torch.isnan(X) + X[nans] = float('inf') + self.min = X.min(0,keepdims=True)[0] + + X[nans] = -float('inf') + self.span = X.max(0,keepdims=True)[0] - self.min + + X[nans] = np.nan + + def transform(self, X): + + assert hasattr(self, 'min') and hasattr(self, 'span'), 'Model not yet fit' + shape = X.shape + X = X.reshape((-1,self.nfeatures)) + t = (X - self.min) / self.span + return t.reshape(shape) + + def inverse_transform(self, X): + + assert hasattr(self, 'min') and hasattr(self, 'span'), 'Model not yet fit' + shape = X.shape + X = X.reshape((-1,self.nfeatures)) + it = X * self.span + self.min + return it.reshape(shape) + + class SciKitTransform(Transform): """ Wrappers around scikit-learn transforms