security_vision
/
Trajectron-plus-plus
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Trajectron-plus-plus/trajectron/model/online/online_mgcvae.py

431 lines
22 KiB
Python
Raw Blame History

import warnings
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from collections import defaultdict, Counter
from trajectron.model.components import *
from trajectron.model.model_utils import *
from trajectron.model.dataset import get_relative_robot_traj
import trajectron.model.dynamics as dynamic_module
from trajectron.model.mgcvae import MultimodalGenerativeCVAE
from trajectron.environment.scene_graph import DirectedEdge
from trajectron.environment.node_type import NodeType
class OnlineMultimodalGenerativeCVAE(MultimodalGenerativeCVAE):
def __init__(self,
env,
node,
model_registrar,
hyperparams,
device):
self.hyperparams = hyperparams
self.node = node
self.node_type = self.node.type
if len(env.scenes) != 1:
raise ValueError("Passed in Environment has number of scenes != 1")
self.robot = env.scenes[0].robot
self.model_registrar = model_registrar
self.device = device
self.node_modules = dict()
self.env = env
self.scene_graph = None
self.state = self.hyperparams['state']
self.pred_state = self.hyperparams['pred_state'][self.node.type]
self.state_length = int(np.sum([len(entity_dims) for entity_dims in self.state[self.node.type].values()]))
if self.hyperparams['incl_robot_node']:
self.robot_state_length = int(
np.sum([len(entity_dims) for entity_dims in self.state[self.robot.type].values()]))
self.pred_state_length = int(np.sum([len(entity_dims) for entity_dims in self.pred_state.values()]))
self.curr_hidden_states = dict()
self.edge_types = Counter()
self.create_graphical_model()
dynamic_class = getattr(dynamic_module, self.hyperparams['dynamic'][self.node_type]['name'])
dyn_limits = hyperparams['dynamic'][self.node_type]['limits']
self.dynamic = dynamic_class(self.env.scenes[0].dt, dyn_limits, device,
self.model_registrar, self.x_size, self.node_type)
def create_graphical_model(self):
"""
Creates or queries all trainable components.
:return: None
"""
self.clear_submodules()
############################
# Everything but Edges #
############################
self.create_node_models()
for name, module in self.node_modules.items():
module.to(self.device)
def update_graph(self, new_scene_graph, new_neighbors, removed_neighbors):
self.scene_graph = new_scene_graph
if self.node in new_neighbors:
for edge_type, new_neighbor_nodes in new_neighbors[self.node].items():
self.add_edge_model(edge_type)
self.edge_types += Counter({edge_type: len(new_neighbor_nodes)})
if self.node in removed_neighbors:
for edge_type, removed_neighbor_nodes in removed_neighbors[self.node].items():
self.remove_edge_model(edge_type)
self.edge_types -= Counter({edge_type: len(removed_neighbor_nodes)})
def get_edge_to(self, other_node):
return DirectedEdge(self.node, other_node)
def add_edge_model(self, edge_type):
if self.hyperparams['edge_encoding']:
if edge_type + '/edge_encoder' not in self.node_modules:
neighbor_state_length = int(
np.sum([len(entity_dims) for entity_dims in
self.state[self._get_other_node_type_from_edge(edge_type)].values()]))
if self.hyperparams['edge_state_combine_method'] == 'pointnet':
self.add_submodule(edge_type + '/pointnet_encoder',
model_if_absent=nn.Sequential(
nn.Linear(self.state_length, 2 * self.state_length),
nn.ReLU(),
nn.Linear(2 * self.state_length, 2 * self.state_length),
nn.ReLU()))
edge_encoder_input_size = 2 * self.state_length + self.state_length
elif self.hyperparams['edge_state_combine_method'] == 'attention':
self.add_submodule(self.node.type + '/edge_attention_combine',
model_if_absent=TemporallyBatchedAdditiveAttention(
encoder_hidden_state_dim=self.state_length,
decoder_hidden_state_dim=self.state_length))
edge_encoder_input_size = self.state_length + neighbor_state_length
else:
edge_encoder_input_size = self.state_length + neighbor_state_length
self.add_submodule(edge_type + '/edge_encoder',
model_if_absent=nn.LSTM(input_size=edge_encoder_input_size,
hidden_size=self.hyperparams['enc_rnn_dim_edge'],
batch_first=True))
def _get_other_node_type_from_edge(self, edge_type_str):
n2_type_str = edge_type_str.split('->')[1]
return NodeType(n2_type_str, self.env.node_type_list.index(n2_type_str) + 1)
def _get_edge_type_from_str(self, edge_type_str):
n1_type_str, n2_type_str = edge_type_str.split('->')
return (NodeType(n1_type_str, self.env.node_type_list.index(n1_type_str) + 1),
NodeType(n2_type_str, self.env.node_type_list.index(n2_type_str) + 1))
def remove_edge_model(self, edge_type):
if self.hyperparams['edge_encoding']:
if len(self.scene_graph.get_neighbors(self.node, self._get_other_node_type_from_edge(edge_type))) == 0:
del self.node_modules[edge_type + '/edge_encoder']
def obtain_encoded_tensors(self,
mode,
inputs,
inputs_st,
inputs_np,
robot_present_and_future,
maps):
x, x_r_t, y_r = None, None, None
batch_size = 1
our_inputs = inputs[self.node]
our_inputs_st = inputs_st[self.node]
initial_dynamics = dict()
initial_dynamics['pos'] = our_inputs[:, 0:2] # TODO: Generalize
initial_dynamics['vel'] = our_inputs[:, 2:4] # TODO: Generalize
self.dynamic.set_initial_condition(initial_dynamics)
#########################################
# Provide basic information to encoders #
#########################################
if self.hyperparams['incl_robot_node'] and self.robot is not None:
robot_present_and_future_st = get_relative_robot_traj(self.env, self.state,
our_inputs, robot_present_and_future,
self.node.type, self.robot.type)
x_r_t = robot_present_and_future_st[..., 0, :]
y_r = robot_present_and_future_st[..., 1:, :]
##################
# Encode History #
##################
node_history_encoded = self.encode_node_history(our_inputs_st)
##############################
# Encode Node Edges per Type #
##############################
total_edge_influence = None
if self.hyperparams['edge_encoding']:
node_edges_encoded = list()
for edge_type in self.edge_types:
connected_nodes_batched = list()
edge_masks_batched = list()
# We get all nodes which are connected to the current node for the current timestep
connected_nodes_batched.append(self.scene_graph.get_neighbors(self.node,
self._get_other_node_type_from_edge(
edge_type)))
if self.hyperparams['dynamic_edges'] == 'yes':
# We get the edge masks for the current node at the current timestep
edge_masks_for_node = self.scene_graph.get_edge_scaling(self.node)
edge_masks_batched.append(torch.tensor(edge_masks_for_node, dtype=torch.float, device=self.device))
# Encode edges for given edge type
encoded_edges_type = self.encode_edge(inputs,
inputs_st,
inputs_np,
edge_type,
connected_nodes_batched,
edge_masks_batched)
node_edges_encoded.append(encoded_edges_type) # List of [bs/nbs, enc_rnn_dim]
#####################
# Encode Node Edges #
#####################
total_edge_influence = self.encode_total_edge_influence(mode,
node_edges_encoded,
node_history_encoded,
batch_size)
self.TD = {'node_history_encoded': node_history_encoded,
'total_edge_influence': total_edge_influence}
################
# Map Encoding #
################
if self.hyperparams['use_map_encoding'] and self.node_type in self.hyperparams['map_encoder']:
if self.node not in maps:
# This means the node was removed (it is only being kept around because of the edge removal filter).
me_params = self.hyperparams['map_encoder'][self.node_type]
self.TD['encoded_map'] = torch.zeros((1, me_params['output_size']))
else:
encoded_map = self.node_modules[self.node_type + '/map_encoder'](maps[self.node] * 2. - 1.,
(mode == ModeKeys.TRAIN))
do = self.hyperparams['map_encoder'][self.node_type]['dropout']
encoded_map = F.dropout(encoded_map, do, training=(mode == ModeKeys.TRAIN))
self.TD['encoded_map'] = encoded_map
######################################
# Concatenate Encoder Outputs into x #
######################################
return self.create_encoder_rep(mode, self.TD, x_r_t, y_r)
def create_encoder_rep(self, mode,
TD,
robot_present_st,
robot_future_st):
# Unpacking TD
node_history_encoded = TD['node_history_encoded']
if self.hyperparams['edge_encoding']:
total_edge_influence = TD['total_edge_influence']
if self.hyperparams['use_map_encoding'] and self.node_type in self.hyperparams['map_encoder']:
encoded_map = TD['encoded_map']
if (self.hyperparams['incl_robot_node']
and self.robot is not None
and robot_future_st is not None
and robot_present_st is not None):
robot_future_encoder = self.encode_robot_future(mode, robot_present_st, robot_future_st)
# Tiling for multiple samples
# This tiling is done because:
# a) we must consider the prediction case where there are many candidate robot future actions,
# b) the edge and history encoders are all the same regardless of which candidate future robot action
# we're evaluating.
node_history_encoded = TD['node_history_encoded'].repeat(robot_future_st.size()[0], 1)
if self.hyperparams['edge_encoding']:
total_edge_influence = TD['total_edge_influence'].repeat(robot_future_st.size()[0], 1)
if self.hyperparams['use_map_encoding'] and self.node_type in self.hyperparams['map_encoder']:
encoded_map = TD['encoded_map'].repeat(robot_future_st.size()[0], 1)
elif self.hyperparams['incl_robot_node'] and self.robot is not None:
# Four times because we're trying to mimic a bi-directional RNN's output (which is c and h from both ends).
robot_future_encoder = torch.zeros([1, 4 * self.hyperparams['enc_rnn_dim_future']], device=self.device)
x_concat_list = list()
# Every node has an edge-influence encoder (which could just be zero).
if self.hyperparams['edge_encoding']:
x_concat_list.append(total_edge_influence) # [bs/nbs, 4*enc_rnn_dim]
# Every node has a history encoder.
x_concat_list.append(node_history_encoded) # [bs/nbs, enc_rnn_dim_history]
if self.hyperparams['incl_robot_node'] and self.robot is not None:
x_concat_list.append(robot_future_encoder) # [bs/nbs, 4*enc_rnn_dim_history]
if self.hyperparams['use_map_encoding'] and self.node_type in self.hyperparams['map_encoder']:
x_concat_list.append(encoded_map) # [bs/nbs, CNN output size]
return torch.cat(x_concat_list, dim=1)
def encode_node_history(self, inputs_st):
new_state = torch.unsqueeze(inputs_st, dim=1) # [bs, 1, state_dim]
if self.node.type + '/node_history_encoder' not in self.curr_hidden_states:
outputs, self.curr_hidden_states[self.node.type + '/node_history_encoder'] = self.node_modules[
self.node.type + '/node_history_encoder'](new_state)
else:
outputs, self.curr_hidden_states[self.node.type + '/node_history_encoder'] = self.node_modules[
self.node.type + '/node_history_encoder'](new_state, self.curr_hidden_states[
self.node.type + '/node_history_encoder'])
return outputs[:, 0, :]
def encode_edge(self, inputs, inputs_st, inputs_np, edge_type, connected_nodes, edge_masks):
edge_type_tuple = self._get_edge_type_from_str(edge_type)
edge_states_list = list() # list of [#of neighbors, max_ht, state_dim]
neighbor_states = list()
orig_rel_state = inputs[self.node].cpu().numpy()
for node in connected_nodes[0]:
neighbor_state_np = inputs_np[node]
# Make State relative to node
_, std = self.env.get_standardize_params(self.state[node.type], node_type=node.type)
std[0:2] = self.env.attention_radius[edge_type_tuple]
# TODO: This all makes the unsafe assumption that the first n dims
# refer to the same quantities even for different agent types!
equal_dims = np.min((neighbor_state_np.shape[-1], orig_rel_state.shape[-1]))
rel_state = np.zeros_like(neighbor_state_np)
rel_state[..., :equal_dims] = orig_rel_state[..., :equal_dims]
neighbor_state_np_st = self.env.standardize(neighbor_state_np,
self.state[node.type],
node_type=node.type,
mean=rel_state,
std=std)
neighbor_state = torch.tensor(neighbor_state_np_st).float().to(self.device)
neighbor_states.append(neighbor_state)
if len(neighbor_states) == 0: # There are no neighbors for edge type # TODO necessary?
neighbor_state_length = int(np.sum([len(entity_dims) for entity_dims in self.state[edge_type[1]].values()]))
edge_states_list.append(torch.zeros((1, 1, neighbor_state_length), device=self.device))
else:
edge_states_list.append(torch.stack(neighbor_states, dim=0))
if self.hyperparams['edge_state_combine_method'] == 'sum':
# Used in Structural-RNN to combine edges as well.
op_applied_edge_states_list = list()
for neighbors_state in edge_states_list:
op_applied_edge_states_list.append(torch.sum(neighbors_state, dim=0))
combined_neighbors = torch.stack(op_applied_edge_states_list, dim=0)
if self.hyperparams['dynamic_edges'] == 'yes':
# Should now be (bs, time, 1)
op_applied_edge_mask_list = list()
for edge_mask in edge_masks:
op_applied_edge_mask_list.append(torch.clamp(torch.sum(edge_mask, dim=0, keepdim=True), max=1.))
combined_edge_masks = torch.stack(op_applied_edge_mask_list, dim=0)
elif self.hyperparams['edge_state_combine_method'] == 'max':
# Used in NLP, e.g. max over word embeddings in a sentence.
op_applied_edge_states_list = list()
for neighbors_state in edge_states_list:
op_applied_edge_states_list.append(torch.max(neighbors_state, dim=0))
combined_neighbors = torch.stack(op_applied_edge_states_list, dim=0)
if self.hyperparams['dynamic_edges'] == 'yes':
# Should now be (bs, time, 1)
op_applied_edge_mask_list = list()
for edge_mask in edge_masks:
op_applied_edge_mask_list.append(torch.clamp(torch.max(edge_mask, dim=0, keepdim=True), max=1.))
combined_edge_masks = torch.stack(op_applied_edge_mask_list, dim=0)
elif self.hyperparams['edge_state_combine_method'] == 'mean':
# Used in NLP, e.g. mean over word embeddings in a sentence.
op_applied_edge_states_list = list()
for neighbors_state in edge_states_list:
op_applied_edge_states_list.append(torch.mean(neighbors_state, dim=0))
combined_neighbors = torch.stack(op_applied_edge_states_list, dim=0)
if self.hyperparams['dynamic_edges'] == 'yes':
# Should now be (bs, time, 1)
op_applied_edge_mask_list = list()
for edge_mask in edge_masks:
op_applied_edge_mask_list.append(torch.clamp(torch.mean(edge_mask, dim=0, keepdim=True), max=1.))
combined_edge_masks = torch.stack(op_applied_edge_mask_list, dim=0)
joint_history = torch.cat([combined_neighbors, torch.unsqueeze(inputs_st[self.node], dim=0)], dim=-1)
if edge_type + '/edge_encoder' not in self.curr_hidden_states:
outputs, self.curr_hidden_states[edge_type + '/edge_encoder'] = self.node_modules[
edge_type + '/edge_encoder'](joint_history)
else:
outputs, self.curr_hidden_states[edge_type + '/edge_encoder'] = self.node_modules[
edge_type + '/edge_encoder'](joint_history, self.curr_hidden_states[edge_type + '/edge_encoder'])
if self.hyperparams['dynamic_edges'] == 'yes':
return outputs[:, 0, :] * combined_edge_masks
else:
return outputs[:, 0, :] # [bs, enc_rnn_dim]
def encoder_forward(self, inputs, inputs_st, inputs_np, robot_present_and_future=None, maps=None):
# Always predicting with the online model.
mode = ModeKeys.PREDICT
self.x = self.obtain_encoded_tensors(mode,
inputs,
inputs_st,
inputs_np,
robot_present_and_future,
maps)
self.n_s_t0 = inputs_st[self.node]
self.latent.p_dist = self.p_z_x(mode, self.x)
# robot_future_st is optional here since you can use the same one from encoder_forward,
# but if it's given then we'll re-run that part of the model (if the node is adjacent to the robot).
def decoder_forward(self, prediction_horizon,
num_samples,
robot_present_and_future=None,
z_mode=False,
gmm_mode=False,
full_dist=False,
all_z_sep=False):
# Always predicting with the online model.
mode = ModeKeys.PREDICT
x_nr_t, y_r = None, None
if (self.hyperparams['incl_robot_node']
and self.robot is not None
and robot_present_and_future is not None):
our_inputs = torch.tensor(self.node.get(np.array([self.node.last_timestep]),
self.state[self.node.type],
padding=0.0),
dtype=torch.float,
device=self.device)
robot_present_and_future_st = get_relative_robot_traj(self.env, self.state,
our_inputs, robot_present_and_future,
self.node.type, self.robot.type)
x_nr_t = robot_present_and_future_st[..., 0, :]
y_r = robot_present_and_future_st[..., 1:, :]
self.x = self.create_encoder_rep(mode, self.TD, x_nr_t, y_r)
self.latent.p_dist = self.p_z_x(mode, self.x)
# Making sure n_s_t0 has the same batch size as x_nr_t
self.n_s_t0 = self.n_s_t0[[0]].repeat(x_nr_t.size()[0], 1)
z, num_samples, num_components = self.latent.sample_p(num_samples,
mode,
most_likely_z=z_mode,
full_dist=full_dist,
all_z_sep=all_z_sep)
y_dist, our_sampled_future = self.p_y_xz(mode, self.x, x_nr_t, y_r, self.n_s_t0, z,
prediction_horizon,
num_samples,
num_components,
gmm_mode)
return y_dist, our_sampled_future
Reference in New Issue View Git Blame Copy Permalink