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

getting hyperparameter tunning with ray tune working. updating default network with optimization and general parameters.

Browse Source
This commit is contained in:
Arec
2021-07-27 14:30:28 -07:00
parent ca80fa19eb
commit 87aa19b86b
6 changed files with 184 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
config/networks.json5
View File

@@ -1,5 +1,5 @@
{
ego_state: {
ego_encoder: {
input_dim: 5, // number of state vars
hidden_n: 0,
hidden_dim: 5,
@@ -30,5 +30,13 @@
hidden_dim: 50,
output_dim: 1, // number of outputs e.g. number of actions, or just one
final_activation: 'sigmoid',
}
},
optim: {
optimizer: 'adam',
lr: 1e-3,
weight_decay: 0.1,
},
train_epochs: 200,
train_batch_size: 32,
loss: 'huber',
}

1
requirements.txt
View File

@@ -5,3 +5,4 @@ pytest
json5
tqdm
tensorboard
ray[tune]

2
src/bc/__init__.py
View File

@@ -1 +1 @@
from src.bc.bc import BehaviorCloningPolicy, train
from src.bc.bc import BehaviorCloningPolicy, train, bc_config

91
src/bc/bc.py
View File

@@ -7,10 +7,12 @@ from torch.utils.tensorboard import SummaryWriter
from src.policies import DeepSetsPolicy
from src.util.transform import MinMaxScaler
from tqdm import tqdm
import json5
from ray import tune
def bc_config(ray_config):
config = {
'ego_state': {'input_dim': 5, 'hidden_n': 0, 'output_dim': 0},
'ego_encoder': {'input_dim': 5, 'hidden_n': 0, 'hidden_dim':0, 'output_dim': 0},
'deepsets': {
'input_dim': 5,
'phi': {
@@ -21,7 +23,7 @@ def bc_config(ray_config):
'rho': {'hidden_n': 0, 'hidden_dim': 10},
'output_dim': 0
},
'path_encoder': {'input_dim': 40, 'hidden_n': 0, 'output_dim': 0},
'path_encoder': {'input_dim': 40, 'hidden_n': 0, 'hidden_dim': 0, 'output_dim': 0},
'head': {
'input_dim': 0, # computed in constructor
'hidden_n': ray_config['head_hidden_n'],
@@ -34,8 +36,8 @@ def bc_config(ray_config):
'lr':ray_config['lr'],
'weight_decay':ray_config['weight_decay']
},
'train_epochs':1000,
'train_batch_size': ray_config['batch_size'],
'train_epochs': 100,
'train_batch_size': ray_config['train_batch_size'],
'loss': ray_config['loss'],
}
@@ -99,7 +101,7 @@ class BehaviorCloningPolicy():
return action
@classmethod
def load_model(cls, config: dict, filestr: str):
def load_model(cls, filestr: str, config: dict = None):
"""
Load a model from a file prefix
Args:
@@ -108,6 +110,9 @@ class BehaviorCloningPolicy():
Returns
model (BehaviorCloningPolicy): loaded model
"""
if not config:
with open(filestr+'_config.json', 'r') as cfg:
config = json5.load(cfg)
transforms = pickle.load(open(filestr+'_transforms.pkl', 'rb'))
model = cls(config, transforms=transforms)
model._policy.load_state_dict(torch.load(filestr+'_model.pt'))
@@ -119,13 +124,17 @@ class BehaviorCloningPolicy():
def parameters(self):
return self._policy.parameters()
def save_model(self, filestr, save_transforms=True):
def save_model(self, filestr, save_config=True, save_transforms=True):
"""
Save transforms and state_dict to a location specificed by filestr
Args:
filestr (str): string prefix to save model to
save_transforms (bool): whether to save transforms
save_config (bool): whether to save the config file (as a json)
save_transforms (bool): whether to save transforms (as a pickle)
"""
if save_config:
with open(filestr+'_config.json', 'w') as cfg:
json5.dump(self._config, cfg)
if save_transforms:
pickle.dump(self._transforms, open(filestr+'_transforms.pkl', 'wb'))
torch.save(self._policy.state_dict(), filestr+'_model.pt')
@@ -148,15 +157,19 @@ def generate_transforms(dataset):
return transforms
def train(train_dataset, cv_dataset, policy, filestr, **kwargs):
def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs):
# hyperparams
train_epochs = 1000
train_batch_size = 64
learning_rate = 1e-3
weight_decay = 0.1
using_ray = kwargs.get('ray', False)
cv_every = 10
# hyperparams
loss_type = config['loss']
train_epochs = config['train_epochs']
train_batch_size = config['train_batch_size']
optimizer_type = config['optim']['optimizer']
learning_rate = config['optim']['lr']
weight_decay = config['optim']['weight_decay']
cv_every = 5
print_epoch_every = 1000
print_cv_every = 1000
checkpoint_every = 100
@@ -176,22 +189,30 @@ def train(train_dataset, cv_dataset, policy, filestr, **kwargs):
policy.policy = policy.policy.type(train_dataset[0]['state'].dtype)
# generate loss function, optimizer
loss_fn = nn.HuberLoss(reduction='sum')
optimizer = torch.optim.Adam(policy.parameters(), lr=learning_rate, weight_decay=weight_decay)
cv_loss_fn = nn.MSELoss(reduction='sum')
if loss_type == 'huber':
loss_fn = nn.HuberLoss(reduction='sum')
elif loss_type == 'mse':
loss_fn = nn.MSELoss(reduction='sum')
else:
raise NotImplementedError
if optimizer_type == 'adam':
optimizer = torch.optim.Adam(policy.parameters(), lr=learning_rate, weight_decay=weight_decay)
else:
raise NotImplementedError
# generate tensorboard writer
writer = SummaryWriter(filestr)
if not using_ray:
writer = SummaryWriter(filestr)
for i in tqdm(range(train_epochs)):
# train
epoch_loss = 0
for (batch_idx, batch) in enumerate(training_loader):
# sample mini-batch and run through policy
pred_action = policy(batch)
if i == 0 and batch_idx==0:
pickle.dump(batch, open(filestr+'_test_batch.pkl', 'wb'))
policy.save_model(filestr)
loss = loss_fn(pred_action, batch['action'])
# compute loss and step optimizer
@@ -202,24 +223,34 @@ def train(train_dataset, cv_dataset, policy, filestr, **kwargs):
epoch_loss += loss.item() / len(train_dataset)
# Write epoch loss
if using_ray:
tune.report(training_loss=epoch_loss, training_iteration=i)
else:
writer.add_scalar('training loss',epoch_loss, i)
writer.add_scalar('training loss',epoch_loss, i)
if i % print_epoch_every == 0:
print('Epoch: {}, Training Loss: {}'.format(i, epoch_loss))
# if i % print_epoch_every == 0:
# print('Epoch: {}, Training Loss: {}'.format(i, epoch_loss))
# measure cv loss
# measure cv loss
if i % cv_every == 0:
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'])
loss = cv_loss_fn(pred_action, batch['action'])
cv_loss += loss.item() / len(cv_dataset)
writer.add_scalar('cv loss', cv_loss, i)
if i % print_cv_every == 0:
print('Epoch: {}, CV Loss: {}'.format(i, cv_loss))
if using_ray:
tune.report(cv_loss=cv_loss, cv_epoch=i)
else:
writer.add_scalar('cv loss', cv_loss, i)
# if i % print_cv_every == 0:
# print('Epoch: {}, CV Loss: {}'.format(i, cv_loss))
# save model checkpoints
if i % checkpoint_every == 0:
policy.save_model(filestr + '_epoch%04i'%(i) )
policy.save_model(filestr)

