purging unused files
This commit is contained in:
Arec Jamgochian
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -1,42 +0,0 @@
|
|||||||
{
|
|
||||||
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',
|
|
||||||
},
|
|
||||||
optim: {
|
|
||||||
optimizer: 'adam',
|
|
||||||
lr: 1e-3,
|
|
||||||
weight_decay: 0.1,
|
|
||||||
},
|
|
||||||
train_epochs: 200,
|
|
||||||
train_batch_size: 32,
|
|
||||||
loss: 'huber',
|
|
||||||
}
|
|
||||||
@@ -1,85 +0,0 @@
|
|||||||
{
|
|
||||||
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: 'id',
|
|
||||||
},
|
|
||||||
},
|
|
||||||
policy_optim: {
|
|
||||||
optimizer: 'adam',
|
|
||||||
lr: 1e-3,
|
|
||||||
weight_decay: 0.1,
|
|
||||||
},
|
|
||||||
value_optim: {
|
|
||||||
optimizer: 'adam',
|
|
||||||
lr: 1e-3,
|
|
||||||
weight_decay: 0.1,
|
|
||||||
},
|
|
||||||
train_epochs: 200,
|
|
||||||
train_batch_size: 32,
|
|
||||||
discount: 0.95,
|
|
||||||
clip_grad_norm: 1.,
|
|
||||||
}
|
|
||||||
@@ -1,203 +0,0 @@
|
|||||||
import json5
|
|
||||||
from functools import partial
|
|
||||||
import os
|
|
||||||
opj = os.path.join
|
|
||||||
|
|
||||||
# set up ray tune
|
|
||||||
import ray
|
|
||||||
from ray import tune
|
|
||||||
from ray.tune import Analysis, ExperimentAnalysis
|
|
||||||
from ray.tune.schedulers import ASHAScheduler
|
|
||||||
from hyperopt import hp
|
|
||||||
from ray.tune.suggest.hyperopt import HyperOptSearch
|
|
||||||
|
|
||||||
# get graphs
|
|
||||||
import intersim
|
|
||||||
from intersim.graphs import ConeVisibilityGraph
|
|
||||||
|
|
||||||
|
|
||||||
from src.main import basestr, main
|
|
||||||
|
|
||||||
def parse_args():
|
|
||||||
"""
|
|
||||||
Parse arguments to main
|
|
||||||
Returns:
|
|
||||||
kwargs: dictionary of arguments:
|
|
||||||
train (bool): whether to run train loop
|
|
||||||
test (bool): whether to run test loop
|
|
||||||
method (str): the method to try for imitation
|
|
||||||
loc (int): the location index of the roundabout
|
|
||||||
config (str): config path
|
|
||||||
seed (int): RNG seed
|
|
||||||
"""
|
|
||||||
import argparse
|
|
||||||
parser = argparse.ArgumentParser(description='Save Expert Trajectories')
|
|
||||||
parser.add_argument('--loc', default=0, type=int,
|
|
||||||
help='location (default 0)')
|
|
||||||
parser.add_argument("--train", help="train model",
|
|
||||||
action="store_true")
|
|
||||||
parser.add_argument("--ray", help="use ray tune to run multiple experiments",
|
|
||||||
action="store_true")
|
|
||||||
parser.add_argument("--test", help="test model",
|
|
||||||
action="store_true")
|
|
||||||
parser.add_argument("--method", help="modeling method",
|
|
||||||
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,
|
|
||||||
help='seed')
|
|
||||||
parser.add_argument('--nframes', default=500, type=int,
|
|
||||||
help='frames for test animation')
|
|
||||||
parser.add_argument('--nsamples', default=200, type=int,
|
|
||||||
help='number of ray samples')
|
|
||||||
parser.add_argument('--graph', action='store_true',
|
|
||||||
help='whether to mask the relative states based on a ConeVisibilityGraph')
|
|
||||||
parser.add_argument('-d', default='./expert_data', type=str,
|
|
||||||
help='data directory')
|
|
||||||
parser.add_argument('-o', default=None, type=str,
|
|
||||||
help='output directory')
|
|
||||||
args = parser.parse_args()
|
|
||||||
kwargs = {
|
|
||||||
'train':args.train,
|
|
||||||
'test':args.test,
|
|
||||||
'method':args.method,
|
|
||||||
'loc':args.loc,
|
|
||||||
'config_path':args.config,
|
|
||||||
'seed':args.seed,
|
|
||||||
'ray':args.ray,
|
|
||||||
'nframes':args.nframes,
|
|
||||||
'nsamples':args.nsamples,
|
|
||||||
'datadir':os.path.abspath(args.d),
|
|
||||||
'graph':None,
|
|
||||||
'outdir': opj('output',args.method,'loc%02i'%(args.loc)),
|
|
||||||
'train_tracks':[0,1,2],
|
|
||||||
'cv_tracks':[3],
|
|
||||||
'test_tracks':[4],
|
|
||||||
}
|
|
||||||
if args.o:
|
|
||||||
kwargs['outdir'] = args.o
|
|
||||||
if args.graph:
|
|
||||||
kwargs['graph'] = ConeVisibilityGraph(r=20, half_angle=120)
|
|
||||||
return kwargs
|
|
||||||
|
|
||||||
def get_full_config(ray_config:dict, method:str)->dict:
|
|
||||||
"""
|
|
||||||
Get full model configuration from ray config and method string
|
|
||||||
Args:
|
|
||||||
ray_config (dict): ray config
|
|
||||||
method (str): method to get full configuration for
|
|
||||||
"""
|
|
||||||
if method == 'bc':
|
|
||||||
from src.bc import bc_config
|
|
||||||
config = bc_config(ray_config)
|
|
||||||
elif method == 'vd':
|
|
||||||
from src.value_dice import vd_config
|
|
||||||
config = vd_config(ray_config)
|
|
||||||
else:
|
|
||||||
raise NotImplementedError
|
|
||||||
return config
|
|
||||||
|
|
||||||
def get_ray_config(method:str)->dict:
|
|
||||||
"""
|
|
||||||
Get configuration for ray based on method.
|
|
||||||
Args:
|
|
||||||
method (str): method to get configuration for
|
|
||||||
Returns:
|
|
||||||
ray_config (dict): configuration for ray
|
|
||||||
"""
|
|
||||||
if method == 'bc':
|
|
||||||
ray_config = {
|
|
||||||
"lr": tune.loguniform(1e-5, 1e-3),
|
|
||||||
"weight_decay": tune.choice([0, 0.1]),
|
|
||||||
"loss": tune.choice(['huber', 'mse']),
|
|
||||||
"train_batch_size": tune.choice([16,32,64]),
|
|
||||||
"deepsets_phi_hidden_n": tune.randint(1,5),
|
|
||||||
"deepsets_phi_hidden_dim": tune.lograndint(8,65),
|
|
||||||
"deepsets_latent_dim": tune.lograndint(8,129),
|
|
||||||
"deepsets_rho_hidden_n": tune.randint(0,3),
|
|
||||||
"deepsets_rho_hidden_dim": tune.lograndint(8,129),
|
|
||||||
"deepsets_output_dim": tune.lograndint(4,129),
|
|
||||||
"head_hidden_n": tune.randint(1,6),
|
|
||||||
"head_hidden_dim": tune.lograndint(16,257),
|
|
||||||
"head_final_activation": tune.choice(['sigmoid', None]),
|
|
||||||
}
|
|
||||||
elif method == 'vd':
|
|
||||||
ray_config = {
|
|
||||||
"policy_lr": tune.loguniform(1e-5, 1e-3),
|
|
||||||
"value_lr": tune.loguniform(1e-5, 1e-3),
|
|
||||||
"policy_weight_decay": tune.choice([0, 0.1]),
|
|
||||||
"value_weight_decay": tune.choice([0, 0.1]),
|
|
||||||
"train_batch_size": tune.choice([16,32,64]),
|
|
||||||
"deepsets_phi_hidden_n": tune.randint(1,5),
|
|
||||||
"deepsets_phi_hidden_dim": tune.lograndint(8,65),
|
|
||||||
"deepsets_latent_dim": tune.lograndint(8,129),
|
|
||||||
"deepsets_rho_hidden_n": tune.randint(0,3),
|
|
||||||
"deepsets_rho_hidden_dim": tune.lograndint(8,129),
|
|
||||||
"deepsets_output_dim": tune.lograndint(4,129),
|
|
||||||
"head_hidden_n": tune.randint(1,6),
|
|
||||||
"head_hidden_dim": tune.lograndint(16,257),
|
|
||||||
"head_final_activation": tune.choice(['sigmoid', None]),
|
|
||||||
"clip_grad_norm": tune.choice([.5, 1., 5., 10.]),
|
|
||||||
"discount": tune.choice([.95, .99])
|
|
||||||
}
|
|
||||||
else:
|
|
||||||
raise NotImplementedError
|
|
||||||
return ray_config
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
|
||||||
kwargs = parse_args()
|
|
||||||
|
|
||||||
# make prefix of output files
|
|
||||||
|
|
||||||
if kwargs['config_path']:
|
|
||||||
# load config
|
|
||||||
with open(kwargs['config_path'], 'r') as cfg:
|
|
||||||
config = json5.load(cfg)
|
|
||||||
if not os.path.isdir(kwargs['outdir']):
|
|
||||||
os.makedirs(kwargs['outdir'])
|
|
||||||
filestr = opj(kwargs['outdir'], basestr(**kwargs))
|
|
||||||
if kwargs['ray']:
|
|
||||||
filestr = kwargs['config_path'].replace('_config.json','')
|
|
||||||
main(config, filestr=filestr, **kwargs)
|
|
||||||
|
|
||||||
elif kwargs['ray'] and kwargs['train']:
|
|
||||||
|
|
||||||
ray.shutdown()
|
|
||||||
ray.init(log_to_driver=False)
|
|
||||||
|
|
||||||
def ray_train(config, datadir=None):
|
|
||||||
full_config = get_full_config(config, kwargs['method'])
|
|
||||||
main(full_config, filestr='exp', **kwargs)
|
|
||||||
|
|
||||||
ray_config = get_ray_config(kwargs['method'])
|
|
||||||
search = HyperOptSearch(ray_config, max_concurrent=8, metric='cv_loss',mode="min",)
|
|
||||||
custom_scheduler = ASHAScheduler(metric='cv_loss', mode="min", grace_period=15)
|
|
||||||
|
|
||||||
analysis = tune.run(
|
|
||||||
ray_train,
|
|
||||||
#config=ray_config,
|
|
||||||
search_alg=search,
|
|
||||||
scheduler=custom_scheduler,
|
|
||||||
local_dir=kwargs['outdir'],
|
|
||||||
#resources_per_trial={"cpu": 2},
|
|
||||||
time_budget_s=120*60,
|
|
||||||
num_samples=kwargs['nsamples'],
|
|
||||||
)
|
|
||||||
elif kwargs['ray'] and kwargs['test']:
|
|
||||||
analysis = Analysis(kwargs['outdir'], default_metric="cv_loss", default_mode="min")
|
|
||||||
config = analysis.get_best_config()
|
|
||||||
filepath = analysis.get_best_logdir()
|
|
||||||
filestr = opj(filepath, 'exp')
|
|
||||||
config_path = filestr+'_config.json'
|
|
||||||
with open(config_path, 'r') as cfg:
|
|
||||||
config = json5.load(cfg)
|
|
||||||
print("Best ray experiment:", filepath)
|
|
||||||
main(config, filestr=filestr, **kwargs)
|
|
||||||
else:
|
|
||||||
raise Exception('No valid config found')
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
@@ -1,9 +0,0 @@
|
|||||||
#!/bin/sh
|
|
||||||
|
|
||||||
python experiments/experiment.py --ray --train -d ./expert_data/base
|
|
||||||
python experiments/experiment.py --ray --test -d ./expert_data/base --nframes 1000
|
|
||||||
python experiments/experiment.py --ray --train -d ./expert_data/reg
|
|
||||||
python experiments/experiment.py --ray --test -d ./expert_data/reg --nframes 1000
|
|
||||||
python experiments/experiment.py --ray --train -d ./expert_data/reg_graph --graph
|
|
||||||
python experiments/experiment.py --ray --test -d ./expert_data/reg_graph --graph --nframes 1000
|
|
||||||
|
|
||||||
@@ -1,5 +0,0 @@
|
|||||||
#!/bin/sh
|
|
||||||
|
|
||||||
# python experiments/experiment.py --method vd --train --ray -d expert_data/reg -o output/vd/loc00/reg --nsamples 400
|
|
||||||
# python experiments/experiment.py --test --ray --method vd -d expert_data/normal -o output/vd/loc00/normal --nframes 1000
|
|
||||||
python experiments/experiment.py --train --method vd --config config/value_dice.json5
|
|
||||||
@@ -1 +0,0 @@
|
|||||||
from src.bc.bc import BehaviorCloningPolicy, train, bc_config
|
|
||||||
191
src/bc/bc.py
191
src/bc/bc.py
@@ -1,191 +0,0 @@
|
|||||||
import torch
|
|
||||||
import torch.nn as nn
|
|
||||||
from torch.utils.data import DataLoader
|
|
||||||
import pickle
|
|
||||||
from torch.utils.tensorboard import SummaryWriter
|
|
||||||
|
|
||||||
from src.policies import IntersimStateNet, IntersimPolicy, generate_transforms
|
|
||||||
from src.util.nn_training import optimizer_factory
|
|
||||||
from tqdm import tqdm
|
|
||||||
import json5
|
|
||||||
from ray import tune
|
|
||||||
|
|
||||||
def bc_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 BehaviorCloningPolicy(IntersimPolicy):
|
|
||||||
"""
|
|
||||||
Class for (continuous) behavior cloning policy
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, config: dict, transforms: dict):
|
|
||||||
"""
|
|
||||||
Initialize BehaviorCloningPolicy
|
|
||||||
Args:
|
|
||||||
config (dict): configuration file to initialize IntersimDeepSetsNet with
|
|
||||||
transforms (dict): dictionary of transforms to apply to different fields
|
|
||||||
"""
|
|
||||||
super(BehaviorCloningPolicy, self).__init__(config, transforms)
|
|
||||||
self._policy = IntersimStateNet(config)
|
|
||||||
|
|
||||||
@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+'_model.pt'))
|
|
||||||
return model
|
|
||||||
|
|
||||||
def eval(self):
|
|
||||||
self._policy.eval()
|
|
||||||
|
|
||||||
def parameters(self):
|
|
||||||
return self._policy.parameters()
|
|
||||||
|
|
||||||
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+'_model.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']
|
|
||||||
|
|
||||||
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']['ego_state'].dtype)
|
|
||||||
|
|
||||||
# generate loss function, optimizer
|
|
||||||
cv_loss_fn = nn.MSELoss(reduction='sum')
|
|
||||||
if loss_type == 'huber':
|
|
||||||
loss_fn = nn.HuberLoss(reduction='sum')
|
|
||||||
elif loss_type == 'mse':
|
|
||||||
loss_fn = nn.MSELoss(reduction='sum')
|
|
||||||
else:
|
|
||||||
raise NotImplementedError
|
|
||||||
optimizer = optimizer_factory(config['optim'], policy.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):
|
|
||||||
|
|
||||||
# sample mini-batch and run through policy
|
|
||||||
pred_action = policy(batch['state'])
|
|
||||||
loss = loss_fn(pred_action, batch['action'])
|
|
||||||
|
|
||||||
# compute loss and step optimizer
|
|
||||||
optimizer.zero_grad()
|
|
||||||
loss.backward()
|
|
||||||
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):
|
|
||||||
pred_action = policy(batch['state'])
|
|
||||||
loss = cv_loss_fn(pred_action, batch['action'])
|
|
||||||
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)
|
|
||||||
@@ -1,96 +0,0 @@
|
|||||||
import torch
|
|
||||||
from torch.utils.data import Dataset
|
|
||||||
import numpy as np
|
|
||||||
from src.data.expert_data import load_expert_data
|
|
||||||
import os
|
|
||||||
opj = os.path.join
|
|
||||||
|
|
||||||
class InteractionDatasetMultiAgent(Dataset):
|
|
||||||
"""
|
|
||||||
Class to handle getting full multi-agent observations and actions
|
|
||||||
"""
|
|
||||||
pass
|
|
||||||
|
|
||||||
class InteractionDatasetSingleAgent(Dataset):
|
|
||||||
"""Class to load states and actions for individual agents."""
|
|
||||||
|
|
||||||
def __init__(self, output_dir='expert_data', loc:int = 0, tracks:list = [0], dtype=torch.float32):
|
|
||||||
"""
|
|
||||||
Args:
|
|
||||||
output_dir (string): Directory with all the images.
|
|
||||||
loc (int): location index
|
|
||||||
tracks (list[int]): track indices
|
|
||||||
"""
|
|
||||||
self.output_dir = output_dir
|
|
||||||
self.loc = loc
|
|
||||||
self.tracks = tracks
|
|
||||||
self.dtype = dtype
|
|
||||||
self.keys = ['ego_state', 'relative_state', 'path', 'action', 'next_ego_state', 'next_relative_state', 'next_path']
|
|
||||||
self._load_dataset()
|
|
||||||
|
|
||||||
def _load_dataset(self):
|
|
||||||
"""
|
|
||||||
Load the full datasets ahead of time
|
|
||||||
"""
|
|
||||||
self.raw_data = {key:[] for key in self.keys}
|
|
||||||
max_nv = 0
|
|
||||||
for track in self.tracks:
|
|
||||||
try:
|
|
||||||
data = load_expert_data(path=self.output_dir, loc=self.loc, track=track)
|
|
||||||
print('Loaded location {} track {}'.format(self.loc,track))
|
|
||||||
except:
|
|
||||||
print('Failed to load location {} track {}'.format(self.loc,track))
|
|
||||||
continue
|
|
||||||
max_nv = max(max_nv, data['relative_state'].shape[1])
|
|
||||||
for key in self.keys:
|
|
||||||
self.raw_data[key].append(data[key])
|
|
||||||
|
|
||||||
# pad second dimension of relative state
|
|
||||||
for i in range(len(self.raw_data['relative_state'])):
|
|
||||||
nv1, nv2, d = self.raw_data['relative_state'][i].shape
|
|
||||||
pad = torch.zeros(nv1, max_nv-nv2, d, dtype=self.dtype) * np.nan
|
|
||||||
self.raw_data['relative_state'][i] = torch.cat((self.raw_data['relative_state'][i], pad), dim=1)
|
|
||||||
self.raw_data['next_relative_state'][i] = torch.cat((self.raw_data['next_relative_state'][i], pad), dim=1)
|
|
||||||
|
|
||||||
# cat lists
|
|
||||||
for key in self.keys:
|
|
||||||
self.raw_data[key] = torch.cat(self.raw_data[key]).type(self.dtype)
|
|
||||||
|
|
||||||
# mandate equal length
|
|
||||||
lengths = [len(self.raw_data[key]) for key in self.keys]
|
|
||||||
assert min(lengths) == max(lengths), 'dataset lengths unequal'
|
|
||||||
|
|
||||||
def __len__(self):
|
|
||||||
return len(self.raw_data['ego_state'])
|
|
||||||
|
|
||||||
def __getitem__(self, idx):
|
|
||||||
"""
|
|
||||||
Sample from the dataset
|
|
||||||
Args:
|
|
||||||
idx: index or indices of B samples
|
|
||||||
Returns:
|
|
||||||
sample (dict): sample dictionary with the following entries:
|
|
||||||
state (dict): state dictionary with the following entries:
|
|
||||||
ego_state (torch.tensor): (B, 5) raw state
|
|
||||||
relative_state (torch.tensor): (B, max_nv, d) relative state (padded with nans)
|
|
||||||
path (torch.tensor): (B, P, 2) tensor of P future path x and y positions
|
|
||||||
action (torch.tensor): (B, 1) actions taken from each state
|
|
||||||
next_stat (dict): next state dictionary with the following entries:
|
|
||||||
ego_state (torch.tensor): (B, 5) raw next state
|
|
||||||
relative_state (torch.tensor): (B, max_nv, d) next relative state (padded with nans)
|
|
||||||
path (torch.tensor): (B, P, 2) tensor of P future next path x and y positions
|
|
||||||
"""
|
|
||||||
#sample = {key:self.raw_data[key][idx] for key in self.keys}
|
|
||||||
sample = {
|
|
||||||
'state':{
|
|
||||||
'ego_state':self.raw_data['ego_state'][idx],
|
|
||||||
'relative_state':self.raw_data['relative_state'][idx],
|
|
||||||
'path':self.raw_data['path'][idx]
|
|
||||||
},
|
|
||||||
'action':self.raw_data['action'][idx],
|
|
||||||
'next_state':{
|
|
||||||
'ego_state':self.raw_data['next_ego_state'][idx],
|
|
||||||
'relative_state':self.raw_data['next_relative_state'][idx],
|
|
||||||
'path':self.raw_data['next_path'][idx]},
|
|
||||||
}
|
|
||||||
return sample
|
|
||||||
@@ -1,210 +0,0 @@
|
|||||||
import torch
|
|
||||||
|
|
||||||
import pickle
|
|
||||||
import gym
|
|
||||||
import numpy as np
|
|
||||||
|
|
||||||
import intersim
|
|
||||||
from intersim.utils import get_map_path, get_svt, SVT_to_stateactions
|
|
||||||
from intersim import collisions
|
|
||||||
from intersim.graphs import ConeVisibilityGraph
|
|
||||||
import os
|
|
||||||
opj = os.path.join
|
|
||||||
|
|
||||||
def generate_expert_data(path: str='expert_data', loc: int = 0, track:int = 0,
|
|
||||||
mask_relstate: bool = False, regularize_actions: bool = False,
|
|
||||||
**kwargs):
|
|
||||||
"""
|
|
||||||
Function to save (joint) states and observations from simulated frame
|
|
||||||
Args:
|
|
||||||
path (str): directory to save data
|
|
||||||
loc (int): location index
|
|
||||||
track (int): track index
|
|
||||||
mask_relstate (bool): whether to mask the relative states from the cone visibility graph
|
|
||||||
regularize_actions (bool): whether to regularize the action selection
|
|
||||||
kwargs: arguments for environment instantiation
|
|
||||||
"""
|
|
||||||
|
|
||||||
action_reg = 0.002 if regularize_actions else 0
|
|
||||||
|
|
||||||
if not os.path.isdir(path):
|
|
||||||
os.makedirs(path)
|
|
||||||
filestr = opj(path,intersim.LOCATIONS[loc]+'_track%03i'%(track))
|
|
||||||
|
|
||||||
svt, svt_path = get_svt(loc=loc, track=track) #base='InteractionSimulator'
|
|
||||||
osm = get_map_path(loc=loc)
|
|
||||||
print('SVT path: {}'.format(svt_path))
|
|
||||||
print('Map path: {}'.format(osm))
|
|
||||||
states, actions = SVT_to_stateactions(svt)
|
|
||||||
|
|
||||||
# animate from environment
|
|
||||||
if mask_relstate:
|
|
||||||
cvg = ConeVisibilityGraph(r=20, half_angle=120)
|
|
||||||
env = gym.make('intersim:intersim-v0', svt=svt, map_path=osm,
|
|
||||||
min_acc=-np.inf, max_acc=np.inf, graph=cvg, mask_relstate=True, **kwargs)
|
|
||||||
else:
|
|
||||||
env = gym.make('intersim:intersim-v0', svt=svt, map_path=osm, **kwargs,
|
|
||||||
min_acc=-np.inf, max_acc=np.inf)
|
|
||||||
|
|
||||||
env.reset()
|
|
||||||
done = False
|
|
||||||
obs, actions_taken, max_devs = [], [], []
|
|
||||||
i = 0
|
|
||||||
while not done and i < len(actions):
|
|
||||||
# check state deviation
|
|
||||||
env_state = env.projected_state
|
|
||||||
nni = ~torch.isnan(env_state[:,0])
|
|
||||||
norms = torch.norm(env_state[nni,:2]-states[i,nni,:2], dim=1)
|
|
||||||
if len(norms)>0:
|
|
||||||
max_devs.append(norms.max())
|
|
||||||
|
|
||||||
# propagate environment
|
|
||||||
ob, r, done, info = env.step(env.target_state(svt.simstate[i+1], mu=action_reg))
|
|
||||||
obs.append(ob)
|
|
||||||
actions_taken.append(info['action_taken'])
|
|
||||||
i += 1
|
|
||||||
|
|
||||||
print("Maximum environment deviation from track: %f m" %(max(max_devs)))
|
|
||||||
|
|
||||||
# check for collisions
|
|
||||||
x = torch.stack([ob['state'] for ob in obs])
|
|
||||||
cols = collisions.check_collisions_trajectory(x, svt.lengths, svt.widths)
|
|
||||||
assert ~torch.any(cols), 'Error: Collisions found at indices {}'.format(cols.nonzero(as_tuple=True))
|
|
||||||
|
|
||||||
# shift actions
|
|
||||||
actions_taken.pop(0)
|
|
||||||
obs.pop(-1)
|
|
||||||
actions = torch.stack(actions_taken)
|
|
||||||
|
|
||||||
# save observations and actions
|
|
||||||
pickle.dump(obs,open(filestr+'_raw_observations.pkl', 'wb'))
|
|
||||||
torch.save(actions, filestr+'_raw_actions.pt')
|
|
||||||
process_expert_observations(obs, actions, filestr)
|
|
||||||
|
|
||||||
def process_expert_observations(obs, actions, filestr, remove_outliers=True, dtype=torch.float32):
|
|
||||||
"""
|
|
||||||
Process the expert observations and save them as torch tensors
|
|
||||||
Args:
|
|
||||||
obs (list[dict]): lost of observations
|
|
||||||
actions (torch.Tensor): (T, nv, a) tensor of actions
|
|
||||||
filestr (str): base filename with which to save out observation tensors
|
|
||||||
remove_outliers (bool): whether to remove datapoints with acceleration above or below 5 m/s/s
|
|
||||||
dtype (torch.Type): type to convert data to
|
|
||||||
"""
|
|
||||||
keys = ['ego_state', 'relative_state', 'path', 'action', 'next_ego_state', 'next_relative_state', 'next_path']
|
|
||||||
data = {key:[] for key in keys}
|
|
||||||
assert len(obs) == len(actions), 'non-matching action and observation lengths'
|
|
||||||
T = len(obs)
|
|
||||||
max_nv = 0
|
|
||||||
for t in range(T-1):
|
|
||||||
nni = ~torch.isnan(obs[t]['state'][:,0]) & ~torch.isnan(obs[t+1]['state'][:,0])
|
|
||||||
max_nv = max(max_nv,nni.count_nonzero())
|
|
||||||
|
|
||||||
# state
|
|
||||||
data['ego_state'].append(obs[t]['state'][nni])
|
|
||||||
data['relative_state'].append(obs[t]['relative_state'].index_select(0,
|
|
||||||
nni.nonzero()[:,0]).index_select(1, nni.nonzero()[:,0]))
|
|
||||||
data['path'].append(torch.stack((obs[t]['paths'][0][nni], obs[t]['paths'][1][nni]), dim=-1))
|
|
||||||
|
|
||||||
# action
|
|
||||||
data['action'].append(actions[t][nni])
|
|
||||||
|
|
||||||
# next state
|
|
||||||
data['next_ego_state'].append(obs[t+1]['state'][nni])
|
|
||||||
data['next_relative_state'].append(obs[t+1]['relative_state'].index_select(0,
|
|
||||||
nni.nonzero()[:,0]).index_select(1, nni.nonzero()[:,0]))
|
|
||||||
data['next_path'].append(torch.stack((obs[t+1]['paths'][0][nni], obs[t+1]['paths'][1][nni]), dim=-1))
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
# pad second dimension of relative state
|
|
||||||
for i in range(len(data['relative_state'])):
|
|
||||||
nv1, nv2, d = data['relative_state'][i].shape
|
|
||||||
pad = torch.zeros(nv1, max_nv-nv2, d, dtype=dtype) * np.nan
|
|
||||||
data['relative_state'][i] = torch.cat((data['relative_state'][i], pad), dim=1)
|
|
||||||
data['next_relative_state'][i] = torch.cat((data['next_relative_state'][i], pad), dim=1)
|
|
||||||
|
|
||||||
# cat lists
|
|
||||||
for key in keys:
|
|
||||||
data[key] = torch.cat(data[key]).type(dtype)
|
|
||||||
|
|
||||||
if remove_outliers:
|
|
||||||
non_outlier_indices = torch.nonzero(torch.abs(data['action'][:,0]) < 5)
|
|
||||||
for key in keys:
|
|
||||||
data[key] = data[key][non_outlier_indices[:,0]]
|
|
||||||
|
|
||||||
# mandate equal length
|
|
||||||
lengths = [len(data[key]) for key in keys]
|
|
||||||
assert min(lengths) == max(lengths), 'dataset lengths unequal'
|
|
||||||
|
|
||||||
# save out data
|
|
||||||
for key in keys:
|
|
||||||
torch.save(data[key], filestr+'_'+key+'.pt')
|
|
||||||
|
|
||||||
def load_expert_data(path='expert_data', loc: int = 0, track:int = 0):
|
|
||||||
"""
|
|
||||||
Load expert data from processed files.
|
|
||||||
Args:
|
|
||||||
path (str): directory to save data
|
|
||||||
loc (int): location index
|
|
||||||
track (int): track index
|
|
||||||
Returns:
|
|
||||||
data (dict[torch.Tensor]): dict of data
|
|
||||||
"""
|
|
||||||
# load observations and actions
|
|
||||||
filestr = opj(path, intersim.LOCATIONS[loc]+'_track%03i'%(track))
|
|
||||||
data = {}
|
|
||||||
keys = ['ego_state', 'relative_state', 'path', 'action', 'next_ego_state', 'next_relative_state', 'next_path']
|
|
||||||
for key in keys:
|
|
||||||
data[key] = torch.load(filestr+'_'+key+'.pt')
|
|
||||||
return data
|
|
||||||
|
|
||||||
def load_expert_data_raw(path='expert_data', loc: int = 0, track:int = 0):
|
|
||||||
"""
|
|
||||||
Load expert data from raw file.
|
|
||||||
Args:
|
|
||||||
path (str): directory to save data
|
|
||||||
loc (int): location index
|
|
||||||
track (int): track index
|
|
||||||
Returns:
|
|
||||||
obs (list[Observations]): list of observations
|
|
||||||
actions (list[torch.tensor]): list of corresponding actions taken in observations
|
|
||||||
"""
|
|
||||||
# load observations and actions
|
|
||||||
filestr = opj(path, intersim.LOCATIONS[loc]+'_track%03i'%(track))
|
|
||||||
obs = pickle.load(open(filestr+'_raw_observations.pkl', 'rb'))
|
|
||||||
actions = torch.load(filestr+'_raw_actions.pt')
|
|
||||||
actions = list(torch.unbind(actions))
|
|
||||||
return obs, actions
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
|
||||||
import argparse
|
|
||||||
parser = argparse.ArgumentParser(description='Save Expert Trajectories')
|
|
||||||
parser.add_argument('--loc', default=0, type=int,
|
|
||||||
help='location (default 0)')
|
|
||||||
parser.add_argument('--track', default=0, type=int,
|
|
||||||
help='track number (default 0)')
|
|
||||||
parser.add_argument('--all-tracks', action='store_true',
|
|
||||||
help='whether to process all tracks at location')
|
|
||||||
parser.add_argument('--graph', action='store_true',
|
|
||||||
help='whether to mask the relative states based on a ConeVisibilityGraph')
|
|
||||||
parser.add_argument('--reg', action='store_true',
|
|
||||||
help='whether to regularize actions in the action targeter')
|
|
||||||
parser.add_argument('-o', default='./expert_data', type=str,
|
|
||||||
help='output folder')
|
|
||||||
args = parser.parse_args()
|
|
||||||
|
|
||||||
kwargs = {
|
|
||||||
'loc':args.loc,
|
|
||||||
'track': args.track,
|
|
||||||
'path':args.o,
|
|
||||||
'mask_relstate':args.graph,
|
|
||||||
'regularize_actions': args.reg
|
|
||||||
}
|
|
||||||
|
|
||||||
if args.all_tracks:
|
|
||||||
for i in range(intersim.MAX_TRACKS):
|
|
||||||
kwargs['track'] = i
|
|
||||||
generate_expert_data(**kwargs)
|
|
||||||
else:
|
|
||||||
generate_expert_data(**kwargs)
|
|
||||||
@@ -1 +0,0 @@
|
|||||||
from src.discriminator.discriminator import *
|
|
||||||
@@ -1,101 +0,0 @@
|
|||||||
import torch
|
|
||||||
|
|
||||||
# imitation.rewards.discrim_nets.DiscrimNetGAIL is composed of self.discriminator (nn.Module),
|
|
||||||
# which gets called with inputs (state, action) when needed.
|
|
||||||
|
|
||||||
class CnnDiscriminator(torch.nn.Module):
|
|
||||||
"""ConvNet similar to stable_baselines3.common.policies.ActorCriticCnnPolicy."""
|
|
||||||
|
|
||||||
def __init__(self, env):
|
|
||||||
super().__init__()
|
|
||||||
|
|
||||||
obs_channels, _, _ = env.observation_space.shape
|
|
||||||
(action_size,) = env.action_space.shape
|
|
||||||
in_channels = obs_channels + action_size
|
|
||||||
|
|
||||||
self.cnn = torch.nn.Sequential(
|
|
||||||
torch.nn.Conv2d(in_channels, 32, kernel_size=(8, 8), stride=(4, 4)), # 5+1 -> 32
|
|
||||||
torch.nn.ReLU(),
|
|
||||||
torch.nn.Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2)), # 32 -> 64
|
|
||||||
torch.nn.ReLU(),
|
|
||||||
torch.nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1)), # 64 -> 64
|
|
||||||
torch.nn.ReLU(),
|
|
||||||
torch.nn.Flatten(start_dim=1, end_dim=-1),
|
|
||||||
torch.nn.LazyLinear(512), # 28224 -> 512
|
|
||||||
torch.nn.ReLU(),
|
|
||||||
torch.nn.LazyLinear(1), # 512 -> 1
|
|
||||||
)
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def _concatenate(state, action):
|
|
||||||
b, _, h, w = state.shape
|
|
||||||
_, a = action.shape
|
|
||||||
act = action.unsqueeze(-1).unsqueeze(-1).expand((b, a, h, w))
|
|
||||||
sa = torch.cat((state, act), -3)
|
|
||||||
return sa
|
|
||||||
|
|
||||||
def forward(self, state, action):
|
|
||||||
sa = self._concatenate(state, action)
|
|
||||||
assert sa.ndim == 4
|
|
||||||
return self.cnn(sa).squeeze(1)
|
|
||||||
|
|
||||||
class CnnDiscriminatorFlatAction(torch.nn.Module):
|
|
||||||
"""ConvNet similar to stable_baselines3.common.policies.ActorCriticCnnPolicy."""
|
|
||||||
|
|
||||||
def __init__(self, env):
|
|
||||||
super().__init__()
|
|
||||||
|
|
||||||
obs_channels, _, _ = env.observation_space.shape
|
|
||||||
(action_size,) = env.action_space.shape
|
|
||||||
in_channels = obs_channels
|
|
||||||
|
|
||||||
self.cnn = torch.nn.Sequential(
|
|
||||||
torch.nn.Conv2d(in_channels, 32, kernel_size=(8, 8), stride=(4, 4)), # in_channels -> 32
|
|
||||||
torch.nn.ReLU(),
|
|
||||||
torch.nn.Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2)), # 32 -> 64
|
|
||||||
torch.nn.ReLU(),
|
|
||||||
torch.nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1)), # 64 -> 64
|
|
||||||
torch.nn.ReLU(),
|
|
||||||
torch.nn.Flatten(start_dim=1, end_dim=-1),
|
|
||||||
torch.nn.LazyLinear(128), # 28224 -> 128
|
|
||||||
)
|
|
||||||
self.decoder = torch.nn.Sequential(
|
|
||||||
torch.nn.LazyLinear(64), #128 + 2 -> 64
|
|
||||||
torch.nn.ReLU(),
|
|
||||||
torch.nn.LazyLinear(64), #64 -> 64
|
|
||||||
torch.nn.ReLU(),
|
|
||||||
torch.nn.LazyLinear(1) #64 -> 1
|
|
||||||
)
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def _concatenate(state, action):
|
|
||||||
b, s= state.shape
|
|
||||||
b, a = action.shape
|
|
||||||
sa = torch.cat((state, action), -1)
|
|
||||||
return sa
|
|
||||||
|
|
||||||
def forward(self, state, action):
|
|
||||||
s = self.cnn(state.float())
|
|
||||||
sa = self._concatenate(s, action)
|
|
||||||
assert sa.ndim == 2
|
|
||||||
return self.decoder(sa).squeeze(1)
|
|
||||||
|
|
||||||
class MlpDiscriminator(torch.nn.Module):
|
|
||||||
"""MLP similar to stable_baselines3.common.policies.ActorCriticPolicy."""
|
|
||||||
|
|
||||||
def __init__(self, env=None):
|
|
||||||
super().__init__()
|
|
||||||
self.flatten = torch.nn.Flatten(start_dim=1, end_dim=-1)
|
|
||||||
self.mlp = torch.nn.Sequential(
|
|
||||||
torch.nn.LazyLinear(64), # 42 -> 64
|
|
||||||
torch.nn.Tanh(),
|
|
||||||
torch.nn.LazyLinear(64), # 64 -> 64
|
|
||||||
torch.nn.Tanh(),
|
|
||||||
torch.nn.LazyLinear(1), # 64 -> 1
|
|
||||||
)
|
|
||||||
|
|
||||||
def forward(self, state, action):
|
|
||||||
flat = self.flatten(state)
|
|
||||||
sa = torch.cat((action, flat), -1)
|
|
||||||
assert sa.ndim == 2
|
|
||||||
return self.mlp(sa).squeeze(1)
|
|
||||||
@@ -1,45 +0,0 @@
|
|||||||
from intersim.envs.intersimple import NRasterized
|
|
||||||
from discriminator import CnnDiscriminator
|
|
||||||
import torch
|
|
||||||
|
|
||||||
def test_image_concatenation():
|
|
||||||
env = NRasterized()
|
|
||||||
disc = CnnDiscriminator(env)
|
|
||||||
s = torch.tensor(env.reset()).unsqueeze(0)
|
|
||||||
a = torch.tensor([[0.5]])
|
|
||||||
sa = disc._concatenate(s, a)
|
|
||||||
|
|
||||||
assert s.shape == (1, 5, 200, 200)
|
|
||||||
assert a.shape == (1, 1)
|
|
||||||
assert sa.shape == (1, 6, 200, 200)
|
|
||||||
assert torch.allclose(sa[:, :5], 1.0 * s)
|
|
||||||
assert (sa[:, 5] == a.unsqueeze(-1)).all()
|
|
||||||
|
|
||||||
def test_image_concatenation3():
|
|
||||||
env = NRasterized()
|
|
||||||
disc = CnnDiscriminator(env)
|
|
||||||
|
|
||||||
s1 = env.reset()
|
|
||||||
a1 = 0.15
|
|
||||||
s2, _, _, _ = env.step(0.9)
|
|
||||||
a2 = 0.25
|
|
||||||
s3, _, _, _ = env.step(-0.9)
|
|
||||||
a3 = 0.35
|
|
||||||
|
|
||||||
s = torch.stack([
|
|
||||||
torch.tensor(s1),
|
|
||||||
torch.tensor(s2),
|
|
||||||
torch.tensor(s3)
|
|
||||||
], axis=0)
|
|
||||||
a = torch.tensor([
|
|
||||||
[a1],
|
|
||||||
[a2],
|
|
||||||
[a3],
|
|
||||||
])
|
|
||||||
sa = disc._concatenate(s, a)
|
|
||||||
|
|
||||||
assert s.shape == (3, 5, 200, 200)
|
|
||||||
assert a.shape == (3, 1)
|
|
||||||
assert sa.shape == (3, 6, 200, 200)
|
|
||||||
assert torch.allclose(sa[:, :5], 1.0 * s)
|
|
||||||
assert (sa[:, 5] == a.unsqueeze(-1)).all()
|
|
||||||
118
src/main.py
118
src/main.py
@@ -1,118 +0,0 @@
|
|||||||
import os
|
|
||||||
import torch
|
|
||||||
import gym
|
|
||||||
import intersim
|
|
||||||
import numpy as np
|
|
||||||
from tqdm import tqdm
|
|
||||||
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):
|
|
||||||
"""
|
|
||||||
Return base prefix for all files relating to a certain experiment
|
|
||||||
Args:
|
|
||||||
kwargs (dict): keyword arguments sent to main training loop
|
|
||||||
Returns:
|
|
||||||
basestr (str): prefix
|
|
||||||
"""
|
|
||||||
return 'base'
|
|
||||||
|
|
||||||
def main(config, method='bc', train=False, test=False, loc=0, datadir='./expert_data', filestr='', **kwargs):
|
|
||||||
"""
|
|
||||||
Main loop for training and testing different imitation models
|
|
||||||
Args:
|
|
||||||
config (dict): configuration dictionary for model
|
|
||||||
train (bool): whether to run train loop
|
|
||||||
test (bool): whether to run test loop
|
|
||||||
method (str): the method to try for imitation
|
|
||||||
loc (int): the location index of the roundabout
|
|
||||||
datadir (str): path to expert data
|
|
||||||
kwargs (dict): remaining kwargs for training loop
|
|
||||||
"""
|
|
||||||
# get/set seed
|
|
||||||
seed = kwargs.get('seed',0)
|
|
||||||
torch.manual_seed(seed)
|
|
||||||
|
|
||||||
# method-based training
|
|
||||||
if method=='bc':
|
|
||||||
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("Method {} not implemented".format(method))
|
|
||||||
|
|
||||||
# default train / cv / test split datasets
|
|
||||||
if train:
|
|
||||||
|
|
||||||
# make policy, train and test datasets, and send to
|
|
||||||
train_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=kwargs['train_tracks'])
|
|
||||||
# generate transform from train_dataset
|
|
||||||
transforms = generate_transforms(train_dataset)
|
|
||||||
policy = policy_class(config, transforms)
|
|
||||||
cv_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=kwargs['cv_tracks'])
|
|
||||||
train_fn(config, policy, train_dataset, cv_dataset, filestr, **kwargs)
|
|
||||||
|
|
||||||
if test:
|
|
||||||
|
|
||||||
# load policy
|
|
||||||
policy = policy_class.load_model(filestr, config)
|
|
||||||
policy.eval()
|
|
||||||
|
|
||||||
# simulate policy
|
|
||||||
simulate_policy(policy, loc=loc, track=kwargs['test_tracks'][0], filestr=filestr, nframes=kwargs['nframes'], graph=kwargs['graph'])
|
|
||||||
|
|
||||||
# run test metrics
|
|
||||||
test_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=kwargs['test_tracks'])
|
|
||||||
writer = SummaryWriter(filestr)
|
|
||||||
info = metrics(filestr, test_dataset, policy)
|
|
||||||
for k, m in info.items():
|
|
||||||
writer.add_scalar('test/{}'.format(k), m, 0)
|
|
||||||
|
|
||||||
|
|
||||||
def simulate_policy(policy, loc=0, track=0, filestr='', nframes=float('inf'), graph=None):
|
|
||||||
"""
|
|
||||||
Simulate a trained policy
|
|
||||||
Args:
|
|
||||||
policy: the policy to simulate, which should return action directly
|
|
||||||
loc (int): location index to test policy
|
|
||||||
track (int): track to test policy
|
|
||||||
filestr (str): path prefix to save simulation to
|
|
||||||
"""
|
|
||||||
# animate from environment
|
|
||||||
basepath = os.path.abspath('./InteractionSimulator')
|
|
||||||
svt, svt_path = get_svt(base=basepath, loc=loc, track=track)
|
|
||||||
osm = get_map_path(base=basepath, loc=loc)
|
|
||||||
if graph:
|
|
||||||
env = gym.make('intersim:intersim-v0', svt=svt, map_path=osm,
|
|
||||||
min_acc=-np.inf, max_acc=np.inf, graph=graph, mask_relstate=True)
|
|
||||||
else:
|
|
||||||
env = gym.make('intersim:intersim-v0', svt=svt, map_path=osm,
|
|
||||||
min_acc=-np.inf, max_acc=np.inf)
|
|
||||||
# env = gym.make('intersim:intersim-v0', loc=loc, track=track,
|
|
||||||
# min_acc=-np.inf, max_acc=np.inf)
|
|
||||||
|
|
||||||
ob, _ = env.reset()
|
|
||||||
env.render()
|
|
||||||
done = False
|
|
||||||
i = 0
|
|
||||||
with tqdm(total=min(nframes, env._svt.Tind)) as pbar:
|
|
||||||
while not done and i < nframes:
|
|
||||||
i += 1
|
|
||||||
|
|
||||||
# get action
|
|
||||||
action = policy(ob)
|
|
||||||
|
|
||||||
# propagate environment
|
|
||||||
ob, r, done, info = env.step(action)
|
|
||||||
env.render()
|
|
||||||
|
|
||||||
pbar.update()
|
|
||||||
|
|
||||||
env.close(filestr=filestr+'_sim')
|
|
||||||
@@ -1,117 +0,0 @@
|
|||||||
import torch
|
|
||||||
from torch import nn
|
|
||||||
|
|
||||||
from src.nets.util import parse_functional
|
|
||||||
|
|
||||||
class DeepSetsModule(nn.Module):
|
|
||||||
def __init__(self, input_dim, phi_hidden_n, phi_hidden_dim, latent_dim, rho_hidden_n, rho_hidden_dim, output_dim):
|
|
||||||
"""
|
|
||||||
Args:
|
|
||||||
input_dim (int): input size of one instance of the set; input size of phi
|
|
||||||
phi_hidden_n (int): number of hidden layers in phi
|
|
||||||
phi_hidden_dim (int): size of hidden layers in phi
|
|
||||||
latent_dim (int): output size of phi network, where sum is taken over instances; input size of rho
|
|
||||||
rho_hidden_n (int): number of hidden layers in rho
|
|
||||||
rho_hidden_dim (int): size of hidden layers in rho
|
|
||||||
output_dim (int): output size of rho
|
|
||||||
"""
|
|
||||||
super(DeepSetsModule, self).__init__()
|
|
||||||
self.input_dim = input_dim
|
|
||||||
self.latent_dim = latent_dim
|
|
||||||
self.phi = Phi(self.input_dim, phi_hidden_n, phi_hidden_dim, self.latent_dim)
|
|
||||||
self.rho = Phi(self.latent_dim, rho_hidden_n, rho_hidden_dim, output_dim)
|
|
||||||
self.output_dim = self.rho.output_dim
|
|
||||||
self.pooling = torch.sum
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def from_config(config):
|
|
||||||
"""
|
|
||||||
Args:
|
|
||||||
config (dict): dictionary with network parameters in the form
|
|
||||||
{
|
|
||||||
"input_dim": 5,
|
|
||||||
"phi": {
|
|
||||||
"hidden_n": 1,
|
|
||||||
"hidden_dim": 10,
|
|
||||||
},
|
|
||||||
"latent_dim": 8,
|
|
||||||
"rho": {
|
|
||||||
"hidden_n": 1,
|
|
||||||
"hidden_dim": 10,
|
|
||||||
},
|
|
||||||
"output_dim" : 1,
|
|
||||||
}
|
|
||||||
Returns:
|
|
||||||
m (nn.Module): deep sets module
|
|
||||||
"""
|
|
||||||
input_dim = config["input_dim"]
|
|
||||||
phi = config["phi"]
|
|
||||||
latent_dim = config["latent_dim"]
|
|
||||||
rho = config["rho"]
|
|
||||||
output_dim = config["output_dim"]
|
|
||||||
m = DeepSetsModule(input_dim, phi["hidden_n"], phi["hidden_dim"], latent_dim, rho["hidden_n"], rho["hidden_dim"], output_dim)
|
|
||||||
return m
|
|
||||||
|
|
||||||
def forward(self, x):
|
|
||||||
"""
|
|
||||||
Args:
|
|
||||||
x (torch.tensor): ([B, ]max_nv, d)
|
|
||||||
Returns:
|
|
||||||
y (torch.tensor): ([B, ]output_dim)
|
|
||||||
"""
|
|
||||||
# mask for selecting only those batches and vehicles where all relative states are not nan
|
|
||||||
# shape (B, max_nv)
|
|
||||||
notnan_mask = torch.all(~torch.isnan(x), dim=-1)
|
|
||||||
# create zero tensor of shape (B, max_nv, latent_dim) to store phi evaluations in
|
|
||||||
latent = torch.zeros([*x.shape[:-1], self.latent_dim], dtype=x.dtype)
|
|
||||||
# evaluate phi for all not NaN entries
|
|
||||||
# x[batch_dynamic_mask] has shape (notnan_mask.sum(), input_dim)
|
|
||||||
latent[notnan_mask] = self.phi(x[notnan_mask])
|
|
||||||
|
|
||||||
# sum over relative state dimension
|
|
||||||
latent = self.pooling(latent, dim=-2)
|
|
||||||
|
|
||||||
# apply rho network
|
|
||||||
y = self.rho(latent)
|
|
||||||
return y
|
|
||||||
|
|
||||||
|
|
||||||
class Phi(nn.Module):
|
|
||||||
def __init__(self, input_dim, hidden_n, hidden_dim, output_dim, final_activation=None):
|
|
||||||
"""
|
|
||||||
Fully connected feedforward network with same size for all hidden layers and ReLU activation
|
|
||||||
|
|
||||||
Args:
|
|
||||||
input_dim (int): input dimension
|
|
||||||
hidden_n (int): number of hidden layers
|
|
||||||
hidden_dim (int): hidden layer dimension
|
|
||||||
output_dim (int): output dimension
|
|
||||||
"""
|
|
||||||
super(Phi, self).__init__()
|
|
||||||
self.input_dim = input_dim
|
|
||||||
self.output_dim = output_dim
|
|
||||||
if hidden_n > 0:
|
|
||||||
self.layers = nn.ModuleList([nn.Linear(self.input_dim, hidden_dim)])
|
|
||||||
for _ in range(hidden_n - 1):
|
|
||||||
self.layers.append(nn.Linear(hidden_dim, hidden_dim))
|
|
||||||
self.layers.append(nn.Linear(hidden_dim, self.output_dim))
|
|
||||||
else:
|
|
||||||
self.layers = nn.ModuleList([nn.Identity()])
|
|
||||||
self.output_dim = self.input_dim
|
|
||||||
self.activation = nn.functional.relu
|
|
||||||
self.final_activation = final_activation if final_activation else lambda x: x
|
|
||||||
|
|
||||||
def forward(self, x):
|
|
||||||
for layer in self.layers[:-1]:
|
|
||||||
x = self.activation(layer(x))
|
|
||||||
x = self.final_activation(self.layers[-1](x))
|
|
||||||
return x
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def from_config(config):
|
|
||||||
args = (config["input_dim"], config["hidden_n"], config["hidden_dim"], config["output_dim"])
|
|
||||||
if "final_activation" in config:
|
|
||||||
kwargs = {"final_activation": parse_functional(config["final_activation"])}
|
|
||||||
else:
|
|
||||||
kwargs = {}
|
|
||||||
return Phi(*args, **kwargs)
|
|
||||||
@@ -1,14 +0,0 @@
|
|||||||
import torch
|
|
||||||
from torch.nn import functional, Identity
|
|
||||||
|
|
||||||
def parse_functional(functional_config):
|
|
||||||
if isinstance(functional_config, str):
|
|
||||||
if functional_config == 'relu':
|
|
||||||
return functional.relu
|
|
||||||
elif functional_config == 'sigmoid':
|
|
||||||
return torch.sigmoid
|
|
||||||
elif functional_config == 'softmax':
|
|
||||||
return functional.softmax
|
|
||||||
elif functional_config == 'id':
|
|
||||||
return Identity()
|
|
||||||
return None
|
|
||||||
@@ -1,2 +0,0 @@
|
|||||||
from src.policies.policy import IntersimPolicy, IntersimStateNet, IntersimStateActionNet, generate_transforms
|
|
||||||
from src.policies.options import OptionsCnnPolicy
|
|
||||||
@@ -1,65 +0,0 @@
|
|||||||
from stable_baselines3.common.policies import ActorCriticPolicy, ActorCriticCnnPolicy
|
|
||||||
from torch.distributions import Categorical
|
|
||||||
|
|
||||||
class OptionsCnnPolicy(ActorCriticPolicy):
|
|
||||||
"""
|
|
||||||
Class for high-level options policy (generator)
|
|
||||||
"""
|
|
||||||
def __init__(self, observation_space, *args, eps=0, **kwargs):
|
|
||||||
super().__init__(observation_space, *args, **kwargs)
|
|
||||||
self.cnn_policy = ActorCriticCnnPolicy(observation_space['obs'], *args, **kwargs)
|
|
||||||
self.eps = eps
|
|
||||||
|
|
||||||
def _prior_distribution(self, s):
|
|
||||||
"""
|
|
||||||
Return prior distribution over high-level options (before masking)
|
|
||||||
Args:
|
|
||||||
s (torch.tensor): observation
|
|
||||||
Returns:
|
|
||||||
values (torch.tensor): values from critic
|
|
||||||
dist (torch.distributions): prior distribution over actions
|
|
||||||
"""
|
|
||||||
latent_pi, latent_vf, latent_sde = self.cnn_policy._get_latent(s)
|
|
||||||
distribution = self.cnn_policy._get_action_dist_from_latent(latent_pi, latent_sde)
|
|
||||||
values = self.cnn_policy.value_net(latent_vf)
|
|
||||||
return values, distribution.distribution
|
|
||||||
|
|
||||||
def forward(self, obs):
|
|
||||||
"""
|
|
||||||
Will mask invalid states before making action selections
|
|
||||||
Args:
|
|
||||||
obs: dict with keys:
|
|
||||||
obs (torch.tensor): (*,o) true observations
|
|
||||||
mask (torch.tensor): (*,m) mask over valid actions
|
|
||||||
Returns:
|
|
||||||
ch (torch.tensor): (*,a) sampled actions
|
|
||||||
values (torch.tensor): (*,) predicted value at observation
|
|
||||||
log_probs (torch.tensor): (*,) log probabilities of selected actions
|
|
||||||
"""
|
|
||||||
s, m = obs['obs'], obs['mask']
|
|
||||||
values, prior = self._prior_distribution(s)
|
|
||||||
posterior = Categorical((prior.probs + self.eps) * m)
|
|
||||||
ch = posterior.sample()
|
|
||||||
return ch, values, posterior.log_prob(ch)
|
|
||||||
|
|
||||||
def _predict(self, obs, deterministic=False):
|
|
||||||
action, _, _ = self.forward(obs)
|
|
||||||
return action
|
|
||||||
|
|
||||||
def evaluate_actions(self, obs, ch):
|
|
||||||
"""
|
|
||||||
Evaluate particular actions
|
|
||||||
Args:
|
|
||||||
obs: dict with keys:
|
|
||||||
obs (torch.tensor): (*,o) true observations
|
|
||||||
mask (torch.tensor): (*,m) masks over valid actions
|
|
||||||
ch (torch.tensor): (*,a) selected actions
|
|
||||||
Returns:
|
|
||||||
values (torch.tensor): (*,) predicted value at observation
|
|
||||||
log_probs (torch.tensor): (*,) log probabilities of selected actions
|
|
||||||
ent (torch.tensor): (*,) entropy of each distribution over actions
|
|
||||||
"""
|
|
||||||
s, m = obs['obs'], obs['mask']
|
|
||||||
values, prior = self._prior_distribution(s)
|
|
||||||
posterior = Categorical((prior.probs + self.eps) * m)
|
|
||||||
return values, posterior.log_prob(ch), posterior.entropy() # additional values used by PPO.train
|
|
||||||
@@ -1,164 +0,0 @@
|
|||||||
|
|
||||||
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):
|
|
||||||
"""
|
|
||||||
Args:
|
|
||||||
config (dict): dictionary for configuring the deep sets policy
|
|
||||||
"""
|
|
||||||
super(IntersimStateNet, 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)
|
|
||||||
|
|
||||||
cat_dim = self.ego_net.output_dim + self.deepsets_net.output_dim + self.path_net.output_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 (torch.tensor): (B, P, 2) tensor of P future path x and 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"])
|
|
||||||
path = self.path_net(sample["path"].reshape((sample["path"].shape[0], -1)))
|
|
||||||
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"])
|
|
||||||
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
|
|
||||||
|
|
||||||
|
|
||||||
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() and key in ob.keys():
|
|
||||||
transformed_ob[key] = self._transforms[key].transform(ob[key])
|
|
||||||
return transformed_ob
|
|
||||||
|
|
||||||
def __call__(self, ob):
|
|
||||||
|
|
||||||
if 'ego_state' in ob.keys():
|
|
||||||
# extract state from dataloader samples
|
|
||||||
pass
|
|
||||||
else:
|
|
||||||
# extract state from observation (using simulator)
|
|
||||||
ob['ego_state'] = ob['state']
|
|
||||||
ob['path'] = torch.stack(ob['paths'],dim=-1)
|
|
||||||
|
|
||||||
ob = self.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(),
|
|
||||||
'ego_state': MinMaxScaler(),
|
|
||||||
'relative_state': MinMaxScaler(reduce_dim=2),
|
|
||||||
'path': MinMaxScaler(reduce_dim=2),
|
|
||||||
}
|
|
||||||
for key in transforms.keys():
|
|
||||||
if key == 'action':
|
|
||||||
transforms[key].fit(dataset[:][key])
|
|
||||||
else:
|
|
||||||
transforms[key].fit(dataset[:]['state'][key])
|
|
||||||
|
|
||||||
return transforms
|
|
||||||
@@ -1,11 +0,0 @@
|
|||||||
import torch
|
|
||||||
|
|
||||||
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
|
|
||||||
@@ -1,140 +0,0 @@
|
|||||||
import torch
|
|
||||||
from torch import nn
|
|
||||||
import numpy as np
|
|
||||||
from sklearn import preprocessing
|
|
||||||
|
|
||||||
class Transform(nn.Module):
|
|
||||||
"""
|
|
||||||
Base class to normalize observations and actions for network.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self):
|
|
||||||
super(Transform, self).__init__()
|
|
||||||
# self.fit(X)
|
|
||||||
|
|
||||||
def fit(self, X):
|
|
||||||
"""
|
|
||||||
Fit transformer to X
|
|
||||||
Args:
|
|
||||||
X (torch.tensor): (B, N) tensor of B data points with N features
|
|
||||||
"""
|
|
||||||
raise NotImplementedError('Please implement fit()')
|
|
||||||
|
|
||||||
def transform(self, X):
|
|
||||||
"""
|
|
||||||
Transform X. fit() has to be called first
|
|
||||||
Args:
|
|
||||||
X (torch.tensor): (B, N) tensor where N has to be the same as during fit()
|
|
||||||
"""
|
|
||||||
raise NotImplementedError('Please implement transform()')
|
|
||||||
|
|
||||||
def inverse_transform(self, X):
|
|
||||||
"""
|
|
||||||
Inverse transformation
|
|
||||||
Args:
|
|
||||||
X (torch.tensor): (B, N) tensor
|
|
||||||
"""
|
|
||||||
raise NotImplementedError('Please implement inverse_transform()')
|
|
||||||
|
|
||||||
def forward(self, X):
|
|
||||||
return self.transform(X)
|
|
||||||
|
|
||||||
class MinMaxScaler(Transform):
|
|
||||||
"""
|
|
||||||
Scale tensor so each feature is in [0, 1]
|
|
||||||
"""
|
|
||||||
def __init__(self, reduce_dim:int=None):
|
|
||||||
"""
|
|
||||||
Initialize SciKitTransform
|
|
||||||
Args:
|
|
||||||
reduce_dim (int): dimension to start calculating featues from
|
|
||||||
e.g. with reduce_dim=2, (A, B, C, D, E) will be reshaped to (A*B, C*D*E)
|
|
||||||
"""
|
|
||||||
self.reduce_dim = reduce_dim
|
|
||||||
super(MinMaxScaler, self).__init__()
|
|
||||||
|
|
||||||
def fit(self, X):
|
|
||||||
nd = X.ndim
|
|
||||||
if self.reduce_dim:
|
|
||||||
self.nfeatures = int(torch.tensor(X.shape[self.reduce_dim:]).prod())
|
|
||||||
else:
|
|
||||||
assert nd==2, 'Invalid ndim'
|
|
||||||
self.nfeatures = X.shape[1]
|
|
||||||
|
|
||||||
X = X.reshape((-1,self.nfeatures))
|
|
||||||
nans = torch.isnan(X)
|
|
||||||
X[nans] = float('inf')
|
|
||||||
self.min = X.min(0,keepdims=True)[0]
|
|
||||||
|
|
||||||
X[nans] = -float('inf')
|
|
||||||
self.span = X.max(0,keepdims=True)[0] - self.min
|
|
||||||
|
|
||||||
X[nans] = np.nan
|
|
||||||
|
|
||||||
def transform(self, X):
|
|
||||||
|
|
||||||
assert hasattr(self, 'min') and hasattr(self, 'span'), 'Model not yet fit'
|
|
||||||
shape = X.shape
|
|
||||||
X = X.reshape((-1,self.nfeatures))
|
|
||||||
t = (X - self.min) / self.span
|
|
||||||
return t.reshape(shape)
|
|
||||||
|
|
||||||
def inverse_transform(self, X):
|
|
||||||
|
|
||||||
assert hasattr(self, 'min') and hasattr(self, 'span'), 'Model not yet fit'
|
|
||||||
shape = X.shape
|
|
||||||
X = X.reshape((-1,self.nfeatures))
|
|
||||||
it = X * self.span + self.min
|
|
||||||
return it.reshape(shape)
|
|
||||||
|
|
||||||
|
|
||||||
class SciKitTransform(Transform):
|
|
||||||
"""
|
|
||||||
Wrappers around scikit-learn transforms
|
|
||||||
"""
|
|
||||||
def __init__(self, tf, reduce_dim:int=None):
|
|
||||||
"""
|
|
||||||
Initialize SciKitTransform
|
|
||||||
Args:
|
|
||||||
tf: transform
|
|
||||||
reduce_dim (int): dimension to start calculating featues from
|
|
||||||
e.g. with reduce_dim=2, (A, B, C, D, E) will be reshaped to (A*B, C*D*E)
|
|
||||||
"""
|
|
||||||
self.tf = tf
|
|
||||||
self.reduce_dim = reduce_dim
|
|
||||||
super(SciKitTransform, self).__init__()
|
|
||||||
|
|
||||||
def fit(self, X):
|
|
||||||
nd = X.ndim
|
|
||||||
if self.reduce_dim:
|
|
||||||
self.nfeatures = int(torch.tensor(X.shape[self.reduce_dim:]).prod())
|
|
||||||
else:
|
|
||||||
assert nd==2, 'Invalid ndim'
|
|
||||||
self.nfeatures = X.shape[1]
|
|
||||||
|
|
||||||
self.tf.fit(X.reshape((-1,self.nfeatures)))
|
|
||||||
|
|
||||||
def transform(self, X):
|
|
||||||
shape = X.shape
|
|
||||||
t = torch.tensor(self.tf.transform(X.reshape((-1,self.nfeatures))), dtype=torch.float)
|
|
||||||
return t.reshape(shape)
|
|
||||||
|
|
||||||
def inverse_transform(self, X):
|
|
||||||
shape = X.shape
|
|
||||||
it = torch.tensor(self.tf.inverse_transform(X.reshape((-1,self.nfeatures))), dtype=torch.float)
|
|
||||||
return it.reshape(shape)
|
|
||||||
|
|
||||||
class SciKitStandardScaler(SciKitTransform):
|
|
||||||
"""
|
|
||||||
Wrapper around scikit-learn's StandardScaler for standardizing each feature individually.
|
|
||||||
"""
|
|
||||||
def __init__(self, **kwargs):
|
|
||||||
super(SciKitStandardScaler, self).__init__(preprocessing.StandardScaler(), **kwargs)
|
|
||||||
|
|
||||||
class SciKitMinMaxScaler(SciKitTransform):
|
|
||||||
"""
|
|
||||||
Wrapper around scikit-learn's MinMaxScaler for scaling features to [0, 1] individually.
|
|
||||||
"""
|
|
||||||
def __init__(self, **kwargs):
|
|
||||||
super(SciKitMinMaxScaler, self).__init__(preprocessing.MinMaxScaler(), **kwargs)
|
|
||||||
|
|
||||||
@@ -1 +0,0 @@
|
|||||||
from src.value_dice.value_dice import ValueDicePolicy, train, vd_config
|
|
||||||
@@ -1,300 +0,0 @@
|
|||||||
import numpy as np
|
|
||||||
|
|
||||||
import torch
|
|
||||||
import torch.nn as nn
|
|
||||||
from torch.utils.data import DataLoader
|
|
||||||
from torch.nn.utils import clip_grad_norm_
|
|
||||||
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
|
|
||||||
|
|
||||||
def vd_config(ray_config):
|
|
||||||
config = {
|
|
||||||
'policy_net': {
|
|
||||||
'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'],
|
|
||||||
},
|
|
||||||
},
|
|
||||||
'value_net': {
|
|
||||||
'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},
|
|
||||||
'action_dim': 1,
|
|
||||||
'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'],
|
|
||||||
},
|
|
||||||
},
|
|
||||||
'policy_optim': {
|
|
||||||
'optimizer':'adam',
|
|
||||||
'lr':ray_config['policy_lr'],
|
|
||||||
'weight_decay':ray_config['policy_weight_decay']
|
|
||||||
},
|
|
||||||
'value_optim': {
|
|
||||||
'optimizer':'adam',
|
|
||||||
'lr':ray_config['value_lr'],
|
|
||||||
'weight_decay':ray_config['value_weight_decay']
|
|
||||||
},
|
|
||||||
'train_epochs': 40,
|
|
||||||
'train_batch_size': ray_config['train_batch_size'],
|
|
||||||
'discount': ray_config['discount'],
|
|
||||||
'clip_grad_norm': ray_config['clip_grad_norm'],
|
|
||||||
}
|
|
||||||
return config
|
|
||||||
|
|
||||||
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
|
|
||||||
|
|
||||||
@value.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())
|
|
||||||
|
|
||||||
def policy_parameters(self):
|
|
||||||
return self.policy.parameters()
|
|
||||||
|
|
||||||
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
|
|
||||||
train_epochs = config['train_epochs']
|
|
||||||
train_batch_size = config['train_batch_size']
|
|
||||||
discount = config['discount']
|
|
||||||
clip_grad_norm = config['clip_grad_norm']
|
|
||||||
|
|
||||||
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
|
|
||||||
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):
|
|
||||||
# get s, a, s', s_0 from batch
|
|
||||||
state = batch['state']
|
|
||||||
action = batch['action']
|
|
||||||
next_state = batch['next_state']
|
|
||||||
initial_state = state
|
|
||||||
|
|
||||||
# append action to state batches
|
|
||||||
# use expert action for s
|
|
||||||
state['action'] = action
|
|
||||||
# run s' and s_0 through policy
|
|
||||||
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__() )
|
|
||||||
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_state)
|
|
||||||
|
|
||||||
# linear loss
|
|
||||||
linear_loss = (1 - discount) * torch.mean(value_init)
|
|
||||||
|
|
||||||
# nonlinear loss
|
|
||||||
value_diff = value - discount * value_next
|
|
||||||
nonlinear_loss = torch.logsumexp(value_diff, dim=0) - np.log(len(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)
|
|
||||||
|
|
||||||
# In original implementation policy is regularized with orthogonal regularization,
|
|
||||||
# value with L2 regularization on gradients
|
|
||||||
policy_loss = -loss
|
|
||||||
value_loss = loss
|
|
||||||
|
|
||||||
# # compute loss and step optimizer
|
|
||||||
# policy_optimizer.zero_grad()
|
|
||||||
# value_optimizer.zero_grad()
|
|
||||||
# policy_loss.backward(retain_graph=True)
|
|
||||||
# value_loss.backward()
|
|
||||||
|
|
||||||
# clip_grad_norm_(policy.policy.parameters(), clip_grad_norm)
|
|
||||||
# clip_grad_norm_(policy.value.parameters(), clip_grad_norm)
|
|
||||||
|
|
||||||
# policy_optimizer.step()
|
|
||||||
# value_optimizer.step()
|
|
||||||
|
|
||||||
if batch_idx % 2 == 0:
|
|
||||||
policy_optimizer.zero_grad()
|
|
||||||
policy_loss.backward()
|
|
||||||
clip_grad_norm_(policy.policy.parameters(), clip_grad_norm)
|
|
||||||
policy_optimizer.step()
|
|
||||||
else:
|
|
||||||
value_optimizer.zero_grad()
|
|
||||||
value_loss.backward()
|
|
||||||
clip_grad_norm_(policy.value.parameters(), clip_grad_norm)
|
|
||||||
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)
|
|
||||||
Reference in New Issue
Block a user