def set_marker_position_yaw(self, pos, yaw):
"""
Set subgoal marker position and orientation
:param pos: position
:param yaw: yaw angle
"""
quat = quatToXYZW(seq='wxyz', orn=euler.euler2quat(0, -np.pi / 2, yaw))
self.marker.set_position(pos)
self.marker_direction.set_position_orientation(pos, quat)
def reset(self):
if self.robot_ids is None:
self._load_model()
self.robot_body.reset_orientation(
quatToXYZW(euler2quat(*self.config["initial_orn"]), 'wxyz'))
self.robot_body.reset_position(self.config["initial_pos"])
self.reset_random_pos()
self.robot_specific_reset()
state = self.calc_state()
return state
def new_orientation_from_dir(orn, next_dir):
initial_dir = np.array([1, 0])
cos = np.dot(initial_dir, normalize(next_dir))
sin = np.cross(initial_dir, normalize(next_dir))
rad = np.arccos(cos)
if sin < 0:
rad = -rad
# print(rad)
next_orn = quatToXYZW(axangle2quat(np.array([0, 0, 1]), rad, is_normalized=True), 'wxyz')
return next_orn
def set_pos_orn_with_z_offset(self, obj, pos, orn=None, offset=None):
"""
Reset position and orientation for the robot or the object
:param obj: an instance of robot or object
:param pos: position
:param orn: orientation
:param offset: z offset
"""
if orn is None:
orn = np.array([0, 0, np.random.uniform(0, np.pi * 2)])
if offset is None:
offset = self.initial_pos_z_offset
obj.set_position_orientation([pos[0], pos[1], pos[2] + offset],
quatToXYZW(euler2quat(*orn), 'wxyz'))
def reset_initial_and_target_pos(self):
collision_links = [-1]
while (-1 in collision_links): # if collision happens restart
pos = [np.random.uniform(4, 5), 0, 0]
self.robots[0].set_position(pos=[pos[0], pos[1], pos[2] + 0.1])
self.robots[0].set_orientation(orn=quatToXYZW(
euler2quat(0, 0, np.random.uniform(0, np.pi * 2)), 'wxyz'))
collision_links = []
for _ in range(self.simulator_loop):
self.simulator_step()
collision_links += [
item[3] for item in p.getContactPoints(
bodyA=self.robots[0].robot_ids[0])
]
collision_links = np.unique(collision_links)
self.initial_pos = pos
self.target_pos = [np.random.uniform(-5, -4), 0, 0]
def reset_initial_and_target_pos(self):
collision_links = [-1]
while -1 in collision_links: # if collision happens, reinitialize
floor, pos = self.scene.get_random_point()
self.robots[0].set_position(pos=[pos[0], pos[1], pos[2] + 0.1])
self.robots[0].set_orientation(orn=quatToXYZW(
euler2quat(0, 0, np.random.uniform(0, np.pi * 2)), 'wxyz'))
collision_links = []
for _ in range(self.simulator_loop):
self.simulator_step()
collision_links += [
item[3] for item in p.getContactPoints(
bodyA=self.robots[0].robot_ids[0])
]
collision_links = np.unique(collision_links)
self.initial_pos = pos
dist = 0.0
while dist < 1.0: # if initial and target positions are < 1 meter away from each other, reinitialize
_, self.target_pos = self.scene.get_random_point_floor(
floor, self.random_height)
dist = l2_distance(self.initial_pos, self.target_pos)
def set_pos_orn_with_z_offset(self, obj, pos, orn=None, offset=None):
"""
Reset position and orientation for the robot or the object
:param obj: an instance of robot or object
:param pos: position
:param orn: orientation
:param offset: z offset
"""
if orn is None:
orn = np.array([0, 0, np.random.uniform(0, np.pi * 2)])
if offset is None:
offset = self.initial_pos_z_offset
is_robot = isinstance(obj, BaseRobot)
body_id = obj.robot_ids[0] if is_robot else obj.body_id
# first set the correct orientation
obj.set_position_orientation(pos, quatToXYZW(euler2quat(*orn), 'wxyz'))
# compute stable z based on this orientation
stable_z = stable_z_on_aabb(body_id, [pos, pos])
# change the z-value of position with stable_z + additional offset
# in case the surface is not perfect smooth (has bumps)
obj.set_position([pos[0], pos[1], stable_z + offset])
def reset_initial_and_target_pos(self):
self.robots[0].set_position(pos=self.initial_pos)
self.robots[0].set_orientation(
orn=quatToXYZW(euler2quat(*self.initial_orn), 'wxyz'))
def main():
step_per_sec = 100
num_directions = 12
obj_count = 0
root_dir = '/cvgl2/u/chengshu/ig_dataset_v5/objects'
s = Simulator(mode='headless',
image_width=512,
image_height=512,
physics_timestep=1 / float(step_per_sec))
p.setGravity(0.0, 0.0, 0.0)
for obj_class_dir in sorted(os.listdir(root_dir)):
obj_class_dir = os.path.join(root_dir, obj_class_dir)
for obj_inst_dir in os.listdir(obj_class_dir):
obj_inst_name = obj_inst_dir
urdf_path = obj_inst_name + '.urdf'
obj_inst_dir = os.path.join(obj_class_dir, obj_inst_dir)
urdf_path = os.path.join(obj_inst_dir, urdf_path)
obj = ArticulatedObject(urdf_path)
s.import_object(obj)
with open(os.path.join(obj_inst_dir, 'misc/bbox.json'),
'r') as bbox_file:
bbox_data = json.load(bbox_file)
bbox_max = np.array(bbox_data['max'])
bbox_min = np.array(bbox_data['min'])
offset = -(bbox_max + bbox_min) / 2.0
z = stable_z_on_aabb(obj.body_id, [[0, 0, 0], [0, 0, 0]])
obj.set_position([offset[0], offset[1], z])
_, extent = get_center_extent(obj.body_id)
max_half_extent = max(extent) / 2.0
px = max_half_extent * 3.0
py = 0.0
pz = extent[2] / 2.0
camera_pose = np.array([px, py, pz])
s.renderer.set_camera(camera_pose, [0, 0, pz], [0, 0, 1])
num_joints = p.getNumJoints(obj.body_id)
if num_joints == 0:
s.reload()
continue
# collect joint low/high limit
joint_low = []
joint_high = []
for j in range(num_joints):
j_low, j_high = p.getJointInfo(obj.body_id, j)[8:10]
joint_low.append(j_low)
joint_high.append(j_high)
# set joints to their lowest limits
for j, j_low in zip(range(num_joints), joint_low):
p.resetJointState(obj.body_id,
j,
targetValue=j_low,
targetVelocity=0.0)
s.sync()
# render the images
joint_low_imgs = []
for i in range(num_directions):
yaw = np.pi * 2.0 / num_directions * i
obj.set_orientation(
quatToXYZW(euler2quat(0.0, 0.0, yaw), 'wxyz'))
s.sync()
rgb, three_d = s.renderer.render(modes=('rgb', '3d'))
depth = -three_d[:, :, 2]
rgb[depth == 0] = 1.0
joint_low_imgs.append(
Image.fromarray((rgb[:, :, :3] * 255).astype(np.uint8)))
# set joints to their highest limits
for j, j_high in zip(range(num_joints), joint_high):
p.resetJointState(obj.body_id,
j,
targetValue=j_high,
targetVelocity=0.0)
s.sync()
# render the images
joint_high_imgs = []
for i in range(num_directions):
yaw = np.pi * 2.0 / num_directions * i
obj.set_orientation(
quatToXYZW(euler2quat(0.0, 0.0, yaw), 'wxyz'))
s.sync()
rgb, three_d = s.renderer.render(modes=('rgb', '3d'))
depth = -three_d[:, :, 2]
rgb[depth == 0] = 1.0
joint_high_imgs.append(
Image.fromarray((rgb[:, :, :3] * 255).astype(np.uint8)))
# concatenate the images
imgs = []
for im1, im2 in zip(joint_low_imgs, joint_high_imgs):
dst = Image.new('RGB', (im1.width + im2.width, im1.height))
dst.paste(im1, (0, 0))
dst.paste(im2, (im1.width, 0))
imgs.append(dst)
gif_path = '{}/visualizations/{}_joint_limit.gif'.format(
obj_inst_dir, obj_inst_name)
# save the gif
imgs[0].save(gif_path,
save_all=True,
append_images=imgs[1:],
optimize=True,
duration=200,
loop=0)
s.reload()
obj_count += 1
print(obj_count, gif_path)