From 5c40de66fa158d4ab82faf19fca34bf7ead75153 Mon Sep 17 00:00:00 2001 From: Johannes Fischer Date: Wed, 4 Aug 2021 19:34:49 +0200 Subject: [PATCH] Implement ValueDICE and some restructuring --- config/networks.json5 | 2 +- config/value_dice.json5 | 84 ++++++++ src/bc/bc.py | 86 +-------- src/main.py | 4 +- src/policies/policy.py | 128 +++++++++++- src/util/nn_training.py | 9 + src/{value-dice => value_dice}/__init__.py | 0 src/value_dice/value_dice.py | 215 +++++++++++++++++++++ 8 files changed, 442 insertions(+), 86 deletions(-) create mode 100644 config/value_dice.json5 create mode 100644 src/util/nn_training.py rename src/{value-dice => value_dice}/__init__.py (100%) create mode 100644 src/value_dice/value_dice.py diff --git a/config/networks.json5 b/config/networks.json5 index ebce18f..ac10d2f 100644 --- a/config/networks.json5 +++ b/config/networks.json5 @@ -36,7 +36,7 @@ lr: 1e-3, weight_decay: 0.1, }, - train_epochs: 200, + train_epochs: 8, train_batch_size: 32, loss: 'huber', } \ No newline at end of file diff --git a/config/value_dice.json5 b/config/value_dice.json5 new file mode 100644 index 0000000..0952bd1 --- /dev/null +++ b/config/value_dice.json5 @@ -0,0 +1,84 @@ +{ + 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: 'sigmoid', + }, + }, + policy_optim: { + optimizer: 'adam', + lr: 1e-3, + weight_decay: 0.1, + }, + value_optim: { + optimizer: 'adam', + lr: 1e-3, + weight_decay: 0.1, + }, + train_epochs: 8, + train_batch_size: 32, + discount: 0.95, +} \ No newline at end of file diff --git a/src/bc/bc.py b/src/bc/bc.py index 687e476..d064ff3 100644 --- a/src/bc/bc.py +++ b/src/bc/bc.py @@ -4,8 +4,9 @@ from torch.utils.data import DataLoader import pickle from torch.utils.tensorboard import SummaryWriter -from src.policies import DeepSetsPolicy +from src.policies import IntersimDeepSetsNet, 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 @@ -46,62 +47,21 @@ def bc_config(ray_config): } return config -class BehaviorCloningPolicy(): +class BehaviorCloningPolicy(IntersimPolicy): """ Class for (continuous) behavior cloning policy """ - def __init__(self, config: dict, transforms: dict={}): + def __init__(self, config: dict, transforms: dict): """ Initialize BehaviorCloningPolicy Args: - config (dict): configuration file to initialize DeepSetsPolicy with + config (dict): configuration file to initialize IntersimDeepSetsNet with transforms (dict): dictionary of transforms to apply to different fields """ - self._config = config - self._transforms = transforms - self._policy = DeepSetsPolicy(config) + super(BehaviorCloningPolicy, self).__init__(config, transforms) + self._policy = IntersimStateNet(config) - @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 __call__(self, ob): - - if 'action' in ob.keys(): - # extract state from dataloader samples - pass - else: - # extract state from observation (using simulator) - ob['path_x'] = ob['paths'][0] - ob['path_y'] = ob['paths'][1] - - # run observation through transforms - for key in ['state', 'relative_state', 'path_x', 'path_y']: - if key in self._transforms.keys(): - ob[key] = self._transforms[key].transform(ob[key]) - - # 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 - @classmethod def load_model(cls, filestr: str, config: dict = None): """ @@ -141,24 +101,6 @@ class BehaviorCloningPolicy(): pickle.dump(self._transforms, open(filestr+'_transforms.pkl', 'wb')) torch.save(self._policy.state_dict(), filestr+'_model.pt') -def generate_transforms(dataset): - """ - Generate transform dictionary from dataset - Args: - dataset (Dataset): dataset of demo observations and actions - """ - transforms = { - 'action': MinMaxScaler(), - 'state': MinMaxScaler(), - 'relative_state': MinMaxScaler(reduce_dim=2), - 'path_x': MinMaxScaler(reduce_dim=2), - 'path_y': MinMaxScaler(reduce_dim=2), - } - for key in transforms.keys(): - transforms[key].fit(dataset[:][key]) - - return transforms - def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs): using_ray = kwargs.get('ray', False) @@ -169,9 +111,6 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs): 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 = 1 print_epoch_every = 1000 @@ -179,12 +118,6 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs): checkpoint_every = 100 cv_batch_size = 256 # doesn't matter - # generate transform from train_dataset - transforms = generate_transforms(train_dataset) - - # initialize policy - policy.transforms = transforms - # training and testing dataloaders training_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True) cv_loader = DataLoader(cv_dataset, batch_size=cv_batch_size, shuffle=True) @@ -200,10 +133,7 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs): 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 + optimizer = optimizer_factory(config['optim'], policy.parameters) # generate tensorboard writer if not using_ray: diff --git a/src/main.py b/src/main.py index 943f8a7..9aa308d 100644 --- a/src/main.py +++ b/src/main.py @@ -47,8 +47,10 @@ def main(config, method='bc', train=False, test=False, loc=0, datadir='./expert_ if train: # make policy, train and test datasets, and send to - policy = policy_class(config) train_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[0,1,2]) + # generate transform from train_dataset + transforms = generate_transforms(train_dataset) + policy = policy_class(config, transforms) cv_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[3]) train_fn(config, policy, train_dataset, cv_dataset, filestr, **kwargs) diff --git a/src/policies/policy.py b/src/policies/policy.py index 524345e..eb3bf72 100644 --- a/src/policies/policy.py +++ b/src/policies/policy.py @@ -4,16 +4,13 @@ from torch import nn from src.nets.deepsets import DeepSetsModule, Phi -class Policy: - pass - -class DeepSetsPolicy(Policy, nn.Module): +class IntersimStateNet(nn.Module): def __init__(self, config): """ Args: config (dict): dictionary for configuring the deep sets policy """ - super(DeepSetsPolicy, self).__init__() + super(IntersimStateNet, self).__init__() ego_config = config['ego_encoder'] deepsets_config = config['deepsets'] pathnet_config = config['path_encoder'] @@ -40,7 +37,7 @@ class DeepSetsPolicy(Policy, nn.Module): Returns: x (torch.tensor): (head_output_dim,) output of common head network """ - ego = self.ego_net(sample["state"]) + ego = self.ego_net(sample["ego_state"]) relative = self.deepsets_net(sample["relative_state"]) # cat path_x, path_y to tensor of dim (B, 2*P) path = torch.cat([sample["path_x"], sample["path_y"]], dim=-1) @@ -48,3 +45,122 @@ class DeepSetsPolicy(Policy, nn.Module): 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"]) + # cat path_x, path_y to tensor of dim (B, 2*P) + path = torch.cat([sample["path_x"], sample["path_y"]], dim=-1) + path = self.path_net(path) + action = sample["action"] + 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(): + transformed_ob[key] = self._transforms[key].transform(ob[key]) + return transformed_ob + + def __call__(self, ob): + + if 'action' in ob.keys(): + # extract state from dataloader samples + pass + else: + # extract state from observation (using simulator) + ob['path_x'] = ob['paths'][0] + ob['path_y'] = ob['paths'][1] + + ob = 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(), + 'state': MinMaxScaler(), + 'relative_state': MinMaxScaler(reduce_dim=2), + 'path_x': MinMaxScaler(reduce_dim=2), + 'path_y': MinMaxScaler(reduce_dim=2), + } + for key in transforms.keys(): + transforms[key].fit(dataset[:][key]) + + return transforms diff --git a/src/util/nn_training.py b/src/util/nn_training.py new file mode 100644 index 0000000..e553776 --- /dev/null +++ b/src/util/nn_training.py @@ -0,0 +1,9 @@ +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/value-dice/__init__.py b/src/value_dice/__init__.py similarity index 100% rename from src/value-dice/__init__.py rename to src/value_dice/__init__.py diff --git a/src/value_dice/value_dice.py b/src/value_dice/value_dice.py new file mode 100644 index 0000000..20f14db --- /dev/null +++ b/src/value_dice/value_dice.py @@ -0,0 +1,215 @@ +import torch +import torch.nn as nn +from torch.utils.data import DataLoader +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 + +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 + + @policy.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()) + + @property + def policy_parameters(self): + return self.policy.parameters() + + @property + 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 + loss_type = config['loss'] + train_epochs = config['train_epochs'] + train_batch_size = config['train_batch_size'] + discount = config['discount'] + + 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'].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 + + ### Linear loss + + # append action to state batches + # use expert action for s + state['action'] = action + # run s' and s_0 through policy + initial_state['action'] = policy(next_state) + next_state['action'] = policy(initial_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_batch) + + value_diff = value - discount * value_next + + linear_loss = (1 - discount) * torch.mean(value_init) + + ### Nonlinear loss + nonlinear_loss = torch.logsumexp(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) + + # TODO: Regularization + policy_loss = -loss #+ ORTHOGONAL_REGULARIZER + value_loss = loss #+ GRADIENT_REGULARIZER + + # compute loss and step optimizer + policy_optimizer.zero_grad() + loss.backward() + policy_optimizer.step() + value_optimizer.zero_grad() + value_loss.backward() + 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)