diff --git a/config/networks.json5 b/config/networks.json5 index 46b4f0d..4471a35 100644 --- a/config/networks.json5 +++ b/config/networks.json5 @@ -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', } \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 8045763..d757a17 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,3 +5,4 @@ pytest json5 tqdm tensorboard +ray[tune] \ No newline at end of file diff --git a/src/bc/__init__.py b/src/bc/__init__.py index 54f5c8d..52cdb6d 100644 --- a/src/bc/__init__.py +++ b/src/bc/__init__.py @@ -1 +1 @@ -from src.bc.bc import BehaviorCloningPolicy, train +from src.bc.bc import BehaviorCloningPolicy, train, bc_config diff --git a/src/bc/bc.py b/src/bc/bc.py index 933e0df..63b0716 100644 --- a/src/bc/bc.py +++ b/src/bc/bc.py @@ -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) diff --git a/src/main.py b/src/main.py index 6cae3f4..64dd944 100644 --- a/src/main.py +++ b/src/main.py @@ -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) \ No newline at end of file + + # 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') + + + + + + \ No newline at end of file diff --git a/src/policies/policy.py b/src/policies/policy.py index dad8f92..524345e 100644 --- a/src/policies/policy.py +++ b/src/policies/policy.py @@ -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