diff --git a/checkpoints/bc-intersimple-setobs2.pt b/checkpoints/bc-intersimple-setobs2.pt deleted file mode 100644 index 944cd70..0000000 Binary files a/checkpoints/bc-intersimple-setobs2.pt and /dev/null differ diff --git a/checkpoints/gail-intersimple-setobs2-03-02-22.pt b/checkpoints/gail-intersimple-setobs2-03-02-22.pt deleted file mode 100644 index 1a94547..0000000 Binary files a/checkpoints/gail-intersimple-setobs2-03-02-22.pt and /dev/null differ diff --git a/checkpoints/gail-options-setobs2-15-02-2022.pt b/checkpoints/gail-options-setobs2-15-02-2022.pt deleted file mode 100644 index 76f70af..0000000 Binary files a/checkpoints/gail-options-setobs2-15-02-2022.pt and /dev/null differ diff --git a/checkpoints/gail-options-setobs2-Feb15_18-49-05.pt b/checkpoints/gail-options-setobs2-Feb15_18-49-05.pt deleted file mode 100644 index 76f70af..0000000 Binary files a/checkpoints/gail-options-setobs2-Feb15_18-49-05.pt and /dev/null differ diff --git a/checkpoints/gail-ppo-intersimple-setobs2.pt b/checkpoints/gail-ppo-intersimple-setobs2.pt deleted file mode 100644 index 800c679..0000000 Binary files a/checkpoints/gail-ppo-intersimple-setobs2.pt and /dev/null differ diff --git a/checkpoints/gail-ppo-options-setobs2-Feb15_22-05-38.pt b/checkpoints/gail-ppo-options-setobs2-Feb15_22-05-38.pt deleted file mode 100644 index a2bd8b6..0000000 Binary files a/checkpoints/gail-ppo-options-setobs2-Feb15_22-05-38.pt and /dev/null differ diff --git a/checkpoints/sgail-options-setobs2-Feb21_13-30-45.pt b/checkpoints/sgail-options-setobs2-Feb21_13-30-45.pt deleted file mode 100644 index 4f02fe6..0000000 Binary files a/checkpoints/sgail-options-setobs2-Feb21_13-30-45.pt and /dev/null differ diff --git a/checkpoints/sgail-ppo-options-setobs2-17-02-2022.pt b/checkpoints/sgail-ppo-options-setobs2-17-02-2022.pt deleted file mode 100644 index 33df960..0000000 Binary files a/checkpoints/sgail-ppo-options-setobs2-17-02-2022.pt and /dev/null differ diff --git a/checkpoints/sgail-ppo-options-setobs2-Feb18_12-53-23.pt b/checkpoints/sgail-ppo-options-setobs2-Feb18_12-53-23.pt deleted file mode 100644 index 76cfb53..0000000 Binary files a/checkpoints/sgail-ppo-options-setobs2-Feb18_12-53-23.pt and /dev/null differ diff --git a/checkpoints/sgail-ppo-options-setobs2-Feb18_16-25-08.pt b/checkpoints/sgail-ppo-options-setobs2-Feb18_16-25-08.pt deleted file mode 100644 index 7013a8d..0000000 Binary files a/checkpoints/sgail-ppo-options-setobs2-Feb18_16-25-08.pt and /dev/null differ diff --git a/checkpoints/wgail-options-setobs2-Feb16_01-06-27.pt b/checkpoints/wgail-options-setobs2-Feb16_01-06-27.pt deleted file mode 100644 index ac6db8c..0000000 Binary files a/checkpoints/wgail-options-setobs2-Feb16_01-06-27.pt and /dev/null differ diff --git a/checkpoints/wgail-ppo-options-setobs2-Feb16_04-02-56.pt b/checkpoints/wgail-ppo-options-setobs2-Feb16_04-02-56.pt deleted file mode 100644 index 255d031..0000000 Binary files a/checkpoints/wgail-ppo-options-setobs2-Feb16_04-02-56.pt and /dev/null differ diff --git a/config/networks.json5 b/config/networks.json5 deleted file mode 100644 index ebce18f..0000000 --- a/config/networks.json5 +++ /dev/null @@ -1,42 +0,0 @@ -{ - ego_encoder: { - input_dim: 5, // number of state vars - hidden_n: 0, - hidden_dim: 5, - output_dim: 5 - }, - deepsets: { - input_dim: 6, // number of relative state vars for others - phi: { - hidden_n: 2, - hidden_dim: 20, - }, - latent_dim: 20, - rho: { - hidden_n: 2, - hidden_dim: 10, - }, - output_dim: 10 - }, - path_encoder: { - input_dim: 40, // 2 * path length for (x,y) coordinates - hidden_n: 0, - hidden_dim: 20, - output_dim: 10, - }, - head: { - input_dim: 0, // computed in policy constructor - hidden_n: 3, - 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/config/value_dice.json5 b/config/value_dice.json5 deleted file mode 100644 index 439eef0..0000000 --- a/config/value_dice.json5 +++ /dev/null @@ -1,85 +0,0 @@ -{ - policy_net: { - ego_encoder: { - input_dim: 5, // number of state vars - hidden_n: 0, - hidden_dim: 5, - output_dim: 5 - }, - deepsets: { - input_dim: 6, // number of relative state vars for others - phi: { - hidden_n: 2, - hidden_dim: 20, - }, - latent_dim: 20, - rho: { - hidden_n: 2, - hidden_dim: 10, - }, - output_dim: 10 - }, - path_encoder: { - input_dim: 40, // 2 * path length for (x,y) coordinates - hidden_n: 0, - hidden_dim: 20, - output_dim: 10, - }, - head: { - input_dim: 0, // computed in policy constructor - hidden_n: 3, - hidden_dim: 50, - output_dim: 1, // number of outputs e.g. number of actions, or just one - final_activation: 'sigmoid', - }, - }, - value_net: { - ego_encoder: { - input_dim: 5, // number of state vars - hidden_n: 0, - hidden_dim: 5, - output_dim: 5 - }, - deepsets: { - input_dim: 6, // number of relative state vars for others - phi: { - hidden_n: 2, - hidden_dim: 20, - }, - latent_dim: 20, - rho: { - hidden_n: 2, - hidden_dim: 10, - }, - output_dim: 10 - }, - path_encoder: { - input_dim: 40, // 2 * path length for (x,y) coordinates - hidden_n: 0, - hidden_dim: 20, - output_dim: 10, - }, - action_dim: 1, // number of actions - head: { - input_dim: 0, // computed in policy constructor - hidden_n: 3, - hidden_dim: 50, - output_dim: 1, // number of outputs e.g. number of actions, or just one - final_activation: 'id', - }, - }, - policy_optim: { - optimizer: 'adam', - lr: 1e-3, - weight_decay: 0.1, - }, - value_optim: { - optimizer: 'adam', - lr: 1e-3, - weight_decay: 0.1, - }, - train_epochs: 200, - train_batch_size: 32, - discount: 0.95, - clip_grad_norm: 1., -} \ No newline at end of file diff --git a/experiments/experiment.py b/experiments/experiment.py deleted file mode 100644 index bcddb06..0000000 --- a/experiments/experiment.py +++ /dev/null @@ -1,203 +0,0 @@ -import json5 -from functools import partial -import os -opj = os.path.join - -# set up ray tune -import ray -from ray import tune -from ray.tune import Analysis, ExperimentAnalysis -from ray.tune.schedulers import ASHAScheduler -from hyperopt import hp -from ray.tune.suggest.hyperopt import HyperOptSearch - -# get graphs -import intersim -from intersim.graphs import ConeVisibilityGraph - - -from src.main import basestr, main - -def parse_args(): - """ - Parse arguments to main - Returns: - kwargs: dictionary of arguments: - train (bool): whether to run train loop - 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') - parser.add_argument('--loc', default=0, type=int, - help='location (default 0)') - parser.add_argument("--train", help="train model", - action="store_true") - parser.add_argument("--ray", 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', 'vd'], default='bc') - parser.add_argument("--config", help="config file path", - default=None, type=str) - parser.add_argument('--seed', default=0, type=int, - help='seed') - parser.add_argument('--nframes', default=500, type=int, - help='frames for test animation') - parser.add_argument('--nsamples', default=200, type=int, - help='number of ray samples') - parser.add_argument('--graph', action='store_true', - help='whether to mask the relative states based on a ConeVisibilityGraph') - parser.add_argument('-d', default='./expert_data', type=str, - help='data directory') - parser.add_argument('-o', default=None, type=str, - help='output directory') - args = parser.parse_args() - kwargs = { - 'train':args.train, - 'test':args.test, - 'method':args.method, - 'loc':args.loc, - 'config_path':args.config, - 'seed':args.seed, - 'ray':args.ray, - 'nframes':args.nframes, - 'nsamples':args.nsamples, - 'datadir':os.path.abspath(args.d), - 'graph':None, - 'outdir': opj('output',args.method,'loc%02i'%(args.loc)), - 'train_tracks':[0,1,2], - 'cv_tracks':[3], - 'test_tracks':[4], - } - if args.o: - kwargs['outdir'] = args.o - if args.graph: - kwargs['graph'] = ConeVisibilityGraph(r=20, half_angle=120) - return 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) - elif method == 'vd': - from src.value_dice import vd_config - config = vd_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.loguniform(1e-5, 1e-3), - "weight_decay": tune.choice([0, 0.1]), - "loss": tune.choice(['huber', 'mse']), - "train_batch_size": tune.choice([16,32,64]), - "deepsets_phi_hidden_n": tune.randint(1,5), - "deepsets_phi_hidden_dim": tune.lograndint(8,65), - "deepsets_latent_dim": tune.lograndint(8,129), - "deepsets_rho_hidden_n": tune.randint(0,3), - "deepsets_rho_hidden_dim": tune.lograndint(8,129), - "deepsets_output_dim": tune.lograndint(4,129), - "head_hidden_n": tune.randint(1,6), - "head_hidden_dim": tune.lograndint(16,257), - "head_final_activation": tune.choice(['sigmoid', None]), - } - elif method == 'vd': - ray_config = { - "policy_lr": tune.loguniform(1e-5, 1e-3), - "value_lr": tune.loguniform(1e-5, 1e-3), - "policy_weight_decay": tune.choice([0, 0.1]), - "value_weight_decay": tune.choice([0, 0.1]), - "train_batch_size": tune.choice([16,32,64]), - "deepsets_phi_hidden_n": tune.randint(1,5), - "deepsets_phi_hidden_dim": tune.lograndint(8,65), - "deepsets_latent_dim": tune.lograndint(8,129), - "deepsets_rho_hidden_n": tune.randint(0,3), - "deepsets_rho_hidden_dim": tune.lograndint(8,129), - "deepsets_output_dim": tune.lograndint(4,129), - "head_hidden_n": tune.randint(1,6), - "head_hidden_dim": tune.lograndint(16,257), - "head_final_activation": tune.choice(['sigmoid', None]), - "clip_grad_norm": tune.choice([.5, 1., 5., 10.]), - "discount": tune.choice([.95, .99]) - } - else: - raise NotImplementedError - return ray_config - -if __name__ == '__main__': - kwargs = parse_args() - - # make prefix of output files - - if kwargs['config_path']: - # load config - with open(kwargs['config_path'], 'r') as cfg: - config = json5.load(cfg) - if not os.path.isdir(kwargs['outdir']): - os.makedirs(kwargs['outdir']) - filestr = opj(kwargs['outdir'], basestr(**kwargs)) - if kwargs['ray']: - filestr = kwargs['config_path'].replace('_config.json','') - main(config, filestr=filestr, **kwargs) - - elif kwargs['ray'] and kwargs['train']: - - ray.shutdown() - ray.init(log_to_driver=False) - - def ray_train(config, datadir=None): - full_config = get_full_config(config, kwargs['method']) - main(full_config, filestr='exp', **kwargs) - - ray_config = get_ray_config(kwargs['method']) - search = HyperOptSearch(ray_config, max_concurrent=8, metric='cv_loss',mode="min",) - custom_scheduler = ASHAScheduler(metric='cv_loss', mode="min", grace_period=15) - - analysis = tune.run( - ray_train, - #config=ray_config, - search_alg=search, - scheduler=custom_scheduler, - local_dir=kwargs['outdir'], - #resources_per_trial={"cpu": 2}, - time_budget_s=120*60, - num_samples=kwargs['nsamples'], - ) - elif kwargs['ray'] and kwargs['test']: - analysis = Analysis(kwargs['outdir'], default_metric="cv_loss", default_mode="min") - config = analysis.get_best_config() - filepath = analysis.get_best_logdir() - filestr = opj(filepath, 'exp') - config_path = filestr+'_config.json' - with open(config_path, 'r') as cfg: - config = json5.load(cfg) - print("Best ray experiment:", filepath) - main(config, filestr=filestr, **kwargs) - else: - raise Exception('No valid config found') - - - - - - diff --git a/experiments/experiments.sh b/experiments/experiments.sh deleted file mode 100755 index 9efc552..0000000 --- a/experiments/experiments.sh +++ /dev/null @@ -1,9 +0,0 @@ -#!/bin/sh - -python experiments/experiment.py --ray --train -d ./expert_data/base -python experiments/experiment.py --ray --test -d ./expert_data/base --nframes 1000 -python experiments/experiment.py --ray --train -d ./expert_data/reg -python experiments/experiment.py --ray --test -d ./expert_data/reg --nframes 1000 -python experiments/experiment.py --ray --train -d ./expert_data/reg_graph --graph -python experiments/experiment.py --ray --test -d ./expert_data/reg_graph --graph --nframes 1000 - diff --git a/experiments/train_vd.sh b/experiments/train_vd.sh deleted file mode 100755 index 13bdc14..0000000 --- a/experiments/train_vd.sh +++ /dev/null @@ -1,5 +0,0 @@ -#!/bin/sh - -# python experiments/experiment.py --method vd --train --ray -d expert_data/reg -o output/vd/loc00/reg --nsamples 400 -# python experiments/experiment.py --test --ray --method vd -d expert_data/normal -o output/vd/loc00/normal --nframes 1000 -python experiments/experiment.py --train --method vd --config config/value_dice.json5 \ No newline at end of file diff --git a/src/bc/__init__.py b/src/bc/__init__.py deleted file mode 100644 index 52cdb6d..0000000 --- a/src/bc/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from src.bc.bc import BehaviorCloningPolicy, train, bc_config diff --git a/src/bc/bc.py b/src/bc/bc.py deleted file mode 100644 index ee3b1d7..0000000 --- a/src/bc/bc.py +++ /dev/null @@ -1,191 +0,0 @@ -import torch -import torch.nn as nn -from torch.utils.data import DataLoader -import pickle -from torch.utils.tensorboard import SummaryWriter - -from src.policies import IntersimStateNet, IntersimPolicy, generate_transforms -from src.util.nn_training import optimizer_factory -from tqdm import tqdm -import json5 -from ray import tune - -def bc_config(ray_config): - config = { - 'ego_encoder': {'input_dim': 5, 'hidden_n': 0, 'hidden_dim':0, 'output_dim': 0}, - 'deepsets': { - 'input_dim': 6, - 'phi': { - 'hidden_n': ray_config['deepsets_phi_hidden_n'], - 'hidden_dim': ray_config['deepsets_phi_hidden_dim'] - }, - 'latent_dim': ray_config['deepsets_latent_dim'], - 'rho': { - 'hidden_n': ray_config['deepsets_rho_hidden_n'], - 'hidden_dim': ray_config['deepsets_rho_hidden_dim'] - }, - 'output_dim': ray_config['deepsets_output_dim'] - }, - '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'], - 'hidden_dim': ray_config['head_hidden_dim'], - 'output_dim': 1, # number of outputs e.g. number of actions, or just one - 'final_activation': ray_config['head_final_activation'], - }, - 'optim': { - 'optimizer':'adam', - 'lr':ray_config['lr'], - 'weight_decay':ray_config['weight_decay'] - }, - 'train_epochs': 40, - 'train_batch_size': ray_config['train_batch_size'], - 'loss': ray_config['loss'], - - } - return config - -class BehaviorCloningPolicy(IntersimPolicy): - """ - Class for (continuous) behavior cloning policy - """ - - def __init__(self, config: dict, transforms: dict): - """ - Initialize BehaviorCloningPolicy - Args: - config (dict): configuration file to initialize IntersimDeepSetsNet with - transforms (dict): dictionary of transforms to apply to different fields - """ - super(BehaviorCloningPolicy, self).__init__(config, transforms) - self._policy = IntersimStateNet(config) - - @classmethod - def load_model(cls, filestr: str, config: dict = None): - """ - 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 - """ - 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')) - return model - - def eval(self): - self._policy.eval() - - def parameters(self): - return self._policy.parameters() - - 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_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') - -def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs): - - using_ray = kwargs.get('ray', False) - if using_ray: - print('using ray') - - # hyperparams - loss_type = config['loss'] - train_epochs = config['train_epochs'] - train_batch_size = config['train_batch_size'] - - cv_every = 1 - print_epoch_every = 1000 - print_cv_every = 5 - checkpoint_every = 100 - cv_batch_size = 256 # doesn't matter - - # 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) - - # change policy dtype - policy.policy = policy.policy.type(train_dataset[0]['state']['ego_state'].dtype) - - # generate loss function, optimizer - 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 - optimizer = optimizer_factory(config['optim'], policy.parameters()) - - # generate tensorboard writer - if not using_ray: - writer = SummaryWriter(filestr) - - for i in tqdm(range(train_epochs)): - - # save model checkpoints - if i % checkpoint_every == 0: - policy.save_model(filestr + '_epoch%04i'%(i) ) - - # train - epoch_loss = 0 - for (batch_idx, batch) in enumerate(training_loader): - - # sample mini-batch and run through policy - pred_action = policy(batch['state']) - loss = loss_fn(pred_action, batch['action']) - - # compute loss and step optimizer - optimizer.zero_grad() - loss.backward() - optimizer.step() - - epoch_loss += loss.item() / len(train_dataset) - - # if i % print_epoch_every == 0: - # print('Epoch: {}, Training Loss: {}'.format(i, epoch_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['state']) - loss = cv_loss_fn(pred_action, batch['action']) - cv_loss += loss.item() / len(cv_dataset) - - - # Write epoch loss - if using_ray: - if i % cv_every == 0: - tune.report(training_loss=epoch_loss, cv_loss=cv_loss, training_iteration=i+1) - else: - tune.report(training_loss=epoch_loss, training_iteration=i+1) - else: - writer.add_scalar('training loss',epoch_loss, i) - if i % cv_every == 0: - writer.add_scalar('cv loss', cv_loss, i) - - # if i % print_cv_every == 0: - # print('Epoch: {}, CV Loss: {}'.format(i, cv_loss)) - - - policy.save_model(filestr) diff --git a/src/data/data_utils.py b/src/data/data_utils.py deleted file mode 100644 index 0bcf9d4..0000000 --- a/src/data/data_utils.py +++ /dev/null @@ -1,96 +0,0 @@ -import torch -from torch.utils.data import Dataset -import numpy as np -from src.data.expert_data import load_expert_data -import os -opj = os.path.join - -class InteractionDatasetMultiAgent(Dataset): - """ - Class to handle getting full multi-agent observations and actions - """ - pass - -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], dtype=torch.float32): - """ - Args: - output_dir (string): Directory with all the images. - loc (int): location index - tracks (list[int]): track indices - """ - self.output_dir = output_dir - self.loc = loc - self.tracks = tracks - self.dtype = dtype - self.keys = ['ego_state', 'relative_state', 'path', 'action', 'next_ego_state', 'next_relative_state', 'next_path'] - self._load_dataset() - - def _load_dataset(self): - """ - Load the full datasets ahead of time - """ - self.raw_data = {key:[] for key in self.keys} - max_nv = 0 - for track in self.tracks: - try: - data = load_expert_data(path=self.output_dir, loc=self.loc, track=track) - print('Loaded location {} track {}'.format(self.loc,track)) - except: - print('Failed to load location {} track {}'.format(self.loc,track)) - continue - max_nv = max(max_nv, data['relative_state'].shape[1]) - for key in self.keys: - self.raw_data[key].append(data[key]) - - # 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, 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['next_relative_state'][i] = torch.cat((self.raw_data['next_relative_state'][i], pad), dim=1) - - # cat lists - for key in self.keys: - self.raw_data[key] = torch.cat(self.raw_data[key]).type(self.dtype) - - # mandate equal length - lengths = [len(self.raw_data[key]) for key in self.keys] - assert min(lengths) == max(lengths), 'dataset lengths unequal' - - def __len__(self): - return len(self.raw_data['ego_state']) - - def __getitem__(self, idx): - """ - Sample from the dataset - Args: - idx: index or indices of B samples - Returns: - sample (dict): sample dictionary with the following entries: - state (dict): state dictionary with the following entries: - ego_state (torch.tensor): (B, 5) raw state - relative_state (torch.tensor): (B, max_nv, d) relative state (padded with nans) - path (torch.tensor): (B, P, 2) tensor of P future path x and y positions - action (torch.tensor): (B, 1) actions taken from each state - next_stat (dict): next state dictionary with the following entries: - ego_state (torch.tensor): (B, 5) raw next state - relative_state (torch.tensor): (B, max_nv, d) next relative state (padded with nans) - path (torch.tensor): (B, P, 2) tensor of P future next path x and y positions - """ - #sample = {key:self.raw_data[key][idx] for key in self.keys} - sample = { - 'state':{ - 'ego_state':self.raw_data['ego_state'][idx], - 'relative_state':self.raw_data['relative_state'][idx], - 'path':self.raw_data['path'][idx] - }, - 'action':self.raw_data['action'][idx], - 'next_state':{ - 'ego_state':self.raw_data['next_ego_state'][idx], - 'relative_state':self.raw_data['next_relative_state'][idx], - 'path':self.raw_data['next_path'][idx]}, - } - return sample \ No newline at end of file diff --git a/src/data/expert_data.py b/src/data/expert_data.py deleted file mode 100644 index 4b4b5ba..0000000 --- a/src/data/expert_data.py +++ /dev/null @@ -1,210 +0,0 @@ -import torch - -import pickle -import gym -import numpy as np - -import intersim -from intersim.utils import get_map_path, get_svt, SVT_to_stateactions -from intersim import collisions -from intersim.graphs import ConeVisibilityGraph -import os -opj = os.path.join - -def generate_expert_data(path: str='expert_data', loc: int = 0, track:int = 0, - mask_relstate: bool = False, regularize_actions: bool = False, - **kwargs): - """ - Function to save (joint) states and observations from simulated frame - Args: - path (str): directory to save data - loc (int): location index - track (int): track index - mask_relstate (bool): whether to mask the relative states from the cone visibility graph - regularize_actions (bool): whether to regularize the action selection - kwargs: arguments for environment instantiation - """ - - action_reg = 0.002 if regularize_actions else 0 - - if not os.path.isdir(path): - os.makedirs(path) - filestr = opj(path,intersim.LOCATIONS[loc]+'_track%03i'%(track)) - - svt, svt_path = get_svt(loc=loc, track=track) #base='InteractionSimulator' - osm = get_map_path(loc=loc) - print('SVT path: {}'.format(svt_path)) - print('Map path: {}'.format(osm)) - states, actions = SVT_to_stateactions(svt) - - # animate from environment - if mask_relstate: - cvg = ConeVisibilityGraph(r=20, half_angle=120) - env = gym.make('intersim:intersim-v0', svt=svt, map_path=osm, - min_acc=-np.inf, max_acc=np.inf, graph=cvg, mask_relstate=True, **kwargs) - else: - env = gym.make('intersim:intersim-v0', svt=svt, map_path=osm, **kwargs, - min_acc=-np.inf, max_acc=np.inf) - - env.reset() - done = False - obs, actions_taken, max_devs = [], [], [] - i = 0 - while not done and i < len(actions): - # check state deviation - env_state = env.projected_state - nni = ~torch.isnan(env_state[:,0]) - norms = torch.norm(env_state[nni,:2]-states[i,nni,:2], dim=1) - if len(norms)>0: - max_devs.append(norms.max()) - - # propagate environment - ob, r, done, info = env.step(env.target_state(svt.simstate[i+1], mu=action_reg)) - obs.append(ob) - actions_taken.append(info['action_taken']) - i += 1 - - print("Maximum environment deviation from track: %f m" %(max(max_devs))) - - # check for collisions - x = torch.stack([ob['state'] for ob in obs]) - cols = collisions.check_collisions_trajectory(x, svt.lengths, svt.widths) - assert ~torch.any(cols), 'Error: Collisions found at indices {}'.format(cols.nonzero(as_tuple=True)) - - # shift actions - actions_taken.pop(0) - obs.pop(-1) - actions = torch.stack(actions_taken) - - # save observations and actions - pickle.dump(obs,open(filestr+'_raw_observations.pkl', 'wb')) - torch.save(actions, filestr+'_raw_actions.pt') - process_expert_observations(obs, actions, filestr) - -def process_expert_observations(obs, actions, filestr, remove_outliers=True, dtype=torch.float32): - """ - Process the expert observations and save them as torch tensors - Args: - obs (list[dict]): lost of observations - actions (torch.Tensor): (T, nv, a) tensor of actions - filestr (str): base filename with which to save out observation tensors - remove_outliers (bool): whether to remove datapoints with acceleration above or below 5 m/s/s - dtype (torch.Type): type to convert data to - """ - keys = ['ego_state', 'relative_state', 'path', 'action', 'next_ego_state', 'next_relative_state', 'next_path'] - data = {key:[] for key in keys} - assert len(obs) == len(actions), 'non-matching action and observation lengths' - T = len(obs) - max_nv = 0 - for t in range(T-1): - nni = ~torch.isnan(obs[t]['state'][:,0]) & ~torch.isnan(obs[t+1]['state'][:,0]) - max_nv = max(max_nv,nni.count_nonzero()) - - # state - data['ego_state'].append(obs[t]['state'][nni]) - data['relative_state'].append(obs[t]['relative_state'].index_select(0, - nni.nonzero()[:,0]).index_select(1, nni.nonzero()[:,0])) - data['path'].append(torch.stack((obs[t]['paths'][0][nni], obs[t]['paths'][1][nni]), dim=-1)) - - # action - data['action'].append(actions[t][nni]) - - # next state - data['next_ego_state'].append(obs[t+1]['state'][nni]) - data['next_relative_state'].append(obs[t+1]['relative_state'].index_select(0, - nni.nonzero()[:,0]).index_select(1, nni.nonzero()[:,0])) - data['next_path'].append(torch.stack((obs[t+1]['paths'][0][nni], obs[t+1]['paths'][1][nni]), dim=-1)) - - - - # pad second dimension of relative state - for i in range(len(data['relative_state'])): - nv1, nv2, d = data['relative_state'][i].shape - pad = torch.zeros(nv1, max_nv-nv2, d, dtype=dtype) * np.nan - data['relative_state'][i] = torch.cat((data['relative_state'][i], pad), dim=1) - data['next_relative_state'][i] = torch.cat((data['next_relative_state'][i], pad), dim=1) - - # cat lists - for key in keys: - data[key] = torch.cat(data[key]).type(dtype) - - if remove_outliers: - non_outlier_indices = torch.nonzero(torch.abs(data['action'][:,0]) < 5) - for key in keys: - data[key] = data[key][non_outlier_indices[:,0]] - - # mandate equal length - lengths = [len(data[key]) for key in keys] - assert min(lengths) == max(lengths), 'dataset lengths unequal' - - # save out data - for key in keys: - torch.save(data[key], filestr+'_'+key+'.pt') - -def load_expert_data(path='expert_data', loc: int = 0, track:int = 0): - """ - Load expert data from processed files. - Args: - path (str): directory to save data - loc (int): location index - track (int): track index - Returns: - data (dict[torch.Tensor]): dict of data - """ - # load observations and actions - filestr = opj(path, intersim.LOCATIONS[loc]+'_track%03i'%(track)) - data = {} - keys = ['ego_state', 'relative_state', 'path', 'action', 'next_ego_state', 'next_relative_state', 'next_path'] - for key in keys: - data[key] = torch.load(filestr+'_'+key+'.pt') - return data - -def load_expert_data_raw(path='expert_data', loc: int = 0, track:int = 0): - """ - Load expert data from raw file. - Args: - path (str): directory to save data - loc (int): location index - track (int): track index - Returns: - obs (list[Observations]): list of observations - actions (list[torch.tensor]): list of corresponding actions taken in observations - """ - # load observations and actions - filestr = opj(path, intersim.LOCATIONS[loc]+'_track%03i'%(track)) - obs = pickle.load(open(filestr+'_raw_observations.pkl', 'rb')) - actions = torch.load(filestr+'_raw_actions.pt') - actions = list(torch.unbind(actions)) - return obs, actions - -if __name__ == '__main__': - import argparse - parser = argparse.ArgumentParser(description='Save Expert Trajectories') - parser.add_argument('--loc', default=0, type=int, - help='location (default 0)') - parser.add_argument('--track', default=0, type=int, - help='track number (default 0)') - parser.add_argument('--all-tracks', action='store_true', - help='whether to process all tracks at location') - parser.add_argument('--graph', action='store_true', - help='whether to mask the relative states based on a ConeVisibilityGraph') - parser.add_argument('--reg', action='store_true', - help='whether to regularize actions in the action targeter') - parser.add_argument('-o', default='./expert_data', type=str, - help='output folder') - args = parser.parse_args() - - kwargs = { - 'loc':args.loc, - 'track': args.track, - 'path':args.o, - 'mask_relstate':args.graph, - 'regularize_actions': args.reg - } - - if args.all_tracks: - for i in range(intersim.MAX_TRACKS): - kwargs['track'] = i - generate_expert_data(**kwargs) - else: - generate_expert_data(**kwargs) \ No newline at end of file diff --git a/src/discriminator/__init__.py b/src/discriminator/__init__.py deleted file mode 100644 index 4ad66b2..0000000 --- a/src/discriminator/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from src.discriminator.discriminator import * \ No newline at end of file diff --git a/src/discriminator/discriminator.py b/src/discriminator/discriminator.py deleted file mode 100644 index 1c9664c..0000000 --- a/src/discriminator/discriminator.py +++ /dev/null @@ -1,101 +0,0 @@ -import torch - -# imitation.rewards.discrim_nets.DiscrimNetGAIL is composed of self.discriminator (nn.Module), -# which gets called with inputs (state, action) when needed. - -class CnnDiscriminator(torch.nn.Module): - """ConvNet similar to stable_baselines3.common.policies.ActorCriticCnnPolicy.""" - - def __init__(self, env): - super().__init__() - - obs_channels, _, _ = env.observation_space.shape - (action_size,) = env.action_space.shape - in_channels = obs_channels + action_size - - self.cnn = torch.nn.Sequential( - torch.nn.Conv2d(in_channels, 32, kernel_size=(8, 8), stride=(4, 4)), # 5+1 -> 32 - torch.nn.ReLU(), - torch.nn.Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2)), # 32 -> 64 - torch.nn.ReLU(), - torch.nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1)), # 64 -> 64 - torch.nn.ReLU(), - torch.nn.Flatten(start_dim=1, end_dim=-1), - torch.nn.LazyLinear(512), # 28224 -> 512 - torch.nn.ReLU(), - torch.nn.LazyLinear(1), # 512 -> 1 - ) - - @staticmethod - def _concatenate(state, action): - b, _, h, w = state.shape - _, a = action.shape - act = action.unsqueeze(-1).unsqueeze(-1).expand((b, a, h, w)) - sa = torch.cat((state, act), -3) - return sa - - def forward(self, state, action): - sa = self._concatenate(state, action) - assert sa.ndim == 4 - return self.cnn(sa).squeeze(1) - -class CnnDiscriminatorFlatAction(torch.nn.Module): - """ConvNet similar to stable_baselines3.common.policies.ActorCriticCnnPolicy.""" - - def __init__(self, env): - super().__init__() - - obs_channels, _, _ = env.observation_space.shape - (action_size,) = env.action_space.shape - in_channels = obs_channels - - self.cnn = torch.nn.Sequential( - torch.nn.Conv2d(in_channels, 32, kernel_size=(8, 8), stride=(4, 4)), # in_channels -> 32 - torch.nn.ReLU(), - torch.nn.Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2)), # 32 -> 64 - torch.nn.ReLU(), - torch.nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1)), # 64 -> 64 - torch.nn.ReLU(), - torch.nn.Flatten(start_dim=1, end_dim=-1), - torch.nn.LazyLinear(128), # 28224 -> 128 - ) - self.decoder = torch.nn.Sequential( - torch.nn.LazyLinear(64), #128 + 2 -> 64 - torch.nn.ReLU(), - torch.nn.LazyLinear(64), #64 -> 64 - torch.nn.ReLU(), - torch.nn.LazyLinear(1) #64 -> 1 - ) - - @staticmethod - def _concatenate(state, action): - b, s= state.shape - b, a = action.shape - sa = torch.cat((state, action), -1) - return sa - - def forward(self, state, action): - s = self.cnn(state.float()) - sa = self._concatenate(s, action) - assert sa.ndim == 2 - return self.decoder(sa).squeeze(1) - -class MlpDiscriminator(torch.nn.Module): - """MLP similar to stable_baselines3.common.policies.ActorCriticPolicy.""" - - def __init__(self, env=None): - super().__init__() - self.flatten = torch.nn.Flatten(start_dim=1, end_dim=-1) - self.mlp = torch.nn.Sequential( - torch.nn.LazyLinear(64), # 42 -> 64 - torch.nn.Tanh(), - torch.nn.LazyLinear(64), # 64 -> 64 - torch.nn.Tanh(), - torch.nn.LazyLinear(1), # 64 -> 1 - ) - - def forward(self, state, action): - flat = self.flatten(state) - sa = torch.cat((action, flat), -1) - assert sa.ndim == 2 - return self.mlp(sa).squeeze(1) diff --git a/src/discriminator/test_discriminator.py b/src/discriminator/test_discriminator.py deleted file mode 100644 index 1de614c..0000000 --- a/src/discriminator/test_discriminator.py +++ /dev/null @@ -1,45 +0,0 @@ -from intersim.envs.intersimple import NRasterized -from discriminator import CnnDiscriminator -import torch - -def test_image_concatenation(): - env = NRasterized() - disc = CnnDiscriminator(env) - s = torch.tensor(env.reset()).unsqueeze(0) - a = torch.tensor([[0.5]]) - sa = disc._concatenate(s, a) - - assert s.shape == (1, 5, 200, 200) - assert a.shape == (1, 1) - assert sa.shape == (1, 6, 200, 200) - assert torch.allclose(sa[:, :5], 1.0 * s) - assert (sa[:, 5] == a.unsqueeze(-1)).all() - -def test_image_concatenation3(): - env = NRasterized() - disc = CnnDiscriminator(env) - - s1 = env.reset() - a1 = 0.15 - s2, _, _, _ = env.step(0.9) - a2 = 0.25 - s3, _, _, _ = env.step(-0.9) - a3 = 0.35 - - s = torch.stack([ - torch.tensor(s1), - torch.tensor(s2), - torch.tensor(s3) - ], axis=0) - a = torch.tensor([ - [a1], - [a2], - [a3], - ]) - sa = disc._concatenate(s, a) - - assert s.shape == (3, 5, 200, 200) - assert a.shape == (3, 1) - assert sa.shape == (3, 6, 200, 200) - assert torch.allclose(sa[:, :5], 1.0 * s) - assert (sa[:, 5] == a.unsqueeze(-1)).all() diff --git a/src/main.py b/src/main.py deleted file mode 100644 index edf255c..0000000 --- a/src/main.py +++ /dev/null @@ -1,118 +0,0 @@ -import os -import torch -import gym -import intersim -import numpy as np -from tqdm import tqdm -from torch.utils.tensorboard import SummaryWriter - -from src import InteractionDatasetSingleAgent, metrics -from intersim.utils import get_map_path, get_svt -from src.policies.policy import generate_transforms - -def basestr(**kwargs): - """ - Return base prefix for all files relating to a certain experiment - Args: - kwargs (dict): keyword arguments sent to main training loop - Returns: - basestr (str): prefix - """ - return 'base' - -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 - loc (int): the location index of the roundabout - datadir (str): path to expert data - kwargs (dict): remaining kwargs for training loop - """ - # get/set seed - seed = kwargs.get('seed',0) - torch.manual_seed(seed) - - # method-based training - if method=='bc': - from src import bc - policy_class = bc.BehaviorCloningPolicy - train_fn = bc.train - elif method=='vd': - from src import value_dice - policy_class = value_dice.ValueDicePolicy - train_fn = value_dice.train - else: - raise NotImplementedError("Method {} not implemented".format(method)) - - # default train / cv / test split datasets - if train: - - # make policy, train and test datasets, and send to - train_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=kwargs['train_tracks']) - # generate transform from train_dataset - transforms = generate_transforms(train_dataset) - policy = policy_class(config, transforms) - cv_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=kwargs['cv_tracks']) - train_fn(config, policy, train_dataset, cv_dataset, filestr, **kwargs) - - if test: - - # load policy - policy = policy_class.load_model(filestr, config) - policy.eval() - - # simulate policy - simulate_policy(policy, loc=loc, track=kwargs['test_tracks'][0], filestr=filestr, nframes=kwargs['nframes'], graph=kwargs['graph']) - - # run test metrics - test_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=kwargs['test_tracks']) - writer = SummaryWriter(filestr) - info = metrics(filestr, test_dataset, policy) - for k, m in info.items(): - writer.add_scalar('test/{}'.format(k), m, 0) - - -def simulate_policy(policy, loc=0, track=0, filestr='', nframes=float('inf'), graph=None): - """ - Simulate a trained policy - Args: - policy: the policy to simulate, which should return action directly - loc (int): location index to test policy - track (int): track to test policy - filestr (str): path prefix to save simulation to - """ - # animate from environment - basepath = os.path.abspath('./InteractionSimulator') - svt, svt_path = get_svt(base=basepath, loc=loc, track=track) - osm = get_map_path(base=basepath, loc=loc) - if graph: - env = gym.make('intersim:intersim-v0', svt=svt, map_path=osm, - min_acc=-np.inf, max_acc=np.inf, graph=graph, mask_relstate=True) - else: - 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, - # min_acc=-np.inf, max_acc=np.inf) - - ob, _ = env.reset() - env.render() - done = False - i = 0 - with tqdm(total=min(nframes, env._svt.Tind)) as pbar: - while not done and i < nframes: - i += 1 - - # get action - action = policy(ob) - - # propagate environment - ob, r, done, info = env.step(action) - env.render() - - pbar.update() - - env.close(filestr=filestr+'_sim') diff --git a/src/models/__init__.py b/src/models/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/nets/__init__py b/src/nets/__init__py deleted file mode 100644 index e69de29..0000000 diff --git a/src/nets/deepsets.py b/src/nets/deepsets.py deleted file mode 100644 index 3b58d16..0000000 --- a/src/nets/deepsets.py +++ /dev/null @@ -1,117 +0,0 @@ -import torch -from torch import nn - -from src.nets.util import parse_functional - -class DeepSetsModule(nn.Module): - def __init__(self, input_dim, phi_hidden_n, phi_hidden_dim, latent_dim, rho_hidden_n, rho_hidden_dim, output_dim): - """ - Args: - input_dim (int): input size of one instance of the set; input size of phi - phi_hidden_n (int): number of hidden layers in phi - phi_hidden_dim (int): size of hidden layers in phi - latent_dim (int): output size of phi network, where sum is taken over instances; input size of rho - rho_hidden_n (int): number of hidden layers in rho - rho_hidden_dim (int): size of hidden layers in rho - output_dim (int): output size of rho - """ - super(DeepSetsModule, self).__init__() - self.input_dim = input_dim - self.latent_dim = latent_dim - self.phi = Phi(self.input_dim, phi_hidden_n, phi_hidden_dim, self.latent_dim) - self.rho = Phi(self.latent_dim, rho_hidden_n, rho_hidden_dim, output_dim) - self.output_dim = self.rho.output_dim - self.pooling = torch.sum - - @staticmethod - def from_config(config): - """ - Args: - config (dict): dictionary with network parameters in the form - { - "input_dim": 5, - "phi": { - "hidden_n": 1, - "hidden_dim": 10, - }, - "latent_dim": 8, - "rho": { - "hidden_n": 1, - "hidden_dim": 10, - }, - "output_dim" : 1, - } - Returns: - m (nn.Module): deep sets module - """ - input_dim = config["input_dim"] - phi = config["phi"] - latent_dim = config["latent_dim"] - rho = config["rho"] - output_dim = config["output_dim"] - m = DeepSetsModule(input_dim, phi["hidden_n"], phi["hidden_dim"], latent_dim, rho["hidden_n"], rho["hidden_dim"], output_dim) - return m - - def forward(self, x): - """ - Args: - x (torch.tensor): ([B, ]max_nv, d) - Returns: - y (torch.tensor): ([B, ]output_dim) - """ - # mask for selecting only those batches and vehicles where all relative states are not nan - # shape (B, max_nv) - notnan_mask = torch.all(~torch.isnan(x), dim=-1) - # create zero tensor of shape (B, max_nv, latent_dim) to store phi evaluations in - latent = torch.zeros([*x.shape[:-1], self.latent_dim], dtype=x.dtype) - # evaluate phi for all not NaN entries - # x[batch_dynamic_mask] has shape (notnan_mask.sum(), input_dim) - latent[notnan_mask] = self.phi(x[notnan_mask]) - - # sum over relative state dimension - latent = self.pooling(latent, dim=-2) - - # apply rho network - y = self.rho(latent) - return y - - -class Phi(nn.Module): - def __init__(self, input_dim, hidden_n, hidden_dim, output_dim, final_activation=None): - """ - Fully connected feedforward network with same size for all hidden layers and ReLU activation - - Args: - input_dim (int): input dimension - hidden_n (int): number of hidden layers - hidden_dim (int): hidden layer dimension - output_dim (int): output dimension - """ - super(Phi, self).__init__() - self.input_dim = input_dim - self.output_dim = output_dim - if hidden_n > 0: - self.layers = nn.ModuleList([nn.Linear(self.input_dim, hidden_dim)]) - for _ in range(hidden_n - 1): - self.layers.append(nn.Linear(hidden_dim, hidden_dim)) - self.layers.append(nn.Linear(hidden_dim, self.output_dim)) - else: - self.layers = nn.ModuleList([nn.Identity()]) - self.output_dim = self.input_dim - self.activation = nn.functional.relu - self.final_activation = final_activation if final_activation else lambda x: x - - def forward(self, x): - for layer in self.layers[:-1]: - x = self.activation(layer(x)) - x = self.final_activation(self.layers[-1](x)) - return x - - @staticmethod - def from_config(config): - args = (config["input_dim"], config["hidden_n"], config["hidden_dim"], config["output_dim"]) - if "final_activation" in config: - kwargs = {"final_activation": parse_functional(config["final_activation"])} - else: - kwargs = {} - return Phi(*args, **kwargs) diff --git a/src/nets/util.py b/src/nets/util.py deleted file mode 100644 index 6634224..0000000 --- a/src/nets/util.py +++ /dev/null @@ -1,14 +0,0 @@ -import torch -from torch.nn import functional, Identity - -def parse_functional(functional_config): - if isinstance(functional_config, str): - if functional_config == 'relu': - return functional.relu - elif functional_config == 'sigmoid': - return torch.sigmoid - elif functional_config == 'softmax': - return functional.softmax - elif functional_config == 'id': - return Identity() - return None \ No newline at end of file diff --git a/src/policies/__init__.py b/src/policies/__init__.py deleted file mode 100644 index 0679815..0000000 --- a/src/policies/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from src.policies.policy import IntersimPolicy, IntersimStateNet, IntersimStateActionNet, generate_transforms -from src.policies.options import OptionsCnnPolicy \ No newline at end of file diff --git a/src/policies/options.py b/src/policies/options.py deleted file mode 100644 index 2b55c12..0000000 --- a/src/policies/options.py +++ /dev/null @@ -1,65 +0,0 @@ -from stable_baselines3.common.policies import ActorCriticPolicy, ActorCriticCnnPolicy -from torch.distributions import Categorical - -class OptionsCnnPolicy(ActorCriticPolicy): - """ - Class for high-level options policy (generator) - """ - def __init__(self, observation_space, *args, eps=0, **kwargs): - super().__init__(observation_space, *args, **kwargs) - self.cnn_policy = ActorCriticCnnPolicy(observation_space['obs'], *args, **kwargs) - self.eps = eps - - def _prior_distribution(self, s): - """ - Return prior distribution over high-level options (before masking) - Args: - s (torch.tensor): observation - Returns: - values (torch.tensor): values from critic - dist (torch.distributions): prior distribution over actions - """ - latent_pi, latent_vf, latent_sde = self.cnn_policy._get_latent(s) - distribution = self.cnn_policy._get_action_dist_from_latent(latent_pi, latent_sde) - values = self.cnn_policy.value_net(latent_vf) - return values, distribution.distribution - - def forward(self, obs): - """ - Will mask invalid states before making action selections - Args: - obs: dict with keys: - obs (torch.tensor): (*,o) true observations - mask (torch.tensor): (*,m) mask over valid actions - Returns: - ch (torch.tensor): (*,a) sampled actions - values (torch.tensor): (*,) predicted value at observation - log_probs (torch.tensor): (*,) log probabilities of selected actions - """ - s, m = obs['obs'], obs['mask'] - values, prior = self._prior_distribution(s) - posterior = Categorical((prior.probs + self.eps) * m) - ch = posterior.sample() - return ch, values, posterior.log_prob(ch) - - def _predict(self, obs, deterministic=False): - action, _, _ = self.forward(obs) - return action - - def evaluate_actions(self, obs, ch): - """ - Evaluate particular actions - Args: - obs: dict with keys: - obs (torch.tensor): (*,o) true observations - mask (torch.tensor): (*,m) masks over valid actions - ch (torch.tensor): (*,a) selected actions - Returns: - values (torch.tensor): (*,) predicted value at observation - log_probs (torch.tensor): (*,) log probabilities of selected actions - ent (torch.tensor): (*,) entropy of each distribution over actions - """ - s, m = obs['obs'], obs['mask'] - values, prior = self._prior_distribution(s) - posterior = Categorical((prior.probs + self.eps) * m) - return values, posterior.log_prob(ch), posterior.entropy() # additional values used by PPO.train diff --git a/src/policies/policy.py b/src/policies/policy.py deleted file mode 100644 index 274d34b..0000000 --- a/src/policies/policy.py +++ /dev/null @@ -1,164 +0,0 @@ - -import torch -from torch import nn - -from src.nets.deepsets import DeepSetsModule, Phi -from src.util.transform import MinMaxScaler - -class IntersimStateNet(nn.Module): - def __init__(self, config): - """ - Args: - config (dict): dictionary for configuring the deep sets policy - """ - super(IntersimStateNet, self).__init__() - ego_config = config['ego_encoder'] - deepsets_config = config['deepsets'] - pathnet_config = config['path_encoder'] - - 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 - head_config = config['head'] - head_config["input_dim"] = cat_dim - self.head = Phi.from_config(head_config) - - def forward(self, sample): - """ - Args: - sample (dict): sample dictionary with the following entries: - state (torch.tensor): (B, 5) raw state - relative_state (torch.tensor): (B, max_nv, d) relative state (padded with nans) - path (torch.tensor): (B, P, 2) tensor of P future path x and y positions - action (torch.tensor): (B, 1) actions taken from each state - Returns: - x (torch.tensor): (head_output_dim,) output of common head network - """ - ego = self.ego_net(sample["ego_state"]) - relative = self.deepsets_net(sample["relative_state"]) - path = self.path_net(sample["path"].reshape((sample["path"].shape[0], -1))) - x = torch.cat([ego, relative, path], dim=-1) - x = self.head(x) - return x - - -class IntersimStateActionNet(nn.Module): - def __init__(self, config): - """ - Args: - config (dict): dictionary for configuring the deep sets policy - """ - super(IntersimStateActionNet, self).__init__() - ego_config = config['ego_encoder'] - deepsets_config = config['deepsets'] - pathnet_config = config['path_encoder'] - - self.ego_net = Phi.from_config(ego_config) - self.deepsets_net = DeepSetsModule.from_config(deepsets_config) - self.path_net = Phi.from_config(pathnet_config) - self.action_dim = config["action_dim"] - - cat_dim = self.ego_net.output_dim + self.deepsets_net.output_dim + self.path_net.output_dim + self.action_dim - # head has number of concatenated features as input - head_config = config['head'] - head_config["input_dim"] = cat_dim - self.head = Phi.from_config(head_config) - - def forward(self, sample): - """ - Args: - sample (dict): sample dictionary with the following entries: - state (torch.tensor): (B, 5) raw state - relative_state (torch.tensor): (B, max_nv, d) relative state (padded with nans) - path_x (torch.tensor): (B, P) tensor of P future path x positions - path_y (torch.tensor): (B, P) tensor of P future path y positions - action (torch.tensor): (B, 1) actions taken from each state - Returns: - x (torch.tensor): (head_output_dim,) output of common head network - """ - ego = self.ego_net(sample["ego_state"]) - relative = self.deepsets_net(sample["relative_state"]) - action = sample["action"] - path = self.path_net(sample["path"].reshape((sample["path"].shape[0], -1))) - x = torch.cat([ego, relative, path, action], dim=-1) - x = self.head(x) - return x - - -class IntersimPolicy(): - """ - Base class for intersim policies - """ - def __init__(self, config, transforms): - super(IntersimPolicy, self).__init__() - self._config = config - self._transforms = transforms - - @property - def transforms(self): - return self._transforms - - @transforms.setter - def transforms(self, transforms): - self._transforms=transforms - - @property - def policy(self): - return self._policy - - @policy.setter - def policy(self, policy): - self._policy = policy - - def transform_observation(self, ob): - # run observation through transforms - transformed_ob = {} - for key in ['ego_state', 'relative_state', 'path', 'action']: - if key in self._transforms.keys() and key in ob.keys(): - transformed_ob[key] = self._transforms[key].transform(ob[key]) - return transformed_ob - - def __call__(self, ob): - - if 'ego_state' in ob.keys(): - # extract state from dataloader samples - pass - else: - # extract state from observation (using simulator) - ob['ego_state'] = ob['state'] - ob['path'] = torch.stack(ob['paths'],dim=-1) - - ob = self.transform_observation(ob) - - # run transformed state through model - action = self._policy(ob) - assert action.ndim == 2, 'action has incorrect shape' - - # untransform action - if 'action' in self._transforms.keys(): - action = self._transforms['action'].inverse_transform(action) - return action - - -def generate_transforms(dataset): - """ - Generate transform dictionary from dataset - Args: - dataset (Dataset): dataset of demo observations and actions - """ - transforms = { - 'action': MinMaxScaler(), - 'ego_state': MinMaxScaler(), - 'relative_state': MinMaxScaler(reduce_dim=2), - 'path': MinMaxScaler(reduce_dim=2), - } - for key in transforms.keys(): - if key == 'action': - transforms[key].fit(dataset[:][key]) - else: - transforms[key].fit(dataset[:]['state'][key]) - - return transforms diff --git a/src/util/nn_training.py b/src/util/nn_training.py deleted file mode 100644 index eea6eea..0000000 --- a/src/util/nn_training.py +++ /dev/null @@ -1,11 +0,0 @@ -import torch - -def optimizer_factory(config, parameters): - optimizer_type = config['optimizer'] - learning_rate = config['lr'] - weight_decay = config['weight_decay'] - if optimizer_type == 'adam': - optimizer = torch.optim.Adam(parameters, lr=learning_rate, weight_decay=weight_decay) - else: - raise NotImplementedError - return optimizer diff --git a/src/util/transform.py b/src/util/transform.py deleted file mode 100644 index 0b244e9..0000000 --- a/src/util/transform.py +++ /dev/null @@ -1,140 +0,0 @@ -import torch -from torch import nn -import numpy as np -from sklearn import preprocessing - -class Transform(nn.Module): - """ - Base class to normalize observations and actions for network. - """ - - def __init__(self): - super(Transform, self).__init__() - # self.fit(X) - - def fit(self, X): - """ - Fit transformer to X - Args: - X (torch.tensor): (B, N) tensor of B data points with N features - """ - raise NotImplementedError('Please implement fit()') - - def transform(self, X): - """ - Transform X. fit() has to be called first - Args: - X (torch.tensor): (B, N) tensor where N has to be the same as during fit() - """ - raise NotImplementedError('Please implement transform()') - - def inverse_transform(self, X): - """ - Inverse transformation - Args: - X (torch.tensor): (B, N) tensor - """ - raise NotImplementedError('Please implement inverse_transform()') - - 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 - """ - def __init__(self, tf, reduce_dim:int=None): - """ - Initialize SciKitTransform - Args: - tf: transform - 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.tf = tf - self.reduce_dim = reduce_dim - super(SciKitTransform, 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] - - self.tf.fit(X.reshape((-1,self.nfeatures))) - - def transform(self, X): - shape = X.shape - t = torch.tensor(self.tf.transform(X.reshape((-1,self.nfeatures))), dtype=torch.float) - return t.reshape(shape) - - def inverse_transform(self, X): - shape = X.shape - it = torch.tensor(self.tf.inverse_transform(X.reshape((-1,self.nfeatures))), dtype=torch.float) - return it.reshape(shape) - -class SciKitStandardScaler(SciKitTransform): - """ - Wrapper around scikit-learn's StandardScaler for standardizing each feature individually. - """ - def __init__(self, **kwargs): - super(SciKitStandardScaler, self).__init__(preprocessing.StandardScaler(), **kwargs) - -class SciKitMinMaxScaler(SciKitTransform): - """ - Wrapper around scikit-learn's MinMaxScaler for scaling features to [0, 1] individually. - """ - def __init__(self, **kwargs): - super(SciKitMinMaxScaler, self).__init__(preprocessing.MinMaxScaler(), **kwargs) - diff --git a/src/value_dice/__init__.py b/src/value_dice/__init__.py deleted file mode 100644 index bb17233..0000000 --- a/src/value_dice/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from src.value_dice.value_dice import ValueDicePolicy, train, vd_config diff --git a/src/value_dice/value_dice.py b/src/value_dice/value_dice.py deleted file mode 100644 index d9e29d6..0000000 --- a/src/value_dice/value_dice.py +++ /dev/null @@ -1,300 +0,0 @@ -import numpy as np - -import torch -import torch.nn as nn -from torch.utils.data import DataLoader -from torch.nn.utils import clip_grad_norm_ -import pickle -import itertools -from torch.utils.tensorboard import SummaryWriter - -from src.policies import IntersimStateNet, IntersimStateActionNet, IntersimPolicy, generate_transforms -from src.util.transform import MinMaxScaler -from src.util.nn_training import optimizer_factory -from tqdm import tqdm -import json5 -from ray import tune - -def vd_config(ray_config): - config = { - 'policy_net': { - 'ego_encoder': {'input_dim': 5, 'hidden_n': 0, 'hidden_dim':0, 'output_dim': 0}, - 'deepsets': { - 'input_dim': 6, - 'phi': { - 'hidden_n': ray_config['deepsets_phi_hidden_n'], - 'hidden_dim': ray_config['deepsets_phi_hidden_dim'] - }, - 'latent_dim': ray_config['deepsets_latent_dim'], - 'rho': { - 'hidden_n': ray_config['deepsets_rho_hidden_n'], - 'hidden_dim': ray_config['deepsets_rho_hidden_dim'] - }, - 'output_dim': ray_config['deepsets_output_dim'] - }, - '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'], - 'hidden_dim': ray_config['head_hidden_dim'], - 'output_dim': 1, # number of outputs e.g. number of actions, or just one - 'final_activation': ray_config['head_final_activation'], - }, - }, - 'value_net': { - 'ego_encoder': {'input_dim': 5, 'hidden_n': 0, 'hidden_dim':0, 'output_dim': 0}, - 'deepsets': { - 'input_dim': 6, - 'phi': { - 'hidden_n': ray_config['deepsets_phi_hidden_n'], - 'hidden_dim': ray_config['deepsets_phi_hidden_dim'] - }, - 'latent_dim': ray_config['deepsets_latent_dim'], - 'rho': { - 'hidden_n': ray_config['deepsets_rho_hidden_n'], - 'hidden_dim': ray_config['deepsets_rho_hidden_dim'] - }, - 'output_dim': ray_config['deepsets_output_dim'] - }, - 'path_encoder': {'input_dim': 40, 'hidden_n': 0, 'hidden_dim': 0, 'output_dim': 0}, - 'action_dim': 1, - 'head': { - 'input_dim': 0, # computed in constructor - 'hidden_n': ray_config['head_hidden_n'], - 'hidden_dim': ray_config['head_hidden_dim'], - 'output_dim': 1, # number of outputs e.g. number of actions, or just one - 'final_activation': ray_config['head_final_activation'], - }, - }, - 'policy_optim': { - 'optimizer':'adam', - 'lr':ray_config['policy_lr'], - 'weight_decay':ray_config['policy_weight_decay'] - }, - 'value_optim': { - 'optimizer':'adam', - 'lr':ray_config['value_lr'], - 'weight_decay':ray_config['value_weight_decay'] - }, - 'train_epochs': 40, - 'train_batch_size': ray_config['train_batch_size'], - 'discount': ray_config['discount'], - 'clip_grad_norm': ray_config['clip_grad_norm'], - } - return config - -class ValueDicePolicy(IntersimPolicy): - """ - Class for value dice policy - """ - - def __init__(self, config: dict, transforms: dict): - """ - Initialize ValueDicePolicy - Args: - config (dict): configuration file to initialize IntersimDeepSetsNet with - transforms (dict): dictionary of transforms to apply to different fields - """ - super(ValueDicePolicy, self).__init__(config, transforms) - self._policy = IntersimStateNet(config['policy_net']) - self._value = IntersimStateActionNet(config['value_net']) - - @property - def value(self): - return self._value - - @value.setter - def value(self, value): - self._value = value - - @classmethod - def load_model(cls, filestr: str, config: dict = None): - """ - 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 - """ - 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+'_policy.pt')) - model._value.load_state_dict(torch.load(filestr+'_value.pt')) - return model - - def parameters(self): - return itertools.chain(self._policy.parameters(), self._value.parameters()) - - def policy_parameters(self): - return self.policy.parameters() - - def value_parameters(self): - return self.value.parameters() - - def eval(self): - self.policy.eval() - self.value.eval() - - 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_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+'_policy.pt') - torch.save(self._value.state_dict(), filestr+'_value.pt') - - - -def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs): - - using_ray = kwargs.get('ray', False) - if using_ray: - print('using ray') - - # hyperparams - train_epochs = config['train_epochs'] - train_batch_size = config['train_batch_size'] - discount = config['discount'] - clip_grad_norm = config['clip_grad_norm'] - - cv_every = 1 - print_epoch_every = 1000 - print_cv_every = 5 - checkpoint_every = 100 - cv_batch_size = 256 # doesn't matter - - # 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) - - # change policy dtype - dtype = train_dataset[0]['state']['ego_state'].dtype - policy.policy = policy.policy.type(dtype) - policy.value = policy.value.type(dtype) - - # define loss function - def f_value_dice_loss(batch): - # get s, a, s', s_0 from batch - state = batch['state'] - action = batch['action'] - next_state = batch['next_state'] - initial_state = state - - # append action to state batches - # use expert action for s - state['action'] = action - # run s' and s_0 through policy - initial_state['action'] = policy(initial_state) - next_state['action'] = policy(next_state) - - # transform state and action before inputting to value network - # (for the policy network this is done in policy.__call__() ) - state = policy.transform_observation(state) - initial_state = policy.transform_observation(initial_state) - next_state = policy.transform_observation(next_state) - - # evaluate value network - value = policy.value(state) - value_init = policy.value(initial_state) - value_next = policy.value(next_state) - - # linear loss - linear_loss = (1 - discount) * torch.mean(value_init) - - # nonlinear loss - value_diff = value - discount * value_next - nonlinear_loss = torch.logsumexp(value_diff, dim=0) - np.log(len(value_diff)) - - loss = nonlinear_loss - linear_loss - return loss - - - policy_optimizer = optimizer_factory(config['policy_optim'], policy.policy_parameters()) - value_optimizer = optimizer_factory(config['value_optim'], policy.value_parameters()) - - # generate tensorboard writer - if not using_ray: - writer = SummaryWriter(filestr) - - for i in tqdm(range(train_epochs)): - - # save model checkpoints - if i % checkpoint_every == 0: - policy.save_model(filestr + '_epoch%04i'%(i) ) - - # train - epoch_loss = 0 - for (batch_idx, batch) in enumerate(training_loader): - - loss = f_value_dice_loss(batch) - - # In original implementation policy is regularized with orthogonal regularization, - # value with L2 regularization on gradients - policy_loss = -loss - value_loss = loss - - # # compute loss and step optimizer - # policy_optimizer.zero_grad() - # value_optimizer.zero_grad() - # policy_loss.backward(retain_graph=True) - # value_loss.backward() - - # clip_grad_norm_(policy.policy.parameters(), clip_grad_norm) - # clip_grad_norm_(policy.value.parameters(), clip_grad_norm) - - # policy_optimizer.step() - # value_optimizer.step() - - if batch_idx % 2 == 0: - policy_optimizer.zero_grad() - policy_loss.backward() - clip_grad_norm_(policy.policy.parameters(), clip_grad_norm) - policy_optimizer.step() - else: - value_optimizer.zero_grad() - value_loss.backward() - clip_grad_norm_(policy.value.parameters(), clip_grad_norm) - value_optimizer.step() - - epoch_loss += loss.item() / len(train_dataset) - - if i % print_epoch_every == 0: - print('Epoch: {}, Training Loss: {}'.format(i, epoch_loss)) - - # measure cv loss - if i % cv_every == 0: - with torch.no_grad(): - cv_loss = 0. - for (batch_idx, batch) in enumerate(cv_loader): - loss = f_value_dice_loss(batch) - cv_loss += loss.item() / len(cv_dataset) - - - # Write epoch loss - if using_ray: - if i % cv_every == 0: - tune.report(training_loss=epoch_loss, cv_loss=cv_loss, training_iteration=i+1) - else: - tune.report(training_loss=epoch_loss, training_iteration=i+1) - else: - writer.add_scalar('training loss',epoch_loss, i) - if i % cv_every == 0: - writer.add_scalar('cv loss', cv_loss, i) - - if i % print_cv_every == 0: - print('Epoch: {}, CV Loss: {}'.format(i, cv_loss)) - - - policy.save_model(filestr)