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)
# calculate the height of the block
pos_z = tower[-1].pose.pos[2] + block.dimensions[2]
else:
pos_xy = ZERO_POS[:2]
pos_z = block.dimensions[2]
block.pose = Pose((pos_xy[0], pos_xy[1], pos_z), (0,0,0,1))
new_tower = copy(tower)
new_tower.append(block)
return new_tower
def find_entropy_tower(blocks, hypotheses, n_samples=250):
"""
Given a set of candidate hypotheses, find the tower that
has the most disagreement between them.
Currently implemeted as a rejection sampling method that
find the highest entropy predictions amongst the models.
"""
best_tower = None
max_entropy = -1.
for _ in range(n_samples):
num_blocks = np.random.randint(2, len(blocks) + 1)
tower = sample_random_tower(blocks[:num_blocks])
tower = [get_rotated_block(b) for b in tower]
tower = [deepcopy(b) for b in tower]
preds = [h(tower) for h in hypotheses]
e = entropy(preds)
if e > max_entropy:
best_tower = tower
max_entropy = e
best_preds = preds
return best_tower
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 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 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])
def sample_unlabeled_data(n_samples, block_set=None, range_n_blocks=(2, 5)):
""" Generate n_samples random towers. For now each sample can also have
random blocks. We should change this later so that the blocks are fixed
(i.e., chosen elsewhere) and we only sample the configuration.
:param n_samples: Number of random towers to consider.
:param block_set (optional): blocks to use in towers. generate new blocks if None
:return: Dict containining numpy arrays of the towers sorted by size.
"""
# initialize a dictionary of lists to store the generated data
sampled_towers = {}
for i in range(range_n_blocks[0], range_n_blocks[1] + 1):
k = f'{i}block'
sampled_towers[k] = {}
sampled_towers[k]['towers'] = []
sampled_towers[k]['labels'] = []
if block_set is not None:
sampled_towers[k]['block_ids'] = []
min_n_blocks = range_n_blocks[0]
max_n_blocks = min(range_n_blocks[1], len(block_set))
# sample random towers and add them to the lists in the dictionary
for ix in range(n_samples):
n_blocks = np.random.randint(min_n_blocks, max_n_blocks + 1)
# get n_blocks, either from scratch or from the block set
if block_set is not None:
blocks = np.random.choice(block_set, n_blocks, replace=False)
else:
blocks = [Object.random(f'obj_{ix}') for ix in range(n_blocks)]
# sample a new tower
tower = sample_random_tower(blocks)
rotated_tower = [get_rotated_block(b) for b in tower]
# and save that tower in the sampled_towers dict
sampled_towers['%dblock' % n_blocks]['towers'].append(
vectorize(rotated_tower))
if block_set is not None:
block_ids = [block.get_id() for block in rotated_tower]
sampled_towers['%dblock' % n_blocks]['block_ids'].append(block_ids)
# convert all the sampled towers to numpy arrays
for k in sampled_towers.keys():
sampled_towers[k]['towers'] = np.array(sampled_towers[k]['towers'])
sampled_towers[k]['labels'] = np.zeros(
(sampled_towers[k]['towers'].shape[0], ))
if block_set is not None:
sampled_towers[k]['block_ids'] = np.array(
sampled_towers[k]['block_ids'])
return sampled_towers
def active(strategy, vis=False):
hypotheses = get_all_hypotheses()
tp = TowerPlanner(stability_mode='contains')
for nx in range(1, MAX_N):
# Generate a random set of 5 blocks.
blocks = [Object.random(f'obj_{ix}') for ix in range(NUM_BLOCKS)]
# Choose a tower to build.
if strategy == 'random':
num_blocks = np.random.randint(2, NUM_BLOCKS + 1)
tower = sample_random_tower(blocks[:num_blocks])
tower = [get_rotated_block(b) for b in tower]
tower = [deepcopy(b) for b in tower]
elif strategy == 'entropy':
tower = find_entropy_tower(blocks, hypotheses)
else:
raise NotImplementedError()
# Check for consistent models.
valid_hypotheses = []
for h in hypotheses:
true = tp.tower_is_stable(tower)
pred = h(tower)
if true == pred:
valid_hypotheses.append(h)
hypotheses = valid_hypotheses
# Visualize the chosen tower and print the updated hypothesis list.
if vis:
TeleportAgent.simulate_tower(tower, vis=True, T=300)
print(hypotheses)
# Check if true model found.
if len(hypotheses) == 1:
break
return nx
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 sample_next_block(n_samples, bases={}, block_set=None):
# if the dataset is empty, we are sampling the 2-block bases of the towers
if bases == {}:
return sample_unlabeled_data(n_samples,
block_set=block_set,
range_n_blocks=(2, 2))
# if the dataset is non-empty, then for each tower in the dataset we need to sample
# a bunch of options for the next block to be placed on top
else:
assert len(
list(bases.keys())
) == 1, 'I want all the towers to be the same height cuz i\'m rushing'
base_n_blocks_key = list(bases.keys())[0]
base_n_blocks = int(base_n_blocks_key.strip('block'))
n_towers = bases[base_n_blocks_key]['towers'].shape[0]
n_new_blocks_per_base = np.ceil(n_samples / n_towers).astype(int)
new_towers = []
new_block_ids = []
for i in range(n_towers):
# pull out some information about the tower we're working on
current_tower = bases[base_n_blocks_key]['towers'][i]
top_block = current_tower[-1]
top_block_dims = top_block[4:7]
top_block_pos = top_block[7:10]
top_of_tower_height = top_block_pos[2] + top_block_dims[2] / 2
# get n_new_blocks_per_base new blocks to add on top of the tower
if block_set is None:
new_top_blocks = [
Object.random(f'obj_{ix}')
for ix in range(n_new_blocks_per_base)
]
else:
# we can't duplicate blocks, so only choose new top blocks from the
# blocks that aren't already in the base
block_ids_already_in_tower = list(
bases[base_n_blocks_key]['block_ids'][i])
remaining_blocks = [
b for b in block_set
if b.name.strip('obj_') not in block_ids_already_in_tower
]
new_top_blocks = sample_with_replacement(
remaining_blocks, k=n_new_blocks_per_base)
# save the block ids for each of the new towers
new_top_block_ids = [
b.name.strip('obj_') for b in new_top_blocks
]
new_block_ids_local = np.zeros(
[n_new_blocks_per_base, base_n_blocks + 1])
new_block_ids_local[:, :-1] = block_ids_already_in_tower
new_block_ids_local[:, -1] = new_top_block_ids
new_block_ids.append(new_block_ids_local)
# get random rotations for each block
orns = sample_with_replacement(QUATERNIONS,
k=n_new_blocks_per_base)
# apply the rotations to each block
rotated_blocks = []
for orn, block in zip(orns, new_top_blocks):
block.pose = Pose(ZERO_POS, orn)
rotated_blocks.append(get_rotated_block(block))
# figure out how far each block can be moved w/ losing contact w/ the block below
dims_xy = np.array([rb.dimensions for rb in rotated_blocks])[:, :2]
# figure out how far each block can be moved w/ losing contact w/ the block below
max_displacements_xy = (top_block_dims[:2] + dims_xy) / 2.
# sample unscaled noise (clip bceause random normal can exceed -1, 1)
noise_xy = np.clip(0.5 * np.random.randn(n_new_blocks_per_base, 2),
-0.95, 0.95)
# and scale the noise by the max allowed displacement
rel_xy = max_displacements_xy * noise_xy
# and get the actual pos by the difference to the top block pos
pos_xy = top_block_pos[:2] + rel_xy
# calculate the height of each block
pos_z = np.array([
rb.dimensions.z / 2 + top_of_tower_height
for rb in rotated_blocks
])
pos_xyz = pos_xyz = np.hstack([pos_xy, pos_z[:, None]])
for pos, orn, block in zip(pos_xyz, orns, new_top_blocks):
block.pose = Pose(Position(*pos), orn)
block.rotation = orn
# create an array to hold all the new towers
new_towers_local = np.zeros([
n_new_blocks_per_base, current_tower.shape[0] + 1,
current_tower.shape[1]
])
# add the existing base of the tower to the array
new_towers_local[:, :-1] = current_tower
# and add the new top block to each tower
new_towers_local[:, -1] = vectorize(new_top_blocks)
new_towers.append(new_towers_local)
# package the new towers into a dict of the appropriate format. include
# block_ids if we are using a block set
new_towers = np.concatenate(new_towers)
new_samples = {
'towers': new_towers,
'labels': np.zeros(new_towers.shape[0])
}
if block_set is not None:
new_samples['block_ids'] = np.concatenate(new_block_ids)
return {f'{base_n_blocks+1}block': new_samples}
def sample_sequential_data(block_set, dataset, n_samples):
""" Generate n_samples random towers. Each tower has the property that its
base (the tower until the last block) is stable. To ensure this, we start
with all the stable towers plus base cases of single block towers.
:param block_set: List of blocks that can be used in the towers.
:param dataset: List of current towers that have been built.
:param n_samples: Number of random towers to consider.
:param max_blocks
:return: Dict containining numpy arrays of the towers sorted by size.
"""
print('Generating data sequentially...')
keys = ['2block', '3block', '4block', '5block']
# initialize a dictionary of lists to store the generated data
sampled_towers = {k: {} for k in keys}
for k in keys:
sampled_towers[k]['towers'] = []
sampled_towers[k]['labels'] = []
if block_set is not None:
sampled_towers[k]['block_ids'] = []
# Gather list of all stable towers. Towers should be of blocks that are rotated in "Block" format.
stable_towers = []
# Get all single block stable towers.
for block in block_set:
for orn in QUATERNIONS:
new_block = deepcopy(block)
new_block.pose = Pose(ZERO_POS, orn)
rot_block = get_rotated_block(new_block)
rot_block.pose = Pose((0., 0., rot_block.dimensions.z / 2.),
(0, 0, 0, 1))
stable_towers.append([rot_block])
# Get all stable towers from the dataset.
for k in keys[:3]:
if dataset is None:
break
tower_tensors = unprocess(
dataset.tower_tensors[k].cpu().numpy().copy())
tower_labels = dataset.tower_labels[k]
for ix, (tower_vec,
tower_label) in enumerate(zip(tower_tensors, tower_labels)):
if tower_label == 1:
block_tower = []
for bx in range(tower_vec.shape[0]):
block = Object.from_vector(tower_vec[bx, :])
if block_set is not None:
block.name = 'obj_' + str(
int(dataset.tower_block_ids[k][ix, bx]))
block_tower.append(block)
stable_towers.append(block_tower)
# maintain block info my block name
# TODO: this will NOT work if using random blocks
block_lookup = {}
for block in block_set:
block_lookup[block.name] = block
# Sample random towers by randomly choosing a stable base then trying to add a block.
for ix in range(n_samples):
# Choose a stable base.
tower_ix = np.random.choice(np.arange(0, len(stable_towers)))
base_tower = stable_towers[tower_ix]
# Choose a block that's not already in the tower.
remaining_blocks = {}
for k in block_lookup:
used = False
for block in base_tower:
if k == block.name:
used = True
if not used:
remaining_blocks[k] = block_lookup[k]
# if we switch block sets during training then the remaining_blocks list
# will be longer than block_set - len(base_tower)
#assert(len(remaining_blocks) == len(block_set) - len(base_tower))
new_block = deepcopy(np.random.choice(list(remaining_blocks.values())))
# Choose an orientation.
orn = QUATERNIONS[np.random.choice(np.arange(0, len(QUATERNIONS)))]
new_block.pose = Pose(ZERO_POS, orn)
rot_block = get_rotated_block(new_block)
# Sample a displacement.
base_dims = np.array(base_tower[-1].dimensions)[:2]
new_dims = np.array(rot_block.dimensions)[:2]
max_displacements_xy = (base_dims + new_dims) / 2.
noise_xy = np.clip(0.5 * np.random.randn(2), -0.95, 0.95)
rel_xy = max_displacements_xy * noise_xy
# Calculate the new pose.
base_pos = np.array(base_tower[-1].pose.pos)[:2]
pos_xy = base_pos + rel_xy
pos_z = np.sum([b.dimensions.z
for b in base_tower]) + rot_block.dimensions.z / 2.
rot_block.pose = Pose((pos_xy[0], pos_xy[1], pos_z), (0, 0, 0, 1))
# Add block to tower.
new_tower = base_tower + [rot_block]
if False:
w = World(new_tower)
env = Environment([w], vis_sim=True, vis_frames=True)
for tx in range(240):
env.step(vis_frames=True)
time.sleep(1 / 240.)
env.disconnect()
# Save that tower in the sampled_towers dict
n_blocks = len(new_tower)
sampled_towers['%dblock' % n_blocks]['towers'].append(
vectorize(new_tower))
# save block id
if block_set is not None:
block_ids = [block.get_id() for block in new_tower]
sampled_towers['%dblock' % n_blocks]['block_ids'].append(block_ids)
# convert all the sampled towers to numpy arrays
for k in keys:
sampled_towers[k]['towers'] = np.array(sampled_towers[k]['towers'])
if sampled_towers[k]['towers'].shape[0] == 0:
sampled_towers[k]['towers'] = sampled_towers[k]['towers'].reshape(
(0, int(k[0]), 21))
sampled_towers[k]['labels'] = np.zeros(
(sampled_towers[k]['towers'].shape[0], ))
if block_set is not None:
sampled_towers[k]['block_ids'] = np.array(
sampled_towers[k]['block_ids'])
return sampled_towers
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
