huangfu/InteractionImitation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Johannes Fischer
2021-07-23 09:30:26 +02:00
parent 9b0aeb182f 563a2cfbd4
commit 5d3b3fa1bc
3 changed files with 72 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
src/bc/bc.py
View File

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

26
src/data_utils.py
View File

@@ -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']) \

51
src/util/transform.py
View File

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