162 lines
5.1 KiB
Python
162 lines
5.1 KiB
Python
import time
|
|
|
|
import numpy as np
|
|
import pybullet as p
|
|
import pybullet_data
|
|
|
|
from sensor_msgs.msg import JointState
|
|
|
|
from model import *
|
|
from utils import *
|
|
|
|
|
|
class PandaArm:
|
|
def __init__(self):
|
|
self.nr_dof = 7
|
|
|
|
self.named_joint_values = {"ready": [0.0, -0.79, 0.0, -2.356, 0.0, 1.57, 0.79]}
|
|
|
|
self.upper_limits = [-7] * self.nr_dof
|
|
self.lower_limits = [7] * self.nr_dof
|
|
self.ranges = [7] * self.nr_dof
|
|
|
|
def load(self, pose):
|
|
self.T_world_base = pose
|
|
self.uid = p.loadURDF(
|
|
"assets/urdfs/panda/panda_arm_hand.urdf",
|
|
basePosition=pose.translation,
|
|
baseOrientation=pose.rotation.as_quat(),
|
|
useFixedBase=True,
|
|
)
|
|
for i, q_i in enumerate(self.named_joint_values["ready"]):
|
|
p.resetJointState(self.uid, i, q_i)
|
|
|
|
def get_state(self):
|
|
joint_states = p.getJointStates(self.uid, range(p.getNumJoints(self.uid)))[:7]
|
|
positions = [state[0] for state in joint_states]
|
|
velocities = [state[1] for state in joint_states]
|
|
return positions, velocities
|
|
|
|
def set_desired_joint_positions(self, positions):
|
|
for i, position in enumerate(positions):
|
|
p.setJointMotorControl2(self.uid, i, p.POSITION_CONTROL, position)
|
|
|
|
def set_desired_joint_velocities(self, velocities):
|
|
for i, velocity in enumerate(velocities):
|
|
p.setJointMotorControl2(
|
|
self.uid, i, p.VELOCITY_CONTROL, targetVelocity=velocity
|
|
)
|
|
|
|
|
|
class PandaGripper:
|
|
def move(self, width):
|
|
for i in [9, 10]:
|
|
p.setJointMotorControl2(
|
|
self.uid,
|
|
i,
|
|
p.POSITION_CONTROL,
|
|
0.5 * width,
|
|
force=10,
|
|
)
|
|
|
|
def read(self):
|
|
return p.getJointState(self.uid, 9)[0] + p.getJointState(self.uid, 10)[0]
|
|
|
|
|
|
class Camera:
|
|
def __init__(self, uid, link_id, width, height, hfov, near, far):
|
|
self.uid = uid
|
|
self.link_id = link_id
|
|
|
|
f = width / (2.0 * np.tan(hfov / 2.0))
|
|
cx, cy = width / 2.0, height / 2.0
|
|
|
|
self.width = width
|
|
self.height = height
|
|
self.K = [f, 0.0, cx, 0.0, f, cy, 0.0, 0.0, 1.0]
|
|
self.near = near
|
|
self.far = far
|
|
|
|
self.proj_mat = p.computeProjectionMatrixFOV(
|
|
np.rad2deg(hfov), width / height, near, far
|
|
)
|
|
|
|
def update(self):
|
|
result = p.getLinkState(self.uid, self.link_id, computeForwardKinematics=1)
|
|
t_world_cam = result[4]
|
|
R_world_cam = Rotation.from_quat(result[5])
|
|
eye = t_world_cam
|
|
target = R_world_cam.apply(np.r_[0.0, 0.0, 1.0]) + t_world_cam
|
|
up = R_world_cam.apply(np.r_[0.0, -1.0, 0.0])
|
|
view_mat = p.computeViewMatrix(eye, target, up)
|
|
|
|
result = p.getCameraImage(
|
|
self.width,
|
|
self.height,
|
|
view_mat,
|
|
self.proj_mat,
|
|
renderer=p.ER_BULLET_HARDWARE_OPENGL,
|
|
)
|
|
color_img, z_buffer = result[2][:, :, :3], result[3]
|
|
|
|
color_img = color_img.astype(np.uint8)
|
|
depth_img = (
|
|
self.far * self.near / (self.far - (self.far - self.near) * z_buffer)
|
|
).astype(np.float32)
|
|
|
|
return color_img, depth_img
|
|
|
|
|
|
class SimPandaEnv(object):
|
|
def __init__(self, gui, publish_state=True):
|
|
self.gui = gui
|
|
self.rate = 240
|
|
p.connect(p.GUI if gui else p.DIRECT)
|
|
p.setAdditionalSearchPath(pybullet_data.getDataPath())
|
|
p.setTimeStep(1.0 / self.rate)
|
|
p.setGravity(0.0, 0.0, -9.8)
|
|
self.step_cnt = 0
|
|
|
|
p.loadURDF("plane.urdf")
|
|
p.loadURDF(
|
|
"table/table.urdf",
|
|
baseOrientation=Rotation.from_rotvec(np.array([0, 0, np.pi / 2])).as_quat(),
|
|
useFixedBase=True,
|
|
)
|
|
p.loadURDF("cube_small.urdf", [0.0, 0.0, 0.8])
|
|
|
|
self.arm = PandaArm()
|
|
self.arm.load(Transform(Rotation.identity(), np.r_[-0.8, 0.0, 0.4]))
|
|
self.gripper = PandaGripper()
|
|
self.gripper.uid = self.arm.uid
|
|
self.model = Model("panda_link0", "panda_ee")
|
|
self.camera = Camera(self.arm.uid, 12, 320, 240, 1.047, 0.2, 2.0)
|
|
|
|
if publish_state:
|
|
self.state_pub = rospy.Publisher("/joint_states", JointState, queue_size=10)
|
|
rospy.Timer(rospy.Duration(1.0 / 30), self._publish_state)
|
|
|
|
def forward(self, dt):
|
|
steps = int(dt * self.rate)
|
|
for _ in range(steps):
|
|
self._step()
|
|
|
|
def _step(self):
|
|
if self.step_cnt % int(self.rate / self.controller.rate) == 0:
|
|
self.controller.update()
|
|
p.stepSimulation()
|
|
time.sleep(1.0 / self.rate)
|
|
|
|
def _publish_state(self, event):
|
|
positions, velocities = self.arm.get_state()
|
|
width = self.gripper.read()
|
|
msg = JointState()
|
|
msg.header.stamp = rospy.Time.now()
|
|
msg.name = ["panda_joint{}".format(i) for i in range(1, 8)] + [
|
|
"panda_finger_joint1",
|
|
"panda_finger_joint2",
|
|
]
|
|
msg.position = np.r_[positions, 0.5 * width, 0.5 * width]
|
|
msg.velocity = velocities
|
|
self.state_pub.publish(msg)
|