pybullet_panda_pick_and_place.py

import itertools
import random
import time

try:
    import numpy as np
    import pybullet as p
    import pybullet_data
except ImportError:
    print("Run `pip install numpy pybullet`")
    quit()


TIMESTEP = 1 / 240  # pybullet is 240hz by default


class PandaPybulletInterface:
    def __init__(self):
        self.robot = p.loadURDF("franka_panda/panda.urdf", useFixedBase=True)
        self.homej = tuple([0, -np.pi / 4, 0, -np.pi / 2, 0, np.pi / 4, np.pi / 4, 0, 0])

        n_joints = p.getNumJoints(self.robot)

        joints = [p.getJointInfo(self.robot, i) for i in range(n_joints)]
        self.joints = [j[0] for j in joints if j[2] != p.JOINT_FIXED]

        self.ee = n_joints - 1  # end-effector

        self.setj(self.homej)

    def get_ee_pose(self):
        ee_to_world = p.getLinkState(self.robot, self.ee)[:2]
        return ee_to_world

    def setj(self, joint_positions):
        for joint, joint_position in zip(self.joints, joint_positions):
            p.resetJointState(self.robot, joint, joint_position)

    def getj(self):
        joint_positions = []
        for joint in self.joints:
            joint_positions.append(p.getJointState(self.robot, joint)[0])
        return joint_positions

    def movej(self, targj, speed=0.01, timeout=5):
        assert len(targj) == len(self.joints)
        for i in itertools.count():
            currj = [p.getJointState(self.robot, i)[0] for i in self.joints]
            currj = np.array(currj)
            diffj = targj - currj
            if all(np.abs(diffj) < 1e-2):
                return

            # Move with constant velocity
            norm = np.linalg.norm(diffj)
            v = diffj / norm if norm > 0 else 0
            stepj = currj + v * speed
            gains = np.ones(len(self.joints))
            p.setJointMotorControlArray(
                bodyIndex=self.robot,
                jointIndices=self.joints,
                controlMode=p.POSITION_CONTROL,
                targetPositions=stepj,
                positionGains=gains,
            )
            yield i

            if i >= (timeout / TIMESTEP):
                print("timeout in joint motor control")
                break

    def solve_ik(self, pose):
        n_joints = p.getNumJoints(self.robot)
        lower_limits = []
        upper_limits = []
        for i in range(n_joints):
            joint_info = p.getJointInfo(self.robot, i)
            lower_limits.append(joint_info[8])
            upper_limits.append(joint_info[9])
        joint_positions = p.calculateInverseKinematics(
            bodyUniqueId=self.robot,
            endEffectorLinkIndex=self.ee,
            targetPosition=pose[0],
            targetOrientation=pose[1],
            lowerLimits=lower_limits,
            upperLimits=upper_limits,
            restPoses=self.homej,
            maxNumIterations=1000,
            residualThreshold=1e-5,
        )
        return joint_positions

    def grasp(self):
        j = self.getj()
        j[-2:] = [0, 0]
        yield from self.movej(j, speed=0.001, timeout=1)

    def ungrasp(self):
        j = self.getj()
        j[-2:] = [0.04, 0.04]
        yield from self.movej(j, speed=0.001, timeout=1)


def pick_and_place(pi, cube):
    for _ in pi.ungrasp():
        p.stepSimulation()
        time.sleep(TIMESTEP)

    ee_to_world = pi.get_ee_pose()  # (position, quaternion)
    cube_to_world = p.getBasePositionAndOrientation(cube)

    # grasp pose
    ee_target1_to_world = cube_to_world[0], ee_to_world[1]

    # pre-grasp pose
    ee_target2_to_world = p.multiplyTransforms(
        *([0, 0, 0.1], [0, 0, 0, 1]),
        *ee_target1_to_world
    )

    j = pi.solve_ik(ee_target2_to_world)
    for _ in pi.movej(j):
        p.stepSimulation()
        time.sleep(TIMESTEP)
    j_pre_grasp = j

    time.sleep(0.2)

    j = pi.solve_ik(ee_target1_to_world)
    for _ in pi.movej(j):
        p.stepSimulation()
        time.sleep(TIMESTEP)
    j_grasp = j

    time.sleep(0.2)

    for _ in pi.grasp():
        p.stepSimulation()
        time.sleep(TIMESTEP)

    j = list(j_pre_grasp)
    j[-2:] = pi.getj()[-2:]
    for _ in pi.movej(j):
        p.stepSimulation()
        time.sleep(TIMESTEP)

    j = list(pi.homej)
    j[-2:] = pi.getj()[-2:]  # keep grasping
    for _ in pi.movej(j):
        p.stepSimulation()
        time.sleep(TIMESTEP)

    time.sleep(0.2)

    j = list(j_pre_grasp)
    j[-2:] = pi.getj()[-2:]
    for _ in pi.movej(j):
        p.stepSimulation()
        time.sleep(TIMESTEP)

    j = list(j_grasp)
    j[-2:] = pi.getj()[-2:]
    for _ in pi.movej(j):
        p.stepSimulation()
        time.sleep(TIMESTEP)

    time.sleep(0.2)

    for _ in pi.ungrasp():
        p.stepSimulation()
        time.sleep(TIMESTEP)

    time.sleep(0.2)

    j = list(j_pre_grasp)
    j[-2:] = pi.getj()[-2:]
    for _ in pi.movej(j):
        p.stepSimulation()
        time.sleep(TIMESTEP)

    j = list(pi.homej)
    j[-2:] = pi.getj()[-2:]
    for _ in pi.movej(j):
        p.stepSimulation()
        time.sleep(TIMESTEP)


def create_cube(pose, color):
    cube = p.loadURDF("cube.urdf", globalScaling=0.05)
    p.changeDynamics(cube, -1, mass=0.1)
    p.resetVisualShapeData(cube, -1, rgbaColor=color)
    p.resetBasePositionAndOrientation(cube, *pose)
    return cube


def main():
    p.connect(p.GUI)
    p.setGravity(0, 0, -9.8)

    p.resetDebugVisualizerCamera(
        cameraDistance=1.5,
        cameraYaw=70,
        cameraPitch=-15,
        cameraTargetPosition=(0, 0, 0.5),
    )
    p.configureDebugVisualizer(p.COV_ENABLE_GUI, False)

    p.setAdditionalSearchPath(pybullet_data.getDataPath())
    p.loadURDF("plane.urdf")

    pi = PandaPybulletInterface()

    cube1 = create_cube(((0.4, 0.2, 0.025), (0, 0, 0, 1)), color=(1, 0, 0, 1))
    cube2 = create_cube(((0.3, -0.1, 0.025), (0, 0, 0, 1)), color=(0, 1, 0, 1))
    cube3 = create_cube(((0.5, 0, 0.025), (0, 0, 0, 1)), color=(0, 0, 1, 1))

    while True:
        cube = random.choice([cube1, cube2, cube3])
        pick_and_place(pi, cube)


if __name__ == "__main__":
    main()

2021-10-13_19-31-43.mp4