def main(args):
NOISE=0.00005
# define real world block params and initial poses
#TODO: Load blocks from pickle file.
if args.use_vision:
with open(args.blocks_file, 'rb') as handle:
blocks = pickle.load(handle)
block_init_xy_poses = None # These will be initialized randomly but updated by the vision system.
else:
# block0 = Object('block0', Dimensions(.0381,.0318,.0635), 1.0, Position(0,0,0), Color(1,0,0))
# block1 = Object('block1', Dimensions(.0381,.0587,.0635), 1.0, Position(0,0,0), Color(0,0,1))
# block2 = Object('block2', Dimensions(.0635,.0381,.0746), 1.0, Position(0,0,0), Color(0,1,0))
block0 = Object('block0', Dimensions(.0381,.0318,.05), 1.0, Position(0,0,0), Color(1,0,0))
block1 = Object('block1', Dimensions(.0381,.0587,.06), 1.0, Position(0,0,0), Color(0,0,1))
block2 = Object('block2', Dimensions(.0635,.0381,.05), 1.0, Position(0,0,0), Color(0,1,0))
blocks = [block0, block1, block2]
block_init_xy_poses = [Pose(Position(0.65,0.3,0), Quaternion(0,0,0,1)),
Pose(Position(0.65,0.15,0), Quaternion(0,0,0,1)),
Pose(Position(0.65,0.0,0), Quaternion(0,0,0,1))]
panda = PandaAgent(blocks, NOISE,
use_platform=False,
block_init_xy_poses=block_init_xy_poses,
teleport=False,
use_vision=args.use_vision,
use_action_server=args.use_action_server,
real=args.real)
# for now hard-code a tower, but in the future will be supplied from
# active data collection or tower found through planning for evaluation
tower_blocks = copy.copy(blocks)
if args.use_vision:
tower_poses = [Pose(Position(0.5,-0.25,0.0725), Quaternion(0,0,0,1)),
Pose(Position(0.5,-0.25,0.18), Quaternion(0,0,0,1)),
Pose(Position(0.5,-0.25,0.28), Quaternion(0,0,0,1))]
else:
tower_poses = [Pose(Position(0.3,0.25,.0318), Quaternion(0,0,0,1)),
Pose(Position(0.3,0.25,.0953), Quaternion(0,0,0,1)),
Pose(Position(0.3,0.25,.1643), Quaternion(0,0,0,1))]
tower = []
for block, pose in zip(tower_blocks, tower_poses):
block.set_pose(pose)
block = get_rotated_block(block) # NOTE: have to do to set rotations field of block
tower.append(block)
# and execute the resulting plan.
if args.use_action_server:
panda.simulate_tower_parallel(tower, base_xy=(0.5, -0.25), real=args.real, vis=True, T=2500)
else:
panda.simulate_tower(tower, base_xy=(0.5, -0.25), real=args.real, vis=True, T=2500)
def test_return_to_start(blocks, n_placements=5, rot_ix=0, block_ix=1):
"""
Let a block fall off the platform and check that we can successfully
pick it up and return it to the starting position.
"""
numpy.random.seed(10)
rot = list(rotation_group())[rot_ix]
for _ in range(n_placements):
# Create agent.
agent = PandaAgent(blocks)
original_pose = agent.pddl_blocks[block_ix].get_base_link_pose()
# Create a random action.
new_dims = numpy.abs(rot.apply(blocks[block_ix].dimensions))
place_pos = new_dims * (-0.5 * numpy.random.rand(3))
x, y, _ = place_pos + numpy.array(agent.platform.get_dimensions()) / 2
action = PlaceAction(pos=None, rot=rot, block=blocks[block_ix])
# Teleport block to platform.
blocks[block_ix].set_pose(
Pose(ZERO_POS, Quaternion(*action.rot.as_quat())))
rotated_block = get_rotated_block(blocks[block_ix])
platform_pose = agent.platform.get_base_link_pose()
platform_tform = pb_robot.geometry.tform_from_pose(platform_pose)
z = agent.platform.get_dimensions(
)[2] / 2 + rotated_block.dimensions[2] / 2 + 1e-5
tform = numpy.array([[1., 0., 0., action.pos[0]],
[0., 1., 0., action.pos[1]], [0., 0., 1., z],
[0., 0., 0., 1.]])
tform[0:3, 0:3] = action.rot.as_matrix()
body_tform = platform_tform @ tform
pose = pb_robot.geometry.pose_from_tform(body_tform)
agent.pddl_blocks[block_ix].set_base_link_pose(pose)
# Execute action.
p.setGravity(0, 0, -10)
for _ in range(500):
p.stepSimulation()
time.sleep(0.01)
check_ungraspable_block(agent)
# Solve PDDL Problem.
pddl_block = agent.pddl_blocks[block_ix]
init = agent._get_initial_pddl_state()
goal_pose = pb_robot.vobj.BodyPose(pddl_block, original_pose)
init += [('Pose', pddl_block, goal_pose),
('Supported', pddl_block, goal_pose, agent.table,
agent.table_pose)]
goal = ('and', ('AtPose', pddl_block, goal_pose), ('On', pddl_block,
agent.table))
# Solve the PDDLStream problem.
print('Init:', init)
print('Goal:', goal)
agent._solve_and_execute_pddl(init, goal)
p.disconnect()
def test_placement_ik(agent, blocks):
"""
To make sure that the platform is in a good position, make sure the
IK is feasible for some grasp position.
"""
get_block_pose = tamp.primitives.get_stable_gen_block(
[agent.table, agent.platform])
get_grasp = tamp.primitives.get_grasp_gen(agent.robot)
get_ik = tamp.primitives.get_ik_fn(agent.robot,
[agent.platform, agent.table])
for r in list(rotation_group()):
r = list(rotation_group())[4]
action = PlaceAction(pos=None, rot=r, block=blocks[0])
blocks[0].set_pose(Pose(ZERO_POS, Quaternion(*action.rot.as_quat())))
rotated_block = get_rotated_block(blocks[0])
x = action.pos[0]
y = action.pos[1]
z = agent.platform.get_dimensions(
)[2] / 2 + rotated_block.dimensions[2] / 2 + 1e-5
tform = numpy.array([[1., 0., 0., x], [0., 1., 0., y], [0., 0., 1., z],
[0., 0., 0., 1.]])
tform[0:3, 0:3] = action.rot.as_matrix()
platform_pose = pb_robot.vobj.BodyPose(
agent.platform, agent.platform.get_base_link_pose())
start_pose = pb_robot.vobj.BodyPose(
agent.pddl_blocks[0], agent.pddl_blocks[0].get_base_link_pose())
placement_pose = get_block_pose(agent.pddl_blocks[0], agent.platform,
platform_pose, tform)[0]
ik_found = False
for grasp in get_grasp(agent.pddl_blocks[0]):
ik_start = get_ik(agent.pddl_blocks[0], start_pose, grasp[0])
ik_placement = get_ik(agent.pddl_blocks[0], placement_pose,
grasp[0])
if ik_start is not None:
print('Y', end='')
else:
print('N', end='')
if ik_placement is not None:
ik_found = True
print('Y', end=' ')
else:
print('N', end=' ')
if ik_found:
print('Found IK.')
else:
print('No IK.')
break
def to_blocks(tower):
blocks = []
for ix in range(tower.shape[0]):
block = Object('block',
dimensions=tower[ix, 4:7].numpy().tolist(),
mass=tower[ix, 0].item(),
com=tower[ix, 1:4].numpy().tolist(),
color=(1, 0, 0))
block.pose = Pose(Position(*tower[ix, 7:10].numpy().tolist()),
Quaternion(*tower[ix, 10:14].numpy().tolist()))
blocks.append(block)
return blocks
def ros_to_tower(msg):
""" Converts a list of ROS BodyInfo messages to a tower of Object variables """
tower = []
for ros_blk in msg:
# Unpack the block information
name = ros_blk.name
dims = (ros_blk.dimensions.x, ros_blk.dimensions.y, ros_blk.dimensions.z)
com = (ros_blk.com.x, ros_blk.com.y, ros_blk.com.z)
mass = ros_blk.mass
color = (ros_blk.color.r, ros_blk.color.g, ros_blk.color.b)
blk = Object(name, dims, mass, com, color)
# Now create a Pose namedtuple
pose = ros_to_pose_tuple(ros_blk.pose)
pos = Position(pose[0][0], pose[0][1], pose[0][2])
rot = Quaternion(pose[1][0], pose[1][1], pose[1][2], pose[1][3])
pose_named = Pose(pos=pos, orn=Quaternion(0,0,0,1))
blk.set_pose(pose_named)
blk.rotation = rot
tower.append(blk)
return tower
def validate_regrasps(agent, blocks, base_xy):
blocks = [blocks[0]]
# For each potential starting rotation:
rotations = all_rotations()
pddl_block = agent.pddl_block_lookup[blocks[0].name]
for rx in range(len(rotations)):
rot = rotations[rx]
q = Quaternion(*rot.as_quat())
init_pos, init_orn = pddl_block.get_base_link_pose()
rot_pose = (init_pos, q)
pddl_block.set_base_link_pose(rot_pose)
stable_z = pb_robot.placements.stable_z(pddl_block, agent.table)
agent.execute()
pddl_block.set_base_link_pose(
((init_pos[0], init_pos[1], stable_z), q))
agent.plan()
pddl_block.set_base_link_pose(
((init_pos[0], init_pos[1], stable_z), q))
for gx, goal_rot in enumerate(rotations):
goal_q = Quaternion(*goal_rot.as_quat())
blocks[0].pose = Pose(ZERO_POS, goal_q)
rb = get_rotated_block(blocks[0])
blocks[0].pose = Pose(Position(0, 0, rb.dimensions.z / 2.), goal_q)
blocks[0].rotation = goal_q
success, stable = agent.simulate_tower(blocks,
real=False,
base_xy=(0.5, -0.3),
vis=True,
T=2500,
save_tower=False)
print('Test:', rx, gx, success)
if not success:
input('Continue?')
def make_platform_world(p_block, action):
""" Given a block, create a world that has a platform to push that block off of.
:param block: The Object which to place on the platform.
"""
platform_table, platform_leg = Object.platform()
p_block.set_pose(Pose(ZERO_POS, Quaternion(*action.rot.as_quat())))
block = get_rotated_block(p_block)
block.set_pose(Pose(pos=Position(x=action.pos.x,
y=action.pos.y,
z=platform_table.get_pose().pos.z+platform_table.dimensions.z/2.+block.dimensions.z/2.),
orn=ZERO_ROT))
return World([platform_table, block, platform_leg])
from copy import copy
from random import choices
from numpy.random import uniform
import numpy as np
import torch
from torch.utils.data import DataLoader
from learning.domains.towers.tower_data import TowerDataset, TowerSampler
from block_utils import all_rotations, get_rotated_block, Quaternion, Pose, ZERO_POS
all_quaternions = [Quaternion(*o.as_quat()) for o in list(all_rotations())]
def random_placement(block, tower, discrete=False):
random_orn = choices(all_quaternions, k=1)[0]
block.pose = Pose(ZERO_POS, random_orn)
block = get_rotated_block(block)
# pick random positions for each block
# figure out how far the block can be moved w/o losing contact w/ the block below
if len(tower) > 0:
# place blocks COM above COM of block below
if discrete:
pos_com_xy = np.add(tower[-1].pose.pos[:2], tower[-1].com[:2])
pos_xy = np.subtract(pos_com_xy, block.com[:2])
# randomly place block on top of block below
else:
max_displacement_xy = np.add(tower[-1].dimensions[:2], block.dimensions[:2])/2.
# randomly sample a displacement
rel_xy = uniform(max_displacement_xy, -max_displacement_xy)
# and get the actual positions of the new block
pos_xy = np.add(tower[-1].pose.pos[:2], rel_xy)
def augment(all_data, K_skip, translate=False, mirror=False, vis_tower=False):
datasets = {}
for k_block in all_data.keys():
num_blocks = int(k_block.strip('block'))
#print('Augmenting %d block towers...' % num_blocks)
data = all_data[f'{num_blocks}block']
# load the tower data
towers = data['towers'][::K_skip, :]
labels = data['labels'][::K_skip]
if 'block_ids' in data.keys() and data['block_ids'].shape != (0,):
block_ids = data['block_ids'][::K_skip, :]
N, K, D = towers.shape
# calculate the number of augmented towers that will be created
N_angles = 4
N_shift = 4 if translate else 1
N_mirror = 3 if mirror else 1
tower_multiplier = N_angles * N_mirror * N_shift
N_towers_to_add = N * tower_multiplier
# and create new arrays to store those towers
augmented_towers = np.zeros((N_towers_to_add, K, D))
augmented_labels = np.zeros(N_towers_to_add)
augmented_block_ids = np.zeros((N_towers_to_add, K))
for ix in range(N):
#if ix % 1000 == 0:
#print(ix)
original_tower = [Object.from_vector(towers[ix, jx, :]) for jx in range(num_blocks)]
rot_towers = []
for kx, z_rot in enumerate([0., np.pi/2., np.pi, 3*np.pi/2]):
rot = R.from_rotvec([0., 0., z_rot])
# rotate each block in the tower and add the new tower to the dataset
rot_poses = rot.apply(np.array([b.pose.pos for b in original_tower]))
rot_tower = []
for bx in range(num_blocks):
rot_block = deepcopy(original_tower[bx])
new_pose = Pose(Position(*rot_poses[bx,:].tolist()),
Quaternion(*rot.as_quat().tolist()))
# new_pose = Pose(Position(*rot.apply(block.pose.pos)),
# Quaternion(*rot.as_quat().tolist()))
rot_block.set_pose(new_pose)
orig_rot = R.from_quat(rot_block.rotation)
rot_block = get_rotated_block(rot_block)
rot_block.rotation = (rot*orig_rot).as_quat().tolist()
rot_tower.append(rot_block)
augmented_towers[tower_multiplier*ix + kx, bx, :] = rot_tower[bx].vectorize()
rot_towers.append(rot_tower)
augmented_labels[tower_multiplier*ix + kx] = labels[ix]
if 'block_ids' in data.keys():
augmented_block_ids[tower_multiplier*ix + kx, :] = block_ids[ix, :]
# translate the base block in the tower and add after the rotated blocks
# if translate:
# dx, dy = np.random.uniform(-0.2, 0.2, 2)
# shifted_tower = augmented_towers[(4+N_shift)*ix + kx, :].copy()
# # Indices 7,8 correspond to the pose.
# # The CoM doesn't need to be shifted because it is relative.
# shifted_tower[:, 7] += dx
# shifted_tower[:, 8] += dy
# augmented_towers[tower_multiplier*ix + N_angles + kx, :, :] = shifted_tower
# augmented_labels[tower_multiplier*ix + N_angles + kx] = labels[ix]
# worlds = [World(original_tower)] + [World(tower) for tower in rot_towers]
# env = Environment(worlds, vis_sim=True, vis_frames=True)
# env.step(vis_frames=True)
# input('Next?')
# env.disconnect()
# flip the mirror the COM about the COG and negate the relative position in x
# and y for each block. Creates a new tower that is the mirror of the original
# tower about the x and y axes
if mirror:
start_index = tower_multiplier*ix
# pull out a vector of all the towers we just added. This will be of shape
# [N_angles x num_blocks x D]
rot_towers = augmented_towers[start_index: start_index+N_angles, ...]
# create a vector that will mirror a tower when we multiply it
# indices 1 and 7 correspond to the x coordinates of COM and relative position
# indices 2 and 8 correspond to the y coordinates of COM and relative position
mirror_in_x = np.ones([1,1,D])
mirror_in_y = np.ones([1,1,D])
mirror_in_x[..., [1,7]] *= -1
mirror_in_y[..., [2,8]] *= -1
# add the mirrored towers to the augmented towers dataset
augmented_towers[start_index+N_angles*1 : start_index+N_angles*2, ...] = rot_towers * mirror_in_x
augmented_towers[start_index+N_angles*2 : start_index+N_angles*3, ...] = rot_towers * mirror_in_y
augmented_labels[start_index:start_index+N_angles*N_mirror] = labels[ix]
if 'block_ids' in data.keys():
augmented_block_ids[start_index:start_index+N_angles*N_mirror, :] = block_ids[ix, :]
if vis_tower:
for i in range(tower_multiplier):
print('VISUALIZE', ix*tower_multiplier+i, N_towers_to_add)
new_tower = [Object.from_vector(augmented_towers[ix*tower_multiplier+i, jx, :]) for jx in range(num_blocks)]
w = World(new_tower)
env = Environment([w], vis_sim=True, vis_frames=True)
for tx in range(60):
env.step(vis_frames=False)
time.sleep(1/240.)
env.disconnect()
datasets[f'{num_blocks}block'] = {'towers': augmented_towers,
'labels': augmented_labels}
if 'block_ids' in data.keys():
datasets[f'{num_blocks}block']['block_ids'] = augmented_block_ids
return datasets