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

moving transforms out of dataset class, will be exclusively in policy classes

Browse Source
This commit is contained in:
Arec
2021-07-20 06:16:54 -07:00
parent 2fb5d5e5b1
commit a15e8c29ff
2 changed files with 12 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
src/data_utils.py
View File

@@ -1,7 +1,6 @@
import torch import torch
from torch.utils.data import Dataset, DataLoader from torch.utils.data import Dataset
import numpy as np import numpy as np
#from torchvision import transforms, utils
from src.expert_data import load_expert_data from src.expert_data import load_expert_data
import os import os
opj = os.path.join opj = os.path.join
@@ -15,24 +14,16 @@ class InteractionDatasetMultiAgent(Dataset):
class InteractionDatasetSingleAgent(Dataset): class InteractionDatasetSingleAgent(Dataset):
"""Class to load states and actions for individual agents.""" """Class to load states and actions for individual agents."""
def __init__(self, output_dir='expert_data', loc:int = 0, tracks:list = [0], transforms={}): def __init__(self, output_dir='expert_data', loc:int = 0, tracks:list = [0]):
""" """
Args: Args:
output_dir (string): Directory with all the images. output_dir (string): Directory with all the images.
loc (int): location index loc (int): location index
tracks (list[int]): track indices tracks (list[int]): track indices
transforms (dict): dictionary of transforms to apply to different variables
""" """
self.output_dir = output_dir self.output_dir = output_dir
self.loc = loc self.loc = loc
self.tracks = tracks self.tracks = tracks
self.transforms = transforms
#self.action_transform = transforms.get('action', None)
#self.state_transform = transforms.get('state', None)
#self.relative_state_transform = transforms.get('relative_state', None)
#self.paths_x_transform = transforms.get('paths_x', None)
#self.paths_y_transform = transform.get('paths_y',None)
self._load_dataset() self._load_dataset()
def _load_dataset(self): def _load_dataset(self):
@@ -95,9 +86,4 @@ class InteractionDatasetSingleAgent(Dataset):
""" """
keys = ['state', 'relative_state', 'path_x', 'path_y', 'action'] keys = ['state', 'relative_state', 'path_x', 'path_y', 'action']
sample = {key:self.raw_data[key][idx] for key in keys} sample = {key:self.raw_data[key][idx] for key in keys}
for key in keys:
if key in self.transforms.keys():
sample[key] = self.transforms[key](sample[key])
return sample return sample

20
src/main.py
View File

@@ -37,31 +37,31 @@ def main(method='bc', train=False, test=False, loc=0, **kwargs):
# method-based training # method-based training
if method=='bc': if method=='bc':
from src import bc from src import bc
policy = bc.BehaviorCloningPolicy(transforms=transforms, **kwargs) policy_class = bc.BehaviorCloningPolicy
load_policy = bc.load_policy load_policy_fn = bc.load_policy
metrics = bc.metrics metrics_fn = bc.metrics
train = bc.train train_fn = bc.train
else: else:
raise NotImplementedError raise NotImplementedError
# default train / cv / test split datasets # default train / cv / test split datasets
if train: if train:
train_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[0,1,2], transforms=transforms, train=True) train_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[0,1,2])
cv_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[3], transforms=transforms) cv_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[3])
train(train_dataset, cv_dataset, policy, filestr=filestr, **kwargs) train_fn(train_dataset, cv_dataset, policy_class, filestr=filestr, **kwargs)
if test: if test:
# load policy # load policy
policy = load_policy(filestr=filestr) policy = load_policy_fn(filestr=filestr)
# simulate policy # simulate policy
test_track = 4 test_track = 4
simulate_policy(policy, loc=loc, track=track, filestr=filestr) simulate_policy(policy, loc=loc, track=track, filestr=filestr)
# run test metrics # run test metrics
# test_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[4], transforms=transforms, train=False) # test_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[4])
# metrics(test_datset, policy) # metrics_fn(test_dataset, policy)
def simulate_policy(policy, loc=0, track=0, filestr=''): def simulate_policy(policy, loc=0, track=0, filestr=''):