Merge branch 'main' of github.com:sisl/InteractionImitation

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(),
+        'action': MinMaxScaler(),
-        'state': SciKitMinMaxScaler(),
+        'state': MinMaxScaler(),
-        'relative_state': SciKitMinMaxScaler(reduce_dim=2),
+        'relative_state': MinMaxScaler(reduce_dim=2),
-        'path_x': SciKitMinMaxScaler(reduce_dim=2),
+        'path_x': MinMaxScaler(reduce_dim=2),
-        'path_y': SciKitMinMaxScaler(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)

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['state'].append(observations[t]['state'][nni])
-                self.raw_data['relative_state'].append(observations[t]['relative_state'][nni.nonzero(),nni.nonzero()].float())
+                self.raw_data['relative_state'].append(observations[t]['relative_state'].index_select(0, 
-                self.raw_data['action'].append(actions[t][nni].float())
+                    nni.nonzero()[:,0]).index_select(1, nni.nonzero()[:,0]))
-                self.raw_data['path_x'].append(observations[t]['paths'][0][nni].float())
+                self.raw_data['action'].append(actions[t][nni])
-                self.raw_data['path_y'].append(observations[t]['paths'][1][nni].float())
+                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['state'] = torch.cat(self.raw_data['state']).type(self.dtype)
-        self.raw_data['action'] = torch.cat(self.raw_data['action'])
+        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'])
+        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'])
+        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']) \

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