diff --git a/experiments/experiment.py b/experiments/experiment.py index 17ae7c5..40235d0 100644 --- a/experiments/experiment.py +++ b/experiments/experiment.py @@ -41,7 +41,7 @@ def parse_args(): parser.add_argument("--test", help="test model", action="store_true") parser.add_argument("--method", help="modeling method", - choices=['bc', 'gail', 'advil'], default='bc') + 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, diff --git a/src/bc/bc.py b/src/bc/bc.py index 5f00916..ee3b1d7 100644 --- a/src/bc/bc.py +++ b/src/bc/bc.py @@ -4,8 +4,7 @@ from torch.utils.data import DataLoader import pickle from torch.utils.tensorboard import SummaryWriter -from src.policies import IntersimDeepSetsNet, IntersimPolicy, generate_transforms -from src.util.transform import MinMaxScaler +from src.policies import IntersimStateNet, IntersimPolicy, generate_transforms from src.util.nn_training import optimizer_factory from tqdm import tqdm import json5 @@ -133,7 +132,7 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs): loss_fn = nn.MSELoss(reduction='sum') else: raise NotImplementedError - optimizer = optimizer_factory(config['optim'], policy.parameters) + 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 551d797..edf255c 100644 --- a/src/main.py +++ b/src/main.py @@ -8,6 +8,7 @@ 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): """ @@ -40,8 +41,12 @@ def main(config, method='bc', train=False, test=False, loc=0, datadir='./expert_ 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 + raise NotImplementedError("Method {} not implemented".format(method)) # default train / cv / test split datasets if train: diff --git a/src/policies/__init__.py b/src/policies/__init__.py index ec9d984..cff41ff 100644 --- a/src/policies/__init__.py +++ b/src/policies/__init__.py @@ -1 +1 @@ -from src.policies.policy import DeepSetsPolicy \ No newline at end of file +from src.policies.policy import IntersimPolicy, IntersimStateNet, IntersimStateActionNet, generate_transforms diff --git a/src/policies/policy.py b/src/policies/policy.py index a4aa132..274d34b 100644 --- a/src/policies/policy.py +++ b/src/policies/policy.py @@ -3,6 +3,7 @@ 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): @@ -80,10 +81,8 @@ class IntersimStateActionNet(nn.Module): """ 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"] + 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 @@ -118,7 +117,7 @@ class IntersimPolicy(): # run observation through transforms transformed_ob = {} for key in ['ego_state', 'relative_state', 'path', 'action']: - if key in self._transforms.keys(): + if key in self._transforms.keys() and key in ob.keys(): transformed_ob[key] = self._transforms[key].transform(ob[key]) return transformed_ob @@ -132,7 +131,7 @@ class IntersimPolicy(): ob['ego_state'] = ob['state'] ob['path'] = torch.stack(ob['paths'],dim=-1) - ob = transform_observation(ob) + ob = self.transform_observation(ob) # run transformed state through model action = self._policy(ob) diff --git a/src/util/nn_training.py b/src/util/nn_training.py index e553776..eea6eea 100644 --- a/src/util/nn_training.py +++ b/src/util/nn_training.py @@ -1,3 +1,5 @@ +import torch + def optimizer_factory(config, parameters): optimizer_type = config['optimizer'] learning_rate = config['lr'] diff --git a/src/value_dice/__init__.py b/src/value_dice/__init__.py index e69de29..bb17233 100644 --- a/src/value_dice/__init__.py +++ b/src/value_dice/__init__.py @@ -0,0 +1 @@ +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 index 20f14db..682f40e 100644 --- a/src/value_dice/value_dice.py +++ b/src/value_dice/value_dice.py @@ -11,6 +11,41 @@ from src.util.nn_training import optimizer_factory from tqdm import tqdm import json5 from ray import tune +def vd_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 ValueDicePolicy(IntersimPolicy): """ @@ -32,7 +67,7 @@ class ValueDicePolicy(IntersimPolicy): def value(self): return self._value - @policy.setter + @value.setter def value(self, value): self._value = value @@ -58,11 +93,9 @@ class ValueDicePolicy(IntersimPolicy): 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() @@ -95,7 +128,6 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs): print('using ray') # hyperparams - loss_type = config['loss'] train_epochs = config['train_epochs'] train_batch_size = config['train_batch_size'] discount = config['discount'] @@ -111,24 +143,24 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs): 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) + 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) + 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) + 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__() ) @@ -137,23 +169,23 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs): 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 - + 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 - nonlinear_loss = torch.logsumexp(value_diff) + # nonlinear loss + value_diff = value - discount * value_next + nonlinear_loss = torch.logsumexp(value_diff, dim=0) 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) + 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: @@ -171,13 +203,13 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs): loss = f_value_dice_loss(batch) - # TODO: Regularization + # TODO: Regularization is done in original source code policy_loss = -loss #+ ORTHOGONAL_REGULARIZER value_loss = loss #+ GRADIENT_REGULARIZER # compute loss and step optimizer policy_optimizer.zero_grad() - loss.backward() + policy_loss.backward(retain_graph=True) policy_optimizer.step() value_optimizer.zero_grad() value_loss.backward()