File size: 11,010 Bytes
d4cbafd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 | import torch
import numpy as np
import collections.abc
from torch.utils.data._utils.collate import default_collate
import dill
container_abcs = collections.abc
def restore(data):
"""
In case we dilled some structures to share between multiple process this function will restore them.
If the data input are not bytes we assume it was not dilled in the first place
:param data: Possibly dilled data structure
:return: Un-dilled data structure
"""
if type(data) is bytes:
return dill.loads(data)
return data
def collate(batch):
if len(batch) == 0:
return batch
elem = batch[0]
if elem is None:
return None
elif isinstance(elem, container_abcs.Sequence):
if len(elem) == 4: # We assume those are the maps, map points, headings and patch_size
scene_map, scene_pts, heading_angle, patch_size = zip(*batch)
if heading_angle[0] is None:
heading_angle = None
else:
heading_angle = torch.Tensor(heading_angle)
map = scene_map[0].get_cropped_maps_from_scene_map_batch(scene_map,
scene_pts=torch.Tensor(scene_pts),
patch_size=patch_size[0],
rotation=heading_angle)
return map
transposed = zip(*batch)
return [collate(samples) for samples in transposed]
elif isinstance(elem, container_abcs.Mapping):
# We have to dill the neighbors structures. Otherwise each tensor is put into
# shared memory separately -> slow, file pointer overhead
# we only do this in multiprocessing
neighbor_dict = {key: [d[key] for d in batch] for key in elem}
return dill.dumps(neighbor_dict) if torch.utils.data.get_worker_info() else neighbor_dict
return default_collate(batch)
def get_relative_robot_traj(env, state, node_traj, robot_traj, node_type, robot_type):
# TODO: We will have to make this more generic if robot_type != node_type
# Make Robot State relative to node
_, std = env.get_standardize_params(state[robot_type], node_type=robot_type)
std[0:2] = env.attention_radius[(node_type, robot_type)]
robot_traj_st = env.standardize(robot_traj,
state[robot_type],
node_type=robot_type,
mean=node_traj,
std=std)
robot_traj_st_t = torch.tensor(robot_traj_st, dtype=torch.float)
return robot_traj_st_t
def get_node_timestep_data(env, scene, t, node, state, pred_state,
edge_types, max_ht, max_ft, hyperparams,
scene_graph=None):
"""
Pre-processes the data for a single batch element: node state over time for a specific time in a specific scene
as well as the neighbour data for it.
:param env: Environment
:param scene: Scene
:param t: Timestep in scene
:param node: Node
:param state: Specification of the node state
:param pred_state: Specification of the prediction state
:param edge_types: List of all Edge Types for which neighbours are pre-processed
:param max_ht: Maximum history timesteps
:param max_ft: Maximum future timesteps (prediction horizon)
:param hyperparams: Model hyperparameters
:param scene_graph: If scene graph was already computed for this scene and time you can pass it here
:return: Batch Element
"""
# Node
timestep_range_x = np.array([t - max_ht, t])
timestep_range_y = np.array([t + 1, t + max_ft])
x = node.get(timestep_range_x, state[node.type])
y = node.get(timestep_range_y, pred_state[node.type])
first_history_index = (max_ht - node.history_points_at(t)).clip(0)
_, std = env.get_standardize_params(state[node.type], node.type)
std[0:2] = env.attention_radius[(node.type, node.type)]
rel_state = np.zeros_like(x[0])
rel_state[0:2] = np.array(x)[-1, 0:2]
x_st = env.standardize(x, state[node.type], node.type, mean=rel_state, std=std)
if list(pred_state[node.type].keys())[0] == 'position': # If we predict position we do it relative to current pos
y_st = env.standardize(y, pred_state[node.type], node.type, mean=rel_state[0:2])
else:
y_st = env.standardize(y, pred_state[node.type], node.type)
x_t = torch.tensor(x, dtype=torch.float)
y_t = torch.tensor(y, dtype=torch.float)
x_st_t = torch.tensor(x_st, dtype=torch.float)
y_st_t = torch.tensor(y_st, dtype=torch.float)
# Neighbors
neighbors_data_st = None
neighbors_edge_value = None
if hyperparams['edge_encoding']:
# Scene Graph
scene_graph = scene.get_scene_graph(t,
env.attention_radius,
hyperparams['edge_addition_filter'],
hyperparams['edge_removal_filter']) if scene_graph is None else scene_graph
neighbors_data_st = dict()
neighbors_edge_value = dict()
for edge_type in edge_types:
neighbors_data_st[edge_type] = list()
# We get all nodes which are connected to the current node for the current timestep
connected_nodes = scene_graph.get_neighbors(node, edge_type[1])
if hyperparams['dynamic_edges'] == 'yes':
# We get the edge masks for the current node at the current timestep
edge_masks = torch.tensor(scene_graph.get_edge_scaling(node), dtype=torch.float)
neighbors_edge_value[edge_type] = edge_masks
for connected_node in connected_nodes:
neighbor_state_np = connected_node.get(np.array([t - max_ht, t]),
state[connected_node.type],
padding=0.0)
# Make State relative to node where neighbor and node have same state
_, std = env.get_standardize_params(state[connected_node.type], node_type=connected_node.type)
std[0:2] = env.attention_radius[edge_type]
equal_dims = np.min((neighbor_state_np.shape[-1], x.shape[-1]))
rel_state = np.zeros_like(neighbor_state_np)
rel_state[:, ..., :equal_dims] = x[-1, ..., :equal_dims]
neighbor_state_np_st = env.standardize(neighbor_state_np,
state[connected_node.type],
node_type=connected_node.type,
mean=rel_state,
std=std)
neighbor_state = torch.tensor(neighbor_state_np_st, dtype=torch.float)
neighbors_data_st[edge_type].append(neighbor_state)
# Robot
robot_traj_st_t = None
timestep_range_r = np.array([t, t + max_ft])
if hyperparams['incl_robot_node']:
x_node = node.get(timestep_range_r, state[node.type])
if scene.non_aug_scene is not None:
robot = scene.get_node_by_id(scene.non_aug_scene.robot.id)
else:
robot = scene.robot
robot_type = robot.type
robot_traj = robot.get(timestep_range_r, state[robot_type], padding=0.0)
robot_traj_st_t = get_relative_robot_traj(env, state, x_node, robot_traj, node.type, robot_type)
# Map
map_tuple = None
if hyperparams['use_map_encoding']:
if node.type in hyperparams['map_encoder']:
if node.non_aug_node is not None:
x = node.non_aug_node.get(np.array([t]), state[node.type])
me_hyp = hyperparams['map_encoder'][node.type]
if 'heading_state_index' in me_hyp:
heading_state_index = me_hyp['heading_state_index']
# We have to rotate the map in the opposit direction of the agent to match them
if type(heading_state_index) is list: # infer from velocity or heading vector
heading_angle = -np.arctan2(x[-1, heading_state_index[1]],
x[-1, heading_state_index[0]]) * 180 / np.pi
else:
heading_angle = -x[-1, heading_state_index] * 180 / np.pi
else:
heading_angle = None
scene_map = scene.map[node.type]
map_point = x[-1, :2]
patch_size = hyperparams['map_encoder'][node.type]['patch_size']
map_tuple = (scene_map, map_point, heading_angle, patch_size)
return (first_history_index, x_t, y_t, x_st_t, y_st_t, neighbors_data_st,
neighbors_edge_value, robot_traj_st_t, map_tuple)
def get_timesteps_data(env, scene, t, node_type, state, pred_state,
edge_types, min_ht, max_ht, min_ft, max_ft, hyperparams):
"""
Puts together the inputs for ALL nodes in a given scene and timestep in it.
:param env: Environment
:param scene: Scene
:param t: Timestep in scene
:param node_type: Node Type of nodes for which the data shall be pre-processed
:param state: Specification of the node state
:param pred_state: Specification of the prediction state
:param edge_types: List of all Edge Types for which neighbors are pre-processed
:param max_ht: Maximum history timesteps
:param max_ft: Maximum future timesteps (prediction horizon)
:param hyperparams: Model hyperparameters
:return:
"""
nodes_per_ts = scene.present_nodes(t,
type=node_type,
min_history_timesteps=min_ht,
min_future_timesteps=max_ft,
return_robot=not hyperparams['incl_robot_node'])
batch = list()
nodes = list()
out_timesteps = list()
for timestep in nodes_per_ts.keys():
scene_graph = scene.get_scene_graph(timestep,
env.attention_radius,
hyperparams['edge_addition_filter'],
hyperparams['edge_removal_filter'])
present_nodes = nodes_per_ts[timestep]
for node in present_nodes:
nodes.append(node)
out_timesteps.append(timestep)
batch.append(get_node_timestep_data(env, scene, timestep, node, state, pred_state,
edge_types, max_ht, max_ft, hyperparams,
scene_graph=scene_graph))
if len(out_timesteps) == 0:
return None
return collate(batch), nodes, out_timesteps |