Skip to content

Instantly share code, notes, and snippets.

@wkentaro
Created October 13, 2021 18:30
Show Gist options
  • Save wkentaro/ff78666ce73323387c7fda1eee451db2 to your computer and use it in GitHub Desktop.
Save wkentaro/ff78666ce73323387c7fda1eee451db2 to your computer and use it in GitHub Desktop.
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()
@wkentaro
Copy link
Author

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment