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

Bugfixes in valuedice

Browse Source
This commit is contained in:
Johannes Fischer
2021-08-04 20:53:03 +02:00
parent 2224e2cd14
commit 40c55478f3
8 changed files with 71 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
src/bc/bc.py
View File

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

7
src/main.py
View File

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

2
src/policies/__init__.py
View File

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

9
src/policies/policy.py
View File

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

2
src/util/nn_training.py
View File

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

1
src/value_dice/__init__.py
View File

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

76
src/value_dice/value_dice.py
View File

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