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

@@ -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/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