def fix_pose(self, name, pose=None, fraction=0.5):
body = self.get_body(name)
if pose is None:
pose = get_pose(body)
else:
set_pose(body, pose)
# TODO: can also drop in simulation
x, y, z = point_from_pose(pose)
roll, pitch, yaw = euler_from_quat(quat_from_pose(pose))
quat = quat_from_euler(Euler(yaw=yaw))
set_quat(body, quat)
surface_name = self.get_supporting(name)
if surface_name is None:
return None, None
if fraction == 0:
new_pose = (Point(x, y, z), quat)
return new_pose, surface_name
surface_aabb = compute_surface_aabb(self, surface_name)
new_z = (1 - fraction) * z + fraction * stable_z_on_aabb(
body, surface_aabb)
point = Point(x, y, new_z)
set_point(body, point)
print('{} error: roll={:.3f}, pitch={:.3f}, z-delta: {:.3f}'.format(
name, roll, pitch, new_z - z))
new_pose = (point, quat)
return new_pose, surface_name
def pose2d_on_surface(world, entity_name, surface_name, pose2d=UNIT_POSE2D):
x, y, yaw = pose2d
body = world.get_body(entity_name)
surface_aabb = compute_surface_aabb(world, surface_name)
z = stable_z_on_aabb(body, surface_aabb)
pose = Pose(Point(x, y, z), Euler(yaw=yaw))
set_pose(body, pose)
return pose
def pose_from_pose2d(self, pose2d, surface):
#assert surface in self.poses_from_surface
#reference_pose = self.poses_from_surface[surface][0]
body = self.world.get_body(self.name)
surface_aabb = compute_surface_aabb(self.world, surface)
world_from_surface = get_surface_reference_pose(
self.world.kitchen, surface)
if DIM == 2:
x, y = pose2d[:DIM]
yaw = np.random.uniform(*CIRCULAR_LIMITS)
else:
x, y, yaw = pose2d
z = stable_z_on_aabb(
body,
surface_aabb) + Z_EPSILON - point_from_pose(world_from_surface)[2]
point = Point(x, y, z)
surface_from_body = Pose(point, Euler(yaw=yaw))
set_pose(body, multiply(world_from_surface, surface_from_body))
return create_relative_pose(self.world, self.name, surface)
def get_supporting(self, obj_name):
# is_placed_on_aabb | is_center_on_aabb
# Only want to generate stable placements, but can operate on initially unstable ones
# TODO: could filter orientation as well
body = self.get_body(obj_name)
supporting = {
surface
for surface in ALL_SURFACES
if is_center_on_aabb(body,
compute_surface_aabb(self, surface),
above_epsilon=5e-2,
below_epsilon=5e-2)
}
if ('range' in supporting) and (len(supporting) == 2):
# TODO: small hack for now
supporting -= {'range'}
if len(supporting) != 1:
print('{} is not supported by a single surface ({})!'.format(
obj_name, supporting))
return None
[surface_name] = supporting
return surface_name
def gen(obj_name, surface_name):
obj_body = world.get_body(obj_name)
surface_body = world.kitchen
if surface_name in ENV_SURFACES:
surface_body = world.environment_bodies[surface_name]
surface_aabb = compute_surface_aabb(world, surface_name)
learned_poses = load_placements(world, surface_name) if learned else [
] # TODO: GROW_PLACEMENT
yaw_range = (-np.pi, np.pi)
#if world.is_real():
# center = -np.pi/4
# half_extent = np.pi / 16
# yaw_range = (center-half_extent, center+half_extent)
while True:
for _ in range(max_attempts):
if surface_name in STOVES:
surface_link = link_from_name(world.kitchen, surface_name)
world_from_surface = get_link_pose(world.kitchen,
surface_link)
z = stable_z_on_aabb(
obj_body,
surface_aabb) - point_from_pose(world_from_surface)[2]
yaw = random.uniform(*yaw_range)
body_pose_surface = Pose(Point(z=z + z_offset),
Euler(yaw=yaw))
body_pose_world = multiply(world_from_surface,
body_pose_surface)
elif learned:
if not learned_poses:
return
surface_pose_world = get_surface_reference_pose(
surface_body, surface_name)
sampled_pose_surface = multiply(
surface_pose_world, random.choice(learned_poses))
[x, y, _] = point_from_pose(sampled_pose_surface)
_, _, yaw = euler_from_quat(
quat_from_pose(sampled_pose_surface))
dx, dy = np.random.normal(
scale=pos_scale, size=2) if pos_scale else np.zeros(2)
# TODO: avoid reloading
z = stable_z_on_aabb(obj_body, surface_aabb)
yaw = random.uniform(*yaw_range)
#yaw = wrap_angle(yaw + np.random.normal(scale=rot_scale))
quat = quat_from_euler(Euler(yaw=yaw))
body_pose_world = ([x + dx, y + dy, z + z_offset], quat)
# TODO: project onto the surface
else:
# TODO: halton sequence
# unit_generator(d, use_halton=True)
body_pose_world = sample_placement_on_aabb(
obj_body, surface_aabb, epsilon=z_offset, percent=2.0)
if body_pose_world is None:
continue # return?
if visibility and not is_visible_by_camera(
world, point_from_pose(body_pose_world)):
continue
# TODO: make sure the surface is open when doing this
robust = True
if robust_radius != 0.:
for theta in np.linspace(0, 5 * np.pi, num=8):
x, y = robust_radius * unit_from_theta(theta)
delta_body = Pose(Point(x, y))
delta_world = multiply(body_pose_world, delta_body)
set_pose(obj_body, delta_world)
if not test_supported(world,
obj_body,
surface_name,
collisions=collisions):
robust = False
break
set_pose(obj_body, body_pose_world)
if robust and test_supported(
world, obj_body, surface_name, collisions=collisions):
rp = create_relative_pose(world, obj_name, surface_name)
yield (rp, )
break
else:
yield None