130
src/main.py
View File

@@ -6,6 +6,7 @@ import json5
import os
opj = os.path.join
from tqdm import tqdm
from functools import partial
from src import InteractionDatasetSingleAgent, metrics
from intersim.utils import get_map_path, get_svt
@@ -20,10 +21,11 @@ def basestr(**kwargs):
"""
return 'base'
def main(method='bc', train=False, test=False, loc=0, config_path=None, **kwargs):
def main(config, method='bc', train=False, test=False, loc=0, datadir='./expert_data', filestr='', **kwargs):
"""
Main loop for training and testing different imitation models
Args:
config (dict): configuration dictionary for model
train (bool): whether to run train loop
test (bool): whether to run test loop
method (str): the method to try for imitation
@@ -35,19 +37,6 @@ def main(method='bc', train=False, test=False, loc=0, config_path=None, **kwargs
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.makedirs(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
@@ -61,9 +50,9 @@ def main(method='bc', train=False, test=False, loc=0, config_path=None, **kwargs
# 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, filestr, **kwargs)
train_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[0])#,1,2])
cv_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[3])
train_fn(config, policy, train_dataset, cv_dataset, filestr, **kwargs)
if test:
@@ -76,7 +65,7 @@ def main(method='bc', train=False, test=False, loc=0, config_path=None, **kwargs
simulate_policy(policy, loc=loc, track=track, filestr=filestr, nframes=500)
# run test metrics
test_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[track])
test_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[track])
metrics(filestr, test_dataset, policy)
@@ -90,8 +79,9 @@ def simulate_policy(policy, loc=0, track=0, filestr='', nframes=float('inf')):
filestr (str): path prefix to save simulation to
"""
# animate from environment
svt, svt_path = get_svt(base='InteractionSimulator', loc=loc, track=track)
osm = get_map_path(base='InteractionSimulator', loc=loc)
basepath = os.path.abspath('./InteractionSimulator')
svt, svt_path = get_svt(base=basepath, loc=loc, track=track)
osm = get_map_path(base=basepath, loc=loc)
env = gym.make('intersim:intersim-v0', svt=svt, map_path=osm,
min_acc=-np.inf, max_acc=np.inf)
# env = gym.make('intersim:intersim-v0', loc=loc, track=track,
@@ -134,12 +124,14 @@ def parse_args():
help='location (default 0)')
parser.add_argument("--train", help="train model",
action="store_true")
parser.add_argument("--all-runs", help="use ray tune to run multiple experiments",
action="store_true")
parser.add_argument("--test", help="test model",
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)
default=None, type=str)
parser.add_argument('--seed', default=0, type=int,
help='seed')
args = parser.parse_args()
@@ -149,11 +141,103 @@ def parse_args():
'method':args.method,
'loc':args.loc,
'config_path':args.config,
'seed':args.seed
'seed':args.seed,
'all_runs':args.all_runs
}
return kwargs
def main_wrapper(**kwargs):
if kwargs['all_runs']:
pass
else:
main(**kwargs)
def get_full_config(ray_config:dict, method:str)->dict:
"""
Get full model configuration from ray config and method string
Args:
ray_config (dict): ray config
method (str): method to get full configuration for
"""
if method == 'bc':
from src.bc import bc_config
config = bc_config(ray_config)
else:
raise NotImplementedError
return config
def get_ray_config(method:str)->dict:
"""
Get configuration for ray based on method.
Args:
method (str): method to get configuration for
Returns:
ray_config (dict): configuration for ray
"""
if method == 'bc':
ray_config = {
"lr": tune.choice([1e-4, 1e-3, 1e-2, 1e-1]),
"weight_decay": tune.choice([0.001, 0.01, 0.1, 0.5, 0.9]),
"loss": tune.choice(['huber', 'mse']),
"train_batch_size": tune.choice([16,32,64]),
"deepsets_phi_hidden_n": tune.choice([1,2,3]),
"deepsets_phi_hidden_dim": tune.choice([16,32,64]),
"deepsets_latent_dim": tune.choice([16,32,64]),
"deepsets_rho_hidden_n": tune.choice([0,1,2]),
"deepsets_rho_hidden_dim": tune.choice([16,32,64]),
"deepsets_output_dim": tune.choice([8,16,32,64]),
"head_hidden_n": tune.choice([0,1,2]),
"head_hidden_dim": tune.choice([16,32,64]),
"head_final_activation": tune.choice(['sigmoid', None]),
}
else:
raise NotImplementedError
return ray_config
if __name__ == '__main__':
kwargs = parse_args()
main(**kwargs)
# make prefix of output files
outdir = opj('output',kwargs['method'],'loc%02i'%(kwargs['loc']))
if kwargs['config_path']:
# load config
with open(kwargs['config_path'], 'r') as cfg:
config = json5.load(cfg)
if not os.path.isdir(outdir):
os.makedirs(outdir)
filestr = opj(outdir, basestr(**kwargs))
main(config, filestr=filestr, **kwargs)
elif kwargs['all_runs'] and kwargs['train']:
def ray_train(config, datadir=None):
full_config = get_full_config(config, kwargs['method'])
main(full_config, filestr='exp', datadir=datadir, ray=True, **kwargs)
# set up ray tune
from ray import tune
from ray.tune.schedulers import ASHAScheduler
datadir = os.path.abspath('./expert_data')
ray_config = get_ray_config(kwargs['method'])
custom_scheduler = ASHAScheduler(
metric='cv_loss',
mode="min",
grace_period=25,
)
analysis = tune.run(
partial(ray_train, datadir=datadir),
config=ray_config,
scheduler=custom_scheduler,
local_dir=outdir,
resources_per_trial={"cpu": 2},
num_samples=20,
)
else:
raise Exception('No valid config found')

8
src/policies/policy.py
View File

@@ -14,13 +14,13 @@ class DeepSetsPolicy(Policy, nn.Module):
config (dict): dictionary for configuring the deep sets policy
"""
super(DeepSetsPolicy, self).__init__()
ego_config = config['ego_state']
ego_config = config['ego_encoder']
deepsets_config = config['deepsets']
pathnet_config = config['path_encoder']
self.ego_net = Phi.from_config(ego_config) if ego_config else lambda x: x
self.deepsets_net = DeepSetsModule.from_config(deepsets_config) if deepsets_config else lambda x: x
self.path_net = Phi.from_config(pathnet_config) if pathnet_config else lambda x: x
self.ego_net = Phi.from_config(ego_config)
self.deepsets_net = DeepSetsModule.from_config(deepsets_config)
self.path_net = Phi.from_config(pathnet_config)
cat_dim = self.ego_net.output_dim + self.deepsets_net.output_dim + self.path_net.output_dim
# head has number of concatenated features as input