def test_to_euler(self):
"""mtx -> euler -> mtx"""
data = random_rotations(13, dtype=torch.float64)
for convention in self._all_euler_angle_conventions():
euler_angles = matrix_to_euler_angles(data, convention)
mdata = euler_angles_to_matrix(euler_angles, convention)
self.assertTrue(torch.allclose(data, mdata))
def test_euler_grad_exists(self):
"""Euler angle calculations are differentiable."""
rotation = random_rotation(dtype=torch.float64, requires_grad=True)
for convention in self._all_euler_angle_conventions():
euler_angles = matrix_to_euler_angles(rotation, convention)
mdata = euler_angles_to_matrix(euler_angles, convention)
[g] = torch.autograd.grad(mdata.sum(), rotation)
self.assertTrue(torch.isfinite(g).all())
def test_from_euler(self):
"""euler -> mtx -> euler"""
n_repetitions = 10
# tolerance is how much we keep the middle angle away from the extreme
# allowed values which make the calculation unstable (Gimbal lock).
tolerance = 0.04
half_pi = math.pi / 2
data = torch.zeros(n_repetitions, 3)
data.uniform_(-math.pi, math.pi)
data[:, 1].uniform_(-half_pi + tolerance, half_pi - tolerance)
for convention in self._tait_bryan_conventions():
matrices = euler_angles_to_matrix(data, convention)
mdata = matrix_to_euler_angles(matrices, convention)
self.assertTrue(torch.allclose(data, mdata))
data[:, 1] += half_pi
for convention in self._proper_euler_conventions():
matrices = euler_angles_to_matrix(data, convention)
mdata = matrix_to_euler_angles(matrices, convention)
self.assertTrue(torch.allclose(data, mdata))
def test_interpolate_rotation():
R1 = euler_angles_to_matrix(torch.tensor([0.0, 0.0, 0.0]),
convention="ZYX")
R2 = euler_angles_to_matrix(torch.tensor([np.pi / 2., 0., 0.]),
convention="ZYX")
R = interp_rotation(R1, R2, interp_factor=0.)
assert torch.allclose(R, R1), (R, R1)
R = interp_rotation(R1, R2, interp_factor=1.)
assert torch.allclose(R, R2), (R, R2)
R = interp_rotation(R2, R2, interp_factor=0.75)
assert torch.allclose(R, R2), (R, R2)
R = interp_rotation(R1, R2, interp_factor=0.5)
expected = euler_angles_to_matrix(torch.tensor([np.pi / 4., 0., 0.]),
convention="ZYX")
assert torch.allclose(R, expected), (R, expected)
R3 = euler_angles_to_matrix(torch.tensor([np.pi, 0., 0.]),
convention="ZYX")
R = interp_rotation(R1, R3, interp_factor=0.5)
angle_distance = torch.arccos(
(torch.trace(torch.matmul(R.transpose(1, 0), R3)) - 1) / 2.)
assert torch.isclose(angle_distance,
torch.tensor(np.pi / 2.)), (angle_distance,
np.pi / 2.)
R4 = euler_angles_to_matrix(torch.tensor([3. * np.pi / 2., 0., 0.]),
convention="ZYX")
R = interp_rotation(R1, R4, interp_factor=0.5)
expected = euler_angles_to_matrix(torch.tensor([-np.pi / 4., 0., 0.]),
convention="ZYX")
assert torch.allclose(R, expected), (R, expected)
def do_fixed_structure_mcmc(grammar,
scene_tree,
num_samples=500,
perturb_in_config_space=False,
verbose=0,
vis_callback=None,
translation_variance=0.1,
rotation_variance=0.1,
do_hit_and_run_postprocess=False):
''' Given a scene tree, resample its continuous variables
(i.e. the node poses) while keeping the root and observed
node poses fixed. Returns a population of trees sampled
from the joint distribution over node poses given the
fixed structure.
Verbose = 0: Print nothing
Verbose = 1: Print updates about accept rate and steps
Verbose = 2: Print NLP output and scoring info
'''
# Strategy sketch:
# - Initialize a random walk at the current scene tree.
# - Repeatedly:
# 1) Apply a random perturbation to the non-observed, non-root
# node poses. If perturb_in_config_space, then this perturbation
# is applied to the root node first, and then that change propagated
# down to its children. Otherwise, every node pose is perturbed
# randomly simultaneously.
# 2) Use an NLP to project the perturbed tree configuration to
# the nearest feasible tree configuration (i.e. satisfyingq
# uniform and joint angle bounds).
# 2b) If do_hit_and_run_postprocess, randomly sample a point uniformly
# between the projected config and the initial config, and project
# that to the nearest good config to get a proposed config.
# 3) Use the Accept/reject with a Metroplis acceptance ratio
# (without proposal probabilities, right now).
# x_proj = projected pose
# x_init = initial pose
# Accept if u ~ [0, 1] <= min(1, p(x_proj)/p(x_init))
#
# TODO: This isn't gonna have a perfect underlying distribution
# unless I account for the Hastings ratio. But given my use of
# nasty nonlinear projection, that seems rough.
current_tree = deepcopy(scene_tree)
# Do steps of random-walk MCMC on those variables.
n_accept = 0
old_score = current_tree.score()
assert torch.isfinite(old_score), "Bad initialization for MCMC."
sample_trees = []
for step_k in range(num_samples):
new_tree = deepcopy(current_tree)
# Update tree from root down.
node_queue = [(current_tree.get_root(), new_tree.get_root())]
while len(node_queue) > 0:
current_parent, new_parent = node_queue.pop()
current_children, current_rules = current_tree.get_children_and_rules(
current_parent)
new_children = new_tree.get_children(new_parent)
## Apply pertubrations to children according to their generating rules.
for current_child, rule, new_child in zip(current_children,
current_rules,
new_children):
if current_child.observed:
# No perturbation to observed nodes necessary.
continue
## The exact perturbation we'll apply
# depends on the relationship to the parent.
xyz_rule, rotation_rule = rule.xyz_rule, rule.rotation_rule
if type(xyz_rule) is WorldFrameBBoxRule:
# Perturb in axes that are not equality constrained
scale = xyz_rule.ub - xyz_rule.lb
perturb = dist.Normal(torch.zeros(3), scale *
translation_variance).sample()
perturb[xyz_rule.xyz_dist.delta_mask] = 0.
new_child.translation = current_child.translation + perturb
elif type(xyz_rule) is WorldFrameBBoxOffsetRule:
# Perturb in axes that are not equality constrained
scale = xyz_rule.ub - xyz_rule.lb
perturb = dist.Normal(torch.zeros(3), scale *
translation_variance).sample()
perturb[xyz_rule.xyz_dist.delta_mask] = 0.
if perturb_in_config_space:
current_offset = current_child.translation - current_parent.translation
new_child.translation = new_parent.translation + current_offset + perturb
else:
new_child.translation = current_child.translation + perturb
elif type(xyz_rule) is WorldFrameGaussianOffsetRule:
# Perturb all axes
scale = xyz_rule.variance
perturb = dist.Normal(torch.zeros(3), scale *
translation_variance).sample()
if perturb_in_config_space:
current_offset = current_child.translation - current_parent.translation
new_child.translation = new_parent.translation + current_offset + perturb
else:
new_child.translation = current_child.translation + perturb
elif type(xyz_rule) is SamePositionRule:
pass
else:
raise NotImplementedError("%s" % xyz_rule)
if type(rotation_rule) is UnconstrainedRotationRule:
# Apply small random rotation
random_small_rotation = euler_angles_to_matrix(
dist.Normal(torch.zeros(3),
torch.ones(3) *
rotation_variance).sample().unsqueeze(0),
convention="ZYX")[0, ...]
new_child.rotation = torch.matmul(current_child.rotation,
random_small_rotation)
elif type(rotation_rule) is SameRotationRule:
pass
elif type(rotation_rule) is UniformBoundedRevoluteJointRule:
# Apply small rotation around axis, unless the rotation is fully constrained
if not np.isclose(rotation_rule.lb, rotation_rule.ub):
scale = rotation_rule.ub - rotation_rule.lb
random_angle = dist.Normal(
torch.zeros(1),
scale * translation_variance).sample()
orig_angle, orig_axis, in_range = recover_relative_angle_axis(
current_parent,
current_child,
target_axis=rotation_rule.axis)
assert in_range
# Add angle to orig angle, and rotate around the joint's actual axis to get
# the new rotation offset.
new_angle_axis = rotation_rule.axis * (orig_angle +
random_angle)
new_R_offset = axis_angle_to_matrix(
new_angle_axis.unsqueeze(0))[0, ...]
if perturb_in_config_space:
new_child.rotation = torch.matmul(
new_parent.rotation, new_R_offset)
else:
new_child.rotation = torch.matmul(
current_child.rotation, new_R_offset)
elif type(rotation_rule) is GaussianChordOffsetRule:
# Apply small rotation around axis
scale = rotation_rule.concentration
random_angle = dist.Normal(torch.zeros(1), scale *
translation_variance).sample()
orig_angle, orig_axis, in_range = recover_relative_angle_axis(
current_parent,
current_child,
target_axis=rotation_rule.axis)
assert in_range
# Add angle to orig angle, and rotate around the joint's actual axis to get
# the new rotation offset.
new_angle_axis = rotation_rule.axis * (orig_angle +
random_angle)
new_R_offset = axis_angle_to_matrix(
new_angle_axis.unsqueeze(0))[0, ...]
if perturb_in_config_space:
new_child.rotation = torch.matmul(
new_parent.rotation, new_R_offset)
else:
new_child.rotation = torch.matmul(
current_child.rotation, new_R_offset)
else:
raise NotImplementedError("%s" % xyz_rule)
# Add children to the node queue.
for current_child, new_child in zip(current_children,
new_children):
node_queue.append((current_child, new_child))
# Now project the tree to the closest feasible tree to that config.
projection_results = optimize_scene_tree_with_nlp(
grammar,
new_tree,
initial_guess_tree=current_tree,
objective="projection",
verbose=verbose > 1)
if projection_results.optim_result.is_success():
new_tree = projection_results.refined_tree
try:
new_score = new_tree.score(verbose=verbose > 1)
except ValueError as e:
logging.warn("Unexpected ValueError: %s", e)
new_score = -torch.tensor(np.inf)
else:
logging.warning("Post-random-step projection failed")
new_score = -torch.tensor(np.inf)
if torch.isfinite(new_score) and do_hit_and_run_postprocess:
# Randomly sample point between this new tree and the
# initial tree.
interp_factor = dist.Uniform(torch.zeros(1),
torch.ones(1)).sample()
current_root = current_tree.get_root()
for original_node, new_node in zip(current_tree.nodes,
new_tree.nodes):
new_node.translation = interp_translation(
original_node.translation, new_node.translation,
interp_factor)
new_node.rotation = interp_rotation(original_node.rotation,
new_node.rotation,
interp_factor)
projection_results = optimize_scene_tree_with_nlp(
grammar,
new_tree,
initial_guess_tree=current_tree,
objective="projection",
verbose=verbose > 1)
if projection_results.optim_result.is_success():
new_tree = projection_results.refined_tree
try:
new_score = new_tree.score(verbose=verbose > 1)
except ValueError as e:
logging.warning(
"Unexpected ValueError after hit-and-run projection: %s",
e)
new_score = -torch.tensor(np.inf)
else:
logging.warning("Post-hit-and-run projection failed")
new_score = -torch.tensor(np.inf)
reject = True
if torch.isfinite(new_score):
# MH acceptance ratio: just score_new / score_old, since proposal is symmetric
alpha = min(1, torch.exp(new_score - old_score))
if verbose:
print("New score %f, old score %f, alpha %f" %
(new_score, old_score, alpha))
if dist.Uniform(0., 1.).sample() <= alpha:
# Accepted. Update our temp pose variables to reflect new "good" state.
n_accept += 1
current_tree = new_tree
sample_trees.append(deepcopy(current_tree))
if vis_callback is not None:
vis_callback(current_tree)
if verbose:
print("%d: Accept rate %f" % (step_k, n_accept / (step_k + 1)))
return sample_trees
