spam.py

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