purging unused files
This commit is contained in:
Arec Jamgochian
@@ -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