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

updating main testing function to use configs and seeds, finishing first pass at behavior cloning policy and training loop. not yet tested

Browse Source
This commit is contained in:
Arec
2021-07-21 08:22:08 -07:00
parent 4350cf8cd5
commit 18b8e0c58f
2 changed files with 104 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

88
src/bc/bc.py
View File

@@ -1,18 +1,28 @@
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.utils.data import DataLoader, RandomSampler
import pickle
from torch.utils.tensorboard import SummaryWriter
from src.policies import
from src.policies import DeepSetsPolicy
from src.util.transform import SciKitMinMaxScaler
import json5
class BehaviorCloningPolicy():
"""
Class for (continuous) behavior cloning policy
"""
def __init__(self, transforms={}, **kwargs):
def __init__(self, config: dict transforms: dict={}):
"""
Initialize BehaviorCloningPolicy
Args:
config (dict): configuration file to initialize DeepSetsPolicy with
transforms (dict): dictionary of transforms to apply to different fields
"""
self._config = config
self._transforms = transforms
self._policy_model = PolicyModel(**kwargs)
self._policy_model = DeepSetsPolicy(config["ego_state"], config["deepsets"], config["path_encoder"], config["head"])
@property
def transforms(self):
@@ -35,31 +45,38 @@ class BehaviorCloningPolicy():
# run observation through transforms
for key in ['state', 'relative_state', 'path_x', 'path_y']:
if key in self._transforms.keys():
ob[key] = self._transforms[key](ob[key])
ob[key] = self._transforms[key].transform(ob[key])
# run transformed state through model
action = self._policy_model(ob)
assert action.ndim == 2, 'action has incorrect shape'
# untransform action
pass
if 'action' in self._transforms.keys():
action = self._transforms['action'].inverse_transform(action)
return action
@classmethod
def load_model(cls, filestr, **kwargs):
def load_model(cls, config: dict, filestr: str):
"""
Load a model from a file prefix
Args:
config (dict): configuration dict to set up model
filestr (str): string prefix to load model from
Returns
model (BehaviorCloningPolicy): loaded model
"""
transforms = pickle.load(open(filestr+'_transforms.pkl', 'rb'))
model = cls(transforms=transforms, **kwargs)
model = cls(config, transforms=transforms)
model._policy_model.load_state_dict(torch.load(filestr+'_model.pt'))
return model
def eval(self):
self._policy_model.eval()
def parameters(self):
return self._policy_model.parameters()
def save_model(self, filestr):
"""
Save transforms and state_dict to a location specificed by filestr
@@ -72,41 +89,72 @@ class BehaviorCloningPolicy():
def generate_transforms(dataset):
"""
Generate transform dictionary from dataset
Args:
dataset (Dataset): dataset of demo observations and actions
"""
pass
transforms = {
'action': SciKitMinMaxScaler()
'state': SciKitMinMaxScaler()
'relative_state': SciKitMinMaxScaler(reduce_dim=2)
'path_x': SciKitMinMaxScaler(reduce_dim=2)
'path_y': SciKitMinMaxScaler(reduce_dim=2)
}
for key in transforms.keys():
transforms[key].fit(dataset[:][key])
def train(train_dataset, cv_dataset, policy_class, filestr=filestr, **kwargs):
return transforms
def train(train_dataset, cv_dataset, policy, filestr, **kwargs):
# hyperparams
train_epochs = 10000
cv_every = 100
train_batch_size = 64
cv_batch_size = 256
cv_batch_size = 256 # doesn't matter
learning_rate = 1e-3
weight_decay=0.1
# generate transform from train_dataset
transforms = generate_transforms(train_dataset)
# initialize policy
policy = BehaviorCloningPolicy(transforms=transforms, **kwargs)
policy.transforms = transforms
# training and testing dataloaders
training_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True)
cv_loader = DataLoader(cv_dataset, batch_size=cv_batch_size, shuffle=True)
# generate loss function, optimizer
loss_fn = nn.HuberLoss(reduction='sum')
optimizer = torch.optim.Adam(policy.parameters(), lr=learning_rate, weight_decay=weight_decay)
for i in train_epochs:
# sample mini-batch and run through policy
pred_action = policy(batch)
epoch_loss = 0
for (batch_idx, batch) in enumerate(training_loader):
# sample mini-batch and run through policy
pred_action = policy(batch)
loss = loss_fn(pred_action, batch['action'])
# compute loss and step optimizer
loss.backwards()
optimizer.step()
# compute loss and step optimizer
optimizer.zero_grad()
loss.backwards()
optimizer.step()
epoch_loss += loss.item() / len(train_dataset)
# measure L2 on every cv epoch
# Write epoch loss
if i % 10 == 0:
print('Epoch: {}, Training Loss: {}'.format(i, epoch_loss))
# measure cv loss
if i % cv_every == 0:
pass
with torch.no_grad():
cv_loss = 0.
for (batch_idx, batch) in enumerate(cv_loader):
pred_action = policy(batch)
loss = loss_fn(pred_action, batch['action'])
cv_loss += loss.item() / len(cv_dataset)
print('Epoch: {}, CV Loss: {}'.format(i, cv_loss))
policy.save_model(filestr)

38
src/main.py
View File

@@ -2,8 +2,11 @@ import torch
import gym
import intersim
import numpy as np
import json5
import os
opj = os.path.join
from src import InteractionDatasetSingleAgent, metrics
def basestr(**kwargs):
@@ -16,7 +19,7 @@ def basestr(**kwargs):
"""
return 'base_'
def main(method='bc', train=False, test=False, loc=0, **kwargs):
def main(method='bc', train=False, test=False, loc=0, config_path=None, **kwargs):
"""
Main loop for training and testing different imitation models
Args:
@@ -24,13 +27,26 @@ def main(method='bc', train=False, test=False, loc=0, **kwargs):
test (bool): whether to run test loop
method (str): the method to try for imitation
loc (int): the location index of the roundabout
kwargs (dict): remaining kwargs for policy and training loop
config_file (str): path to config file
kwargs (dict): remaining kwargs for training loop
"""
# get/set seed
seed = kwargs.get('seed',0)
torch.manual_seed(seed)
# make prefix of output files
outdir = opj('output',method,'loc%02i'%(loc))
if not os.path.isdir(outdir):
os.mkdir(outdir)
filestr = opj(outdir, basestr(**kwargs))
# load config
if config_path:
with open(config_path, 'r') as cfg:
config = json5.load(cfg)
else:
raise Exception('No config path specified')
# method-based training
if method=='bc':
from src import bc
@@ -41,14 +57,17 @@ def main(method='bc', train=False, test=False, loc=0, **kwargs):
# default train / cv / test split datasets
if train:
# make policy, train and test datasets, and send to
policy = policy_class(config)
train_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[0,1,2])
cv_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[3])
train_fn(train_dataset, cv_dataset, policy_class, filestr=filestr, **kwargs)
train_fn(train_dataset, cv_dataset, policy, filestr, **kwargs)
if test:
# load policy
policy = policy_class.load_model(filestr=filestr, **kwargs)
policy = policy_class.load_model(config, filestr)
policy.eval()
# simulate policy
@@ -77,6 +96,7 @@ def simulate_policy(policy, loc=0, track=0, filestr=''):
env.render()
done = False
while not done:
# get action
action = policy(ob)
@@ -95,6 +115,8 @@ def parse_args():
test (bool): whether to run test loop
method (str): the method to try for imitation
loc (int): the location index of the roundabout
config (str): config path
seed (int): RNG seed
"""
import argparse
parser = argparse.ArgumentParser(description='Save Expert Trajectories')
@@ -106,6 +128,10 @@ def parse_args():
action="store_true")
parser.add_argument("--method", help="modeling method",
choices=['bc', 'gail', 'advil'], default='bc')
parser.add_argument("--config", help="config file path",
default='config/networks.json5', type=str)
parser.add_argument('--seed', default=0, type=int,
help='seed')
parser.add_argument()
parser.add_argument()
args = parser.parse_args()
@@ -113,7 +139,9 @@ def parse_args():
'train'=args.train,
'test'=args.test,
'method'=args.method,
'loc'=args.loc
'loc'=args.loc,
'config'=args.config,
'seed'=args.seed
}
return kwargs