Bugfixes in valuedice

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,

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.policies import IntersimStateNet, IntersimPolicy, generate_transforms
-from src.util.transform import MinMaxScaler
 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:

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:

src/policies/__init__.py

@@ -1 +1 @@
-from src.policies.policy import DeepSetsPolicy
+from src.policies.policy import IntersimPolicy, IntersimStateNet, IntersimStateActionNet, generate_transforms

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)

src/util/nn_training.py

@@ -1,3 +1,5 @@
+import torch
+
 def optimizer_factory(config, parameters):
     optimizer_type = config['optimizer']
     learning_rate = config['lr']

src/value_dice/__init__.py

@@ -0,0 +1 @@
+from src.value_dice.value_dice import ValueDicePolicy, train, vd_config

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)
+        initial_state['action'] = policy(initial_state)
-        next_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 = (policy.value(state))
-        value_init = policy.value(initial_state)  
+        value_init = (policy.value(initial_state))
-        value_next = policy.value(next_batch)
+        value_next = (policy.value(next_state))
-
-        value_diff = value - discount * value_next
         
+        # linear loss
         linear_loss = (1 - discount) * torch.mean(value_init)
 
-        ### Nonlinear loss
+        # nonlinear loss
-        nonlinear_loss = torch.logsumexp(value_diff)
+        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)
+    policy_optimizer = optimizer_factory(config['policy_optim'], policy.policy_parameters())
-    value_optimizer = optimizer_factory(config['value_optim'], policy.value_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()