huangfu/InteractionImitation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implement ValueDICE and some restructuring

Browse Source
This commit is contained in:
Johannes Fischer
2021-08-04 19:34:49 +02:00
parent cba42c6e4d
commit 5c40de66fa
8 changed files with 442 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
config/networks.json5
View File

@@ -36,7 +36,7 @@
lr: 1e-3,
weight_decay: 0.1,
},
train_epochs: 200,
train_epochs: 8,
train_batch_size: 32,
loss: 'huber',
}

84
config/value_dice.json5 Normal file
View File

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

86
src/bc/bc.py
View File

@@ -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,61 +47,20 @@ 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)
@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
super(BehaviorCloningPolicy, self).__init__(config, transforms)
self._policy = IntersimStateNet(config)
@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:

4
src/main.py
View File

@@ -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)

128
src/policies/policy.py
View File

@@ -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

9
src/util/nn_training.py Normal file
View File

@@ -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

0
src/value-dice/__init__.py → src/value_dice/__init__.py
View File

215
src/value_dice/value_dice.py Normal file
View File

@@ -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)