Implement ValueDICE and some restructuring

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',
 }

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,
+}

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
+        super(BehaviorCloningPolicy, self).__init__(config, transforms)
-        self._transforms = transforms
+        self._policy = IntersimStateNet(config)
-        self._policy = DeepSetsPolicy(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 = optimizer_factory(config['optim'], policy.parameters)
-        optimizer = torch.optim.Adam(policy.parameters(), lr=learning_rate, weight_decay=weight_decay)
-    else:
-        raise NotImplementedError
 
     # generate tensorboard writer
     if not using_ray:

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)
     

src/policies/policy.py

@@ -4,16 +4,13 @@ from torch import nn
 
 from src.nets.deepsets import DeepSetsModule, Phi
 
-class Policy:
+class IntersimStateNet(nn.Module):
-    pass
-
-class DeepSetsPolicy(Policy, 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

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

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)