Skip to content

Instantly share code, notes, and snippets.

@wkentaro
Created July 1, 2020 23:59
Show Gist options
  • Save wkentaro/97a51c3e4669df080cae796de7275f16 to your computer and use it in GitHub Desktop.
Save wkentaro/97a51c3e4669df080cae796de7275f16 to your computer and use it in GitHub Desktop.
import contrib
import numpy as np
import pickpp
env = contrib.RlightPickAndPlaceEnv(
version="v3", num_instances=1, train=False
).env
env.launch()
env._pyrep.register_cinematic("cam_cinematic_placeholder")
for _ in range(10):
obs = env.reset()
mask = obs["mask"] == obs["objects"][0, 0]
pcd = pickpp.geometry.pointcloud_from_depth(
obs["depth"], *obs["intrinsic"]
)
pcd_masked = pcd[mask]
pos, qua = obs["extrinsic"][:3], obs["extrinsic"][3:]
T_cam2robot = pickpp.geometry.transformation_matrix(qua, pos)
pcd_masked = pickpp.geometry.transform_points(pcd_masked, T_cam2robot)
pos, qua = obs["robot"][:3], obs["robot"][3:]
T_ee2robot = pickpp.geometry.transformation_matrix(qua, pos)
pcd_masked = pickpp.geometry.transform_points(
pcd_masked, np.linalg.inv(T_ee2robot)
)
normals = pickpp.geometry.normals_from_pointcloud(pcd_masked)
index = np.random.randint(len(pcd_masked))
position = pcd_masked[index]
normal = normals[index]
quaternion = pickpp.geometry.quaternion_from_vec2vec([0, 0, 1], normal)
coord = pickpp.geometry.Coordinate(
position=position, quaternion=quaternion
)
coord.translate([0, 0, -0.05], wrt="local")
coord.euler = [coord.euler[0], coord.euler[1], np.pi]
T_robot2world = env._robot.get_transform()
coord.transform(T_ee2robot, wrt="world")
try:
with env._robot.ik_constraint():
path = env._robot.get_path(
position=coord.position,
euler=coord.euler,
ignore_collisions=True,
relative_to=env._robot,
)
while not path.step():
env._pyrep.step()
env._go_down_grasp()
env._go_up(ignore_collisions=True)
env._place(pickpp.extra.pyrep.Shape("target_box"))
except Exception:
pass
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment