def get_map_feats(feature_map,
batch_shape,
batch_size,
last_agent_positions_grid,
A,
phi=None):
"""Interpolate the positions into the feature map.
:param feature_map: (B, H, W, F), feature map
:param batch_shape: first several dimensions of positions. could be inferred but we require it for sanity checks.
:param batch_size: size of batch. could be inferred but we require it for sanity checks.
:param last_agent_positions_grid: batch_shape + (A, D)
:param A: number of agents. could be inferred but we require it for sanity checks.
:param F: feature map dimension. could be inferred but we require it for sanity checks.
:param phi: DEPRECATED
:returns:
:rtype:
"""
assert (tensoru.rank(feature_map) == 4)
assert (tensoru.rank(last_agent_positions_grid) >= 3)
F = tensoru.size(feature_map, -1)
# (B, batch_shape[1:]*A, d)
last_agent_positions_grid_r_xy = tf.reshape(last_agent_positions_grid,
(batch_shape[0], -1, 2))
last_agent_positions_grid_r_ij = last_agent_positions_grid_r_xy[..., ::-1]
# (B, batch_shape[1:]*A, F)
# N.B. the indexing order! Our points are stored in xy-format, with x corresponding to the W dimension of the feature grid.
map_feat = tfu.interpolate_bilinear(feature_map,
last_agent_positions_grid_r_ij,
indexing='ij')
# (BSize, A, F)
map_feats = tf.reshape(map_feat, (batch_size, A, F))
return map_feats
def prepare(self, batch_shape):
if len(batch_shape) == 2: self.batch_str = 'bk'
elif len(batch_shape) == 1: self.batch_str = 'b'
else: raise ValueError("Unhandled batch size")
if len(batch_shape) == 1:
self.current_local2world = self.original_local2world = self.world2local.invert(
)
# Input spaces: Car frames. Output space: grid frame.
self.current_local2grid = self.original_local2grid = self.world2grid * self.current_local2world
else:
R = tensoru.repeat_expand_dims(self.world2local.R,
axis=1,
n=len(batch_shape) - 1)
R = tf.tile(R, (1, ) + batch_shape[1:] + (1, ) *
(tensoru.rank(self.world2local.R) - 1))
t = tensoru.repeat_expand_dims(self.world2local.t,
axis=1,
n=len(batch_shape) - 1)
t = tf.tile(t, (1, ) + batch_shape[1:] + (1, ) *
(tensoru.rank(self.world2local.t) - 1))
# Use this object to track the car frames during the rollout.
# Input space: world coordinates. Output spaces: car frame coordinates.
world2local = similarityu.SimilarityTransform(
R=R, t=t, scale=self.world2local.scale)
# "current_local2<X>" frames track the frame as it rolls out. The "local" frames are fixed, however.
self.current_local2world = self.original_local2world = world2local.invert(
)
# Input spaces: Car frames. Output space: grid frame.
self.current_local2grid = self.original_local2grid = self.world2grid * self.current_local2world
def signed_distance_transform(binary_image,
normalize=True,
clip_top=1,
clip_bottom=-10,
dtype=np.float64):
"""
:param binary_image: np.ndarray with boolean type and two dimensions.
:param normalize: whether to normalize the result to [0, 1]
:param clip_top: if normalizing, positive distance at which to clip
:param clip_bottom: if normalizing, negative distance at which to clip
:param dtype: dtype of the
:returns:
:rtype:
"""
assert (binary_image.dtype is np.dtype(np.bool))
assert (tensoru.rank(binary_image) == 2)
dt = morph.distance_transform_edt
sdt = (dt(binary_image) - dt(1 - binary_image)).astype(dtype)
if not normalize:
return sdt
else:
assert (clip_top > clip_bottom)
assert (clip_top >= 1)
sdt[sdt > clip_top] = clip_top
sdt[sdt < clip_bottom] = clip_bottom
return (sdt - clip_bottom) / (clip_top - clip_bottom)
def one_hotify_light_strings(light_strings):
"""Featurize light strings in the tensorflow graph. By putting it in the graph, we don't have to remember how they're featurized
when it comes to inference time.
:param light_strings:
:returns:
:rtype:
"""
light_strings = tf.strings.upper(light_strings)
assert (tensoru.rank(light_strings) == 1)
eye = tf.eye(5, dtype=tf.float64)[..., None]
c = lambda x: tf.cast(x, tf.float64)
none = c(tf.equal(light_strings, 'NONE'))[None]
green = c(tf.equal(light_strings, 'GREEN'))[None]
yellow = c(tf.equal(light_strings, 'YELLOW'))[None]
intersection = c(tf.equal(light_strings, 'INTERSECTION'))[None]
red = c(tf.equal(light_strings, 'RED'))[None]
# Ensure there's exactly one catgoirzation for each.
with tf.control_dependencies([
tf.compat.v1.assert_equal(
none + green + yellow + intersection + red, tf.ones_like(none))
]):
# This puts them in separate bins.
feats = tf.transpose(eye[0] * none + eye[1] * green +
eye[2] * intersection + eye[3] * yellow +
eye[4] * red)
# This puts them all in the same bin.
# feats = tf.transpose(eye[0] * none + eye[0] * green + eye[0] * intersection + eye[0] * yellow + eye[4] * red)
return feats
def _Rtheta(cls, theta, lib):
Rs = lib.stack([(lib.cos(theta), -lib.sin(theta)),
(lib.sin(theta), lib.cos(theta))],
axis=0)
# (x, y, ...) -> (..., x, y)
if tensoru.rank(Rs) > 2:
Rs = tensoru.rotate_left(tensoru.rotate_left(Rs, lib=lib), lib=lib)
return Rs
def __init__(self, R, t, scale, lib=tf):
"""
Implements x' = f(x; R, t, scale) = scale * R * x + t for batched data.
t represents the origin of the transformed coordinates.
R represents the pre-translation rotation matrix.
:param R: (B, A, 2, 2) or (B, 2, 2) or (2, 2)
:param t: (B, A, 2) or (B, 2) or (2,)
:param scale: scalar
:returns:
:rtype:
"""
assert (lib in (np, tf))
self.rr = tensoru.rank(R)
self.tr = tensoru.rank(t)
# self.theta = lib.atan2(R[..., 1, 0], R[..., 0, 0])
joint_rank = (self.rr, self.tr)
assert (joint_rank in [(5, 4), (4, 3), (3, 2), (2, 1)])
if self.rr == 5: self._ein = 'bkaij'
elif self.rr == 4: self._ein = 'baij'
elif self.rr == 3: self._ein = 'bij'
elif self.rr == 2: self._ein = 'ij'
else: raise ValueError("Unhandled Rotation matrix dimensionality!")
# Maintain some tiled versions for broadcasted adding.
self._ts = {}
if self.tr == 3:
self._ts['baj'] = self.t
self._ts['batj'] = tensoru.repeat_expand_dims(self.t, 1, axis=-2)
self._ts['bkaj'] = tensoru.swap_axes(self._ts['batj'], 1, 2)
self._ts['bkatj'] = tensoru.repeat_expand_dims(self._ts['batj'],
1,
axis=1)
self._ts['baNj'] = self.t[..., None, :]
elif self.tr == 4:
self._ts['bkaj'] = self.t
self._ts['batj'] = tensoru.swap_axes(self._ts['bkaj'], 1, 2)
self._ts['bkaNj'] = self.t[..., None, :]
self._ts['bkaNj'] = self.t[..., None, :]
else:
pass
def get_whisker_map_feats(feature_map,
batch_shape,
batch_size,
cars2grid,
template_cars,
A,
phi=None):
"""Interpolate the positions into the feature map.
:param feature_map: (B, H, W, F), feature map
:param batch_shape: first several dimensions of positions. could be inferred but we require it for sanity checks.
:param batch_size: size of batch. could be inferred but we require it for sanity checks.
:param last_agent_positions_cars: batch_shape + (A, D)
:param cars2grid: SimilarityTransform from car frames to grid
:param template: (N, D) template of positions in local frame at which to interpolate
:param A: number of agents. could be inferred but we require it for sanity checks.
:param F: feature map dimension. could be inferred but we require it for sanity checks.
:param phi: DEPRECATED
:returns:
:rtype:
"""
assert (tensoru.rank(feature_map) == 4)
F = tensoru.size(feature_map, -1)
if len(batch_shape) == 2:
points_ein = 'bkaNj'
template_cars = template_cars[None]
elif len(batch_shape) == 1:
points_ein = 'baNj'
else:
raise ValueError
n_whiskers = tensoru.size(template_cars, -2)
whiskers_grid = cars2grid.apply(template_cars, points_ein=points_ein)
whiskers_grid_r_xy = tf.reshape(whiskers_grid, (batch_shape[0], -1, 2))
# (B, batch_shape[1:]*A, F)
# N.B. the indexing order! Our points are stored in xy-format, with x corresponding to the W dimension of the feature grid.
map_feat = tfu.interpolate_bilinear(feature_map,
whiskers_grid_r_xy,
indexing='xy')
# (B, ..., A, n_whiskers, F)
map_feat_r = tf.reshape(map_feat, batch_shape + (A, n_whiskers, F))
# (B*..., A, n_whiskers*F)
map_feats = tf.reshape(map_feat_r, (batch_size, A, n_whiskers * F))
return map_feats, whiskers_grid
def apply(self,
points,
points_ein=None,
dtype=None,
name=None,
translate=True):
input_rank = tensoru.rank(points)
if dtype is None: dtype = self.lib.float64
assert (dtype in (self.lib.float64, self.lib.float32, self.lib.int32))
if points_ein is None:
if input_rank == 5:
points_ein = 'bkatj'
# Note that we can't handle 'bkaj' input implicitly... i.e. force the second batching to be the second axis.
elif input_rank == 4:
points_ein = 'batj'
elif input_rank == 3:
points_ein = 'baj'
elif input_rank == 2:
points_ein = 'bj'
elif input_rank == 1:
points_ein = 'j'
else:
raise ValueError(
"Unhandled points dimensionality: {}".format(input_rank))
# Ensure input only uses certain characters
assert (len(set(points_ein) - set('bkatjN')) == 0)
einstr = '{},{}->{}'.format(self._ein, points_ein,
points_ein.replace('j', 'i'))
points_R = self.scale * self.lib.einsum(einstr, self.R, points)
# Add the translation.
if translate:
if self.tr == 1: points_Rt = points_R + self.t
else: points_Rt = points_R + self._ts[points_ein]
else:
points_Rt = points_R
if dtype == self.lib.int32:
return self.lib.cast(self.lib.round(points_Rt), self.lib.int32)
else:
return points_Rt
def plot_joint_trajectory(joint_traj,
key='joint_trajectories',
fig=None,
ax=None,
render=True,
limit=100,
**kwargs):
"""
:param joint_traj: K x A x T x d
:param key:
:param fig:
:param ax:
:param render:
:param kwargs:
:returns:
:rtype:
"""
assert (tensoru.rank(joint_traj) == 4)
A = tensoru.size(joint_traj, 1)
for a in range(A):
render_a = (a == A - 1) and render
color = kwargs.get('color', COLORS[a])
if isinstance(joint_traj, tf.Tensor):
single_traj = joint_traj[:, a].numpy()
else:
single_traj = joint_traj[:, a]
kwargs.pop('color', None)
fig, ax = plot_trajectory(key,
single_traj,
color=color,
render=render_a,
fig=fig,
ax=ax,
axis=[0, limit, limit, 0],
**kwargs)
assert (fig is not None)
return fig, ax
def plot_joint_trajectory(joint_traj,
limit,
scale=1,
agents=None,
fig=None,
ax=None,
**kwargs):
assert (tensoru.rank(joint_traj) == 4)
if agents is None:
agents = range(tensoru.size(joint_traj, 1))
for a in agents:
if "color" in kwargs:
color = kwargs.pop("color")
else:
color = cm.get_cmap("tab10").colors[a]
if isinstance(joint_traj, tf.Tensor):
single_traj = joint_traj[:, a].numpy().copy()
else:
single_traj = joint_traj[:, a].copy()
if scale > 0:
single_traj *= scale
single_traj[..., -1] *= -1
fig, ax = plot_trajectory(single_traj,
color=color,
fig=fig,
ax=ax,
axis=limit,
**kwargs)
assert (fig is not None)
return fig, ax
def inverse(self, S, phi):
"""Implements the inverse component of the bijection.
:param S: output states. Must be in __CAR FRAMES__.
:param phi: context
:returns:
:rtype:
"""
r = tensoru.rank(S)
assert (r in (4, 5))
if r == 4:
B, A, T, D = tensoru.shape(S)
S_0 = phi.S_past_car_frames[..., -1, :]
S_m1 = phi.S_past_car_frames[..., -2, :]
batch_shape = (B, )
elif r == 5:
B, K, A, T, D = tensoru.shape(S)
S_0 = tensoru.expand_and_tile_axis(phi.S_past_car_frames[...,
-1, :],
N=K,
axis=1)
S_m1 = tensoru.expand_and_tile_axis(phi.S_past_car_frames[...,
-2, :],
N=K,
axis=1)
batch_shape = (B, K)
else:
raise ValueError("Unhandled rank of data to invert.")
S_history = [S_m1, S_0]
Z_history = []
m_history = []
mu_history = []
sigel_history = []
sigma_history = []
metadata_history = []
phi.prepare(batch_shape)
self._prepare(batch_shape, phi)
for t_idx in range(T):
m_t, sigel_t, sigma_t = self.step_generate(S_history, phi)
mu_t = m_t + 2 * S_history[-1] - S_history[-2]
# S first, then compute Z.
S_t = S[..., t_idx, :]
# expm(X) expm(-X) = I -> (expm(X))^{-1} = expm(-X). Avoids computing any inverses explicitly.
sigma_t_inv = tf.linalg.expm(-1 * sigel_t)
Z_t = tf.einsum('...ij,...j->...i', sigma_t_inv, S_t - mu_t)
phi.update_frames(S_t_car_frames=S_t,
S_tm1_car_frames=S_history[-1])
m_history.append(m_t)
mu_history.append(mu_t)
sigel_history.append(sigel_t)
sigma_history.append(sigma_t)
S_history.append(S_t)
Z_history.append(Z_t)
metadata_history.append(self.current_metadata)
roll = interface.ESPRollout(S_car_frames_list=S_history[2:],
Z_list=Z_history,
m_list=m_history,
mu_list=mu_history,
sigma_list=sigma_history,
sigel_list=sigel_history,
metadata_list=metadata_history,
phi=phi)
return roll
def forward(self, Z, phi):
"""Implements the forward component of the bijection. Implemented in __CAR FRAMES__.
:param Z: latent states
:param phi: context
:returns:
:rtype:
"""
r = tensoru.rank(Z)
with contextlib.ExitStack() as stack:
if getattr(self, 'debug_eager', False):
tape = tf.GradientTape(persistent=True)
stack.enter_context(tape)
tape.watch(phi.S_past_car_frames)
tape.watch(phi.overhead_features)
assert (r in (4, 5))
if r == 4:
B, A, T, D = tensoru.shape(Z)
S_0 = phi.S_past_car_frames[..., -1, :]
S_m1 = phi.S_past_car_frames[..., -2, :]
batch_shape = (B, )
elif r == 5:
B, K, A, T, D = tensoru.shape(Z)
# (B, A, T, D) -> (B, K, A, T, D)
S_0 = tensoru.expand_and_tile_axis(
phi.S_past_car_frames[..., -1, :], N=K, axis=1)
S_m1 = tensoru.expand_and_tile_axis(
phi.S_past_car_frames[..., -2, :], N=K, axis=1)
batch_shape = (B, K)
else:
raise ValueError("Unhandled rank of latents to warp.")
S_history = [S_m1, S_0]
Z_history = []
m_history = []
mu_history = []
sigel_history = []
sigma_history = []
metadata_history = []
phi.prepare(batch_shape)
self._prepare(batch_shape, phi)
for t_idx in range(T):
m_t, sigel_t, sigma_t = self.step_generate(S_history, phi)
mu_t = m_t + 2 * S_history[-1] - S_history[-2]
# Z first, then compute S.
Z_t = Z[..., t_idx, :]
S_t = mu_t + tf.einsum('...ij,...j->...i', sigma_t, Z_t)
phi.update_frames(S_t_car_frames=S_t,
S_tm1_car_frames=S_history[-1])
m_history.append(m_t)
mu_history.append(mu_t)
sigel_history.append(sigel_t)
sigma_history.append(sigma_t)
S_history.append(S_t)
Z_history.append(Z_t)
metadata_history.append(self.current_metadata)
if getattr(self, 'debug_eager', False): pdb.set_trace()
roll = interface.ESPRollout(S_car_frames_list=S_history[2:],
Z_list=Z_history,
m_list=m_history,
mu_list=mu_history,
sigma_list=sigma_history,
sigel_list=sigel_history,
metadata_list=metadata_history,
phi=phi)
return roll
def _traj_to_s_T(self, trajectories):
assert (tensoru.rank(trajectories) == 5)
s_T = tf.squeeze(trajectories[..., -1, :])
assert (tensoru.shape(s_T) == self.full_shape)
return s_T
